morningman commented on code in PR #69581: URL: https://github.com/apache/doris/pull/69581#discussion_r4397407803
########## be/src/format_v2/lance/lance_reader_helper.cpp: ########## @@ -0,0 +1,849 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "format_v2/lance/lance_reader_helper.h" + +#include <arrow/array.h> +#include <arrow/builder.h> +#include <arrow/c/bridge.h> +#include <arrow/compute/cast.h> +#include <arrow/extension_type.h> +#include <arrow/type.h> +#include <arrow/util/key_value_metadata.h> +#include <fmt/format.h> +#include <lance/lance.h> + +#include <bit> +#include <limits> +#include <optional> +#include <unordered_set> + +#include "common/config.h" +#include "common/logging.h" +#include "core/data_type/data_type_array.h" +#include "core/data_type/data_type_factory.hpp" +#include "core/data_type/data_type_map.h" +#include "core/data_type/data_type_nothing.h" +#include "core/data_type/data_type_nullable.h" +#include "core/data_type/data_type_struct.h" +#include "core/data_type_serde/arrow_validation.h" +#include "exec/common/endian.h" + +namespace doris::format::lance { + +RuntimeProfile::Counter* add_lance_counter(RuntimeProfile* profile, const char* name, + TUnit::type unit, const char* group) { + // These counters label sections; they do not represent time or a sum of their children. + ADD_CHILD_COUNTER_WITH_LEVEL(profile, group, TUnit::NONE, LANCE_READER_PROFILE, 1); + return ADD_CHILD_COUNTER_WITH_LEVEL(profile, name, unit, group, 1); +} + +namespace { + +constexpr std::string_view ARROW_EXTENSION_NAME = "ARROW:extension:name"; +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, +}; + +int arrow_time_precision(arrow::TimeUnit::type unit) { + switch (unit) { + case arrow::TimeUnit::SECOND: + return 0; + case arrow::TimeUnit::MILLI: + return 3; + case arrow::TimeUnit::MICRO: + case arrow::TimeUnit::NANO: + return 6; + } + return 6; +} + +// Extract and validate extension names from metadata and registered Arrow 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()); +} + +// Map Arrow fields recursively, preserving Null leaves in complex types. +Status arrow_field_to_doris_type(const std::shared_ptr<arrow::Field>& field, + DataTypePtr* doris_type) { + const auto nullable_primitive = [&](PrimitiveType type, int precision = 0, int scale = 0, + int len = -1) { + *doris_type = + DataTypeFactory::instance().create_data_type(type, true, precision, scale, len); + return Status::OK(); + }; + + LanceExtensionKind extension_kind; + std::shared_ptr<arrow::DataType> arrow_type; + RETURN_IF_ERROR(get_lance_extension(field, &extension_kind, &arrow_type)); + switch (extension_kind) { + case LanceExtensionKind::JSON: + return nullable_primitive(TYPE_JSONB); + case LanceExtensionKind::BFLOAT16: + return nullable_primitive(TYPE_FLOAT); + case LanceExtensionKind::NONE: + break; + } + if (arrow_type->id() == arrow::Type::DICTIONARY) { + return Status::NotSupported("unsupported Lance Arrow dictionary type for field '{}': {}", + field->name(), arrow_type->ToString()); + } + + switch (arrow_type->id()) { + case arrow::Type::NA: + // Required Null leaves can bypass NullableSerDe through FE schema reconstruction. + // Reject them before Arrow NA buffers reach the Boolean-backed physical SerDe. + if (!field->nullable()) { + return Status::NotSupported("non-nullable Lance Arrow Null field: {}", field->name()); + } + return nullable_primitive(TYPE_NULL); + case arrow::Type::BOOL: + return nullable_primitive(TYPE_BOOLEAN); + case arrow::Type::INT8: + return nullable_primitive(TYPE_TINYINT); + case arrow::Type::UINT8: + case arrow::Type::INT16: + return nullable_primitive(TYPE_SMALLINT); + case arrow::Type::UINT16: + case arrow::Type::INT32: + return nullable_primitive(TYPE_INT); + case arrow::Type::UINT32: + case arrow::Type::INT64: + return nullable_primitive(TYPE_BIGINT); + case arrow::Type::UINT64: + return nullable_primitive(TYPE_LARGEINT); + case arrow::Type::HALF_FLOAT: + case arrow::Type::FLOAT: + return nullable_primitive(TYPE_FLOAT); + case arrow::Type::DOUBLE: + return nullable_primitive(TYPE_DOUBLE); + case arrow::Type::STRING: + case arrow::Type::LARGE_STRING: + return nullable_primitive(TYPE_STRING); + case arrow::Type::BINARY: + case arrow::Type::LARGE_BINARY: + return nullable_primitive(TYPE_VARBINARY, 0, 0, std::numeric_limits<int32_t>::max()); + case arrow::Type::FIXED_SIZE_BINARY: { + const auto binary = std::static_pointer_cast<arrow::FixedSizeBinaryType>(arrow_type); + return nullable_primitive(TYPE_VARBINARY, 0, 0, binary->byte_width()); + } + case arrow::Type::DATE32: + case arrow::Type::DATE64: + return nullable_primitive(TYPE_DATEV2); + case arrow::Type::TIME32: + case arrow::Type::TIME64: { + const auto time = std::static_pointer_cast<arrow::TimeType>(arrow_type); + return nullable_primitive(TYPE_TIMEV2, 0, arrow_time_precision(time->unit())); + } + case arrow::Type::TIMESTAMP: { + const auto timestamp = std::static_pointer_cast<arrow::TimestampType>(arrow_type); + const auto timestamp_type = + timestamp->timezone().empty() ? TYPE_DATETIMEV2 : TYPE_TIMESTAMPTZ; + return nullable_primitive(timestamp_type, 0, arrow_time_precision(timestamp->unit())); + } + case arrow::Type::DURATION: + return nullable_primitive(TYPE_BIGINT); + case arrow::Type::DECIMAL128: + case arrow::Type::DECIMAL256: { + const auto decimal = std::static_pointer_cast<arrow::DecimalType>(arrow_type); + const int precision = decimal->precision(); + const int scale = decimal->scale(); + if (precision <= 0 || precision > arrow::Decimal256Type::kMaxPrecision || scale < 0 || + scale > precision) { + return Status::NotSupported( + "unsupported Lance Arrow decimal type for field '{}': precision={}, scale={}", + field->name(), precision, scale); + } + const PrimitiveType doris_decimal_type = precision <= 9 ? TYPE_DECIMAL32 + : precision <= 18 ? TYPE_DECIMAL64 + : precision <= 38 ? TYPE_DECIMAL128I + : TYPE_DECIMAL256; + return nullable_primitive(doris_decimal_type, precision, scale); + } + case arrow::Type::LIST: + case arrow::Type::LARGE_LIST: + case arrow::Type::FIXED_SIZE_LIST: { + const auto list = std::static_pointer_cast<arrow::BaseListType>(arrow_type); + DataTypePtr value_type; + RETURN_IF_ERROR(arrow_field_to_doris_type(list->value_field(), &value_type)); + *doris_type = make_nullable(std::make_shared<DataTypeArray>(value_type)); + return Status::OK(); + } + case arrow::Type::MAP: { + const auto map = std::static_pointer_cast<arrow::MapType>(arrow_type); + DataTypePtr key_type; + DataTypePtr item_type; + RETURN_IF_ERROR(arrow_field_to_doris_type(map->key_field(), &key_type)); + RETURN_IF_ERROR(arrow_field_to_doris_type(map->item_field(), &item_type)); + *doris_type = make_nullable(std::make_shared<DataTypeMap>(key_type, item_type)); + return Status::OK(); + } + case arrow::Type::STRUCT: { + const auto struct_type = std::static_pointer_cast<arrow::StructType>(arrow_type); + // Lance decodes STRUCT<>, but a Doris ColumnStruct takes its row count from its first + // child, so a zero-field struct cannot be materialized at any nesting level. + if (struct_type->num_fields() == 0) { + return Status::NotSupported("unsupported Lance Arrow zero-field struct: {}", + field->name()); + } + DataTypes field_types; + Strings field_names; + field_types.reserve(struct_type->num_fields()); + field_names.reserve(struct_type->num_fields()); + for (const auto& child : struct_type->fields()) { + DataTypePtr field_type; + RETURN_IF_ERROR(arrow_field_to_doris_type(child, &field_type)); + field_types.emplace_back(std::move(field_type)); + field_names.emplace_back(child->name()); + } + *doris_type = make_nullable(std::make_shared<DataTypeStruct>(field_types, field_names)); + return Status::OK(); + } + default: + return Status::NotSupported("unsupported Lance Arrow type: {}", arrow_type->ToString()); + } +} + +// Widen little-endian Lance BFloat16 values to Arrow Float32 without precision loss. +Status convert_bfloat16_array(const std::shared_ptr<arrow::Array>& array, + arrow::MemoryPool* memory_pool, + std::shared_ptr<arrow::Array>* normalized) { + DORIS_CHECK(array != nullptr); + DORIS_CHECK(memory_pool != 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(memory_pool); + 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(); +} + +// Narrow a Decimal256 array whose declared precision fits Decimal128. The safe cast keeps every +// in-precision value and its scale, and rejects a value that exceeds the declared precision +// instead of truncating it: unscaled 10^18 stored in a decimal256(18, 2) field fails with +// "Decimal value does not fit in precision 18". +Status narrow_decimal256_array(std::string_view field_name, + const std::shared_ptr<arrow::Array>& array, + const std::shared_ptr<arrow::DataType>& decimal128_type, + arrow::MemoryPool* memory_pool, + std::shared_ptr<arrow::Array>* normalized) { + DORIS_CHECK(array != nullptr); + DORIS_CHECK(memory_pool != nullptr); + DORIS_CHECK(normalized != nullptr); + arrow::compute::ExecContext exec_context(memory_pool); + auto narrowed = arrow::compute::Cast(*array, decimal128_type, + arrow::compute::CastOptions::Safe(), &exec_context); + if (!narrowed.ok()) { + return Status::InvalidArgument("narrow Lance Decimal256 field '{}' to {} failed: {}", + field_name, decimal128_type->ToString(), + narrowed.status().message()); + } + *normalized = std::move(narrowed).ValueUnsafe(); + return Status::OK(); +} + +// Rebuild a nested Arrow type after one or more child arrays changed physical type. +Status set_lance_nested_type(std::string_view field_name, + const std::shared_ptr<arrow::DataType>& source_type, + const arrow::FieldVector& child_fields, + std::shared_ptr<arrow::ArrayData>* data) { + switch (source_type->id()) { + case arrow::Type::LIST: + (*data)->type = arrow::list(child_fields[0]); + break; + case arrow::Type::LARGE_LIST: + (*data)->type = arrow::large_list(child_fields[0]); + break; + case arrow::Type::FIXED_SIZE_LIST: + (*data)->type = arrow::fixed_size_list( + child_fields[0], + std::static_pointer_cast<arrow::FixedSizeListType>(source_type)->list_size()); + break; + case arrow::Type::STRUCT: + (*data)->type = arrow::struct_(child_fields); + break; + case arrow::Type::MAP: { + const auto map_type = std::static_pointer_cast<arrow::MapType>(source_type); + auto normalized_type = arrow::MapType::Make(child_fields[0], map_type->keys_sorted()); + if (!normalized_type.ok()) { + return Status::InvalidArgument("normalize Lance map field '{}' failed: {}", field_name, + normalized_type.status().message()); + } + (*data)->type = std::move(normalized_type).ValueUnsafe(); + break; + } + default: + return Status::InvalidArgument("Lance field '{}' has unexpected child-bearing type {}", + field_name, source_type->ToString()); + } + return Status::OK(); +} + +// Materialize the visible range into an offset-zero Arrow array for Doris SerDes. +Status compact_lance_array(const std::shared_ptr<arrow::Array>& array, + arrow::MemoryPool* memory_pool, + std::shared_ptr<arrow::Array>* compacted) { + DORIS_CHECK(memory_pool != nullptr); + const auto validation = array->ValidateFull(); + if (!validation.ok()) { + return Status::InvalidArgument("validate sliced Lance array failed: {}", + validation.message()); + } + auto builder_result = arrow::MakeBuilder(array->type(), memory_pool); + if (!builder_result.ok()) { + return Status::InternalError("create sliced Lance array builder failed: {}", + builder_result.status().message()); + } + auto builder = std::move(builder_result).ValueUnsafe(); + auto arrow_status = builder->Reserve(array->length()); + if (!arrow_status.ok()) { + return Status::InternalError("reserve sliced Lance array builder failed: {}", + arrow_status.message()); + } + arrow_status = builder->AppendArraySlice(*array->data(), 0, array->length()); + if (!arrow_status.ok()) { + return Status::InternalError("copy sliced Lance array failed: {}", arrow_status.message()); + } + arrow_status = builder->Finish(compacted); + if (!arrow_status.ok()) { + return Status::InternalError("finish sliced Lance array copy failed: {}", + arrow_status.message()); + } + if ((*compacted)->offset() != 0) { + return Status::InternalError("compacted Lance array retained offset {}", + (*compacted)->offset()); + } + return Status::OK(); +} + +template <typename OffsetType> +bool lance_list_offsets_are_tight(const arrow::ArrayData& data) { + const auto* offsets = data.GetValues<OffsetType>(1); + return offsets[0] == 0 && offsets[data.length] == data.child_data[0]->length; +} + +// Whether Doris SerDes can read this array in place: with enable_arrow_input_validation (the +// default) they reject a non-zero offset at any depth ("non-zero array offset is not supported"), +// and the nested SerDes expect each child to hold exactly the elements its parent references. The +// array's own offset and offsets are not enough. Lance's legacy (v1) plain decoder returns the +// Boolean range [start, end) as a bitmap with bit offset start % 8, and its struct and list readers +// wrap that child unchanged: with batch_size 3, rows 3..5 of STRUCT<b: BOOLEAN> arrive as a +// zero-offset struct whose child has offset 3, and a LIST<BOOLEAN> batch whose first value is +// value k of the writer batch has values at offset k % 8. +bool lance_array_layout_is_compact(const arrow::ArrayData& data) { + if (data.offset != 0) { + return false; + } + switch (data.type->id()) { + case arrow::Type::LIST: + case arrow::Type::MAP: + return lance_list_offsets_are_tight<int32_t>(data) && + lance_array_layout_is_compact(*data.child_data[0]); + case arrow::Type::LARGE_LIST: + return lance_list_offsets_are_tight<int64_t>(data) && + lance_array_layout_is_compact(*data.child_data[0]); + case arrow::Type::FIXED_SIZE_LIST: { + const auto list_size = static_cast<const arrow::FixedSizeListType&>(*data.type).list_size(); + return data.child_data[0]->length == data.length * list_size && + lance_array_layout_is_compact(*data.child_data[0]); + } + case arrow::Type::STRUCT: + for (const auto& child : data.child_data) { + if (child->length != data.length || !lance_array_layout_is_compact(*child)) { + return false; + } + } + return true; + default: + return true; + } +} + +// Compact a sliced variable-offset parent and all of its visible descendants. +template <typename ArrayType> +Status compact_lance_offset_array(const std::shared_ptr<arrow::Array>& array, + arrow::MemoryPool* memory_pool, + std::shared_ptr<arrow::Array>* compacted) { + const auto offset_array = std::dynamic_pointer_cast<ArrayType>(array); + if (offset_array == nullptr) { + return Status::InvalidArgument("invalid sliced Lance offset array: {}", + array->type()->ToString()); + } + const auto child_begin = static_cast<int64_t>(offset_array->value_offset(0)); + const auto child_end = static_cast<int64_t>(offset_array->value_offset(offset_array->length())); + const auto& values = offset_array->values(); + if (child_begin < 0 || child_end < child_begin || child_end > values->length()) { + return Status::InvalidArgument("invalid sliced Lance offsets [{}, {}) for child length {}", + child_begin, child_end, values->length()); + } + if (lance_array_layout_is_compact(*array->data())) { + *compacted = array; + return Status::OK(); + } + return compact_lance_array(array, memory_pool, compacted); +} + +// Compact sliced arrays and nested children only when Doris cannot consume their current layout. +// A nested array is passed through only when its whole subtree is compact, because a tight struct +// or list can still hold an offset child (see lance_array_layout_is_compact()). +Status compact_lance_array_if_needed(const std::shared_ptr<arrow::Array>& array, + arrow::MemoryPool* memory_pool, + std::shared_ptr<arrow::Array>* compacted) { + switch (array->type_id()) { + case arrow::Type::LIST: + return compact_lance_offset_array<arrow::ListArray>(array, memory_pool, compacted); + case arrow::Type::LARGE_LIST: + return compact_lance_offset_array<arrow::LargeListArray>(array, memory_pool, compacted); + case arrow::Type::MAP: + return compact_lance_offset_array<arrow::MapArray>(array, memory_pool, compacted); + case arrow::Type::FIXED_SIZE_LIST: { + const auto list = std::dynamic_pointer_cast<arrow::FixedSizeListArray>(array); + if (list == nullptr) { + return Status::InvalidArgument("invalid sliced Lance fixed-size list array: {}", + array->type()->ToString()); + } + const auto child_begin = list->value_offset(0); + const auto child_length = list->length() * list->value_length(); + const auto& values = list->values(); + if (child_begin < 0 || child_length < 0 || child_begin > values->length() - child_length) { + return Status::InvalidArgument( + "invalid sliced Lance fixed-size list range [{}, {}) for child length {}", + child_begin, child_begin + child_length, values->length()); + } + if (lance_array_layout_is_compact(*array->data())) { + *compacted = array; + return Status::OK(); + } + return compact_lance_array(array, memory_pool, compacted); + } + case arrow::Type::STRUCT: { + const auto struct_array = std::dynamic_pointer_cast<arrow::StructArray>(array); + if (struct_array == nullptr) { + return Status::InvalidArgument("invalid sliced Lance struct array: {}", + array->type()->ToString()); + } + if (lance_array_layout_is_compact(*array->data())) { + *compacted = array; + return Status::OK(); + } + return compact_lance_array(array, memory_pool, compacted); + } + default: + if (array->offset() == 0) { + *compacted = array; + return Status::OK(); + } + return compact_lance_array(array, memory_pool, compacted); + } +} + +} // namespace + +Status LanceArrowArrayNormalizer::create(const std::shared_ptr<arrow::Field>& field, + LanceArrowArrayNormalizer* normalizer) { + DORIS_CHECK(field != nullptr); + DORIS_CHECK(normalizer != nullptr); + + LanceExtensionKind extension_kind; + std::shared_ptr<arrow::DataType> storage_type; + RETURN_IF_ERROR(get_lance_extension(field, &extension_kind, &storage_type)); + + LanceArrowArrayNormalizer result; + result._field_name = field->name(); + result._storage_type = std::move(storage_type); + result._unwrap_registered_extension = field->type()->id() == arrow::Type::EXTENSION; + result._convert_bfloat16 = extension_kind == LanceExtensionKind::BFLOAT16; + // Lance keeps the 256-bit carrier of pa.decimal256(18, 2), while the FE and + // arrow_field_to_doris_type() type the column DECIMAL64 by precision. The DECIMAL32/64/128I + // SerDes read only Decimal128Array, so a Decimal256Array reaching them made their dynamic_cast + // return null and the BE crash on the first value, also for nested leaves. Precision 39..76 + // stays Decimal256 and maps to DECIMAL256. + if (result._storage_type->id() == arrow::Type::DECIMAL256) { + const auto& decimal = static_cast<const arrow::Decimal256Type&>(*result._storage_type); + if (decimal.precision() <= arrow::Decimal128Type::kMaxPrecision) { + result._decimal128_type = arrow::decimal128(decimal.precision(), decimal.scale()); + } + } + result._requires_special_handling = result._unwrap_registered_extension || + result._convert_bfloat16 || + result._decimal128_type != nullptr; + result._child_normalizers.reserve(result._storage_type->num_fields()); + for (const auto& child_field : result._storage_type->fields()) { + LanceArrowArrayNormalizer child_normalizer; + RETURN_IF_ERROR(create(child_field, &child_normalizer)); + result._requires_special_handling |= child_normalizer._requires_special_handling; + result._child_normalizers.emplace_back(std::move(child_normalizer)); + } + *normalizer = std::move(result); + return Status::OK(); +} + +Status LanceArrowArrayNormalizer::normalize_for_doris( + const std::shared_ptr<arrow::Array>& array, arrow::MemoryPool* memory_pool, + std::shared_ptr<arrow::Array>* normalized) const { + DORIS_CHECK(array != nullptr); + DORIS_CHECK(memory_pool != nullptr); + DORIS_CHECK(normalized != nullptr); + + // The schema binding already resolved extension metadata and nested special types. For common + // arrays, runtime work is limited to checking whether the visible slice needs compaction. + if (!_requires_special_handling) { + return compact_lance_array_if_needed(array, memory_pool, normalized); + } + + auto storage_array = array; + if (_unwrap_registered_extension && 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: {}", + array->type()->ToString()); + } + storage_array = extension_array->storage(); + } + // Extension and BFloat16 arrays are first converted to their physical storage representation. + 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 (_convert_bfloat16) { + return convert_bfloat16_array(storage_array, memory_pool, normalized); + } + if (_decimal128_type != nullptr) { + // The cast reads only the visible slice and returns an offset-zero array, so no + // compaction follows. For a nested leaf, the parent's call swaps in this array in the + // child loop below and rebuilds its own type through set_lance_nested_type(). + return narrow_decimal256_array(_field_name, storage_array, _decimal128_type, memory_pool, + normalized); + } + + // Normalize extension children before compacting the parent. Arrow cannot create a builder + // for list/struct types whose child is still an ExtensionType (Arrow 24 returns + // NotImplemented), while the rebuilt physical type can be compacted normally. + const auto& child_data = storage_array->data()->child_data; + if (_child_normalizers.size() != child_data.size()) { + return Status::InvalidArgument( + "Lance field '{}' has {} child fields but its Arrow array has {} children", + _field_name, _child_normalizers.size(), child_data.size()); + } + + arrow::FieldVector normalized_fields; + std::shared_ptr<arrow::ArrayData> normalized_data; + for (size_t child_idx = 0; child_idx < _child_normalizers.size(); ++child_idx) { + const auto& child_normalizer = _child_normalizers[child_idx]; + if (!child_normalizer._requires_special_handling) { + continue; + } + auto child_array = arrow::MakeArray(child_data[child_idx]); + std::shared_ptr<arrow::Array> normalized_child; + RETURN_IF_ERROR( + child_normalizer.normalize_for_doris(child_array, memory_pool, &normalized_child)); Review Comment: Confirmed, and fixed in `1c74dbab92c`. The Lance normalizer narrows a nested Decimal256 leaf by casting the whole raw child array before any SerDe runs, so the safe cast checked the values under a NULL list, fixed-size-list, map or struct row and, for a sliced parent, the values of rows outside the slice; the SerDe-level masking from the previous round only applies after normalization. Before narrowing, the normalizer now marks NULL every child slot that no non-NULL row of the parent references, level by level down to the leaf, so only values a reader can see are checked, and a visible value outside the declared precision is still rejected with the field name. A map is masked through its values: Arrow allows no NULL in a map's entries struct or keys (`MapArray::SetData()` aborts on one), so a Decimal256 key under a NULL map row is still checked. BFloat16 widening and extension unwrapping never inspect a value, so only the subtrees that narrow a decimal are masked. Covered by `Lanc eTableReaderTypeTest.NarrowsDecimal256PayloadsHiddenByNullParents` (NULL fixed-size-list, list, struct, map and list-of-struct rows, and a sliced list) and by a nested visible value added to `LanceTableReaderTypeTest.RejectsDecimal256ValuesOutsideDeclaredPrecision`; the masking helper has its own cases in `ArrowParentNullsTest` (`MasksTheSlotsOfNullListRows`, `MasksTheSlotsOfNullFixedSizeListAndStructRows`, `MasksTheItemsOfNullMapRows`, `MasksTheChildSlotsOutsideASlicedParent`, `ReturnsTheChildUnchangedWhenEverySlotIsVisible`, `RejectsSlotsOutsideTheChildAndUnknownParents`). Building such a payload needs hand-made buffers (pyarrow's converters reject a value outside the declared precision), so the test builds the batch shape directly through the converter, like the DATE32 test of the previous round; the Lance writer's pass-through of fixed-size-list child data was checked for that thread. ########## be/src/core/data_type_serde/arrow_parent_nulls.cpp: ########## @@ -0,0 +1,82 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "core/data_type_serde/arrow_parent_nulls.h" + +#include <arrow/array/array_base.h> +#include <arrow/array/data.h> +#include <arrow/array/util.h> +#include <arrow/buffer.h> +#include <arrow/type.h> +#include <arrow/util/bit_util.h> +#include <arrow/util/bitmap_ops.h> + +#include "common/check.h" +#include "common/logging.h" + +namespace doris { + +Status mask_arrow_children_hidden_by_null_parents( + const arrow::Array& parent, int64_t start, int64_t end, + const std::shared_ptr<arrow::Array>& child, + const std::function<std::pair<int64_t, int64_t>(int64_t)>& child_slots, + std::shared_ptr<arrow::Array>* masked) { + DORIS_CHECK(child != nullptr); + DORIS_CHECK(masked != nullptr); + *masked = child; + // An Arrow Null child has no validity buffer and is NULL in every slot already. + if (parent.null_count() == 0 || child->type_id() == arrow::Type::NA) { + return Status::OK(); + } + const auto& data = *child->data(); + // Validity bits are addressed from the start of the buffer, before the array offset. + const int64_t bit_length = data.offset + data.length; + std::shared_ptr<arrow::Buffer> validity; + for (int64_t row = start; row < end; ++row) { + if (!parent.IsNull(row)) { + continue; + } + if (validity == nullptr) { + auto allocated = arrow::AllocateBitmap(bit_length); Review Comment: Fixed in `1c74dbab92c`: the masked validity bitmap is now allocated from `ExecEnv::arrow_memory_pool()`, the tracked `doris::ArrowMemoryPool` that the Lance converter and the ADBC reader already use, so it is charged to the query and passes its limit check; unit tests that never initialize ExecEnv fall back to Arrow's default pool, as `LanceRecordBatchConverter` does. `ArrowParentNullsTest.AllocatesTheMaskFromTheDorisArrowPool` puts a proxy pool into ExecEnv and checks that the bitmap is charged to it and released with the masked array, and `ArrowParentNullsTest.ReportsMaskAllocationFailure` checks that a pool refusing the bitmap fails the read with a Status. The scale is smaller than described, though. The bitmap is one bit per child slot and lives only until the child SerDe returns, while the same call appends at least two tracked bytes per child slot to the Doris column (value plus null map; the segment of a NULL list row is read as NULL elements), so the 100-million-child exam ple already charges about 200 MB to the query in that call, against 12.5 MB here. Columns are converted one after another and each mask is released once its child is read, so at most one bitmap per nesting level (two for a map) is alive at a time, not one per nested column. -- 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]
