diff --git a/be/src/format/arrow/arrow_array_normalizer.cpp b/be/src/format/arrow/arrow_array_normalizer.cpp index 4b4483cb922a2b..4f8aa4df38cc73 100644 --- a/be/src/format/arrow/arrow_array_normalizer.cpp +++ b/be/src/format/arrow/arrow_array_normalizer.cpp @@ -18,9 +18,13 @@ #include "format/arrow/arrow_array_normalizer.h" #include +#include +#include +#include #include #include +#include #include #include "common/check.h" @@ -30,22 +34,164 @@ namespace doris { namespace { -// The accepted counterpart of an encoding-only variant. Null when there is none. -std::shared_ptr target_type_for(const arrow::DataType& type) { - switch (type.id()) { +// The accepted counterpart of an encoding-only variant. Null when neither this type nor any child +// needs normalization. +std::shared_ptr target_type_for(const std::shared_ptr& type) { + switch (type->id()) { case arrow::Type::LARGE_STRING: case arrow::Type::STRING_VIEW: return arrow::utf8(); case arrow::Type::LARGE_BINARY: case arrow::Type::BINARY_VIEW: return arrow::binary(); + case arrow::Type::DICTIONARY: { + const auto& dictionary = static_cast(*type); + auto nested_target = target_type_for(dictionary.value_type()); + return nested_target != nullptr ? nested_target : dictionary.value_type(); + } + case arrow::Type::RUN_END_ENCODED: { + const auto& encoded = static_cast(*type); + auto nested_target = target_type_for(encoded.value_type()); + return nested_target != nullptr ? nested_target : encoded.value_type(); + } + case arrow::Type::LIST: + case arrow::Type::LARGE_LIST: + case arrow::Type::FIXED_SIZE_LIST: { + const auto& list = static_cast(*type); + auto child_target = target_type_for(list.value_type()); + if (child_target == nullptr) { + return nullptr; + } + auto child = list.value_field()->WithType(std::move(child_target)); + if (type->id() == arrow::Type::LIST) { + return arrow::list(std::move(child)); + } + if (type->id() == arrow::Type::LARGE_LIST) { + return arrow::large_list(std::move(child)); + } + const auto& fixed = static_cast(*type); + return arrow::fixed_size_list(std::move(child), fixed.list_size()); + } + case arrow::Type::STRUCT: { + arrow::FieldVector fields; + fields.reserve(type->num_fields()); + bool changed = false; + for (const auto& field : type->fields()) { + auto child_target = target_type_for(field->type()); + changed |= child_target != nullptr; + fields.push_back(child_target != nullptr ? field->WithType(std::move(child_target)) + : field); + } + return changed ? arrow::struct_(fields) : nullptr; + } + case arrow::Type::MAP: { + const auto& map = static_cast(*type); + auto key_target = target_type_for(map.key_type()); + auto item_target = target_type_for(map.item_type()); + if (key_target == nullptr && item_target == nullptr) { + return nullptr; + } + auto key = key_target != nullptr ? map.key_field()->WithType(std::move(key_target)) + : map.key_field(); + auto item = item_target != nullptr ? map.item_field()->WithType(std::move(item_target)) + : map.item_field(); + return std::make_shared(std::move(key), std::move(item), map.keys_sorted()); + } + default: + return nullptr; + } +} + +bool contains_list_view(const arrow::DataType& type) { + switch (type.id()) { + case arrow::Type::LIST_VIEW: + case arrow::Type::LARGE_LIST_VIEW: + return true; case arrow::Type::DICTIONARY: - return static_cast(type).value_type(); + return contains_list_view(*static_cast(type).value_type()); case arrow::Type::RUN_END_ENCODED: - return static_cast(type).value_type(); + return contains_list_view(*static_cast(type).value_type()); + case arrow::Type::EXTENSION: + return contains_list_view(*static_cast(type).storage_type()); default: - return nullptr; + for (const auto& field : type.fields()) { + if (contains_list_view(*field->type())) { + return true; + } + } + return false; + } +} + +template +arrow::Result> canonicalize_list_view(const ViewArray& source, + arrow::MemoryPool* pool) { + // Imported C stream arrays are not guaranteed to have validated ranges; copying an invalid + // range before this check could read beyond the child array. + auto validation = source.ValidateFull(); + if (!validation.ok()) { + return validation; + } + + // Preflight the full expansion before any builder mutation. Nested view builders and Arrow's + // NullBuilder can otherwise overflow signed lengths before their parent reports capacity. + if (contains_list_view(*source.value_type())) { + return arrow::Status::Invalid("nested list view canonicalization is not supported"); + } + int64_t logical_value_count = 0; + constexpr int64_t max_value_count = std::numeric_limits::max(); + for (int64_t i = 0; i < source.length(); ++i) { + if (source.IsNull(i)) { + continue; + } + const int64_t value_length = source.value_length(i); + if (value_length > max_value_count - logical_value_count) { + return arrow::Status::CapacityError("list view logical values exceed output capacity"); + } + logical_value_count += value_length; + } + + auto child_builder_result = arrow::MakeBuilder(source.value_type(), pool); + if (!child_builder_result.ok()) { + return child_builder_result.status(); + } + std::shared_ptr child_builder(child_builder_result.MoveValueUnsafe()); + ListBuilder builder(pool, child_builder); + auto reserve_status = builder.Reserve(source.length()); + if (!reserve_status.ok()) { + return reserve_status; + } + reserve_status = child_builder->Reserve(logical_value_count); + if (!reserve_status.ok()) { + return reserve_status; } + + arrow::ArraySpan values(*source.values()->data()); + for (int64_t i = 0; i < source.length(); ++i) { + if (source.IsNull(i)) { + auto append_status = builder.AppendNull(); + if (!append_status.ok()) { + return append_status; + } + continue; + } + auto append_status = builder.Append(); + if (!append_status.ok()) { + return append_status; + } + append_status = child_builder->AppendArraySlice(values, source.value_offset(i), + source.value_length(i)); + if (!append_status.ok()) { + return append_status; + } + } + + std::shared_ptr out; + auto finish_status = builder.Finish(&out); + if (!finish_status.ok()) { + return finish_status; + } + return std::static_pointer_cast(out); } } // namespace @@ -59,6 +205,8 @@ bool is_serde_acceptable_arrow_type(const arrow::DataType& type) { case arrow::Type::BINARY_VIEW: case arrow::Type::DICTIONARY: case arrow::Type::RUN_END_ENCODED: + case arrow::Type::LIST_VIEW: + case arrow::Type::LARGE_LIST_VIEW: return false; // No Doris column can hold these, so they must not reach a serde either. case arrow::Type::INTERVAL_MONTHS: @@ -69,13 +217,19 @@ bool is_serde_acceptable_arrow_type(const arrow::DataType& type) { case arrow::Type::DENSE_UNION: return false; default: + for (const auto& field : type.fields()) { + if (!is_serde_acceptable_arrow_type(*field->type())) { + return false; + } + } return true; } } -Status normalize_arrow_array(const std::shared_ptr& arr, +Status normalize_arrow_array(const std::shared_ptr& arr, arrow::MemoryPool* pool, std::shared_ptr* out) { DORIS_CHECK(arr != nullptr); + DORIS_CHECK(pool != nullptr); DORIS_CHECK(out != nullptr); std::shared_ptr current = arr; @@ -89,14 +243,41 @@ Status normalize_arrow_array(const std::shared_ptr& arr, return Status::OK(); } - auto target = target_type_for(type); + // List views may share or reorder value ranges, so rebuild canonical offsets instead of + // exposing their buffers to a serde that requires contiguous list values. + if (type.id() == arrow::Type::LIST_VIEW) { + auto converted = canonicalize_list_view( + static_cast(*current), pool); + if (!converted.ok()) { + return Status::InternalError("ADBC: failed to normalize arrow type '{}': {}", + type.ToString(), converted.status().ToString()); + } + current = converted.MoveValueUnsafe(); + continue; + } + if (type.id() == arrow::Type::LARGE_LIST_VIEW) { + auto converted = canonicalize_list_view( + static_cast(*current), pool); + if (!converted.ok()) { + return Status::InternalError("ADBC: failed to normalize arrow type '{}': {}", + type.ToString(), converted.status().ToString()); + } + current = converted.MoveValueUnsafe(); + continue; + } + + auto target = target_type_for(current->type()); if (target == nullptr) { return Status::NotSupported( "ADBC: arrow type '{}' cannot be materialized into a Doris column", type.ToString()); } - auto casted = arrow::compute::Cast(*current, target); + arrow::compute::ExecContext exec_context(pool); + auto casted = arrow::compute::Cast(*current, target, arrow::compute::CastOptions::Safe(), + &exec_context); if (!casted.ok()) { return Status::InternalError("ADBC: failed to normalize arrow type '{}' to '{}': {}", type.ToString(), target->ToString(), diff --git a/be/src/format/arrow/arrow_array_normalizer.h b/be/src/format/arrow/arrow_array_normalizer.h index eed9ab02293fbb..b2a6a49fa51d41 100644 --- a/be/src/format/arrow/arrow_array_normalizer.h +++ b/be/src/format/arrow/arrow_array_normalizer.h @@ -30,15 +30,17 @@ namespace doris { /// emits string_view, Go-based drivers may emit large_* and dictionary. This normalizes them into /// a shape the serdes accept. /// -/// Only top-level types are normalized. A nested type whose child is an unaccepted variant (say -/// list) passes through and fails inside the serde, loudly rather than silently. +/// Encoding-only descendants are normalized recursively so a complex serde never mistakes encoded +/// child buffers for logical values. Nested list views remain unsupported because their shared +/// ranges cannot be canonicalized safely by an outer Arrow cast. /// Whether the serdes take this Arrow type as-is. bool is_serde_acceptable_arrow_type(const arrow::DataType& type); -/// Normalizes `arr` into a serde-acceptable shape. Returns it unchanged (no copy) when it already -/// is one, and fails with the offending type named when no accepted shape exists. -Status normalize_arrow_array(const std::shared_ptr& arr, +/// Normalizes `arr` into a serde-acceptable shape using `pool` for every conversion. Returns it +/// unchanged (no copy) when it already is one, and fails with the offending type named when no +/// accepted shape exists. +Status normalize_arrow_array(const std::shared_ptr& arr, arrow::MemoryPool* pool, std::shared_ptr* out); } // namespace doris diff --git a/be/src/format_v2/column_mapper.cpp b/be/src/format_v2/column_mapper.cpp index 341ecd45395a07..e80f7aba0f4722 100644 --- a/be/src/format_v2/column_mapper.cpp +++ b/be/src/format_v2/column_mapper.cpp @@ -1916,6 +1916,32 @@ static void attach_timestamp_semantics(const ColumnMapping& mapping, LocalColumn attach_timestamp_semantics(child_mapping, &*child_it); } } +static Status apply_projected_file_definition_to_mapping(const ColumnDefinition& projected_field, + ColumnMapping* mapping) { + DORIS_CHECK(mapping != nullptr); + mapping->file_type = projected_field.type; + mapping->projected_file_children = projected_field.children; + for (auto& child_mapping : mapping->child_mappings) { + if (!child_mapping.file_local_id.has_value()) { + continue; + } + const auto child_it = + std::ranges::find_if(projected_field.children, [&](const ColumnDefinition& child) { + return child.file_local_id() == *child_mapping.file_local_id; + }); + if (child_it == projected_field.children.end()) { + return Status::InternalError( + "Projected file type for '{}' is missing mapped child id {}", + mapping->file_column_name, *child_mapping.file_local_id); + } + // A full root projection changes every descendant's runtime shape too. Keep the recursive + // mapping in sync so a formerly pruned child cannot be mistaken for a trivial direct column. + RETURN_IF_ERROR(apply_projected_file_definition_to_mapping(*child_it, &child_mapping)); + } + mapping->is_trivial = mapping_can_use_file_column_directly(*mapping); + return Status::OK(); +} + // Update the mapping's file type according to the projection, and determine whether the projection // is trivial (i.e. the projected file type is the same as the table type, so no need to // rematerialize the complex value back to table layout after reading from file). @@ -1934,10 +1960,7 @@ static Status apply_projection_to_mapping_file_type(const LocalColumnIndex& proj field.children = mapping->original_file_children; ColumnDefinition projected_field; RETURN_IF_ERROR(project_column_definition(field, projection, &projected_field)); - mapping->file_type = std::move(projected_field.type); - mapping->projected_file_children = std::move(projected_field.children); - mapping->is_trivial = mapping_can_use_file_column_directly(*mapping); - return Status::OK(); + return apply_projected_file_definition_to_mapping(projected_field, mapping); } static const ColumnDefinition* find_file_child_by_name( diff --git a/be/src/format_v2/table/adbc_reader.cpp b/be/src/format_v2/table/adbc_reader.cpp index 98847c40603f2b..0ab86359baecc6 100644 --- a/be/src/format_v2/table/adbc_reader.cpp +++ b/be/src/format_v2/table/adbc_reader.cpp @@ -40,13 +40,18 @@ #include "core/data_type/data_type_map.h" #include "core/data_type/data_type_nullable.h" #include "core/data_type/data_type_struct.h" +#include "core/data_type/primitive_type.h" #include "core/data_type_serde/data_type_serde.h" #include "format/arrow/arrow_array_normalizer.h" +#include "format/parquet/arrow_memory_pool.h" +#include "format_v2/column_mapper.h" #include "format_v2/materialized_reader_util.h" #include "runtime/descriptors.h" +#include "runtime/exec_env.h" #include "runtime/file_scan_profile.h" #include "runtime/runtime_state.h" #include "util/adbc_driver_registry.h" +#include "util/string_util.h" #include "util/timezone_utils.h" #include "util/url_coding.h" @@ -105,6 +110,136 @@ Status validate_adbc_range(const TFileRangeDesc& range) { return Status::OK(); } +Status arrow_type_mismatch(const std::string& path, const arrow::DataType& arrow_type, + const DataTypePtr& doris_type) { + return Status::InvalidArgument( + "ADBC Arrow type mismatch at '{}': runtime type '{}' does not match cached Doris " + "type '{}'", + path, arrow_type.ToString(), doris_type->get_name()); +} + +int timestamp_scale(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 -1; +} + +Status validate_arrow_type(const arrow::DataType& arrow_type, const DataTypePtr& doris_type, + const std::string& path) { + DORIS_CHECK(doris_type != nullptr); + const auto nested_doris_type = remove_nullable(doris_type); + const auto primitive_type = nested_doris_type->get_primitive_type(); + const auto mismatch = [&]() { return arrow_type_mismatch(path, arrow_type, doris_type); }; + + switch (arrow_type.id()) { + case arrow::Type::BOOL: + return primitive_type == TYPE_BOOLEAN ? Status::OK() : mismatch(); + case arrow::Type::INT8: + return primitive_type == TYPE_TINYINT ? Status::OK() : mismatch(); + case arrow::Type::UINT8: + return primitive_type == TYPE_SMALLINT ? Status::OK() : mismatch(); + case arrow::Type::INT16: + return primitive_type == TYPE_SMALLINT ? Status::OK() : mismatch(); + case arrow::Type::UINT16: + return primitive_type == TYPE_INT ? Status::OK() : mismatch(); + case arrow::Type::INT32: + return primitive_type == TYPE_INT ? Status::OK() : mismatch(); + case arrow::Type::UINT32: + return primitive_type == TYPE_BIGINT ? Status::OK() : mismatch(); + case arrow::Type::INT64: + return primitive_type == TYPE_BIGINT ? Status::OK() : mismatch(); + case arrow::Type::UINT64: + return primitive_type == TYPE_LARGEINT ? Status::OK() : mismatch(); + case arrow::Type::HALF_FLOAT: + case arrow::Type::FLOAT: + return primitive_type == TYPE_FLOAT ? Status::OK() : mismatch(); + case arrow::Type::DOUBLE: + return primitive_type == TYPE_DOUBLE ? Status::OK() : mismatch(); + case arrow::Type::STRING: + case arrow::Type::BINARY: + case arrow::Type::FIXED_SIZE_BINARY: + return primitive_type == TYPE_STRING ? Status::OK() : mismatch(); + case arrow::Type::DATE32: + return primitive_type == TYPE_DATEV2 ? Status::OK() : mismatch(); + case arrow::Type::DATE64: + return primitive_type == TYPE_DATETIMEV2 && nested_doris_type->get_scale() == 3 + ? Status::OK() + : mismatch(); + case arrow::Type::TIMESTAMP: { + const auto& timestamp = static_cast(arrow_type); + const bool zoned = !timestamp.timezone().empty(); + const auto expected = zoned ? TYPE_TIMESTAMPTZ : TYPE_DATETIMEV2; + return primitive_type == expected && + nested_doris_type->get_scale() == timestamp_scale(timestamp.unit()) + ? Status::OK() + : mismatch(); + } + case arrow::Type::DECIMAL128: + case arrow::Type::DECIMAL256: { + const auto& decimal = static_cast(arrow_type); + return is_decimal(primitive_type) && + nested_doris_type->get_precision() == decimal.precision() && + nested_doris_type->get_scale() == decimal.scale() + ? Status::OK() + : mismatch(); + } + case arrow::Type::LIST: + case arrow::Type::LARGE_LIST: + case arrow::Type::FIXED_SIZE_LIST: { + if (primitive_type != TYPE_ARRAY) { + return mismatch(); + } + const auto& list = static_cast(arrow_type); + const auto& array = assert_cast(*nested_doris_type); + return validate_arrow_type(*list.value_type(), array.get_nested_type(), path + ".element"); + } + case arrow::Type::STRUCT: { + if (primitive_type != TYPE_STRUCT) { + return mismatch(); + } + const auto& arrow_struct = static_cast(arrow_type); + const auto& doris_struct = assert_cast(*nested_doris_type); + if (cast_set(arrow_struct.num_fields()) != doris_struct.get_elements().size()) { + return mismatch(); + } + for (int field_idx = 0; field_idx < arrow_struct.num_fields(); ++field_idx) { + const auto& arrow_field = arrow_struct.field(field_idx); + const auto& doris_name = doris_struct.get_element_name(field_idx); + if (to_lower(arrow_field->name()) != to_lower(doris_name)) { + return Status::InvalidArgument( + "ADBC Arrow field mismatch at '{}': runtime field '{}' does not match " + "cached field '{}' at ordinal {}", + path, arrow_field->name(), doris_name, field_idx); + } + RETURN_IF_ERROR(validate_arrow_type(*arrow_field->type(), + doris_struct.get_element(field_idx), + path + "." + doris_name)); + } + return Status::OK(); + } + case arrow::Type::MAP: { + if (primitive_type != TYPE_MAP) { + return mismatch(); + } + const auto& arrow_map = static_cast(arrow_type); + const auto& doris_map = assert_cast(*nested_doris_type); + RETURN_IF_ERROR(validate_arrow_type(*arrow_map.key_type(), doris_map.get_key_type(), + path + ".key")); + return validate_arrow_type(*arrow_map.item_type(), doris_map.get_value_type(), + path + ".value"); + } + default: + return mismatch(); + } +} + // Drivers allocate the strings inside AdbcError, so every populated error has to be released. class AdbcErrorGuard { public: @@ -498,6 +633,13 @@ Status AdbcFileReader::get_schema(std::vector* file_schema) co return Status::OK(); } +std::unique_ptr AdbcFileReader::create_column_mapper( + TableColumnMapperOptions options) const { + // ADBC streams return complete Arrow roots, so TableReader must own every nested projection and + // reorder after the full source shape has been materialized. + return std::make_unique(std::move(options)); +} + Status AdbcFileReader::open(std::shared_ptr request) { SCOPED_TIMER(_total_time); SCOPED_TIMER(_open_stream_time); @@ -594,6 +736,13 @@ Status AdbcFileReader::_materialize_record_batch(const arrow::RecordBatch& batch return Status::OK(); } + ArrowMemoryPool<> local_arrow_pool; + arrow::MemoryPool* arrow_pool = ExecEnv::GetInstance()->arrow_memory_pool(); + if (arrow_pool == nullptr) { + // Embedded and unit-test runtimes may omit ExecEnv memory initialization; keep conversions + // on Doris' tracked allocator instead of falling back to Arrow's untracked default pool. + arrow_pool = &local_arrow_pool; + } std::vector materialized_columns(file_block->columns(), false); for (int arrow_idx = 0; arrow_idx < batch.num_columns(); ++arrow_idx) { const std::string& column_name = batch.schema()->field(arrow_idx)->name(); @@ -608,7 +757,7 @@ Status AdbcFileReader::_materialize_record_batch(const arrow::RecordBatch& batch std::shared_ptr array; { SCOPED_TIMER(_normalize_time); - RETURN_IF_ERROR(normalize_arrow_array(batch.column(arrow_idx), &array)); + RETURN_IF_ERROR(normalize_arrow_array(batch.column(arrow_idx), arrow_pool, &array)); } RETURN_IF_ERROR(_materialize_arrow_column(column_name, array, batch.num_rows(), file_id_it->second, block_position_it->second, @@ -643,9 +792,15 @@ Status AdbcFileReader::_materialize_arrow_column(const std::string& column_name, return Status::InternalError("ADBC block position {} out of range, block columns {}", block_position.value(), file_block->columns()); } - auto columns_guard = file_block->mutate_columns_scoped(); - auto& columns = columns_guard.mutable_columns(); - const auto& target_type = columns_guard.get_datatype_by_position(block_position.value()); + const auto& target_type = file_block->get_by_position(block_position.value()).type; + + // The cached FE target remains authoritative when a source schema changes mid-cache-window; + // otherwise SerDes without a null map can silently turn source nulls into default values. + if (array->null_count() > 0 && !target_type->is_nullable()) { + return Status::InternalError( + "ADBC Arrow column '{}' contains {} null rows for non-nullable Doris type {}", + column_name, array->null_count(), target_type->get_name()); + } // An all-null column arrives with a type that says nothing about the column. // @@ -665,10 +820,18 @@ Status AdbcFileReader::_materialize_arrow_column(const std::string& column_name, // Only for a nullable target: substituting defaults into a NOT NULL column would turn a source // that wrongly sent nulls into silently wrong data, so that keeps failing in the serde. if (array->null_count() == array->length() && target_type->is_nullable()) { + auto columns_guard = file_block->mutate_columns_scoped(); + auto& columns = columns_guard.mutable_columns(); columns[block_position.value()]->insert_many_defaults(cast_set(num_rows)); return Status::OK(); } + // Validate the whole logical shape before acquiring mutable block columns. Struct SerDes map + // children by ordinal, so a stale cache entry can otherwise produce plausible but wrong values. + RETURN_IF_ERROR(validate_arrow_type(*array->type(), target_type, column_name)); + + auto columns_guard = file_block->mutate_columns_scoped(); + auto& columns = columns_guard.mutable_columns(); try { RETURN_IF_ERROR(target_type->get_serde()->read_column_from_arrow( *columns[block_position.value()], array.get(), 0, num_rows, _ctz)); diff --git a/be/src/format_v2/table/adbc_reader.h b/be/src/format_v2/table/adbc_reader.h index 8e184364e99d72..39addf8eec2a97 100644 --- a/be/src/format_v2/table/adbc_reader.h +++ b/be/src/format_v2/table/adbc_reader.h @@ -79,6 +79,8 @@ class AdbcFileReader final : public FileReader { Status init(RuntimeState* state) override; Status get_schema(std::vector* file_schema) const override; + std::unique_ptr create_column_mapper( + TableColumnMapperOptions options) const override; Status open(std::shared_ptr request) override; Status get_block(Block* file_block, size_t* rows, bool* eof) override; Status close() override; diff --git a/be/test/format/arrow/arrow_array_normalizer_test.cpp b/be/test/format/arrow/arrow_array_normalizer_test.cpp index d83385981fbecf..75b73b99dbe92a 100644 --- a/be/test/format/arrow/arrow_array_normalizer_test.cpp +++ b/be/test/format/arrow/arrow_array_normalizer_test.cpp @@ -19,9 +19,13 @@ #include #include +#include #include #include +#include +#include +#include #include #include #include @@ -40,6 +44,59 @@ std::shared_ptr build_large_string(const std::vector& return out; } +class AllocateBeforeFreeMemoryPool : public arrow::MemoryPool { +public: + explicit AllocateBeforeFreeMemoryPool(arrow::MemoryPool* delegate) : _delegate(delegate) {} + + arrow::Status Allocate(int64_t size, int64_t alignment, uint8_t** out) override { + auto status = _delegate->Allocate(size, alignment, out); + if (status.ok()) { + _record_allocation(size); + } + return status; + } + + arrow::Status Reallocate(int64_t old_size, int64_t new_size, int64_t alignment, + uint8_t** ptr) override { + uint8_t* replacement = nullptr; + auto status = _delegate->Allocate(new_size, alignment, &replacement); + if (!status.ok()) { + return status; + } + std::memcpy(replacement, *ptr, static_cast(std::min(old_size, new_size))); + _record_allocation(new_size); + _delegate->Free(*ptr, old_size, alignment); + _bytes_allocated -= old_size; + *ptr = replacement; + return arrow::Status::OK(); + } + + void Free(uint8_t* buffer, int64_t size, int64_t alignment) override { + _delegate->Free(buffer, size, alignment); + _bytes_allocated -= size; + } + + int64_t bytes_allocated() const override { return _bytes_allocated; } + int64_t max_memory() const override { return _max_memory; } + int64_t total_bytes_allocated() const override { return _total_bytes_allocated; } + int64_t num_allocations() const override { return _num_allocations; } + std::string backend_name() const override { return "allocate-before-free-test"; } + +private: + void _record_allocation(int64_t size) { + _bytes_allocated += size; + _max_memory = std::max(_max_memory, _bytes_allocated); + _total_bytes_allocated += size; + ++_num_allocations; + } + + arrow::MemoryPool* _delegate; + int64_t _bytes_allocated = 0; + int64_t _max_memory = 0; + int64_t _total_bytes_allocated = 0; + int64_t _num_allocations = 0; +}; + } // namespace // Doris' own Arrow is already in the accepted shape; normalizing must not copy it. @@ -52,7 +109,7 @@ TEST(ArrowArrayNormalizerTest, PlainStringIsAcceptedUnchanged) { EXPECT_TRUE(is_serde_acceptable_arrow_type(*in->type())); std::shared_ptr out; - ASSERT_TRUE(normalize_arrow_array(in, &out).ok()); + ASSERT_TRUE(normalize_arrow_array(in, arrow::default_memory_pool(), &out).ok()); EXPECT_EQ(in.get(), out.get()); } @@ -62,7 +119,7 @@ TEST(ArrowArrayNormalizerTest, LargeStringIsConvertedToString) { EXPECT_FALSE(is_serde_acceptable_arrow_type(*in->type())); std::shared_ptr out; - ASSERT_TRUE(normalize_arrow_array(in, &out).ok()); + ASSERT_TRUE(normalize_arrow_array(in, arrow::default_memory_pool(), &out).ok()); ASSERT_EQ(out->type_id(), arrow::Type::STRING); ASSERT_EQ(out->length(), 3); auto sa = std::static_pointer_cast(out); @@ -87,7 +144,7 @@ TEST(ArrowArrayNormalizerTest, DictionaryIsDecodedToValueType) { EXPECT_FALSE(is_serde_acceptable_arrow_type(*in->type())); std::shared_ptr out; - ASSERT_TRUE(normalize_arrow_array(in, &out).ok()); + ASSERT_TRUE(normalize_arrow_array(in, arrow::default_memory_pool(), &out).ok()); ASSERT_EQ(out->type_id(), arrow::Type::STRING); auto sa = std::static_pointer_cast(out); ASSERT_EQ(sa->length(), 3); @@ -105,7 +162,7 @@ TEST(ArrowArrayNormalizerTest, NullsArePreservedAcrossConversion) { ASSERT_TRUE(b.Finish(&in).ok()); std::shared_ptr out; - ASSERT_TRUE(normalize_arrow_array(in, &out).ok()); + ASSERT_TRUE(normalize_arrow_array(in, arrow::default_memory_pool(), &out).ok()); ASSERT_EQ(out->length(), 2); EXPECT_FALSE(out->IsNull(0)); EXPECT_TRUE(out->IsNull(1)); @@ -124,7 +181,7 @@ TEST(ArrowArrayNormalizerTest, DictionaryOfLargeStringIsFullyNormalized) { auto in = std::make_shared(dict_type, idx, dict); std::shared_ptr out; - ASSERT_TRUE(normalize_arrow_array(in, &out).ok()); + ASSERT_TRUE(normalize_arrow_array(in, arrow::default_memory_pool(), &out).ok()); ASSERT_EQ(out->type_id(), arrow::Type::STRING); auto sa = std::static_pointer_cast(out); ASSERT_EQ(sa->length(), 2); @@ -132,13 +189,311 @@ TEST(ArrowArrayNormalizerTest, DictionaryOfLargeStringIsFullyNormalized) { EXPECT_EQ(sa->GetString(1), "x"); } +TEST(ArrowArrayNormalizerTest, ListViewIsConvertedToListInLogicalOrder) { + arrow::Int32Builder offsets_builder; + ASSERT_TRUE(offsets_builder.AppendValues({2, 0, 1}).ok()); + std::shared_ptr offsets; + ASSERT_TRUE(offsets_builder.Finish(&offsets).ok()); + + arrow::Int32Builder sizes_builder; + ASSERT_TRUE(sizes_builder.AppendValues({2, 1, 2}).ok()); + std::shared_ptr sizes; + ASSERT_TRUE(sizes_builder.Finish(&sizes).ok()); + + arrow::Int32Builder values_builder; + ASSERT_TRUE(values_builder.AppendValues({1, 2, 3, 4}).ok()); + std::shared_ptr values; + ASSERT_TRUE(values_builder.Finish(&values).ok()); + + auto in = arrow::ListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + EXPECT_FALSE(is_serde_acceptable_arrow_type(*in->type())); + + std::shared_ptr out; + ASSERT_TRUE(normalize_arrow_array(in, arrow::default_memory_pool(), &out).ok()); + ASSERT_EQ(out->type_id(), arrow::Type::LIST); + auto list = std::static_pointer_cast(out); + ASSERT_EQ(list->length(), 3); + EXPECT_EQ(list->value_slice(0)->ToString(), "[\n 3,\n 4\n]"); + EXPECT_EQ(list->value_slice(1)->ToString(), "[\n 1\n]"); + EXPECT_EQ(list->value_slice(2)->ToString(), "[\n 2,\n 3\n]"); +} + +TEST(ArrowArrayNormalizerTest, LargeListViewIsConvertedToLargeList) { + arrow::Int64Builder offsets_builder; + ASSERT_TRUE(offsets_builder.AppendValues({1, 0}).ok()); + std::shared_ptr offsets; + ASSERT_TRUE(offsets_builder.Finish(&offsets).ok()); + + arrow::Int64Builder sizes_builder; + ASSERT_TRUE(sizes_builder.AppendValues({2, 1}).ok()); + std::shared_ptr sizes; + ASSERT_TRUE(sizes_builder.Finish(&sizes).ok()); + + arrow::Int32Builder values_builder; + ASSERT_TRUE(values_builder.AppendValues({10, 20, 30}).ok()); + std::shared_ptr values; + ASSERT_TRUE(values_builder.Finish(&values).ok()); + + auto in = arrow::LargeListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + EXPECT_FALSE(is_serde_acceptable_arrow_type(*in->type())); + + std::shared_ptr out; + ASSERT_TRUE(normalize_arrow_array(in, arrow::default_memory_pool(), &out).ok()); + ASSERT_EQ(out->type_id(), arrow::Type::LARGE_LIST); + auto list = std::static_pointer_cast(out); + ASSERT_EQ(list->length(), 2); + EXPECT_EQ(list->value_slice(0)->ToString(), "[\n 20,\n 30\n]"); + EXPECT_EQ(list->value_slice(1)->ToString(), "[\n 10\n]"); +} + +TEST(ArrowArrayNormalizerTest, ListViewCanonicalizationUsesCallerMemoryPool) { + arrow::Int32Builder offsets_builder; + ASSERT_TRUE(offsets_builder.AppendValues({0, 0, 0}).ok()); + std::shared_ptr offsets; + ASSERT_TRUE(offsets_builder.Finish(&offsets).ok()); + + arrow::Int32Builder sizes_builder; + ASSERT_TRUE(sizes_builder.AppendValues({3, 3, 3}).ok()); + std::shared_ptr sizes; + ASSERT_TRUE(sizes_builder.Finish(&sizes).ok()); + + arrow::Int32Builder values_builder; + ASSERT_TRUE(values_builder.AppendValues({10, 20, 30}).ok()); + std::shared_ptr values; + ASSERT_TRUE(values_builder.Finish(&values).ok()); + + auto in = arrow::ListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + arrow::ProxyMemoryPool pool(arrow::default_memory_pool()); + std::shared_ptr out; + ASSERT_TRUE(normalize_arrow_array(in, &pool, &out).ok()); + EXPECT_GT(pool.bytes_allocated(), 0); + out.reset(); + EXPECT_EQ(pool.bytes_allocated(), 0); +} + +TEST(ArrowArrayNormalizerTest, ListViewPreReservesExpandedChildAllocation) { + constexpr int32_t kRows = 1025; + constexpr int32_t kValuesPerRow = 1000; + arrow::Int32Builder offsets_builder; + ASSERT_TRUE(offsets_builder.AppendValues(std::vector(kRows, 0)).ok()); + std::shared_ptr offsets; + ASSERT_TRUE(offsets_builder.Finish(&offsets).ok()); + + arrow::Int32Builder sizes_builder; + ASSERT_TRUE(sizes_builder.AppendValues(std::vector(kRows, kValuesPerRow)).ok()); + std::shared_ptr sizes; + ASSERT_TRUE(sizes_builder.Finish(&sizes).ok()); + + arrow::Int32Builder values_builder; + ASSERT_TRUE(values_builder.AppendValues(std::vector(kValuesPerRow, 7)).ok()); + std::shared_ptr values; + ASSERT_TRUE(values_builder.Finish(&values).ok()); + + auto in = arrow::ListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + AllocateBeforeFreeMemoryPool pool(arrow::default_memory_pool()); + std::shared_ptr out; + ASSERT_TRUE(normalize_arrow_array(in, &pool, &out).ok()); + EXPECT_LE(pool.max_memory(), pool.bytes_allocated() + 256 * 1024); +} + +TEST(ArrowArrayNormalizerTest, NestedListViewOverflowIsRejectedBeforeAllocation) { + const int64_t range_length = std::numeric_limits::max() - 1; + arrow::Int64Builder inner_offsets_builder; + ASSERT_TRUE(inner_offsets_builder.AppendValues({0, 0}).ok()); + std::shared_ptr inner_offsets; + ASSERT_TRUE(inner_offsets_builder.Finish(&inner_offsets).ok()); + arrow::Int64Builder inner_sizes_builder; + ASSERT_TRUE(inner_sizes_builder.AppendValues({range_length, range_length}).ok()); + std::shared_ptr inner_sizes; + ASSERT_TRUE(inner_sizes_builder.Finish(&inner_sizes).ok()); + auto null_values = std::make_shared(std::numeric_limits::max()); + auto inner = arrow::LargeListViewArray::FromArrays(*inner_offsets, *inner_sizes, *null_values) + .ValueOrDie(); + + arrow::Int64Builder outer_offsets_builder; + ASSERT_TRUE(outer_offsets_builder.Append(0).ok()); + std::shared_ptr outer_offsets; + ASSERT_TRUE(outer_offsets_builder.Finish(&outer_offsets).ok()); + arrow::Int64Builder outer_sizes_builder; + ASSERT_TRUE(outer_sizes_builder.Append(2).ok()); + std::shared_ptr outer_sizes; + ASSERT_TRUE(outer_sizes_builder.Finish(&outer_sizes).ok()); + auto outer = arrow::LargeListViewArray::FromArrays(*outer_offsets, *outer_sizes, *inner) + .ValueOrDie(); + + arrow::ProxyMemoryPool pool(arrow::default_memory_pool()); + std::shared_ptr out; + EXPECT_FALSE(normalize_arrow_array(outer, &pool, &out).ok()); + EXPECT_EQ(pool.max_memory(), 0); +} + +TEST(ArrowArrayNormalizerTest, NullableListViewSlicePreservesValidity) { + arrow::Int32Builder offsets_builder; + ASSERT_TRUE(offsets_builder.Append(0).ok()); + ASSERT_TRUE(offsets_builder.AppendNull().ok()); + ASSERT_TRUE(offsets_builder.Append(1).ok()); + std::shared_ptr offsets; + ASSERT_TRUE(offsets_builder.Finish(&offsets).ok()); + + arrow::Int32Builder sizes_builder; + ASSERT_TRUE(sizes_builder.AppendValues({1, 0, 1}).ok()); + std::shared_ptr sizes; + ASSERT_TRUE(sizes_builder.Finish(&sizes).ok()); + + arrow::Int32Builder values_builder; + ASSERT_TRUE(values_builder.AppendValues({10, 20}).ok()); + std::shared_ptr values; + ASSERT_TRUE(values_builder.Finish(&values).ok()); + + auto full = arrow::ListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + auto in = full->Slice(1, 2); + ASSERT_EQ(in->offset(), 1); + + std::shared_ptr out; + ASSERT_TRUE(normalize_arrow_array(in, arrow::default_memory_pool(), &out).ok()); + ASSERT_EQ(out->type_id(), arrow::Type::LIST); + EXPECT_TRUE(out->IsNull(0)); + EXPECT_TRUE(out->IsValid(1)); + auto list = std::static_pointer_cast(out); + EXPECT_EQ(list->value_slice(1)->ToString(), "[\n 20\n]"); +} + +TEST(ArrowArrayNormalizerTest, NullableLargeListViewSlicePreservesValidity) { + arrow::Int64Builder offsets_builder; + ASSERT_TRUE(offsets_builder.Append(0).ok()); + ASSERT_TRUE(offsets_builder.AppendNull().ok()); + ASSERT_TRUE(offsets_builder.Append(1).ok()); + std::shared_ptr offsets; + ASSERT_TRUE(offsets_builder.Finish(&offsets).ok()); + + arrow::Int64Builder sizes_builder; + ASSERT_TRUE(sizes_builder.AppendValues({1, 0, 1}).ok()); + std::shared_ptr sizes; + ASSERT_TRUE(sizes_builder.Finish(&sizes).ok()); + + arrow::Int32Builder values_builder; + ASSERT_TRUE(values_builder.AppendValues({10, 20}).ok()); + std::shared_ptr values; + ASSERT_TRUE(values_builder.Finish(&values).ok()); + + auto full = arrow::LargeListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + auto in = full->Slice(1, 2); + ASSERT_EQ(in->offset(), 1); + + std::shared_ptr out; + ASSERT_TRUE(normalize_arrow_array(in, arrow::default_memory_pool(), &out).ok()); + ASSERT_EQ(out->type_id(), arrow::Type::LARGE_LIST); + EXPECT_TRUE(out->IsNull(0)); + EXPECT_TRUE(out->IsValid(1)); + auto list = std::static_pointer_cast(out); + EXPECT_EQ(list->value_slice(1)->ToString(), "[\n 20\n]"); +} + +TEST(ArrowArrayNormalizerTest, InvalidListViewRangeIsRejectedBeforeCopy) { + arrow::Int32Builder offsets_builder; + ASSERT_TRUE(offsets_builder.AppendValues({0, 2}).ok()); + std::shared_ptr offsets; + ASSERT_TRUE(offsets_builder.Finish(&offsets).ok()); + + arrow::Int32Builder sizes_builder; + ASSERT_TRUE(sizes_builder.AppendValues({1, 1}).ok()); + std::shared_ptr sizes; + ASSERT_TRUE(sizes_builder.Finish(&sizes).ok()); + + arrow::Int32Builder values_builder; + ASSERT_TRUE(values_builder.AppendValues({10, 20}).ok()); + std::shared_ptr values; + ASSERT_TRUE(values_builder.Finish(&values).ok()); + + auto in = arrow::ListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + std::shared_ptr out; + EXPECT_FALSE(normalize_arrow_array(in, arrow::default_memory_pool(), &out).ok()); +} + +TEST(ArrowArrayNormalizerTest, InvalidLargeListViewRangeIsRejectedBeforeCopy) { + arrow::Int64Builder offsets_builder; + ASSERT_TRUE(offsets_builder.AppendValues({0, 2}).ok()); + std::shared_ptr offsets; + ASSERT_TRUE(offsets_builder.Finish(&offsets).ok()); + + arrow::Int64Builder sizes_builder; + ASSERT_TRUE(sizes_builder.AppendValues({1, 1}).ok()); + std::shared_ptr sizes; + ASSERT_TRUE(sizes_builder.Finish(&sizes).ok()); + + arrow::Int32Builder values_builder; + ASSERT_TRUE(values_builder.AppendValues({10, 20}).ok()); + std::shared_ptr values; + ASSERT_TRUE(values_builder.Finish(&values).ok()); + + auto in = arrow::LargeListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + std::shared_ptr out; + EXPECT_FALSE(normalize_arrow_array(in, arrow::default_memory_pool(), &out).ok()); +} + +void expect_dictionary_child_list_view_is_decoded(bool large_offsets) { + arrow::Int8Builder dictionary_builder; + ASSERT_TRUE(dictionary_builder.AppendValues({42, 43}).ok()); + std::shared_ptr dictionary; + ASSERT_TRUE(dictionary_builder.Finish(&dictionary).ok()); + + arrow::Int8Builder indices_builder; + ASSERT_TRUE(indices_builder.AppendValues({0, 1}).ok()); + std::shared_ptr indices; + ASSERT_TRUE(indices_builder.Finish(&indices).ok()); + auto dictionary_type = arrow::dictionary(arrow::int8(), arrow::int8()); + auto values = std::make_shared(dictionary_type, indices, dictionary); + + std::shared_ptr input; + if (large_offsets) { + arrow::Int64Builder offsets_builder; + ASSERT_TRUE(offsets_builder.Append(0).ok()); + std::shared_ptr offsets; + ASSERT_TRUE(offsets_builder.Finish(&offsets).ok()); + arrow::Int64Builder sizes_builder; + ASSERT_TRUE(sizes_builder.Append(2).ok()); + std::shared_ptr sizes; + ASSERT_TRUE(sizes_builder.Finish(&sizes).ok()); + input = arrow::LargeListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + } else { + arrow::Int32Builder offsets_builder; + ASSERT_TRUE(offsets_builder.Append(0).ok()); + std::shared_ptr offsets; + ASSERT_TRUE(offsets_builder.Finish(&offsets).ok()); + arrow::Int32Builder sizes_builder; + ASSERT_TRUE(sizes_builder.Append(2).ok()); + std::shared_ptr sizes; + ASSERT_TRUE(sizes_builder.Finish(&sizes).ok()); + input = arrow::ListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + } + + std::shared_ptr output; + ASSERT_TRUE(normalize_arrow_array(input, arrow::default_memory_pool(), &output).ok()); + ASSERT_TRUE(is_serde_acceptable_arrow_type(*output->type())); + const auto decoded_values = + large_offsets ? std::static_pointer_cast(output)->values() + : std::static_pointer_cast(output)->values(); + ASSERT_EQ(decoded_values->type_id(), arrow::Type::INT8); + const auto decoded = std::static_pointer_cast(decoded_values); + EXPECT_EQ(decoded->Value(0), 42); + EXPECT_EQ(decoded->Value(1), 43); +} + +TEST(ArrowArrayNormalizerTest, ListViewDictionaryChildIsDecoded) { + expect_dictionary_child_list_view_is_decoded(false); +} + +TEST(ArrowArrayNormalizerTest, LargeListViewDictionaryChildIsDecoded) { + expect_dictionary_child_list_view_is_decoded(true); +} + // An unsupported type must name itself, otherwise the offending column cannot be found in prod. TEST(ArrowArrayNormalizerTest, UnsupportedTypeFailsLoudWithTypeName) { auto in = arrow::MakeArrayOfNull(arrow::month_interval(), 1).ValueOrDie(); EXPECT_FALSE(is_serde_acceptable_arrow_type(*in->type())); std::shared_ptr out; - Status st = normalize_arrow_array(in, &out); + Status st = normalize_arrow_array(in, arrow::default_memory_pool(), &out); EXPECT_FALSE(st.ok()); EXPECT_NE(st.to_string().find("interval"), std::string::npos); } diff --git a/be/test/format_v2/table/adbc_reader_test.cpp b/be/test/format_v2/table/adbc_reader_test.cpp index 0d463641cba455..00e7f68554b0c2 100644 --- a/be/test/format_v2/table/adbc_reader_test.cpp +++ b/be/test/format_v2/table/adbc_reader_test.cpp @@ -35,11 +35,17 @@ #include "common/object_pool.h" #include "core/assert_cast.h" #include "core/block/block.h" +#include "core/column/column_array.h" #include "core/column/column_nullable.h" #include "core/column/column_string.h" +#include "core/column/column_struct.h" #include "core/column/column_vector.h" +#include "core/data_type/data_type_array.h" +#include "core/data_type/data_type_nullable.h" #include "core/data_type/data_type_number.h" #include "core/data_type/data_type_string.h" +#include "core/data_type/data_type_struct.h" +#include "format_v2/column_mapper.h" #include "format_v2/file_reader.h" #include "gen_cpp/PlanNodes_types.h" #include "io/file_factory.h" @@ -118,6 +124,33 @@ std::vector string_slot(ObjectPool* pool, DescriptorTbl** desc_ return (*desc_tbl)->get_tuple_descriptor(0)->slots(); } +std::vector int_array_slot(ObjectPool* pool, DescriptorTbl** desc_tbl) { + DescriptorTblBuilder builder(pool); + builder.declare_tuple() << std::make_tuple( + std::make_shared(std::make_shared()), + std::string("c_array")); + *desc_tbl = builder.build(); + return (*desc_tbl)->get_tuple_descriptor(0)->slots(); +} + +DataTypePtr nullable_int32_type() { + return make_nullable(std::make_shared()); +} + +DataTypePtr array_struct_type(const Strings& names) { + DataTypes fields(names.size(), nullable_int32_type()); + return make_nullable(std::make_shared( + make_nullable(std::make_shared(std::move(fields), names)))); +} + +std::vector array_struct_slot(ObjectPool* pool, DescriptorTbl** desc_tbl) { + DescriptorTblBuilder builder(pool); + builder.declare_tuple() << std::make_tuple(array_struct_type({"a", "b"}), + std::string("c_array")); + *desc_tbl = builder.build(); + return (*desc_tbl)->get_tuple_descriptor(0)->slots(); +} + std::vector sqlite_slots(ObjectPool* pool, DescriptorTbl** desc_tbl) { DescriptorTblBuilder builder(pool); // SQLite stores integers as 64-bit and reals as doubles; the ADBC driver reports them as such. @@ -146,6 +179,143 @@ std::shared_ptr make_large_string_batch() { return make_named_large_string_batch("c_str"); } +std::shared_ptr make_list_view_batch() { + arrow::Int32Builder offsets_builder; + EXPECT_TRUE(offsets_builder.AppendValues({2, 0}).ok()); + std::shared_ptr offsets; + EXPECT_TRUE(offsets_builder.Finish(&offsets).ok()); + + arrow::Int32Builder sizes_builder; + EXPECT_TRUE(sizes_builder.AppendValues({1, 2}).ok()); + std::shared_ptr sizes; + EXPECT_TRUE(sizes_builder.Finish(&sizes).ok()); + + arrow::Int32Builder values_builder; + EXPECT_TRUE(values_builder.AppendValues({10, 20, 30}).ok()); + std::shared_ptr values; + EXPECT_TRUE(values_builder.Finish(&values).ok()); + + auto array = arrow::ListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + auto schema = arrow::schema({arrow::field("c_array", array->type())}); + return arrow::RecordBatch::Make(schema, 2, {array}); +} + +std::shared_ptr make_nullable_list_view_batch(bool large_offsets) { + std::shared_ptr offsets; + if (large_offsets) { + arrow::Int64Builder offsets_builder; + EXPECT_TRUE(offsets_builder.Append(0).ok()); + EXPECT_TRUE(offsets_builder.AppendNull().ok()); + EXPECT_TRUE(offsets_builder.Finish(&offsets).ok()); + } else { + arrow::Int32Builder offsets_builder; + EXPECT_TRUE(offsets_builder.Append(0).ok()); + EXPECT_TRUE(offsets_builder.AppendNull().ok()); + EXPECT_TRUE(offsets_builder.Finish(&offsets).ok()); + } + + std::shared_ptr sizes; + if (large_offsets) { + arrow::Int64Builder sizes_builder; + EXPECT_TRUE(sizes_builder.AppendValues({1, 0}).ok()); + EXPECT_TRUE(sizes_builder.Finish(&sizes).ok()); + } else { + arrow::Int32Builder sizes_builder; + EXPECT_TRUE(sizes_builder.AppendValues({1, 0}).ok()); + EXPECT_TRUE(sizes_builder.Finish(&sizes).ok()); + } + + arrow::Int32Builder values_builder; + EXPECT_TRUE(values_builder.Append(10).ok()); + std::shared_ptr values; + EXPECT_TRUE(values_builder.Finish(&values).ok()); + + std::shared_ptr array; + if (large_offsets) { + array = arrow::LargeListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + } else { + array = arrow::ListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + } + auto schema = arrow::schema({arrow::field("c_array", array->type())}); + return arrow::RecordBatch::Make(schema, 2, {array}); +} + +std::shared_ptr make_float_list_view_batch(bool large_offsets) { + std::shared_ptr offsets; + std::shared_ptr sizes; + if (large_offsets) { + arrow::Int64Builder offsets_builder; + EXPECT_TRUE(offsets_builder.Append(0).ok()); + EXPECT_TRUE(offsets_builder.Finish(&offsets).ok()); + arrow::Int64Builder sizes_builder; + EXPECT_TRUE(sizes_builder.Append(1).ok()); + EXPECT_TRUE(sizes_builder.Finish(&sizes).ok()); + } else { + arrow::Int32Builder offsets_builder; + EXPECT_TRUE(offsets_builder.Append(0).ok()); + EXPECT_TRUE(offsets_builder.Finish(&offsets).ok()); + arrow::Int32Builder sizes_builder; + EXPECT_TRUE(sizes_builder.Append(1).ok()); + EXPECT_TRUE(sizes_builder.Finish(&sizes).ok()); + } + + arrow::FloatBuilder values_builder; + EXPECT_TRUE(values_builder.Append(1.0F).ok()); + std::shared_ptr values; + EXPECT_TRUE(values_builder.Finish(&values).ok()); + + std::shared_ptr array; + if (large_offsets) { + array = arrow::LargeListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + } else { + array = arrow::ListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + } + return arrow::RecordBatch::Make(arrow::schema({arrow::field("c_array", array->type())}), 1, + {array}); +} + +std::shared_ptr make_struct_list_view_batch(const Strings& names = {"a", "b"}) { + arrow::Int32Builder offsets_builder; + EXPECT_TRUE(offsets_builder.Append(0).ok()); + std::shared_ptr offsets; + EXPECT_TRUE(offsets_builder.Finish(&offsets).ok()); + arrow::Int32Builder sizes_builder; + EXPECT_TRUE(sizes_builder.Append(2).ok()); + std::shared_ptr sizes; + EXPECT_TRUE(sizes_builder.Finish(&sizes).ok()); + + arrow::Int32Builder a_builder; + EXPECT_TRUE(a_builder.AppendValues({11, 12}).ok()); + std::shared_ptr a; + EXPECT_TRUE(a_builder.Finish(&a).ok()); + arrow::Int32Builder b_builder; + EXPECT_TRUE(b_builder.AppendValues({21, 22}).ok()); + std::shared_ptr b; + EXPECT_TRUE(b_builder.Finish(&b).ok()); + auto values = arrow::StructArray::Make({a, b}, {arrow::field(names[0], arrow::int32()), + arrow::field(names[1], arrow::int32())}) + .ValueOrDie(); + auto array = arrow::ListViewArray::FromArrays(*offsets, *sizes, *values).ValueOrDie(); + return arrow::RecordBatch::Make(arrow::schema({arrow::field("c_array", array->type())}), 1, + {array}); +} + +ColumnDefinition projected_array_struct_b() { + ColumnDefinition b {.identifier = Field::create_field("b"), + .name = "b", + .type = nullable_int32_type()}; + auto struct_type = + make_nullable(std::make_shared(DataTypes {b.type}, Strings {"b"})); + ColumnDefinition element {.identifier = Field::create_field("element"), + .name = "element", + .type = struct_type, + .children = {b}}; + return {.identifier = Field::create_field("c_array"), + .name = "c_array", + .type = make_nullable(std::make_shared(struct_type)), + .children = {element}}; +} + std::unique_ptr create_reader(RuntimeProfile* profile, const TFileRangeDesc& range, const std::vector& slots, AdbcStreamFactory factory) { @@ -301,6 +471,242 @@ TEST(AdbcReaderTest, NormalizesLargeStringBeforeMaterializing) { EXPECT_EQ(*close_count, 1); } +TEST(AdbcReaderTest, NormalizesListViewBeforeMaterializing) { + ObjectPool pool; + DescriptorTbl* desc_tbl = nullptr; + const auto slots = int_array_slot(&pool, &desc_tbl); + RuntimeState state; + RuntimeProfile profile("adbc_reader_list_view_test"); + auto close_count = std::make_shared(0); + + auto reader = create_reader( + &profile, fake_adbc_range(), slots, + [close_count](const TFileRangeDesc&, std::unique_ptr* out) { + *out = std::make_unique( + std::vector> {make_list_view_batch()}, + close_count); + return Status::OK(); + }); + ASSERT_TRUE(reader->init(&state).ok()); + + std::vector schema; + ASSERT_TRUE(reader->get_schema(&schema).ok()); + auto request = std::make_shared(); + FileScanRequestBuilder builder(request.get()); + ASSERT_TRUE(builder.add_non_predicate_column(LocalColumnId(0)).ok()); + ASSERT_TRUE(reader->open(request).ok()); + + auto block = make_request_block(schema, {0}); + size_t rows = 0; + bool eof = false; + const auto status = reader->get_block(&block, &rows, &eof); + ASSERT_TRUE(status.ok()) << status.to_string(); + ASSERT_EQ(rows, 2); + const auto& nullable = assert_cast(*block.get_by_position(0).column); + const auto& array = assert_cast(nullable.get_nested_column()); + ASSERT_EQ(array.get_offsets(), ColumnArray::Offsets64({1, 3})); + const auto& elements = assert_cast(array.get_data()); + const auto& data = assert_cast(elements.get_nested_column()).get_data(); + EXPECT_EQ(data, ColumnInt32::Container({30, 10, 20})); + + ASSERT_TRUE(reader->close().ok()); + EXPECT_EQ(*close_count, 1); +} + +void assert_required_array_rejects_null_list_view(bool large_offsets) { + ObjectPool pool; + DescriptorTbl* desc_tbl = nullptr; + const auto slots = int_array_slot(&pool, &desc_tbl); + RuntimeState state; + RuntimeProfile profile(large_offsets ? "adbc_reader_required_large_list_view_test" + : "adbc_reader_required_list_view_test"); + auto close_count = std::make_shared(0); + + auto reader = create_reader( + &profile, fake_adbc_range(), slots, + [close_count, large_offsets](const TFileRangeDesc&, std::unique_ptr* out) { + *out = std::make_unique( + std::vector> { + make_nullable_list_view_batch(large_offsets)}, + close_count); + return Status::OK(); + }); + ASSERT_TRUE(reader->init(&state).ok()); + + std::vector schema; + ASSERT_TRUE(reader->get_schema(&schema).ok()); + auto request = std::make_shared(); + FileScanRequestBuilder builder(request.get()); + ASSERT_TRUE(builder.add_non_predicate_column(LocalColumnId(0)).ok()); + ASSERT_TRUE(reader->open(request).ok()); + + const auto required_type = remove_nullable(schema[0].type); + ASSERT_FALSE(required_type->is_nullable()); + Block block; + block.insert({required_type->create_column(), required_type, schema[0].name}); + size_t rows = 0; + bool eof = false; + const auto status = reader->get_block(&block, &rows, &eof); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("c_array"), std::string::npos) << status.to_string(); + EXPECT_NE(status.to_string().find("non-nullable"), std::string::npos) << status.to_string(); + + ASSERT_TRUE(reader->close().ok()); + EXPECT_EQ(*close_count, 1); +} + +TEST(AdbcReaderTest, RejectsNullListViewForRequiredArray) { + assert_required_array_rejects_null_list_view(false); +} + +TEST(AdbcReaderTest, RejectsNullLargeListViewForRequiredArray) { + assert_required_array_rejects_null_list_view(true); +} + +void assert_int_array_rejects_float_list_view(bool large_offsets) { + ObjectPool pool; + DescriptorTbl* desc_tbl = nullptr; + const auto slots = int_array_slot(&pool, &desc_tbl); + RuntimeState state; + RuntimeProfile profile(large_offsets ? "adbc_reader_large_list_view_schema_drift_test" + : "adbc_reader_list_view_schema_drift_test"); + auto close_count = std::make_shared(0); + + auto reader = create_reader( + &profile, fake_adbc_range(), slots, + [close_count, large_offsets](const TFileRangeDesc&, std::unique_ptr* out) { + *out = std::make_unique( + std::vector> { + make_float_list_view_batch(large_offsets)}, + close_count); + return Status::OK(); + }); + ASSERT_TRUE(reader->init(&state).ok()); + + std::vector schema; + ASSERT_TRUE(reader->get_schema(&schema).ok()); + auto request = std::make_shared(); + FileScanRequestBuilder builder(request.get()); + ASSERT_TRUE(builder.add_non_predicate_column(LocalColumnId(0)).ok()); + ASSERT_TRUE(reader->open(request).ok()); + + auto block = make_request_block(schema, {0}); + size_t rows = 0; + bool eof = false; + const auto status = reader->get_block(&block, &rows, &eof); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("c_array"), std::string::npos) << status.to_string(); + EXPECT_NE(status.to_string().find("float"), std::string::npos) << status.to_string(); + EXPECT_NE(status.to_string().find("INT"), std::string::npos) << status.to_string(); + + ASSERT_TRUE(reader->close().ok()); + EXPECT_EQ(*close_count, 1); +} + +TEST(AdbcReaderTest, RejectsFloatListViewForCachedIntArray) { + assert_int_array_rejects_float_list_view(false); +} + +TEST(AdbcReaderTest, RejectsFloatLargeListViewForCachedIntArray) { + assert_int_array_rejects_float_list_view(true); +} + +TEST(AdbcReaderTest, CreatesMaterializedColumnMapper) { + ObjectPool pool; + DescriptorTbl* desc_tbl = nullptr; + const auto slots = int_array_slot(&pool, &desc_tbl); + RuntimeProfile profile("adbc_reader_materialized_mapper_test"); + auto reader = create_reader(&profile, fake_adbc_range(), slots, {}); + + auto mapper = reader->create_column_mapper({.mode = TableColumnMappingMode::BY_NAME}); + + ASSERT_NE(dynamic_cast(mapper.get()), nullptr); +} + +TEST(AdbcReaderTest, RejectsReorderedSameTypedStructFields) { + ObjectPool pool; + DescriptorTbl* desc_tbl = nullptr; + const auto slots = array_struct_slot(&pool, &desc_tbl); + RuntimeState state; + RuntimeProfile profile("adbc_reader_reordered_struct_test"); + auto close_count = std::make_shared(0); + auto reader = + create_reader(&profile, fake_adbc_range(), slots, + [close_count](const TFileRangeDesc&, std::unique_ptr* out) { + *out = std::make_unique( + std::vector> { + make_struct_list_view_batch({"b", "a"})}, + close_count); + return Status::OK(); + }); + ASSERT_TRUE(reader->init(&state).ok()); + std::vector schema; + ASSERT_TRUE(reader->get_schema(&schema).ok()); + auto request = std::make_shared(); + FileScanRequestBuilder builder(request.get()); + ASSERT_TRUE(builder.add_non_predicate_column(LocalColumnId(0)).ok()); + ASSERT_TRUE(reader->open(request).ok()); + + auto block = make_request_block(schema, {0}); + size_t rows = 0; + bool eof = false; + const auto status = reader->get_block(&block, &rows, &eof); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("runtime field 'b'"), std::string::npos) + << status.to_string(); + EXPECT_NE(status.to_string().find("cached field 'a'"), std::string::npos) << status.to_string(); + + ASSERT_TRUE(reader->close().ok()); + EXPECT_EQ(*close_count, 1); +} + +TEST(AdbcReaderTest, TableReaderProjectsAListViewStructChildByName) { + ObjectPool pool; + DescriptorTbl* desc_tbl = nullptr; + const auto slots = array_struct_slot(&pool, &desc_tbl); + RuntimeState state; + RuntimeProfile profile("adbc_reader_nested_projection_test"); + auto close_count = std::make_shared(0); + AdbcReader reader([close_count](const TFileRangeDesc&, std::unique_ptr* out) { + *out = std::make_unique( + std::vector> {make_struct_list_view_batch()}, + close_count); + return Status::OK(); + }); + const auto projected = projected_array_struct_b(); + TFileScanRangeParams scan_params; + ASSERT_TRUE(reader.init({.projected_columns = {projected}, + .conjuncts = {}, + .format = FileFormat::ARROW, + .scan_params = &scan_params, + .io_ctx = std::make_shared(), + .runtime_state = &state, + .scanner_profile = &profile, + .file_slot_descs = &slots}) + .ok()); + SplitReadOptions split; + split.current_range = fake_adbc_range(); + split.current_split_format = FileFormat::ARROW; + ASSERT_TRUE(reader.prepare_split(split).ok()); + + Block block; + block.insert({projected.type->create_column(), projected.type, projected.name}); + bool eos = false; + const auto status = reader.get_block(&block, &eos); + ASSERT_TRUE(status.ok()) << status.to_string(); + ASSERT_EQ(block.rows(), 1); + const auto& root = assert_cast(*block.get_by_position(0).column); + const auto& array = assert_cast(root.get_nested_column()); + const auto& element = assert_cast(array.get_data()); + const auto& row = assert_cast(element.get_nested_column()); + const auto& b = assert_cast(row.get_column(0)); + const auto& values = assert_cast(b.get_nested_column()).get_data(); + EXPECT_EQ(values, ColumnInt32::Container({21, 22})); + + ASSERT_TRUE(reader.close().ok()); + EXPECT_EQ(*close_count, 1); +} + // An all-null int64 array standing in for a TEXT column, which is what a source that infers Arrow // types from the values it returns sends when a filter leaves only nulls. std::shared_ptr make_all_null_int64_batch(const std::string& column_name) {