From 6344250fd9ea7540c54765ce67c98c9698cfb997 Mon Sep 17 00:00:00 2001 From: Kam Cheung Ting Date: Mon, 26 Oct 2026 10:07:10 +0000 Subject: [PATCH] feat(data): declare stable physical Arrow schema on task streams ProjectionContext::Make now takes the reader's physical input Arrow schema and builds its reusable projection state from it. The context owns deep copies of the input and output Arrow schemas: each selected output field becomes a deep copy of its physical counterpart, preserving the physical layout, nested list offset widths, metadata, and type-specific flags such as sorted map keys, while only top-level names and nullability come from the projected schema. The borrowed logical-schema pointers and their accessors are removed, so schema inputs are only used during construction. FileScanTaskReader::Open resolves the physical schema before emitting any batch so pass-through, filtered, and projected batches all match the schema reported by the stream, including files that mix list and large_list layouts. Tests cover both projection backends (nanoarrow and Arrow compute): mixed list layouts, sorted-map flags under nullability changes, zero-column projections, and 36 task-reader scenarios that destroy the reader before consuming the stream and validate retained batches after the stream is released. --- src/iceberg/arrow_c_data_util.cc | 54 +++- src/iceberg/arrow_c_data_util_internal.h | 25 +- src/iceberg/data/file_scan_task_reader.cc | 61 ++-- src/iceberg/data/file_scan_task_reader.h | 4 +- src/iceberg/test/arrow_c_data_util_test.cc | 132 ++++++++- .../test/file_scan_task_reader_test.cc | 274 +++++++++++++++++- 6 files changed, 460 insertions(+), 90 deletions(-) diff --git a/src/iceberg/arrow_c_data_util.cc b/src/iceberg/arrow_c_data_util.cc index a77012a69d..331812c871 100644 --- a/src/iceberg/arrow_c_data_util.cc +++ b/src/iceberg/arrow_c_data_util.cc @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -221,9 +222,7 @@ Status AppendValue(const ArrowSchema& input_schema, const ArrowArray& input_arra } // namespace ProjectionContext::ProjectionContext(ProjectionContext&& other) noexcept - : input_schema_(std::exchange(other.input_schema_, nullptr)), - output_schema_(std::exchange(other.output_schema_, nullptr)), - selected_field_indices_(std::move(other.selected_field_indices_)), + : selected_field_indices_(std::move(other.selected_field_indices_)), input_arrow_schema_(other.input_arrow_schema_), output_arrow_schema_(other.output_arrow_schema_), project_batch_function_(std::exchange(other.project_batch_function_, nullptr)), @@ -244,8 +243,6 @@ ProjectionContext& ProjectionContext::operator=(ProjectionContext&& other) noexc output_arrow_schema_.release(&output_arrow_schema_); } - input_schema_ = std::exchange(other.input_schema_, nullptr); - output_schema_ = std::exchange(other.output_schema_, nullptr); selected_field_indices_ = std::move(other.selected_field_indices_); input_arrow_schema_ = other.input_arrow_schema_; other.input_arrow_schema_.release = nullptr; @@ -267,26 +264,53 @@ ProjectionContext::~ProjectionContext() { Result ProjectionContext::Make( const Schema& input_schema, const Schema& output_schema, + const ArrowSchema& physical_input_schema, ProjectionContext::ProjectBatchFunction project_batch_function) { + ArrowSchemaView schema_view; + ArrowError error; + ArrowErrorInit(&error); + ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR( + ArrowSchemaViewInit(&schema_view, &physical_input_schema, &error), error); + ICEBERG_PRECHECK(schema_view.type == NANOARROW_TYPE_STRUCT, + "Physical input schema must be a struct, but is {}", + ArrowTypeString(schema_view.type)); + ICEBERG_PRECHECK( + std::cmp_equal(physical_input_schema.n_children, input_schema.fields().size()), + "Physical input schema must have {} fields to match the logical schema, but has {}", + input_schema.fields().size(), physical_input_schema.n_children); + ICEBERG_ASSIGN_OR_RAISE( auto selected_field_indices, BuildSelectedFieldIndices(input_schema.fields(), output_schema.fields())); - ProjectionContext context; - context.input_schema_ = &input_schema; - context.output_schema_ = &output_schema; - context.selected_field_indices_ = std::move(selected_field_indices); - ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(input_schema, &context.input_arrow_schema_)); - ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(output_schema, &context.output_arrow_schema_)); context.project_batch_function_ = project_batch_function; + ICEBERG_NANOARROW_RETURN_UNEXPECTED( + ArrowSchemaDeepCopy(&physical_input_schema, &context.input_arrow_schema_)); + ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(output_schema, &context.output_arrow_schema_)); + // Replace each logical output field with the physical layout of its selected input + // field. The deep copy also carries the physical metadata (e.g. field ids and + // extension annotations); only the top-level name and nullability come from + // output_schema. + const auto& output_fields = output_schema.fields(); + for (size_t i = 0; i < output_fields.size(); ++i) { + auto* output_field = context.output_arrow_schema_.children[i]; + output_field->release(output_field); + ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowSchemaDeepCopy( + physical_input_schema.children[selected_field_indices[i]], output_field)); + ICEBERG_NANOARROW_RETURN_UNEXPECTED( + ArrowSchemaSetName(output_field, std::string(output_fields[i].name()).c_str())); + // Preserve type-specific flags such as sorted map keys when setting nullability. + if (output_fields[i].optional()) { + output_field->flags |= ARROW_FLAG_NULLABLE; + } else { + output_field->flags &= ~ARROW_FLAG_NULLABLE; + } + } + context.selected_field_indices_ = std::move(selected_field_indices); return context; } -const Schema& ProjectionContext::input_schema() const { return *input_schema_; } - -const Schema& ProjectionContext::output_schema() const { return *output_schema_; } - const ArrowSchema& ProjectionContext::input_arrow_schema() const { return input_arrow_schema_; } diff --git a/src/iceberg/arrow_c_data_util_internal.h b/src/iceberg/arrow_c_data_util_internal.h index e02db29a64..45fdb82586 100644 --- a/src/iceberg/arrow_c_data_util_internal.h +++ b/src/iceberg/arrow_c_data_util_internal.h @@ -61,19 +61,26 @@ class ICEBERG_EXPORT ProjectionContext { /// \brief Resolve the registered ProjectBatch implementation. static auto ResolveProjectBatchFunction() -> ProjectBatchFunction; - /// \brief Build reusable projection state for a validated schema pair. + /// \brief Build reusable projection state using the reader's physical Arrow layout. /// /// \param input_schema Schema that describes every input batch. /// \param output_schema Final schema and column order requested by the caller. + /// \param physical_input_schema Physical fields corresponding to input_schema in + /// the same order. Each selected output field becomes a deep copy of its physical + /// counterpart, preserving the physical layout, nested list offset widths, + /// metadata, and type-specific flags; only the top-level name and nullability come + /// from output_schema. /// \param project_batch_function Optional implementation returned by /// ProjectionContext::ResolveProjectBatchFunction, or nullptr to use the nanoarrow /// path. /// \note It validates that output_schema selects or reorders complete top-level fields - /// by field id. Nested pruning and type changes are rejected. The input_schema and - /// output_schema passed to Make must outlive the context. ProjectBatch may lazily - /// initialize backend cache; do not share one context across concurrent calls. + /// by field id. Nested pruning and type changes are rejected. Schema inputs are only + /// used during construction; the context owns its physical Arrow schemas. + /// ProjectBatch may lazily initialize backend cache; do not share one context across + /// concurrent calls. static Result Make(const Schema& input_schema, const Schema& output_schema, + const ArrowSchema& physical_input_schema, ProjectBatchFunction project_batch_function); ProjectionContext(ProjectionContext&&) noexcept; @@ -83,10 +90,6 @@ class ICEBERG_EXPORT ProjectionContext { ProjectionContext(const ProjectionContext&) = delete; ProjectionContext& operator=(const ProjectionContext&) = delete; - const Schema& input_schema() const; - - const Schema& output_schema() const; - const ArrowSchema& input_arrow_schema() const; const ArrowSchema& output_arrow_schema() const; @@ -100,10 +103,6 @@ class ICEBERG_EXPORT ProjectionContext { private: ProjectionContext() = default; - // Raw schema pointers are borrowed from caller-owned schemas. FileScanTaskReader - // keeps those schema objects alive in the same stream source that owns this context. - const Schema* input_schema_ = nullptr; - const Schema* output_schema_ = nullptr; std::vector selected_field_indices_; ArrowSchema input_arrow_schema_{}; ArrowSchema output_arrow_schema_{}; @@ -237,7 +236,7 @@ ICEBERG_EXPORT Result MakeArrowArrayStream( /// \param row_indices Zero-based row positions to copy from `input_batch`. /// \param projection Reusable schema/projection state created by /// ProjectionContext::Make. -/// \return A newly owned ArrowArray matching `projection.output_schema()`. +/// \return A newly owned ArrowArray matching `projection.output_arrow_schema()`. ICEBERG_EXPORT Result ProjectBatch(ArrowArray* input_batch, std::span row_indices, ProjectionContext& projection); diff --git a/src/iceberg/data/file_scan_task_reader.cc b/src/iceberg/data/file_scan_task_reader.cc index 8c55e31b2b..b25c6cdc64 100644 --- a/src/iceberg/data/file_scan_task_reader.cc +++ b/src/iceberg/data/file_scan_task_reader.cc @@ -25,13 +25,15 @@ #include #include +#include + #include "iceberg/arrow_c_data_guard_internal.h" #include "iceberg/arrow_c_data_util_internal.h" #include "iceberg/data/delete_filter.h" #include "iceberg/file_reader.h" #include "iceberg/manifest/manifest_entry.h" +#include "iceberg/nanoarrow_status_internal.h" #include "iceberg/schema.h" -#include "iceberg/schema_internal.h" #include "iceberg/table_scan.h" #include "iceberg/util/macros.h" @@ -68,22 +70,12 @@ class MergeOnReadStreamSource { public: MergeOnReadStreamSource(std::unique_ptr reader, std::unique_ptr delete_filter, - std::shared_ptr<::iceberg::Schema> required_schema, - std::shared_ptr<::iceberg::Schema> projected_schema, - ProjectionContext projection_context) + bool project_all_columns, ProjectionContext projection_context) : reader_(std::move(reader)), delete_filter_(std::move(delete_filter)), - required_schema_(std::move(required_schema)), - projected_schema_(std::move(projected_schema)), - project_all_rows_(required_schema_->SameSchema(*projected_schema_)), + project_all_columns_(project_all_columns), projection_context_(std::move(projection_context)) {} - ~MergeOnReadStreamSource() { - if (cached_schema_.has_value() && cached_schema_->release != nullptr) { - cached_schema_->release(&cached_schema_.value()); - } - } - Status Close() { if (reader_ == nullptr) { return {}; @@ -92,12 +84,6 @@ class MergeOnReadStreamSource { } Result> Next() { - if (!cached_schema_.has_value()) { - // File readers expose one stable Arrow schema for every batch in the stream. - ICEBERG_ASSIGN_OR_RAISE(cached_schema_, reader_->Schema()); - } - ArrowSchema& input_arrow_schema = cached_schema_.value(); - while (true) { ICEBERG_ASSIGN_OR_RAISE(auto next_batch, reader_->Next()); if (!next_batch.has_value()) { @@ -107,13 +93,14 @@ class MergeOnReadStreamSource { ArrowArray input_batch = std::move(next_batch.value()); internal::ArrowArrayGuard input_batch_guard(&input_batch); - ICEBERG_ASSIGN_OR_RAISE( - auto alive, delete_filter_->ComputeAliveRows(input_arrow_schema, input_batch)); + ICEBERG_ASSIGN_OR_RAISE(auto alive, + delete_filter_->ComputeAliveRows( + projection_context_.input_arrow_schema(), input_batch)); if (alive.empty()) { continue; } - if (alive.alive_count() == input_batch.length && project_all_rows_) { + if (alive.alive_count() == input_batch.length && project_all_columns_) { // Transfer ownership to the stream result; the local guard must not release it. ArrowArray output_batch = input_batch; input_batch.release = nullptr; @@ -125,19 +112,17 @@ class MergeOnReadStreamSource { } Result Schema() { - ArrowSchema schema; - ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(*projected_schema_, &schema)); + ArrowSchema schema{}; + ICEBERG_NANOARROW_RETURN_UNEXPECTED( + ArrowSchemaDeepCopy(&projection_context_.output_arrow_schema(), &schema)); return schema; } private: std::unique_ptr reader_; std::unique_ptr delete_filter_; - std::shared_ptr<::iceberg::Schema> required_schema_; - std::shared_ptr<::iceberg::Schema> projected_schema_; - bool project_all_rows_ = false; + bool project_all_columns_ = false; ProjectionContext projection_context_; - std::optional cached_schema_; }; } // namespace @@ -209,21 +194,25 @@ class FileScanTaskReader::Impl { io_, field_lookup_, has_position_deletes, delete_counter_)); auto required_schema = delete_filter->RequiredSchema(); - auto project_batch_function = ProjectionContext::ResolveProjectBatchFunction(); - // ProjectionContext borrows schemas that are kept in MergeOnReadStreamSource. - ICEBERG_ASSIGN_OR_RAISE(auto projection_context, - ProjectionContext::Make(*required_schema, *projected_schema_, - project_batch_function)); - auto options = MakeReaderOptions(task, io_, required_schema, task.residual_filter(), name_mapping_, properties_, data_file->first_row_id, data_file->data_sequence_number); ICEBERG_ASSIGN_OR_RAISE(auto reader, ReaderFactoryRegistry::Open(data_file->file_format, options)); + // Resolve the physical layout before emitting any batch. Both pass-through and + // projected batches must use the same types, including nested list offset widths. + ICEBERG_ASSIGN_OR_RAISE(auto input_arrow_schema, reader->Schema()); + internal::ArrowSchemaGuard input_schema_guard(&input_arrow_schema); + ICEBERG_ASSIGN_OR_RAISE( + auto projection_context, + ProjectionContext::Make(*required_schema, *projected_schema_, input_arrow_schema, + ProjectionContext::ResolveProjectBatchFunction())); + + const bool project_all_columns = required_schema->SameSchema(*projected_schema_); auto mor_reader = std::make_unique( - std::move(reader), std::move(delete_filter), std::move(required_schema), - projected_schema_, std::move(projection_context)); + std::move(reader), std::move(delete_filter), project_all_columns, + std::move(projection_context)); return MakeArrowArrayStream(std::move(mor_reader)); } diff --git a/src/iceberg/data/file_scan_task_reader.h b/src/iceberg/data/file_scan_task_reader.h index 33a4b82fff..03b0f0c9f0 100644 --- a/src/iceberg/data/file_scan_task_reader.h +++ b/src/iceberg/data/file_scan_task_reader.h @@ -38,7 +38,9 @@ namespace iceberg { /// /// FileScanTaskReader chooses the copy-on-write path for tasks without deletes and /// the merge-on-read path for tasks with v2 position or equality deletes. The returned -/// stream always exposes `projected_schema`. +/// stream exposes the fields of `projected_schema` in one stable physical Arrow layout, +/// which may depend on the file format and reader properties (for example, `large_list` +/// instead of `list`). Every batch matches the schema reported by the stream. /// /// TODO(gangwu): Add a mode that emits a `_deleted` column instead of filtering rows. /// TODO(gangwu): Use evaluator to apply residual expression filters. diff --git a/src/iceberg/test/arrow_c_data_util_test.cc b/src/iceberg/test/arrow_c_data_util_test.cc index c93a928a3d..1d0567e9bf 100644 --- a/src/iceberg/test/arrow_c_data_util_test.cc +++ b/src/iceberg/test/arrow_c_data_util_test.cc @@ -44,12 +44,8 @@ namespace iceberg::internal { namespace { -std::shared_ptr<::arrow::RecordBatch> MakeBatch(const Schema& schema, - std::string_view json) { - ArrowSchema c_schema; - EXPECT_THAT(ToArrowSchema(schema, &c_schema), IsOk()); - // ImportSchema takes ownership of c_schema and calls release. - auto arrow_schema = ::arrow::ImportSchema(&c_schema).ValueOrDie(); +std::shared_ptr<::arrow::RecordBatch> MakeBatch( + const std::shared_ptr<::arrow::Schema>& arrow_schema, std::string_view json) { auto struct_type = ::arrow::struct_(arrow_schema->fields()); return ::arrow::RecordBatch::FromStructArray( ::arrow::json::ArrayFromJSONString(struct_type, std::string(json)) @@ -57,6 +53,13 @@ std::shared_ptr<::arrow::RecordBatch> MakeBatch(const Schema& schema, .ValueOrDie(); } +std::shared_ptr<::arrow::RecordBatch> MakeBatch(const Schema& schema, + std::string_view json) { + ArrowSchema c_schema; + EXPECT_THAT(ToArrowSchema(schema, &c_schema), IsOk()); + return MakeBatch(::arrow::ImportSchema(&c_schema).ValueOrDie(), json); +} + ProjectionContext::ProjectBatchFunction ArrowComputeFunction() { arrow::RegisterAll(); auto function = ProjectionContext::ResolveProjectBatchFunction(); @@ -74,15 +77,17 @@ std::shared_ptr<::arrow::RecordBatch> RunProjectBatch( ArrowSchemaGuard schema_guard(&c_schema); ArrowArrayGuard array_guard(&c_array); - auto projection = - ProjectionContext::Make(required_schema, projected_schema, project_batch_function); + auto projection = ProjectionContext::Make(required_schema, projected_schema, c_schema, + project_batch_function); EXPECT_THAT(projection, IsOk()); + // The context must own its Arrow schemas independently of the caller. + c_schema.release(&c_schema); auto result = ProjectBatch(&c_array, alive_indices, projection.value()); EXPECT_THAT(result, IsOk()); ArrowSchema out_c_schema; - EXPECT_THAT(ToArrowSchema(projected_schema, &out_c_schema), IsOk()); + EXPECT_EQ(ArrowSchemaDeepCopy(&projection->output_arrow_schema(), &out_c_schema), 0); auto arrow_out_schema = ::arrow::ImportSchema(&out_c_schema).ValueOrDie(); ArrowArray out_array = std::exchange(result.value(), ArrowArray{}); @@ -171,6 +176,87 @@ TEST(ProjectBatchTest, EmptyRowSelection) { ExpectProjectBatch(*batch, alive, *schema, *schema, R"([])"); } +TEST(ProjectBatchTest, EmptyProjectionPreservesRowCount) { + auto required = MakeFullSchema(); + Schema projected(std::vector{}); + auto batch = MakeBatch(*required, R"([[1,"a",1.0],[2,"b",2.0],[3,"c",3.0]])"); + + for (auto backend : + {ProjectionContext::ProjectBatchFunction{nullptr}, ArrowComputeFunction()}) { + SCOPED_TRACE(backend == nullptr ? "nanoarrow" : "Arrow compute"); + for (const auto& alive : {std::vector{}, std::vector{0, 2}}) { + auto actual = RunProjectBatch(*batch, alive, *required, projected, backend); + EXPECT_EQ(actual->num_columns(), 0); + EXPECT_EQ(actual->num_rows(), alive.size()); + EXPECT_TRUE(actual->ValidateFull().ok()) << actual->ValidateFull().ToString(); + } + } +} + +TEST(ProjectBatchTest, PreservesPhysicalListLayoutsWithBothBackends) { + auto nested_list = std::make_shared(SchemaField::MakeOptional( + 3, "element", + std::make_shared(SchemaField::MakeOptional(4, "element", int32())))); + Schema required( + std::vector{SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeOptional(2, "numbers", nested_list), + SchemaField::MakeOptional(5, "unused", string())}); + Schema projected(std::vector{required.fields()[1].WithName("renamed"), + required.fields()[0]}); + ArrowSchema c_schema; + ASSERT_THAT(ToArrowSchema(required, &c_schema), IsOk()); + auto physical_schema = ::arrow::ImportSchema(&c_schema).ValueOrDie(); + auto fields = physical_schema->fields(); + // Keep the inner list at 32-bit offsets while the outer list uses 64-bit offsets. + fields[1] = fields[1]->WithType(::arrow::large_list(fields[1]->type()->field(0))); + physical_schema = ::arrow::schema(std::move(fields)); + auto batch = MakeBatch(physical_schema, + R"([[1, [[10, null], []], "a"], [2, [], "b"], + [3, [[30]], "c"], [4, null, "d"]])"); + auto expected_schema = ::arrow::schema( + {physical_schema->field(1)->WithName("renamed"), physical_schema->field(0)}); + auto expected = + MakeBatch(expected_schema, R"([[[[10, null], []], 1], [[], 2], [null, 4]])"); + + for (auto backend : + {ProjectionContext::ProjectBatchFunction{nullptr}, ArrowComputeFunction()}) { + SCOPED_TRACE(backend == nullptr ? "nanoarrow" : "Arrow compute"); + const std::vector alive{0, 1, 3}; + auto actual = RunProjectBatch(*batch, alive, required, projected, backend); + ASSERT_TRUE(actual->schema()->Equals(*expected_schema)); + ASSERT_TRUE(actual->ValidateFull().ok()) << actual->ValidateFull().ToString(); + EXPECT_TRUE(actual->Equals(*expected)) << actual->ToString(); + } +} + +TEST(ProjectBatchTest, PreservesPhysicalMapFlagsWhenChangingNullability) { + auto map_type = map(SchemaField::MakeRequired(2, "key", int32()), + SchemaField::MakeOptional(3, "value", string())); + Schema required( + std::vector{SchemaField::MakeRequired(1, "values", map_type)}); + Schema projected( + std::vector{SchemaField::MakeOptional(1, "renamed", map_type)}); + auto canonical = MakeBatch(required, "[]")->schema(); + const auto& canonical_map = + static_cast(*canonical->field(0)->type()); + auto physical_field = canonical->field(0)->WithType(std::make_shared<::arrow::MapType>( + canonical_map.key_field(), canonical_map.item_field(), /*keys_sorted=*/true)); + auto batch = + MakeBatch(::arrow::schema({physical_field}), R"([[[[1, "a"], [2, "b"]]], [[]]])"); + auto expected_schema = + ::arrow::schema({physical_field->WithName("renamed")->WithNullable(true)}); + auto expected = MakeBatch(expected_schema, R"([[[[1, "a"], [2, "b"]]]])"); + + for (auto backend : + {ProjectionContext::ProjectBatchFunction{nullptr}, ArrowComputeFunction()}) { + SCOPED_TRACE(backend == nullptr ? "nanoarrow" : "Arrow compute"); + auto actual = RunProjectBatch(*batch, {0}, required, projected, backend); + EXPECT_TRUE(actual->schema()->Equals(*expected_schema)); + EXPECT_TRUE(actual->ValidateFull().ok()) << actual->ValidateFull().ToString(); + EXPECT_TRUE(actual->Equals(*expected)) << actual->ToString(); + } +} + TEST(ProjectBatchTest, ProjectionRejectsNestedPruning) { auto input_schema = Schema(std::vector{ SchemaField::MakeOptional(1, "person", @@ -186,9 +272,33 @@ TEST(ProjectBatchTest, ProjectionRejectsNestedPruning) { })), }); - auto projection = ProjectionContext::Make(input_schema, output_schema, nullptr); + ArrowSchema physical_schema; + ASSERT_THAT(ToArrowSchema(input_schema, &physical_schema), IsOk()); + ArrowSchemaGuard schema_guard(&physical_schema); + EXPECT_THAT( + ProjectionContext::Make(input_schema, output_schema, physical_schema, nullptr), + IsError(ErrorKind::kInvalidArgument)); +} + +TEST(ProjectBatchTest, PhysicalSchemaMustMatchLogicalFieldCount) { + auto required = MakeFullSchema(); + Schema projected(std::vector{required->fields()[0]}); + ArrowSchema physical_schema; + ASSERT_THAT(ToArrowSchema(projected, &physical_schema), IsOk()); + ArrowSchemaGuard schema_guard(&physical_schema); + + EXPECT_THAT(ProjectionContext::Make(*required, projected, physical_schema, nullptr), + HasErrorMessage("Physical input schema must have 3 fields")); +} + +TEST(ProjectBatchTest, PhysicalSchemaMustBeAStruct) { + Schema schema(std::vector{}); + ArrowSchema physical_schema; + ASSERT_EQ(ArrowSchemaInitFromType(&physical_schema, NANOARROW_TYPE_INT32), 0); + ArrowSchemaGuard schema_guard(&physical_schema); - EXPECT_THAT(projection, IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT(ProjectionContext::Make(schema, schema, physical_schema, nullptr), + HasErrorMessage("Physical input schema must be a struct")); } } // namespace iceberg::internal diff --git a/src/iceberg/test/file_scan_task_reader_test.cc b/src/iceberg/test/file_scan_task_reader_test.cc index 05416bd7d7..3e936ee90c 100644 --- a/src/iceberg/test/file_scan_task_reader_test.cc +++ b/src/iceberg/test/file_scan_task_reader_test.cc @@ -33,6 +33,7 @@ #include #include #include +#include #include #include #include @@ -264,11 +265,8 @@ class FileScanTaskReaderTest : public TempFileTestBase { return path; } - Result> MakeDataFile(std::shared_ptr schema, - const std::string& json_data, - int64_t record_count = 3) { - ICEBERG_ASSIGN_OR_RAISE(auto path, - CreateParquetDataFile(std::move(schema), json_data)); + Result> MakeDataFileForPath(const std::string& path, + int64_t record_count) { ICEBERG_ASSIGN_OR_RAISE(auto input_file, file_io_->NewInputFile(path)); ICEBERG_ASSIGN_OR_RAISE(auto size, input_file->Size()); return std::make_shared(DataFile{ @@ -280,21 +278,21 @@ class FileScanTaskReaderTest : public TempFileTestBase { }); } + Result> MakeDataFile(std::shared_ptr schema, + const std::string& json_data, + int64_t record_count = 3) { + ICEBERG_ASSIGN_OR_RAISE(auto path, + CreateParquetDataFile(std::move(schema), json_data)); + return MakeDataFileForPath(path, record_count); + } + Result> MakeDataFile( std::shared_ptr schema, const std::vector& json_batches, int64_t record_count, int64_t max_row_group_length) { ICEBERG_ASSIGN_OR_RAISE( auto path, CreateParquetDataFile(std::move(schema), json_batches, max_row_group_length)); - ICEBERG_ASSIGN_OR_RAISE(auto input_file, file_io_->NewInputFile(path)); - ICEBERG_ASSIGN_OR_RAISE(auto size, input_file->Size()); - return std::make_shared(DataFile{ - .content = DataFile::Content::kData, - .file_path = path, - .file_format = FileFormatType::kParquet, - .record_count = record_count, - .file_size_in_bytes = size, - }); + return MakeDataFileForPath(path, record_count); } Result> MakePositionDeleteFile( @@ -644,6 +642,42 @@ TEST_F(FileScanTaskReaderTest, OpenWithPositionDeletesFiltersRowsAndPrunesPos) { ASSERT_NO_FATAL_FAILURE(VerifyStream(&stream, R"([[1, "Foo"], [3, "Baz"]])")); } +TEST_F(FileScanTaskReaderTest, EmptyProjectionPreservesLiveRowCount) { + ICEBERG_UNWRAP_OR_FAIL( + auto data_file, + MakeDataFile(table_schema_, + R"([[1, "Foo", "blue"], [2, "Bar", "red"], [3, "Baz", "green"]])")); + ICEBERG_UNWRAP_OR_FAIL( + auto deletion, MakePositionDeleteFile(CreateNewTempFilePathWithSuffix(".parquet"), + {1}, data_file->file_path)); + ICEBERG_UNWRAP_OR_FAIL( + auto reader, + FileScanTaskReader::Make({ + .io = file_io_, + .table_schema = table_schema_, + .projected_schema = std::make_shared(std::vector{}), + })); + ICEBERG_UNWRAP_OR_FAIL(auto opened_stream, + reader->Open(FileScanTask(data_file, {deletion}))); + nanoarrow::UniqueArrayStream stream(&opened_stream); + nanoarrow::UniqueSchema schema; + ASSERT_EQ(stream->get_schema(stream.get(), schema.get()), 0) + << stream->get_last_error(stream.get()); + EXPECT_EQ(schema->n_children, 0); + + nanoarrow::UniqueArray batch; + ASSERT_EQ(stream->get_next(stream.get(), batch.get()), 0) + << stream->get_last_error(stream.get()); + ASSERT_NE(batch->release, nullptr); + EXPECT_EQ(batch->n_children, 0); + EXPECT_EQ(batch->length, 2); + + nanoarrow::UniqueArray end; + ASSERT_EQ(stream->get_next(stream.get(), end.get()), 0) + << stream->get_last_error(stream.get()); + EXPECT_EQ(end->release, nullptr); +} + TEST_F(FileScanTaskReaderTest, OpenWithDeletionVectorFiltersRows) { ICEBERG_UNWRAP_OR_FAIL( auto data_file, @@ -793,6 +827,218 @@ TEST_F(FileScanTaskReaderTest, OpenWithEqualityDeletesKeepsInputBatchWhenAllRows VerifyStream(&stream, R"([[1, "Foo"], [2, "Bar"], [3, "Baz"]])")); } +enum class PhysicalListLayout { kList, kLargeList, kStoredMixed }; +enum class DeleteBatchSelection { + kFirstRow, + kSecondBatchRow, + kNoOp, + kFirstBatch, + kAll, + kEmptySplit +}; + +class FileScanTaskPhysicalSchemaTest + : public FileScanTaskReaderTest, + public ::testing::WithParamInterface< + std::tuple> { + public: + static std::string ParamName(const ::testing::TestParamInfo& info) { + const auto& [deletes, prune, layout] = info.param; + constexpr std::array names{"DeleteFirst", "DeleteSecond", "NoOp", + "DeleteFirstBatch", "DeleteAll", "EmptySplit"}; + std::string name = names[static_cast(deletes)]; + name += prune ? "Pruned" : "AllColumns"; + name += + layout == PhysicalListLayout::kList + ? "List" + : (layout == PhysicalListLayout::kLargeList ? "LargeList" : "StoredMixed"); + return name; + } + + protected: + Result> MakeDataFileWithStoredSchema( + const std::shared_ptr<::arrow::Schema>& schema, const std::string& json) { + ICEBERG_ASSIGN_OR_RAISE(auto batch, MakeRecordBatch(schema, json)); + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto table, + ::arrow::Table::FromRecordBatches(schema, {batch})); + auto path = CreateNewTempFilePathWithSuffix(".parquet"); + ICEBERG_ASSIGN_OR_RAISE(auto output, arrow::OpenArrowOutputStream(file_io_, path)); + // store_schema() embeds the ARROW:schema metadata so the file keeps the exact + // list/large_list layout of `schema` instead of the reader's configured type. + auto writer_properties = + ::parquet::ArrowWriterProperties::Builder().store_schema()->build(); + ICEBERG_ARROW_RETURN_NOT_OK(::parquet::arrow::WriteTable( + *table, ::arrow::default_memory_pool(), output, table->num_rows(), + ::parquet::default_writer_properties(), writer_properties)); + ICEBERG_ARROW_RETURN_NOT_OK(output->Close()); + return MakeDataFileForPath(path, table->num_rows()); + } +}; + +TEST_P(FileScanTaskPhysicalSchemaTest, StablePhysicalSchemaAcrossDeleteBatches) { + const auto [deletes, prune_columns, layout] = GetParam(); + const bool use_large_list = layout == PhysicalListLayout::kLargeList; + table_schema_ = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeOptional( + 2, "numbers", + std::make_shared(SchemaField::MakeOptional(3, "element", int32()))), + SchemaField::MakeOptional( + 4, "nested", + std::make_shared(std::vector{SchemaField::MakeOptional( + 5, "tags", + std::make_shared( + SchemaField::MakeOptional(6, "element", int32())))}))}); + projected_schema_ = prune_columns + ? std::make_shared(std::vector{ + table_schema_->fields()[2], table_schema_->fields()[1]}) + : table_schema_; + ICEBERG_UNWRAP_OR_FAIL(auto expected_schema, MakeArrowSchema(*table_schema_)); + if (layout != PhysicalListLayout::kList) { + auto fields = expected_schema->fields(); + fields[1] = fields[1]->WithType(::arrow::large_list(fields[1]->type()->field(0))); + if (use_large_list) { + auto tags = fields[2]->type()->field(0); + fields[2] = fields[2]->WithType(::arrow::struct_( + {tags->WithType(::arrow::large_list(tags->type()->field(0)))})); + } + expected_schema = ::arrow::schema(std::move(fields)); + } + const std::string json = + R"([[1, [10, 11], {"tags": [1, null]}], [2, null, {"tags": []}], + [3, [], null], [4, [40], {"tags": null}]])"; + std::string delete_json; + std::vector expected_json{ + R"([[1, [10, 11], {"tags": [1, null]}], [2, null, {"tags": []}]])", + R"([[3, [], null], [4, [40], {"tags": null}]])"}; + switch (deletes) { + case DeleteBatchSelection::kFirstRow: + delete_json = "[[1]]"; + expected_json[0] = R"([[2, null, {"tags": []}]])"; + break; + case DeleteBatchSelection::kSecondBatchRow: + delete_json = "[[3]]"; + expected_json[1] = R"([[4, [40], {"tags": null}]])"; + break; + case DeleteBatchSelection::kNoOp: + delete_json = "[[99]]"; + break; + case DeleteBatchSelection::kFirstBatch: + delete_json = "[[1], [2]]"; + expected_json.erase(expected_json.begin()); + break; + case DeleteBatchSelection::kAll: + delete_json = "[[1], [2], [3], [4]]"; + expected_json.clear(); + break; + case DeleteBatchSelection::kEmptySplit: + delete_json = "[[99]]"; + expected_json.clear(); + break; + } + ICEBERG_UNWRAP_OR_FAIL(auto data_file, + layout == PhysicalListLayout::kStoredMixed + ? MakeDataFileWithStoredSchema(expected_schema, json) + : MakeDataFile(table_schema_, json, /*record_count=*/4)); + auto delete_schema = + std::make_shared(std::vector{table_schema_->fields()[0]}); + ICEBERG_UNWRAP_OR_FAIL( + auto eq_delete, MakeEqualityDeleteFile(CreateNewTempFilePathWithSuffix(".parquet"), + delete_schema, delete_json, {1})); + // A split covering only the first Parquet magic byte selects no row groups. + const int64_t length = + deletes == DeleteBatchSelection::kEmptySplit ? 1 : data_file->file_size_in_bytes; + ICEBERG_UNWRAP_OR_FAIL(auto task, + FileScanTask::MakeSplit(data_file, 0, length, {eq_delete})); + + ReaderProperties properties; + properties.Set(ReaderProperties::kBatchSize, int64_t{2}); + properties.Set(ReaderProperties::kArrowUseLargeList, use_large_list); + FileScanTaskReader::Options options{.io = file_io_, + .table_schema = table_schema_, + .projected_schema = projected_schema_, + .properties = properties.configs()}; + ICEBERG_UNWRAP_OR_FAIL(auto reader, FileScanTaskReader::Make(std::move(options))); + ICEBERG_UNWRAP_OR_FAIL(auto opened_stream, reader->Open(*task)); + nanoarrow::UniqueArrayStream stream(&opened_stream); + // The stream owns everything it needs after the task reader is destroyed. + reader.reset(); + nanoarrow::UniqueSchema retained_schema; + ASSERT_EQ(stream->get_schema(stream.get(), retained_schema.get()), 0) + << stream->get_last_error(stream.get()); + auto output_schema = + prune_columns + ? ::arrow::schema({expected_schema->field(2), expected_schema->field(1)}) + : expected_schema; + auto verify_schema = [&]() { + nanoarrow::UniqueSchema declared_schema; + ASSERT_EQ(stream->get_schema(stream.get(), declared_schema.get()), 0) + << stream->get_last_error(stream.get()); + auto import_result = ::arrow::ImportSchema(declared_schema.get()); + ASSERT_TRUE(import_result.ok()) << import_result.status().ToString(); + auto actual = import_result.ValueOrDie(); + ASSERT_TRUE(actual->Equals(*output_schema)) << actual->ToString() << "\nExpected:\n" + << output_schema->ToString(); + }; + // Check the declaration before importing arrays: a list/large_list mismatch can + // otherwise interpret 64-bit offsets as 32-bit offsets without an import error. + // Each callback result is released independently of the stream's cached schema. + ASSERT_NO_FATAL_FAILURE(verify_schema()); + ASSERT_NO_FATAL_FAILURE(verify_schema()); + + std::vector> retained_batches; + std::vector> expected_batches; + for (const auto& json : expected_json) { + nanoarrow::UniqueArray array; + ASSERT_EQ(stream->get_next(stream.get(), array.get()), 0) + << stream->get_last_error(stream.get()); + ASSERT_NE(array->release, nullptr); + ASSERT_NO_FATAL_FAILURE(verify_schema()); + auto result = ::arrow::ImportRecordBatch(array.get(), output_schema); + ASSERT_TRUE(result.ok()) << result.status().ToString(); + auto actual = result.ValueOrDie(); + ASSERT_NE(actual, nullptr); + ASSERT_TRUE(actual->ValidateFull().ok()) << actual->ValidateFull().ToString(); + ASSERT_TRUE(actual->schema()->Equals(*output_schema)); + ICEBERG_UNWRAP_OR_FAIL(auto expected, MakeRecordBatch(expected_schema, json)); + if (prune_columns) { + expected = expected->SelectColumns({2, 1}).ValueOrDie(); + } + EXPECT_TRUE(actual->Equals(*expected)) << "Actual:\n" + << actual->ToString() << "\nExpected:\n" + << expected->ToString(); + retained_batches.push_back(std::move(actual)); + expected_batches.push_back(std::move(expected)); + } + nanoarrow::UniqueArray end; + ASSERT_EQ(stream->get_next(stream.get(), end.get()), 0) + << stream->get_last_error(stream.get()); + EXPECT_EQ(end->release, nullptr); + ASSERT_NO_FATAL_FAILURE(verify_schema()); + + // Returned schemas and batches remain valid after the stream's context is released. + stream.reset(); + auto retained_result = ::arrow::ImportSchema(retained_schema.get()); + ASSERT_TRUE(retained_result.ok()) << retained_result.status().ToString(); + EXPECT_TRUE(retained_result.ValueOrDie()->Equals(*output_schema)); + for (size_t i = 0; i < retained_batches.size(); ++i) { + ASSERT_TRUE(retained_batches[i]->ValidateFull().ok()); + EXPECT_TRUE(retained_batches[i]->Equals(*expected_batches[i])); + } +} + +INSTANTIATE_TEST_SUITE_P( + DeleteBatches, FileScanTaskPhysicalSchemaTest, + ::testing::Combine( + ::testing::Values(DeleteBatchSelection::kFirstRow, + DeleteBatchSelection::kSecondBatchRow, + DeleteBatchSelection::kNoOp, DeleteBatchSelection::kFirstBatch, + DeleteBatchSelection::kAll, DeleteBatchSelection::kEmptySplit), + ::testing::Bool(), + ::testing::Values(PhysicalListLayout::kList, PhysicalListLayout::kLargeList, + PhysicalListLayout::kStoredMixed)), + FileScanTaskPhysicalSchemaTest::ParamName); + TEST_F(FileScanTaskReaderTest, OpenWithSchemasResolvesDroppedEqualityField) { auto current_schema = std::make_shared( std::vector{SchemaField::MakeRequired(1, "id", int32()),