github-actions[bot] commented on code in PR #67325:
URL: https://github.com/apache/doris/pull/67325#discussion_r3932951847
##########
be/src/format_v2/table/lance_reader.cpp:
##########
@@ -61,6 +65,250 @@ constexpr std::string_view DISTANCE_COLUMN = "_distance";
constexpr std::string_view ROW_ID_COLUMN = "_rowid";
constexpr std::string_view ARROW_EXTENSION_NAME = "ARROW:extension:name";
constexpr const char* LANCE_READER_PROFILE = "LanceReader";
+constexpr std::string_view ARROW_JSON_EXTENSION = "arrow.json";
+constexpr std::string_view LANCE_JSON_EXTENSION = "lance.json";
+constexpr std::string_view LANCE_BFLOAT16_EXTENSION = "lance.bfloat16";
+
+enum class LanceExtensionKind {
+ NONE,
+ JSON,
+ BFLOAT16,
+};
+
+// Extracts and validates extension names from metadata and registered types.
+Status get_lance_extension(const std::shared_ptr<arrow::Field>& field,
+ LanceExtensionKind* extension_kind,
+ std::shared_ptr<arrow::DataType>* storage_type) {
+ DORIS_CHECK(field != nullptr);
+ DORIS_CHECK(extension_kind != nullptr);
+ DORIS_CHECK(storage_type != nullptr);
+ *extension_kind = LanceExtensionKind::NONE;
+ *storage_type = field->type();
+
+ std::optional<std::string> extension_name;
+ if (field->HasMetadata()) {
+ const auto metadata_name =
field->metadata()->Get(ARROW_EXTENSION_NAME);
+ if (metadata_name.ok() && !metadata_name.ValueUnsafe().empty()) {
+ extension_name = metadata_name.ValueUnsafe();
+ }
+ }
+ if (field->type()->id() == arrow::Type::EXTENSION) {
+ const auto extension_type =
std::dynamic_pointer_cast<arrow::ExtensionType>(field->type());
+ if (extension_type == nullptr) {
+ return Status::InvalidArgument("invalid Arrow extension type for
Lance field '{}'",
+ field->name());
+ }
+ if (extension_name.has_value() && *extension_name !=
extension_type->extension_name()) {
+ return Status::InvalidArgument(
+ "conflicting Arrow extension names for Lance field '{}':
'{}' and '{}'",
+ field->name(), *extension_name,
extension_type->extension_name());
+ }
+ extension_name = extension_type->extension_name();
+ *storage_type = extension_type->storage_type();
+ }
+ if (!extension_name.has_value()) {
+ return Status::OK();
+ }
+
+ if (*extension_name == ARROW_JSON_EXTENSION) {
+ if ((*storage_type)->id() != arrow::Type::STRING &&
+ (*storage_type)->id() != arrow::Type::LARGE_STRING) {
+ return Status::NotSupported(
+ "Arrow JSON extension for Lance field '{}' requires UTF8
storage, got {}",
+ field->name(), (*storage_type)->ToString());
+ }
+ *extension_kind = LanceExtensionKind::JSON;
+ return Status::OK();
+ }
+ if (*extension_name == LANCE_JSON_EXTENSION) {
+ if ((*storage_type)->id() != arrow::Type::LARGE_BINARY) {
+ return Status::NotSupported(
+ "Lance JSON extension for field '{}' requires LARGE_BINARY
storage, got {}",
+ field->name(), (*storage_type)->ToString());
+ }
+ *extension_kind = LanceExtensionKind::JSON;
+ return Status::OK();
+ }
+ if (*extension_name == LANCE_BFLOAT16_EXTENSION) {
+ if ((*storage_type)->id() != arrow::Type::FIXED_SIZE_BINARY ||
+
std::static_pointer_cast<arrow::FixedSizeBinaryType>(*storage_type)->byte_width()
!=
+ 2) {
+ return Status::NotSupported(
+ "Lance BFloat16 extension for field '{}' requires
FIXED_SIZE_BINARY(2) "
+ "storage, got {}",
+ field->name(), (*storage_type)->ToString());
+ }
+ *extension_kind = LanceExtensionKind::BFLOAT16;
+ return Status::OK();
+ }
+ return Status::NotSupported("unsupported Lance Arrow extension type '{}'
for field '{}'",
+ *extension_name, field->name());
+}
+
+// Determines whether a field subtree contains BFloat16 values or extension
arrays to normalize.
+Status field_requires_lance_normalization(const std::shared_ptr<arrow::Field>&
field,
+ bool* requires_normalization) {
+ DORIS_CHECK(field != nullptr);
+ DORIS_CHECK(requires_normalization != nullptr);
+
+ LanceExtensionKind extension_kind;
+ std::shared_ptr<arrow::DataType> storage_type;
+ RETURN_IF_ERROR(get_lance_extension(field, &extension_kind,
&storage_type));
+ bool required = extension_kind == LanceExtensionKind::BFLOAT16 ||
+ field->type()->id() == arrow::Type::EXTENSION;
+ for (const auto& child : storage_type->fields()) {
+ bool child_required = false;
+ RETURN_IF_ERROR(field_requires_lance_normalization(child,
&child_required));
+ required |= child_required;
+ }
+ *requires_normalization = required;
+ return Status::OK();
+}
+
+// Widens little-endian Lance BFloat16 values to Arrow Float32 without
precision loss.
+Status convert_bfloat16_array(const std::shared_ptr<arrow::Array>& array,
+ std::shared_ptr<arrow::Array>* normalized) {
+ DORIS_CHECK(array != nullptr);
+ DORIS_CHECK(normalized != nullptr);
+ const auto fixed_binary =
std::dynamic_pointer_cast<arrow::FixedSizeBinaryArray>(array);
+ if (fixed_binary == nullptr || fixed_binary->byte_width() != 2) {
+ return Status::InvalidArgument("invalid Lance BFloat16 array storage:
{}",
+ array->type()->ToString());
+ }
+ if (config::enable_arrow_input_validation) {
+ check_arrow_fixed_width_buffer(*fixed_binary, sizeof(uint16_t));
+ }
+
+ arrow::FloatBuilder builder;
+ auto arrow_status = builder.Reserve(fixed_binary->length());
+ if (!arrow_status.ok()) {
+ return Status::InternalError("reserve Lance BFloat16 output failed:
{}",
+ arrow_status.message());
+ }
+ for (int64_t row = 0; row < fixed_binary->length(); ++row) {
+ if (fixed_binary->IsNull(row)) {
+ arrow_status = builder.AppendNull();
+ } else {
+ const auto bits =
LittleEndian::Load16(fixed_binary->GetValue(row));
+ arrow_status =
builder.Append(std::bit_cast<float>(static_cast<uint32_t>(bits) << 16));
+ }
+ if (!arrow_status.ok()) {
+ return Status::InternalError("append Lance BFloat16 value failed:
{}",
+ arrow_status.message());
+ }
+ }
+ std::shared_ptr<arrow::FloatArray> result;
+ arrow_status = builder.Finish(&result);
+ if (!arrow_status.ok()) {
+ return Status::InternalError("finish Lance BFloat16 conversion failed:
{}",
+ arrow_status.message());
+ }
+ *normalized = std::move(result);
+ return Status::OK();
+}
+
+// Normalizes nested BFloat16 arrays while preserving offsets and null bitmaps.
+Status normalize_lance_arrow_array(const std::shared_ptr<arrow::Field>& field,
+ const std::shared_ptr<arrow::Array>& array,
+ std::shared_ptr<arrow::Array>* normalized) {
+ DORIS_CHECK(field != nullptr);
+ DORIS_CHECK(array != nullptr);
+ DORIS_CHECK(normalized != nullptr);
+
+ LanceExtensionKind extension_kind;
+ std::shared_ptr<arrow::DataType> storage_type;
+ RETURN_IF_ERROR(get_lance_extension(field, &extension_kind,
&storage_type));
+
+ auto storage_array = array;
+ if (array->type_id() == arrow::Type::EXTENSION) {
+ const auto extension_array =
std::dynamic_pointer_cast<arrow::ExtensionArray>(array);
+ if (extension_array == nullptr) {
+ return Status::InvalidArgument("invalid Arrow extension array for
Lance field '{}'",
+ field->name());
+ }
+ storage_array = extension_array->storage();
+ }
+ if (storage_array->type_id() != storage_type->id()) {
+ return Status::InvalidArgument(
+ "Lance field '{}' storage type {} does not match array type
{}", field->name(),
+ storage_type->ToString(), storage_array->type()->ToString());
+ }
+ if (extension_kind == LanceExtensionKind::BFLOAT16) {
+ return convert_bfloat16_array(storage_array, normalized);
+ }
+
+ const auto& child_fields = storage_type->fields();
+ const auto& child_data = storage_array->data()->child_data;
+ if (child_fields.empty()) {
+ *normalized = std::move(storage_array);
+ return Status::OK();
+ }
+ if (child_fields.size() != child_data.size()) {
+ return Status::InvalidArgument(
+ "Lance field '{}' has {} child fields but its Arrow array has
{} children",
+ field->name(), child_fields.size(), child_data.size());
+ }
+
+ arrow::FieldVector normalized_fields;
+ std::shared_ptr<arrow::ArrayData> normalized_data;
+ for (size_t child_idx = 0; child_idx < child_fields.size(); ++child_idx) {
+ bool child_requires_normalization = false;
+
RETURN_IF_ERROR(field_requires_lance_normalization(child_fields[child_idx],
+
&child_requires_normalization));
+ if (!child_requires_normalization) {
+ continue;
+ }
+ auto child_array = arrow::MakeArray(child_data[child_idx]);
Review Comment:
[P2] Normalize only the visible child range of sliced offset containers. For
a stable-row-ID Lance dataset, an unfiltered pushed LIMIT cannot become a
pre-scan range, so pinned Lance keeps `GlobalLimitExec` over the full
projection and DataFusion returns a `RecordBatch` slice. Arrow
List/LargeList/Map slices retain the complete values/entries child, and this
line recurses into that unsliced child; `convert_bfloat16_array` then widens
every backing value. `LIMIT 1` can therefore allocate Float32 storage for a
whole batch's nested BFloat16 values per fragment (about 267 MiB for an 8160 x
8192 child), even though only one row is visible. Please normalize only the
child interval covered by the visible parent offsets and rebase the parent,
with a stable-row-ID LIMIT or explicitly sliced offset-container regression.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
