From c6e1875927cf2580216923a14ee7b2076ed012a9 Mon Sep 17 00:00:00 2001 From: Xin Huang Date: Fri, 18 Sep 2026 15:45:02 -0700 Subject: [PATCH 1/9] =?UTF-8?q?=E2=9C=A8=20feat(schema):=20Add=20write-def?= =?UTF-8?q?ault=20and=20strict=20projection=20options?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend Project with ProjectionOptions so the same field-id matching serves both paths: reads fill missing fields with initial-default and allow type promotion, writes fill with write-default and require exact types. --- src/iceberg/schema_util.cc | 43 +++++++++++------ src/iceberg/schema_util.h | 23 +++++++++- src/iceberg/test/schema_util_test.cc | 69 ++++++++++++++++++++++++++++ 3 files changed, 121 insertions(+), 14 deletions(-) diff --git a/src/iceberg/schema_util.cc b/src/iceberg/schema_util.cc index b9f32346a9..32082f5590 100644 --- a/src/iceberg/schema_util.cc +++ b/src/iceberg/schema_util.cc @@ -35,7 +35,8 @@ namespace iceberg { namespace { -Status ValidateSchemaEvolution(const Type& expected_type, const Type& source_type) { +Status ValidateSchemaEvolution(const Type& expected_type, const Type& source_type, + bool allow_type_promotion) { if (expected_type.is_nested()) { // Nested type requires identical type ids but their sub-fields are checked // recursively and individually. @@ -52,6 +53,10 @@ Status ValidateSchemaEvolution(const Type& expected_type, const Type& source_typ if (source_type.type_id() == TypeId::kUnknown && expected_type.is_primitive()) { return {}; } + if (!allow_type_promotion) { + return NotSupported("Cannot project {} as {}: type promotion is not allowed", + source_type, expected_type); + } switch (expected_type.type_id()) { case TypeId::kLong: { @@ -83,11 +88,11 @@ Status ValidateSchemaEvolution(const Type& expected_type, const Type& source_typ } Result ProjectNested(const Type& expected_type, const Type& source_type, - bool prune_source); + bool prune_source, const ProjectionOptions& options); Result ProjectField(const SchemaField& expected_field, const SchemaField& source_field, size_t source_index, - bool prune_source) { + bool prune_source, const ProjectionOptions& options) { FieldProjection projection; if (expected_field.type()->type_id() == TypeId::kUnknown) { @@ -112,10 +117,12 @@ Result ProjectField(const SchemaField& expected_field, if (expected_field.type()->is_nested()) { ICEBERG_ASSIGN_OR_RAISE( projection, - ProjectNested(*expected_field.type(), *source_field.type(), prune_source)); + ProjectNested(*expected_field.type(), *source_field.type(), prune_source, + options)); } else { - ICEBERG_RETURN_UNEXPECTED( - ValidateSchemaEvolution(*expected_field.type(), *source_field.type())); + ICEBERG_RETURN_UNEXPECTED(ValidateSchemaEvolution(*expected_field.type(), + *source_field.type(), + options.allow_type_promotion)); } // If `prune_source` is false, all fields will be read so the local index is exactly @@ -127,7 +134,7 @@ Result ProjectField(const SchemaField& expected_field, } Result ProjectNested(const Type& expected_type, const Type& source_type, - bool prune_source) { + bool prune_source, const ProjectionOptions& options) { if (!expected_type.is_nested()) { return InvalidSchema("Expected a nested type, but got {}", expected_type); } @@ -169,17 +176,25 @@ Result ProjectNested(const Type& expected_type, const Type& sou if (auto iter = source_field_map.find(field_id); iter != source_field_map.cend()) { ICEBERG_ASSIGN_OR_RAISE(child_projection, ProjectField(expected_field, *iter->second.field, - iter->second.local_index, prune_source)); + iter->second.local_index, prune_source, + options)); } else if (MetadataColumns::IsMetadataColumn(field_id)) { child_projection.kind = FieldProjection::Kind::kMetadata; - } else if (expected_field.initial_default() != nullptr) { - // Rows written before the field existed assume its `initial-default` value. + } else if (const auto& default_value = + options.default_policy == ProjectionOptions::DefaultPolicy::kWrite + ? expected_field.write_default() + : expected_field.initial_default(); + default_value != nullptr) { child_projection.kind = FieldProjection::Kind::kDefault; - child_projection.from = *expected_field.initial_default(); + child_projection.from = *default_value; } else if (expected_field.optional()) { child_projection.kind = FieldProjection::Kind::kNull; } else { // TODO(gangwu): support constant value + if (options.default_policy == ProjectionOptions::DefaultPolicy::kWrite) { + return InvalidSchema("Missing required field: {} without a write-default", + expected_field.ToString()); + } return InvalidSchema("Missing required field: {}", expected_field.ToString()); } result.children.emplace_back(std::move(child_projection)); @@ -195,9 +210,11 @@ Result ProjectNested(const Type& expected_type, const Type& sou } // namespace Result Project(const Schema& expected_schema, - const Schema& source_schema, bool prune_source) { + const Schema& source_schema, bool prune_source, + const ProjectionOptions& options) { ICEBERG_ASSIGN_OR_RAISE(auto field_projection, - ProjectNested(expected_schema, source_schema, prune_source)); + ProjectNested(expected_schema, source_schema, prune_source, + options)); return SchemaProjection{std::move(field_projection.children)}; } diff --git a/src/iceberg/schema_util.h b/src/iceberg/schema_util.h index 421d00e0ad..7223abee0e 100644 --- a/src/iceberg/schema_util.h +++ b/src/iceberg/schema_util.h @@ -81,6 +81,25 @@ struct ICEBERG_EXPORT SchemaProjection { std::vector fields; }; +/// \brief Options to control schema projection behavior. +struct ICEBERG_EXPORT ProjectionOptions { + /// \brief Which default value fills a field that is missing in the source schema. + enum class DefaultPolicy { + /// \brief Use the v3 `initial-default`; for reading data written before the field + /// existed. + kInitial, + /// \brief Use the v3 `write-default`; for aligning data to the write schema. + kWrite, + }; + + /// \brief The default value policy for missing fields. + DefaultPolicy default_policy = DefaultPolicy::kInitial; + /// \brief Whether type promotion (e.g. int to long) is allowed for projected fields. + /// Must be disabled when projecting for the write path, where the source values are + /// stored as-is and therefore must exactly match the expected type. + bool allow_type_promotion = true; +}; + /// \brief Project the expected schema on top of the source schema. /// /// \param expected_schema The expected schema. @@ -90,10 +109,12 @@ struct ICEBERG_EXPORT SchemaProjection { /// column pruning, so `prune_source` is set to true in this case such that the `from` /// field in `FieldProjection` exactly reflects the position (relative to its nesting /// level) to get the column value from the reader. +/// \param options Options controlling default values and type promotion. /// \return The projection result. ICEBERG_EXPORT Result Project(const Schema& expected_schema, const Schema& source_schema, - bool prune_source); + bool prune_source, + const ProjectionOptions& options = {}); ICEBERG_EXPORT std::string_view ToString(FieldProjection::Kind kind); ICEBERG_EXPORT std::string ToString(const FieldProjection& projection); diff --git a/src/iceberg/test/schema_util_test.cc b/src/iceberg/test/schema_util_test.cc index 9a3eff8870..1138704ced 100644 --- a/src/iceberg/test/schema_util_test.cc +++ b/src/iceberg/test/schema_util_test.cc @@ -232,6 +232,75 @@ TEST(SchemaUtilTest, ProjectPresentFieldIgnoresInitialDefault) { AssertProjectedField(projection_result->fields[0], 0); } +TEST(SchemaUtilTest, ProjectMissingFieldWithWriteDefault) { + // The write path fills a missing field with `write-default`, not `initial-default`. + Schema source_schema = CreateFlatSchema(); + Schema expected_schema({ + SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int64()), + SchemaField(/*field_id=*/10, "extra", iceberg::int32(), /*optional=*/false, + /*doc=*/{}, std::make_shared(Literal::Int(42)), + std::make_shared(Literal::Int(7))), + }); + + auto projection_result = + Project(expected_schema, source_schema, /*prune_source=*/false, + ProjectionOptions{ + .default_policy = ProjectionOptions::DefaultPolicy::kWrite}); + ASSERT_THAT(projection_result, IsOk()); + + const auto& projection = *projection_result; + ASSERT_EQ(projection.fields.size(), 2); + AssertProjectedField(projection.fields[0], 0); + ASSERT_EQ(projection.fields[1].kind, FieldProjection::Kind::kDefault); + ASSERT_EQ(std::get(projection.fields[1].from), Literal::Int(7)); +} + +TEST(SchemaUtilTest, ProjectWritePolicyIgnoresInitialDefault) { + // With the write default source, a field having only `initial-default` behaves as if + // it had no default: optional becomes null and required is rejected. + Schema source_schema = CreateFlatSchema(); + Schema expected_schema({ + SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int64()), + SchemaField(/*field_id=*/10, "extra", iceberg::string(), /*optional=*/true, + /*doc=*/{}, std::make_shared(Literal::String("n/a"))), + }); + + auto projection_result = + Project(expected_schema, source_schema, /*prune_source=*/false, + ProjectionOptions{ + .default_policy = ProjectionOptions::DefaultPolicy::kWrite}); + ASSERT_THAT(projection_result, IsOk()); + ASSERT_EQ(projection_result->fields[1].kind, FieldProjection::Kind::kNull); +} + +TEST(SchemaUtilTest, ProjectWritePolicyRejectsMissingRequired) { + Schema source_schema = CreateFlatSchema(); + Schema expected_schema({ + SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int64()), + SchemaField::MakeRequired(/*field_id=*/10, "extra", iceberg::string()), + }); + + auto projection_result = + Project(expected_schema, source_schema, /*prune_source=*/false, + ProjectionOptions{ + .default_policy = ProjectionOptions::DefaultPolicy::kWrite}); + ASSERT_THAT(projection_result, IsError(ErrorKind::kInvalidSchema)); + ASSERT_THAT(projection_result, HasErrorMessage("without a write-default")); +} + +TEST(SchemaUtilTest, ProjectWithoutTypePromotion) { + Schema source_schema( + {SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int32())}); + Schema expected_schema( + {SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int64())}); + + auto projection_result = + Project(expected_schema, source_schema, /*prune_source=*/false, + ProjectionOptions{.allow_type_promotion = false}); + ASSERT_THAT(projection_result, IsError(ErrorKind::kNotSupported)); + ASSERT_THAT(projection_result, HasErrorMessage("type promotion is not allowed")); +} + TEST(SchemaUtilTest, ProjectMetadataColumn) { Schema source_schema = CreateFlatSchema(); Schema expected_schema({ From 17690805852edb90f0937545352044d30a04b18b Mon Sep 17 00:00:00 2001 From: Xin Huang Date: Fri, 18 Sep 2026 15:45:28 -0700 Subject: [PATCH 2/9] =?UTF-8?q?=E2=9C=A8=20feat(arrow):=20Align=20batches?= =?UTF-8?q?=20to=20write=20schemas?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add arrow::AlignBatchForWrite which projects an input batch onto the write schema via Project with write-default options: present columns are reordered by field id (recursively for structs), missing columns are materialized from write-default or null, and unknown input columns are rejected rather than silently dropped. --- src/iceberg/arrow/arrow_c_data_util.cc | 151 +++++++++++++++++++++ src/iceberg/arrow_c_data_util_internal.h | 21 +++ src/iceberg/test/arrow_c_data_util_test.cc | 96 +++++++++++++ 3 files changed, 268 insertions(+) diff --git a/src/iceberg/arrow/arrow_c_data_util.cc b/src/iceberg/arrow/arrow_c_data_util.cc index 5c8ea5bbf3..521ee8062d 100644 --- a/src/iceberg/arrow/arrow_c_data_util.cc +++ b/src/iceberg/arrow/arrow_c_data_util.cc @@ -22,20 +22,31 @@ #include #include #include +#include #include +#include #include #include #include #include #include +#include #include #include "iceberg/arrow/arrow_status_internal.h" +#include "iceberg/arrow/literal_util_internal.h" #include "iceberg/arrow_c_data_guard_internal.h" #include "iceberg/arrow_c_data_util_internal.h" #include "iceberg/nanoarrow_status_internal.h" #include "iceberg/result.h" +#include "iceberg/schema.h" +#include "iceberg/schema_field.h" +#include "iceberg/schema_internal.h" +#include "iceberg/schema_util.h" +#include "iceberg/type.h" +#include "iceberg/util/checked_cast.h" +#include "iceberg/util/formatter.h" #include "iceberg/util/macros.h" namespace iceberg { @@ -128,4 +139,144 @@ void RegisterArrowProjectBatch() { }); } +namespace arrow { +namespace { + +Result> ImportIcebergSchema(const Schema& schema) { + ArrowSchema c_schema; + ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(schema, &c_schema)); + internal::ArrowSchemaGuard c_schema_guard(&c_schema); + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto arrow_schema, ::arrow::ImportSchema(&c_schema)); + c_schema_guard.Release(); + return arrow_schema; +} + +// Writers must not silently drop data: every input field, at any struct nesting level, +// has to exist in the write schema. +Status RejectUnknownInputFields(const StructType& input_struct, + const StructType& write_struct) { + for (const auto& input_field : input_struct.fields()) { + ICEBERG_ASSIGN_OR_RAISE(auto write_field, + write_struct.GetFieldById(input_field.field_id())); + if (!write_field.has_value()) { + return InvalidSchema("Input field id {} is not present in write schema", + input_field.field_id()); + } + if (input_field.type()->type_id() == TypeId::kStruct && + write_field->get().type()->type_id() == TypeId::kStruct) { + ICEBERG_RETURN_UNEXPECTED(RejectUnknownInputFields( + internal::checked_cast(*input_field.type()), + internal::checked_cast(*write_field->get().type()))); + } + } + return {}; +} + +Result>> MaterializeProjection( + std::span projections, + const std::vector>& input_arrays, + const std::vector>& write_arrow_fields, int64_t length, + ::arrow::MemoryPool* pool) { + std::vector> result; + result.reserve(projections.size()); + for (size_t index = 0; index < projections.size(); ++index) { + const auto& projection = projections[index]; + const auto& write_arrow_type = write_arrow_fields[index]->type(); + switch (projection.kind) { + case FieldProjection::Kind::kProjected: { + const auto& input_array = input_arrays[std::get(projection.from)]; + if (projection.children.empty()) { + result.emplace_back(input_array); + break; + } + auto input_struct = std::dynamic_pointer_cast<::arrow::StructArray>(input_array); + auto write_struct_type = + std::dynamic_pointer_cast<::arrow::StructType>(write_arrow_type); + if (input_struct == nullptr || write_struct_type == nullptr) { + return InvalidArrowData("Expected struct array for projected field {}", + index); + } + ICEBERG_ASSIGN_OR_RAISE( + auto children, + MaterializeProjection(projection.children, input_struct->fields(), + write_struct_type->fields(), input_struct->length(), + pool)); + if (children.empty()) { + result.emplace_back(std::make_shared<::arrow::StructArray>( + write_struct_type, input_struct->length(), std::move(children), + input_struct->null_bitmap(), input_struct->null_count(), + input_struct->offset())); + break; + } + ICEBERG_ARROW_ASSIGN_OR_RETURN( + auto struct_array, + ::arrow::StructArray::Make(std::move(children), write_struct_type->fields(), + input_struct->null_bitmap(), + input_struct->null_count(), + input_struct->offset())); + result.emplace_back(std::move(struct_array)); + break; + } + case FieldProjection::Kind::kDefault: { + ICEBERG_ASSIGN_OR_RAISE( + auto default_array, + MakeDefaultArray(std::get(projection.from), write_arrow_type, + length, pool)); + result.emplace_back(std::move(default_array)); + break; + } + case FieldProjection::Kind::kNull: { + ICEBERG_ARROW_ASSIGN_OR_RETURN( + auto null_array, ::arrow::MakeArrayOfNull(write_arrow_type, length, pool)); + result.emplace_back(std::move(null_array)); + break; + } + default: + return NotSupported("Projection kind {} is not supported for write alignment", + ToString(projection.kind)); + } + } + return result; +} + +} // namespace + +Result AlignBatchForWrite(ArrowArray* input_batch, const Schema& input_schema, + const Schema& write_schema, + ::arrow::MemoryPool* pool) { + ICEBERG_PRECHECK(input_batch != nullptr, "input_batch must not be null"); + ICEBERG_PRECHECK(pool != nullptr, "pool must not be null"); + internal::ArrowArrayGuard input_guard(input_batch); + + ICEBERG_ASSIGN_OR_RAISE(auto input_arrow_schema, ImportIcebergSchema(input_schema)); + ICEBERG_ARROW_ASSIGN_OR_RETURN( + auto input_record_batch, + ::arrow::ImportRecordBatch(input_batch, input_arrow_schema)); + input_guard.Release(); + + ICEBERG_ASSIGN_OR_RAISE(auto write_arrow_schema, ImportIcebergSchema(write_schema)); + + ICEBERG_RETURN_UNEXPECTED(RejectUnknownInputFields(input_schema, write_schema)); + ICEBERG_ASSIGN_OR_RAISE( + auto projection, + Project(write_schema, input_schema, /*prune_source=*/false, + ProjectionOptions{ + .default_policy = ProjectionOptions::DefaultPolicy::kWrite, + .allow_type_promotion = false})); + ICEBERG_ASSIGN_OR_RAISE( + auto columns, + MaterializeProjection(projection.fields, input_record_batch->columns(), + write_arrow_schema->fields(), + input_record_batch->num_rows(), pool)); + auto output_record_batch = ::arrow::RecordBatch::Make( + write_arrow_schema, input_record_batch->num_rows(), std::move(columns)); + + ArrowArray output{}; + ICEBERG_ARROW_RETURN_NOT_OK(::arrow::ExportRecordBatch(*output_record_batch, &output)); + internal::ArrowArrayGuard output_guard(&output); + return std::exchange(output, ArrowArray{}); +} + +} // namespace arrow + } // namespace iceberg diff --git a/src/iceberg/arrow_c_data_util_internal.h b/src/iceberg/arrow_c_data_util_internal.h index e02db29a64..974eac768b 100644 --- a/src/iceberg/arrow_c_data_util_internal.h +++ b/src/iceberg/arrow_c_data_util_internal.h @@ -35,6 +35,10 @@ #include "iceberg/result.h" #include "iceberg/type_fwd.h" +namespace arrow { +class MemoryPool; +} + namespace iceberg { /// \brief Cached state for ProjectBatch over one input/output schema pair. @@ -242,4 +246,21 @@ ICEBERG_EXPORT Result ProjectBatch(ArrowArray* input_batch, std::span row_indices, ProjectionContext& projection); +namespace arrow { + +/// \brief Align an Arrow batch to an Iceberg write schema. +/// +/// Fields are matched by Iceberg field id. Missing fields are materialized from their +/// `write-default`, or as null when optional. Missing required fields without a +/// `write-default` are rejected. Existing values, including explicit nulls, are +/// preserved. Nested structs are aligned recursively. +/// +/// `input_batch` is consumed even when alignment fails. +ICEBERG_EXPORT Result AlignBatchForWrite(ArrowArray* input_batch, + const Schema& input_schema, + const Schema& write_schema, + ::arrow::MemoryPool* pool); + +} // namespace arrow + } // namespace iceberg diff --git a/src/iceberg/test/arrow_c_data_util_test.cc b/src/iceberg/test/arrow_c_data_util_test.cc index c93a928a3d..cd8d794d32 100644 --- a/src/iceberg/test/arrow_c_data_util_test.cc +++ b/src/iceberg/test/arrow_c_data_util_test.cc @@ -27,18 +27,22 @@ #include #include +#include #include #include #include #include "iceberg/arrow/arrow_register.h" +#include "iceberg/arrow/arrow_status_internal.h" #include "iceberg/arrow_c_data_guard_internal.h" #include "iceberg/arrow_c_data_util_internal.h" +#include "iceberg/expression/literal.h" #include "iceberg/schema.h" #include "iceberg/schema_field.h" #include "iceberg/schema_internal.h" #include "iceberg/test/matchers.h" #include "iceberg/type.h" +#include "iceberg/util/macros.h" namespace iceberg::internal { @@ -119,6 +123,32 @@ std::shared_ptr MakeFullSchema() { SchemaField::MakeOptional(3, "score", float64())}); } +Result> AlignForWrite(const Schema& input_schema, + const Schema& write_schema, + std::string_view input_json) { + auto input = MakeBatch(input_schema, input_json); + ArrowArray c_array; + ICEBERG_ARROW_RETURN_NOT_OK(::arrow::ExportRecordBatch(*input, &c_array)); + ICEBERG_ASSIGN_OR_RAISE(auto aligned, + arrow::AlignBatchForWrite(&c_array, input_schema, write_schema, + ::arrow::default_memory_pool())); + + ArrowSchema c_schema; + ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(write_schema, &c_schema)); + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto arrow_schema, ::arrow::ImportSchema(&c_schema)); + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto result, + ::arrow::ImportRecordBatch(&aligned, arrow_schema)); + return result; +} + +SchemaField DefaultedField(int32_t id, std::string name, std::shared_ptr type, + bool optional, Literal initial_default, + Literal write_default) { + return SchemaField(id, std::move(name), std::move(type), optional, {}, + std::make_shared(std::move(initial_default)), + std::make_shared(std::move(write_default))); +} + } // namespace TEST(ProjectBatchTest, ProjectSelectedRowsWithoutColumnProjection) { @@ -191,4 +221,70 @@ TEST(ProjectBatchTest, ProjectionRejectsNestedPruning) { EXPECT_THAT(projection, IsError(ErrorKind::kInvalidArgument)); } +TEST(AlignBatchForWriteTest, ReordersAndMaterializesWriteDefaultsAndNulls) { + Schema input_schema({SchemaField::MakeOptional(2, "name", string()), + SchemaField::MakeRequired(1, "id", int32())}); + Schema write_schema({ + SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeOptional(2, "name", string()), + DefaultedField(3, "score", int32(), /*optional=*/false, Literal::Int(42), + Literal::Int(7)), + SchemaField::MakeOptional(4, "comment", string()), + }); + + ICEBERG_UNWRAP_OR_FAIL(auto actual, AlignForWrite(input_schema, write_schema, + R"([["alice",1],[null,2]])")); + auto expected = MakeBatch(write_schema, R"([[1,"alice",7,null],[2,null,7,null]])"); + EXPECT_TRUE(actual->Equals(*expected)) << "actual:\n" << actual->ToString(); +} + +TEST(AlignBatchForWriteTest, RejectsMissingRequiredFieldWithoutWriteDefault) { + Schema input_schema({SchemaField::MakeRequired(1, "id", int32())}); + Schema write_schema({SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeRequired(2, "required", string())}); + + auto result = AlignForWrite(input_schema, write_schema, R"([[1]])"); + + EXPECT_THAT(result, IsError(ErrorKind::kInvalidSchema)); + EXPECT_THAT(result, HasErrorMessage("without a write-default")); +} + +TEST(AlignBatchForWriteTest, RecursivelyAlignsNestedStructs) { + auto input_person = std::make_shared( + std::vector{SchemaField::MakeRequired(2, "name", string())}); + auto write_person = std::make_shared(std::vector{ + SchemaField::MakeRequired(2, "name", string()), + DefaultedField(3, "age", int32(), /*optional=*/false, Literal::Int(10), + Literal::Int(18)), + }); + Schema input_schema({SchemaField::MakeOptional(1, "person", std::move(input_person))}); + Schema write_schema({SchemaField::MakeOptional(1, "person", std::move(write_person))}); + + ICEBERG_UNWRAP_OR_FAIL(auto actual, AlignForWrite(input_schema, write_schema, + R"([[{"name":"alice"}],[null]])")); + auto expected = MakeBatch(write_schema, R"([[{"name":"alice","age":18}],[null]])"); + EXPECT_TRUE(actual->Equals(*expected)) << "actual:\n" << actual->ToString(); +} + +TEST(AlignBatchForWriteTest, RejectsIncompatibleTypes) { + Schema input_schema({SchemaField::MakeRequired(1, "id", int32())}); + Schema write_schema({SchemaField::MakeRequired(1, "id", int64())}); + + auto result = AlignForWrite(input_schema, write_schema, R"([[1]])"); + + EXPECT_THAT(result, IsError(ErrorKind::kNotSupported)); + EXPECT_THAT(result, HasErrorMessage("type promotion is not allowed")); +} + +TEST(AlignBatchForWriteTest, RejectsInputFieldsOutsideWriteSchema) { + Schema input_schema({SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeOptional(2, "extra", string())}); + Schema write_schema({SchemaField::MakeRequired(1, "id", int32())}); + + auto result = AlignForWrite(input_schema, write_schema, R"([[1,"extra"]])"); + + EXPECT_THAT(result, IsError(ErrorKind::kInvalidSchema)); + EXPECT_THAT(result, HasErrorMessage("Input field id 2")); +} + } // namespace iceberg::internal From 8e01c0eaf0ebf212280e966bb724b5a33e9c40a1 Mon Sep 17 00:00:00 2001 From: Xin Huang Date: Fri, 18 Sep 2026 15:45:28 -0700 Subject: [PATCH 3/9] =?UTF-8?q?=E2=9C=A8=20feat(data):=20Apply=20write=20d?= =?UTF-8?q?efaults=20in=20writers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DataWriter aligns batches to the write schema when the input schema differs, via a new Write(data, input_schema) overload; schemas that already match keep the zero-conversion Write(data) path. Adds Parquet and Avro end-to-end tests where a column added with write-default reads back the default for rows written without it. --- src/iceberg/avro/avro_writer.cc | 13 ++++ src/iceberg/avro/avro_writer.h | 2 + src/iceberg/data/data_writer.cc | 10 ++- src/iceberg/data/data_writer.h | 2 + src/iceberg/file_writer.h | 8 ++- src/iceberg/parquet/parquet_writer.cc | 12 ++++ src/iceberg/parquet/parquet_writer.h | 2 + src/iceberg/test/data_writer_test.cc | 13 ++++ src/iceberg/test/default_value_test.cc | 96 ++++++++++++++++++++++++++ 9 files changed, 156 insertions(+), 2 deletions(-) diff --git a/src/iceberg/avro/avro_writer.cc b/src/iceberg/avro/avro_writer.cc index 01d8cb3b14..48558abe3a 100644 --- a/src/iceberg/avro/avro_writer.cc +++ b/src/iceberg/avro/avro_writer.cc @@ -32,6 +32,7 @@ #include "iceberg/arrow/arrow_io_internal.h" #include "iceberg/arrow/arrow_status_internal.h" +#include "iceberg/arrow_c_data_util_internal.h" #include "iceberg/avro/avro_data_util_internal.h" #include "iceberg/avro/avro_direct_encoder_internal.h" #include "iceberg/avro/avro_metrics.h" @@ -241,6 +242,14 @@ class AvroWriter::Impl { return {}; } + Status Write(ArrowArray* data, const Schema& input_schema) { + ICEBERG_ASSIGN_OR_RAISE( + auto aligned_data, + arrow::AlignBatchForWrite(data, input_schema, *write_schema_, + ::arrow::default_memory_pool())); + return Write(&aligned_data); + } + Status Close() { if (!backend_->Closed()) { backend_->Close(); @@ -290,6 +299,10 @@ AvroWriter::~AvroWriter() = default; Status AvroWriter::Write(ArrowArray* data) { return impl_->Write(data); } +Status AvroWriter::Write(ArrowArray* data, const Schema& input_schema) { + return impl_->Write(data, input_schema); +} + Status AvroWriter::Open(const WriterOptions& options) { impl_ = std::make_unique(); return impl_->Open(options); diff --git a/src/iceberg/avro/avro_writer.h b/src/iceberg/avro/avro_writer.h index af68fac440..87b770df8d 100644 --- a/src/iceberg/avro/avro_writer.h +++ b/src/iceberg/avro/avro_writer.h @@ -40,6 +40,8 @@ class ICEBERG_BUNDLE_EXPORT AvroWriter : public Writer { Status Write(ArrowArray* data) final; + Status Write(ArrowArray* data, const Schema& input_schema) final; + Result metrics() final; Result length() final; diff --git a/src/iceberg/data/data_writer.cc b/src/iceberg/data/data_writer.cc index 512f77cc96..c4820e9bb1 100644 --- a/src/iceberg/data/data_writer.cc +++ b/src/iceberg/data/data_writer.cc @@ -26,6 +26,7 @@ #include "iceberg/file_writer.h" #include "iceberg/manifest/manifest_entry.h" #include "iceberg/partition_spec.h" +#include "iceberg/schema.h" #include "iceberg/util/macros.h" namespace iceberg { @@ -33,6 +34,8 @@ namespace iceberg { class DataWriter::Impl { public: static Result> Make(DataWriterOptions options) { + ICEBERG_PRECHECK(options.schema != nullptr, "Write schema must not be null"); + ICEBERG_PRECHECK(options.input_schema != nullptr, "Input schema must not be null"); WriterOptions writer_options{ .path = options.path, .schema = options.schema, @@ -46,7 +49,12 @@ class DataWriter::Impl { return std::unique_ptr(new Impl(std::move(options), std::move(writer))); } - Status Write(ArrowArray* data) { return writer_->Write(data); } + Status Write(ArrowArray* data) { + if (*options_.input_schema == *options_.schema) { + return writer_->Write(data); + } + return writer_->Write(data, *options_.input_schema); + } Result Length() const { return writer_->length(); } diff --git a/src/iceberg/data/data_writer.h b/src/iceberg/data/data_writer.h index 9a2002f0db..3104048eb7 100644 --- a/src/iceberg/data/data_writer.h +++ b/src/iceberg/data/data_writer.h @@ -42,6 +42,8 @@ namespace iceberg { struct ICEBERG_DATA_EXPORT DataWriterOptions { std::string path; std::shared_ptr schema; + /// Required schema of Arrow batches passed to Write. + std::shared_ptr input_schema; std::shared_ptr spec; PartitionValues partition; FileFormatType format = FileFormatType::kParquet; diff --git a/src/iceberg/file_writer.h b/src/iceberg/file_writer.h index 4b1045a409..d6fcb30f03 100644 --- a/src/iceberg/file_writer.h +++ b/src/iceberg/file_writer.h @@ -75,7 +75,7 @@ class ICEBERG_EXPORT WriterProperties : public ConfigBase { struct ICEBERG_EXPORT WriterOptions { /// \brief The path to the file to write. std::string path; - /// \brief The schema of the data to write. + /// \brief The target schema to write to the file. std::shared_ptr schema; /// \brief FileIO instance to create the file. std::shared_ptr io; @@ -107,6 +107,12 @@ class ICEBERG_EXPORT Writer { /// \note Ownership of the data is transferred to the writer. virtual Status Write(ArrowArray* data) = 0; + /// \brief Write Arrow data described by an input schema that may differ from the + /// writer schema. + /// + /// \note Ownership of the data is transferred to the writer. + virtual Status Write(ArrowArray* data, const Schema& input_schema) = 0; + /// \brief Get the file statistics. /// Only valid after the file is closed. virtual Result metrics() = 0; diff --git a/src/iceberg/parquet/parquet_writer.cc b/src/iceberg/parquet/parquet_writer.cc index 08c3941383..c01ccbe8d7 100644 --- a/src/iceberg/parquet/parquet_writer.cc +++ b/src/iceberg/parquet/parquet_writer.cc @@ -39,6 +39,7 @@ #include "iceberg/arrow/arrow_io_internal.h" #include "iceberg/arrow/arrow_status_internal.h" +#include "iceberg/arrow_c_data_util_internal.h" #include "iceberg/parquet/parquet_metrics_internal.h" #include "iceberg/parquet/parquet_schema_util_internal.h" #include "iceberg/schema_internal.h" @@ -295,6 +296,13 @@ class ParquetWriter::Impl { return {}; } + Status Write(ArrowArray* array, const Schema& input_schema) { + ICEBERG_ASSIGN_OR_RAISE( + auto aligned_array, + arrow::AlignBatchForWrite(array, input_schema, *schema_, pool_)); + return Write(&aligned_array); + } + // Close the writer and release resources Status Close() { if (writer_ == nullptr) { @@ -371,6 +379,10 @@ Status ParquetWriter::Open(const WriterOptions& options) { Status ParquetWriter::Write(ArrowArray* array) { return impl_->Write(array); } +Status ParquetWriter::Write(ArrowArray* array, const Schema& input_schema) { + return impl_->Write(array, input_schema); +} + Status ParquetWriter::Close() { return impl_->Close(); } Result ParquetWriter::metrics() { return impl_->metrics(); } diff --git a/src/iceberg/parquet/parquet_writer.h b/src/iceberg/parquet/parquet_writer.h index 562ffb8013..bcd5b65c44 100644 --- a/src/iceberg/parquet/parquet_writer.h +++ b/src/iceberg/parquet/parquet_writer.h @@ -40,6 +40,8 @@ class ICEBERG_BUNDLE_EXPORT ParquetWriter : public Writer { Status Write(ArrowArray* array) final; + Status Write(ArrowArray* array, const Schema& input_schema) final; + Result metrics() final; Result length() final; diff --git a/src/iceberg/test/data_writer_test.cc b/src/iceberg/test/data_writer_test.cc index d95361675e..71832d8b43 100644 --- a/src/iceberg/test/data_writer_test.cc +++ b/src/iceberg/test/data_writer_test.cc @@ -71,6 +71,7 @@ class DataWriterTest : public ::testing::Test { return DataWriterOptions{ .path = "test_data.parquet", .schema = schema_, + .input_schema = schema_, .spec = partition_spec_, .partition = std::move(partition), .format = FileFormatType::kParquet, @@ -138,11 +139,22 @@ class DataWriterFormatTest : public DataWriterTest, public ::testing::WithParamInterface> {}; +TEST_F(DataWriterTest, RejectsMissingInputSchema) { + auto options = MakeDefaultOptions(); + options.input_schema.reset(); + + auto result = DataWriter::Make(options); + + EXPECT_THAT(result, IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT(result, HasErrorMessage("Input schema must not be null")); +} + TEST_P(DataWriterFormatTest, CreateWithFormat) { auto [format, path] = GetParam(); DataWriterOptions options{ .path = path, .schema = schema_, + .input_schema = schema_, .spec = partition_spec_, .partition = PartitionValues{}, .format = format, @@ -162,6 +174,7 @@ TEST_P(DataWriterFormatTest, WriteRowLineage) { DataWriterOptions options{ .path = path, .schema = schema, + .input_schema = schema, .spec = partition_spec_, .partition = PartitionValues{}, .format = format, diff --git a/src/iceberg/test/default_value_test.cc b/src/iceberg/test/default_value_test.cc index f85f2e0d22..ce76f7dfc3 100644 --- a/src/iceberg/test/default_value_test.cc +++ b/src/iceberg/test/default_value_test.cc @@ -23,6 +23,8 @@ #include #include #include +#include +#include #include #include @@ -38,11 +40,14 @@ #include "iceberg/arrow/arrow_io_internal.h" #include "iceberg/avro/avro_register.h" +#include "iceberg/data/data_writer.h" #include "iceberg/expression/literal.h" #include "iceberg/file_format.h" #include "iceberg/file_reader.h" #include "iceberg/file_writer.h" #include "iceberg/parquet/parquet_register.h" +#include "iceberg/partition_spec.h" +#include "iceberg/row/partition_values.h" #include "iceberg/schema.h" #include "iceberg/schema_field.h" #include "iceberg/schema_internal.h" @@ -72,6 +77,12 @@ struct DefaultValueEndToEndParam { bool avro_skip_datum = true; }; +struct WriteDefaultEndToEndParam { + std::string name; + FileFormatType format; + std::string path; +}; + class DefaultValueEndToEndTest : public UpdateTestBase, public ::testing::WithParamInterface { @@ -86,6 +97,37 @@ class DefaultValueEndToEndTest } }; +class WriteDefaultEndToEndTest + : public ::testing::TestWithParam { + protected: + static void SetUpTestSuite() { + parquet::RegisterAll(); + avro::RegisterAll(); + } + + void SetUp() override { file_io_ = arrow::ArrowFileSystemFileIO::MakeMockFileIO(); } + + std::shared_ptr<::arrow::Array> CreateArray(const Schema& schema, + std::string_view json) { + ArrowSchema arrow_c_schema; + ICEBERG_THROW_NOT_OK(ToArrowSchema(schema, &arrow_c_schema)); + auto arrow_schema = ::arrow::ImportType(&arrow_c_schema).ValueOrDie(); + return ::arrow::json::ArrayFromJSONString(::arrow::struct_(arrow_schema->fields()), + std::string(json)) + .ValueOrDie(); + } + + static std::unordered_map FormatProperties( + FileFormatType format) { + if (format == FileFormatType::kParquet) { + return {{"write.parquet.compression-codec", "uncompressed"}}; + } + return {}; + } + + std::shared_ptr file_io_; +}; + TEST_P(DefaultValueEndToEndTest, WriteEvolveReadFillsInitialDefault) { const auto& param = GetParam(); ICEBERG_UNWRAP_OR_FAIL(auto original_schema, table_->schema()); @@ -186,6 +228,50 @@ TEST_P(DefaultValueEndToEndTest, WriteEvolveReadFillsInitialDefault) { ASSERT_FALSE(next_batch.has_value()); } +TEST_P(WriteDefaultEndToEndTest, MissingColumnUsesWriteDefault) { + const auto& param = GetParam(); + auto input_schema = std::make_shared( + std::vector{SchemaField::MakeRequired(1, "id", int32())}); + auto write_schema = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), + SchemaField(2, "added", int32(), /*optional=*/false, /*doc=*/{}, + std::make_shared(Literal::Int(42)), + std::make_shared(Literal::Int(7))), + }); + DataWriterOptions options{ + .path = param.path, + .schema = write_schema, + .input_schema = input_schema, + .spec = PartitionSpec::Unpartitioned(), + .partition = PartitionValues{}, + .format = param.format, + .io = file_io_, + .properties = FormatProperties(param.format), + }; + + ICEBERG_UNWRAP_OR_FAIL(auto writer, DataWriter::Make(options)); + auto input = CreateArray(*input_schema, R"([[1], [2]])"); + ArrowArray arrow_array; + ASSERT_TRUE(::arrow::ExportArray(*input, &arrow_array).ok()); + ASSERT_THAT(writer->Write(&arrow_array), IsOk()); + ASSERT_THAT(writer->Close(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL( + auto reader, ReaderFactoryRegistry::Open( + param.format, + {.path = param.path, .io = file_io_, .projection = write_schema})); + ICEBERG_UNWRAP_OR_FAIL(auto batch, reader->Next()); + ASSERT_TRUE(batch.has_value()); + + auto expected = CreateArray(*write_schema, R"([[1, 7], [2, 7]])"); + auto actual = ::arrow::ImportArray(&batch.value(), expected->type()).ValueOrDie(); + + // This file is written after the column exists, so the missing input column uses + // write-default (7), not initial-default (42). + ASSERT_TRUE(actual->Equals(expected)) + << "actual: " << actual->ToString() << "\nexpected: " << expected->ToString(); +} + namespace { // Two-row column of `json` values at `type`, for the simple cases. @@ -261,6 +347,16 @@ INSTANTIATE_TEST_SUITE_P( return info.param.name; }); +INSTANTIATE_TEST_SUITE_P( + Formats, WriteDefaultEndToEndTest, + ::testing::Values(WriteDefaultEndToEndParam{"parquet", FileFormatType::kParquet, + "write-default.parquet"}, + WriteDefaultEndToEndParam{"avro", FileFormatType::kAvro, + "write-default.avro"}), + [](const ::testing::TestParamInfo& info) { + return info.param.name; + }); + } // namespace } // namespace iceberg From 744fdd897d782f40d97e2e61626f5eb872e41773 Mon Sep 17 00:00:00 2001 From: Xin Huang Date: Sun, 11 Oct 2026 21:49:35 -0700 Subject: [PATCH 4/9] =?UTF-8?q?=F0=9F=90=9B=20fix(data):=20Harden=20write?= =?UTF-8?q?=20schema=20alignment?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Preserve sliced struct offsets and align collection children by field id. - Materialize unknown inputs as typed nulls. - Configure optional input schemas without extending the writer interface. --- src/iceberg/arrow/arrow_c_data_util.cc | 106 ++++++------ src/iceberg/arrow_c_data_util_internal.h | 2 +- src/iceberg/avro/avro_writer.cc | 26 +-- src/iceberg/avro/avro_writer.h | 2 - src/iceberg/data/data_writer.cc | 10 +- src/iceberg/data/data_writer.h | 4 +- src/iceberg/file_writer.h | 8 +- src/iceberg/parquet/parquet_writer.cc | 25 +-- src/iceberg/parquet/parquet_writer.h | 2 - src/iceberg/schema_util.cc | 41 ++--- src/iceberg/test/arrow_c_data_util_test.cc | 177 ++++++++++++++++++++- src/iceberg/test/data_writer_test.cc | 14 +- src/iceberg/test/default_value_test.cc | 91 ++++++----- src/iceberg/test/schema_util_test.cc | 50 ++++-- 14 files changed, 377 insertions(+), 181 deletions(-) diff --git a/src/iceberg/arrow/arrow_c_data_util.cc b/src/iceberg/arrow/arrow_c_data_util.cc index 521ee8062d..90964f2dff 100644 --- a/src/iceberg/arrow/arrow_c_data_util.cc +++ b/src/iceberg/arrow/arrow_c_data_util.cc @@ -25,8 +25,10 @@ #include #include -#include +#include #include +#include +#include #include #include #include @@ -151,22 +153,21 @@ Result> ImportIcebergSchema(const Schema& schem return arrow_schema; } -// Writers must not silently drop data: every input field, at any struct nesting level, -// has to exist in the write schema. -Status RejectUnknownInputFields(const StructType& input_struct, - const StructType& write_struct) { - for (const auto& input_field : input_struct.fields()) { +// Writers must not silently drop data at any nesting level. +Status RejectUnknownInputFields(const NestedType& input_type, + const NestedType& write_type) { + for (const auto& input_field : input_type.fields()) { ICEBERG_ASSIGN_OR_RAISE(auto write_field, - write_struct.GetFieldById(input_field.field_id())); + write_type.GetFieldById(input_field.field_id())); if (!write_field.has_value()) { return InvalidSchema("Input field id {} is not present in write schema", input_field.field_id()); } - if (input_field.type()->type_id() == TypeId::kStruct && - write_field->get().type()->type_id() == TypeId::kStruct) { + if (input_field.type()->is_nested() && + input_field.type()->type_id() == write_field->get().type()->type_id()) { ICEBERG_RETURN_UNEXPECTED(RejectUnknownInputFields( - internal::checked_cast(*input_field.type()), - internal::checked_cast(*write_field->get().type()))); + internal::checked_cast(*input_field.type()), + internal::checked_cast(*write_field->get().type()))); } } return {}; @@ -175,8 +176,8 @@ Status RejectUnknownInputFields(const StructType& input_struct, Result>> MaterializeProjection( std::span projections, const std::vector>& input_arrays, - const std::vector>& write_arrow_fields, int64_t length, - ::arrow::MemoryPool* pool) { + const std::vector>& write_arrow_fields, + int64_t length, ::arrow::MemoryPool* pool) { std::vector> result; result.reserve(projections.size()); for (size_t index = 0; index < projections.size(); ++index) { @@ -189,39 +190,54 @@ Result>> MaterializeProjection( result.emplace_back(input_array); break; } - auto input_struct = std::dynamic_pointer_cast<::arrow::StructArray>(input_array); - auto write_struct_type = - std::dynamic_pointer_cast<::arrow::StructType>(write_arrow_type); - if (input_struct == nullptr || write_struct_type == nullptr) { - return InvalidArrowData("Expected struct array for projected field {}", - index); + auto data = input_array->data()->Copy(); + data->type = write_arrow_type; + auto child_data = data; + auto child_type = write_arrow_type; + int64_t child_length = data->offset + data->length; + switch (write_arrow_type->id()) { + case ::arrow::Type::STRUCT: + break; + case ::arrow::Type::LIST: + child_length = data->child_data[0]->length; + break; + case ::arrow::Type::MAP: + // Arrow maps have an entries struct between the map and its key/value. + child_data = data->child_data[0]->Copy(); + child_type = write_arrow_type->field(0)->type(); + child_length = child_data->offset + child_data->length; + break; + default: + return InvalidArrowData("Expected nested array for projected field {}", + index); } + std::vector> input_children; + input_children.reserve(child_data->child_data.size()); + for (const auto& child : child_data->child_data) { + input_children.emplace_back(::arrow::MakeArray(child)); + } + // Raw child data retains the parent's physical row positions. + // StructArray::fields() would apply the parent slice, causing its offset to be + // applied twice. ICEBERG_ASSIGN_OR_RAISE( auto children, - MaterializeProjection(projection.children, input_struct->fields(), - write_struct_type->fields(), input_struct->length(), - pool)); - if (children.empty()) { - result.emplace_back(std::make_shared<::arrow::StructArray>( - write_struct_type, input_struct->length(), std::move(children), - input_struct->null_bitmap(), input_struct->null_count(), - input_struct->offset())); - break; + MaterializeProjection(projection.children, input_children, + child_type->fields(), child_length, pool)); + child_data->type = child_type; + child_data->child_data.clear(); + for (const auto& child : children) { + child_data->child_data.emplace_back(child->data()); } - ICEBERG_ARROW_ASSIGN_OR_RETURN( - auto struct_array, - ::arrow::StructArray::Make(std::move(children), write_struct_type->fields(), - input_struct->null_bitmap(), - input_struct->null_count(), - input_struct->offset())); - result.emplace_back(std::move(struct_array)); + if (write_arrow_type->id() == ::arrow::Type::MAP) { + data->child_data = {std::move(child_data)}; + } + result.emplace_back(::arrow::MakeArray(std::move(data))); break; } case FieldProjection::Kind::kDefault: { - ICEBERG_ASSIGN_OR_RAISE( - auto default_array, - MakeDefaultArray(std::get(projection.from), write_arrow_type, - length, pool)); + ICEBERG_ASSIGN_OR_RAISE(auto default_array, + MakeDefaultArray(std::get(projection.from), + write_arrow_type, length, pool)); result.emplace_back(std::move(default_array)); break; } @@ -259,15 +275,15 @@ Result AlignBatchForWrite(ArrowArray* input_batch, const Schema& inp ICEBERG_RETURN_UNEXPECTED(RejectUnknownInputFields(input_schema, write_schema)); ICEBERG_ASSIGN_OR_RAISE( auto projection, - Project(write_schema, input_schema, /*prune_source=*/false, - ProjectionOptions{ - .default_policy = ProjectionOptions::DefaultPolicy::kWrite, - .allow_type_promotion = false})); + Project( + write_schema, input_schema, /*prune_source=*/false, + ProjectionOptions{.default_policy = ProjectionOptions::DefaultPolicy::kWrite, + .allow_type_promotion = false})); ICEBERG_ASSIGN_OR_RAISE( auto columns, MaterializeProjection(projection.fields, input_record_batch->columns(), - write_arrow_schema->fields(), - input_record_batch->num_rows(), pool)); + write_arrow_schema->fields(), input_record_batch->num_rows(), + pool)); auto output_record_batch = ::arrow::RecordBatch::Make( write_arrow_schema, input_record_batch->num_rows(), std::move(columns)); diff --git a/src/iceberg/arrow_c_data_util_internal.h b/src/iceberg/arrow_c_data_util_internal.h index 974eac768b..fe2d25d668 100644 --- a/src/iceberg/arrow_c_data_util_internal.h +++ b/src/iceberg/arrow_c_data_util_internal.h @@ -253,7 +253,7 @@ namespace arrow { /// Fields are matched by Iceberg field id. Missing fields are materialized from their /// `write-default`, or as null when optional. Missing required fields without a /// `write-default` are rejected. Existing values, including explicit nulls, are -/// preserved. Nested structs are aligned recursively. +/// preserved. Structs, lists, and maps are aligned recursively. /// /// `input_batch` is consumed even when alignment fails. ICEBERG_EXPORT Result AlignBatchForWrite(ArrowArray* input_batch, diff --git a/src/iceberg/avro/avro_writer.cc b/src/iceberg/avro/avro_writer.cc index 48558abe3a..cedb112028 100644 --- a/src/iceberg/avro/avro_writer.cc +++ b/src/iceberg/avro/avro_writer.cc @@ -32,6 +32,7 @@ #include "iceberg/arrow/arrow_io_internal.h" #include "iceberg/arrow/arrow_status_internal.h" +#include "iceberg/arrow_c_data_guard_internal.h" #include "iceberg/arrow_c_data_util_internal.h" #include "iceberg/avro/avro_data_util_internal.h" #include "iceberg/avro/avro_direct_encoder_internal.h" @@ -182,6 +183,10 @@ class AvroWriter::Impl { public: Status Open(const WriterOptions& options) { write_schema_ = options.schema; + input_schema_ = options.input_schema; + if (input_schema_ != nullptr && *input_schema_ == *write_schema_) { + input_schema_.reset(); + } ::avro::NodePtr root; ICEBERG_RETURN_UNEXPECTED(ToAvroNodeVisitor{}.Visit(*write_schema_, &root)); @@ -230,6 +235,14 @@ class AvroWriter::Impl { } Status Write(ArrowArray* data) { + ArrowArray aligned_data{}; + internal::ArrowArrayGuard aligned_guard(&aligned_data); + if (input_schema_ != nullptr) { + ICEBERG_ASSIGN_OR_RAISE( + aligned_data, arrow::AlignBatchForWrite(data, *input_schema_, *write_schema_, + ::arrow::default_memory_pool())); + data = &aligned_data; + } ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batch, ::arrow::ImportRecordBatch(data, arrow_schema_)); @@ -242,14 +255,6 @@ class AvroWriter::Impl { return {}; } - Status Write(ArrowArray* data, const Schema& input_schema) { - ICEBERG_ASSIGN_OR_RAISE( - auto aligned_data, - arrow::AlignBatchForWrite(data, input_schema, *write_schema_, - ::arrow::default_memory_pool())); - return Write(&aligned_data); - } - Status Close() { if (!backend_->Closed()) { backend_->Close(); @@ -281,6 +286,7 @@ class AvroWriter::Impl { private: // The schema to write. std::shared_ptr<::iceberg::Schema> write_schema_; + std::shared_ptr<::iceberg::Schema> input_schema_; // The avro schema to write. std::shared_ptr<::avro::ValidSchema> avro_schema_; // Arrow output stream of the Avro file to write @@ -299,10 +305,6 @@ AvroWriter::~AvroWriter() = default; Status AvroWriter::Write(ArrowArray* data) { return impl_->Write(data); } -Status AvroWriter::Write(ArrowArray* data, const Schema& input_schema) { - return impl_->Write(data, input_schema); -} - Status AvroWriter::Open(const WriterOptions& options) { impl_ = std::make_unique(); return impl_->Open(options); diff --git a/src/iceberg/avro/avro_writer.h b/src/iceberg/avro/avro_writer.h index 87b770df8d..af68fac440 100644 --- a/src/iceberg/avro/avro_writer.h +++ b/src/iceberg/avro/avro_writer.h @@ -40,8 +40,6 @@ class ICEBERG_BUNDLE_EXPORT AvroWriter : public Writer { Status Write(ArrowArray* data) final; - Status Write(ArrowArray* data, const Schema& input_schema) final; - Result metrics() final; Result length() final; diff --git a/src/iceberg/data/data_writer.cc b/src/iceberg/data/data_writer.cc index c4820e9bb1..a61cad4a6b 100644 --- a/src/iceberg/data/data_writer.cc +++ b/src/iceberg/data/data_writer.cc @@ -26,7 +26,6 @@ #include "iceberg/file_writer.h" #include "iceberg/manifest/manifest_entry.h" #include "iceberg/partition_spec.h" -#include "iceberg/schema.h" #include "iceberg/util/macros.h" namespace iceberg { @@ -35,12 +34,12 @@ class DataWriter::Impl { public: static Result> Make(DataWriterOptions options) { ICEBERG_PRECHECK(options.schema != nullptr, "Write schema must not be null"); - ICEBERG_PRECHECK(options.input_schema != nullptr, "Input schema must not be null"); WriterOptions writer_options{ .path = options.path, .schema = options.schema, .io = options.io, .properties = WriterProperties::FromMap(options.properties), + .input_schema = options.input_schema, }; ICEBERG_ASSIGN_OR_RAISE(auto writer, @@ -49,12 +48,7 @@ class DataWriter::Impl { return std::unique_ptr(new Impl(std::move(options), std::move(writer))); } - Status Write(ArrowArray* data) { - if (*options_.input_schema == *options_.schema) { - return writer_->Write(data); - } - return writer_->Write(data, *options_.input_schema); - } + Status Write(ArrowArray* data) { return writer_->Write(data); } Result Length() const { return writer_->length(); } diff --git a/src/iceberg/data/data_writer.h b/src/iceberg/data/data_writer.h index 3104048eb7..39d9f50873 100644 --- a/src/iceberg/data/data_writer.h +++ b/src/iceberg/data/data_writer.h @@ -42,14 +42,14 @@ namespace iceberg { struct ICEBERG_DATA_EXPORT DataWriterOptions { std::string path; std::shared_ptr schema; - /// Required schema of Arrow batches passed to Write. - std::shared_ptr input_schema; std::shared_ptr spec; PartitionValues partition; FileFormatType format = FileFormatType::kParquet; std::shared_ptr io; std::optional sort_order_id; std::unordered_map properties; + /// \brief Schema of input batches; when omitted, batches use the target schema. + std::shared_ptr input_schema; }; /// \brief Writer for Iceberg data files. diff --git a/src/iceberg/file_writer.h b/src/iceberg/file_writer.h index d6fcb30f03..d4aea3431a 100644 --- a/src/iceberg/file_writer.h +++ b/src/iceberg/file_writer.h @@ -85,6 +85,8 @@ struct ICEBERG_EXPORT WriterOptions { std::shared_ptr metrics_config = MetricsConfig::Default(); /// \brief Format-specific or implementation-specific properties. WriterProperties properties; + /// \brief Schema of input batches; when omitted, batches use the target schema. + std::shared_ptr input_schema; }; /// \brief Base writer class to write data from different file formats. @@ -107,12 +109,6 @@ class ICEBERG_EXPORT Writer { /// \note Ownership of the data is transferred to the writer. virtual Status Write(ArrowArray* data) = 0; - /// \brief Write Arrow data described by an input schema that may differ from the - /// writer schema. - /// - /// \note Ownership of the data is transferred to the writer. - virtual Status Write(ArrowArray* data, const Schema& input_schema) = 0; - /// \brief Get the file statistics. /// Only valid after the file is closed. virtual Result metrics() = 0; diff --git a/src/iceberg/parquet/parquet_writer.cc b/src/iceberg/parquet/parquet_writer.cc index c01ccbe8d7..176d8248fb 100644 --- a/src/iceberg/parquet/parquet_writer.cc +++ b/src/iceberg/parquet/parquet_writer.cc @@ -39,9 +39,11 @@ #include "iceberg/arrow/arrow_io_internal.h" #include "iceberg/arrow/arrow_status_internal.h" +#include "iceberg/arrow_c_data_guard_internal.h" #include "iceberg/arrow_c_data_util_internal.h" #include "iceberg/parquet/parquet_metrics_internal.h" #include "iceberg/parquet/parquet_schema_util_internal.h" +#include "iceberg/schema.h" #include "iceberg/schema_internal.h" #include "iceberg/type.h" #include "iceberg/util/macros.h" @@ -240,6 +242,10 @@ class ParquetWriter::Impl { public: Status Open(const WriterOptions& options) { schema_ = options.schema; + input_schema_ = options.input_schema; + if (input_schema_ != nullptr && *input_schema_ == *schema_) { + input_schema_.reset(); + } ICEBERG_ASSIGN_OR_RAISE(auto compression, ParseCompression(options.properties)); ICEBERG_ASSIGN_OR_RAISE(auto compression_level, ParseCodecLevel(options.properties)); @@ -282,6 +288,13 @@ class ParquetWriter::Impl { } Status Write(ArrowArray* array) { + ArrowArray aligned_array{}; + internal::ArrowArrayGuard aligned_guard(&aligned_array); + if (input_schema_ != nullptr) { + ICEBERG_ASSIGN_OR_RAISE(aligned_array, arrow::AlignBatchForWrite( + array, *input_schema_, *schema_, pool_)); + array = &aligned_array; + } ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batch, ::arrow::ImportRecordBatch(array, arrow_schema_)); @@ -296,13 +309,6 @@ class ParquetWriter::Impl { return {}; } - Status Write(ArrowArray* array, const Schema& input_schema) { - ICEBERG_ASSIGN_OR_RAISE( - auto aligned_array, - arrow::AlignBatchForWrite(array, input_schema, *schema_, pool_)); - return Write(&aligned_array); - } - // Close the writer and release resources Status Close() { if (writer_ == nullptr) { @@ -350,6 +356,7 @@ class ParquetWriter::Impl { ::arrow::MemoryPool* pool_ = ::arrow::default_memory_pool(); // Schema to write from the Iceberg table. std::shared_ptr schema_; + std::shared_ptr input_schema_; // Schema to write from the Parquet file. std::shared_ptr<::arrow::Schema> arrow_schema_; // Parquet schema descriptor generated from the Arrow schema. @@ -379,10 +386,6 @@ Status ParquetWriter::Open(const WriterOptions& options) { Status ParquetWriter::Write(ArrowArray* array) { return impl_->Write(array); } -Status ParquetWriter::Write(ArrowArray* array, const Schema& input_schema) { - return impl_->Write(array, input_schema); -} - Status ParquetWriter::Close() { return impl_->Close(); } Result ParquetWriter::metrics() { return impl_->metrics(); } diff --git a/src/iceberg/parquet/parquet_writer.h b/src/iceberg/parquet/parquet_writer.h index bcd5b65c44..562ffb8013 100644 --- a/src/iceberg/parquet/parquet_writer.h +++ b/src/iceberg/parquet/parquet_writer.h @@ -40,8 +40,6 @@ class ICEBERG_BUNDLE_EXPORT ParquetWriter : public Writer { Status Write(ArrowArray* array) final; - Status Write(ArrowArray* array, const Schema& input_schema) final; - Result metrics() final; Result length() final; diff --git a/src/iceberg/schema_util.cc b/src/iceberg/schema_util.cc index 32082f5590..e6894d923e 100644 --- a/src/iceberg/schema_util.cc +++ b/src/iceberg/schema_util.cc @@ -88,11 +88,13 @@ Status ValidateSchemaEvolution(const Type& expected_type, const Type& source_typ } Result ProjectNested(const Type& expected_type, const Type& source_type, - bool prune_source, const ProjectionOptions& options); + bool prune_source, + const ProjectionOptions& options); Result ProjectField(const SchemaField& expected_field, const SchemaField& source_field, size_t source_index, - bool prune_source, const ProjectionOptions& options) { + bool prune_source, + const ProjectionOptions& options) { FieldProjection projection; if (expected_field.type()->type_id() == TypeId::kUnknown) { @@ -109,20 +111,22 @@ Result ProjectField(const SchemaField& expected_field, expected_field.field_id()); } if (source_field.type()->type_id() == TypeId::kUnknown && - expected_field.type()->is_nested()) { + (expected_field.type()->is_nested() || + (options.default_policy == ProjectionOptions::DefaultPolicy::kWrite && + expected_field.type()->is_primitive()))) { + // The write path must materialize nulls with the expected type rather than copy + // the source's unknown array into a differently typed column. projection.kind = FieldProjection::Kind::kNull; return projection; } if (expected_field.type()->is_nested()) { - ICEBERG_ASSIGN_OR_RAISE( - projection, - ProjectNested(*expected_field.type(), *source_field.type(), prune_source, - options)); + ICEBERG_ASSIGN_OR_RAISE(projection, + ProjectNested(*expected_field.type(), *source_field.type(), + prune_source, options)); } else { - ICEBERG_RETURN_UNEXPECTED(ValidateSchemaEvolution(*expected_field.type(), - *source_field.type(), - options.allow_type_promotion)); + ICEBERG_RETURN_UNEXPECTED(ValidateSchemaEvolution( + *expected_field.type(), *source_field.type(), options.allow_type_promotion)); } // If `prune_source` is false, all fields will be read so the local index is exactly @@ -134,7 +138,8 @@ Result ProjectField(const SchemaField& expected_field, } Result ProjectNested(const Type& expected_type, const Type& source_type, - bool prune_source, const ProjectionOptions& options) { + bool prune_source, + const ProjectionOptions& options) { if (!expected_type.is_nested()) { return InvalidSchema("Expected a nested type, but got {}", expected_type); } @@ -174,10 +179,10 @@ Result ProjectNested(const Type& expected_type, const Type& sou FieldProjection child_projection; if (auto iter = source_field_map.find(field_id); iter != source_field_map.cend()) { - ICEBERG_ASSIGN_OR_RAISE(child_projection, - ProjectField(expected_field, *iter->second.field, - iter->second.local_index, prune_source, - options)); + ICEBERG_ASSIGN_OR_RAISE( + child_projection, + ProjectField(expected_field, *iter->second.field, iter->second.local_index, + prune_source, options)); } else if (MetadataColumns::IsMetadataColumn(field_id)) { child_projection.kind = FieldProjection::Kind::kMetadata; } else if (const auto& default_value = @@ -212,9 +217,9 @@ Result ProjectNested(const Type& expected_type, const Type& sou Result Project(const Schema& expected_schema, const Schema& source_schema, bool prune_source, const ProjectionOptions& options) { - ICEBERG_ASSIGN_OR_RAISE(auto field_projection, - ProjectNested(expected_schema, source_schema, prune_source, - options)); + ICEBERG_ASSIGN_OR_RAISE( + auto field_projection, + ProjectNested(expected_schema, source_schema, prune_source, options)); return SchemaProjection{std::move(field_projection.children)}; } diff --git a/src/iceberg/test/arrow_c_data_util_test.cc b/src/iceberg/test/arrow_c_data_util_test.cc index cd8d794d32..9c12de052e 100644 --- a/src/iceberg/test/arrow_c_data_util_test.cc +++ b/src/iceberg/test/arrow_c_data_util_test.cc @@ -21,6 +21,8 @@ /// Verifies ProjectBatch behavior across registered implementations. #include +#include +#include #include #include #include @@ -37,6 +39,7 @@ #include "iceberg/arrow_c_data_guard_internal.h" #include "iceberg/arrow_c_data_util_internal.h" #include "iceberg/expression/literal.h" +#include "iceberg/result.h" #include "iceberg/schema.h" #include "iceberg/schema_field.h" #include "iceberg/schema_internal.h" @@ -123,12 +126,11 @@ std::shared_ptr MakeFullSchema() { SchemaField::MakeOptional(3, "score", float64())}); } -Result> AlignForWrite(const Schema& input_schema, - const Schema& write_schema, - std::string_view input_json) { - auto input = MakeBatch(input_schema, input_json); +Result> AlignForWrite( + const Schema& input_schema, const Schema& write_schema, + const ::arrow::RecordBatch& input) { ArrowArray c_array; - ICEBERG_ARROW_RETURN_NOT_OK(::arrow::ExportRecordBatch(*input, &c_array)); + ICEBERG_ARROW_RETURN_NOT_OK(::arrow::ExportRecordBatch(input, &c_array)); ICEBERG_ASSIGN_OR_RAISE(auto aligned, arrow::AlignBatchForWrite(&c_array, input_schema, write_schema, ::arrow::default_memory_pool())); @@ -138,15 +140,26 @@ Result> AlignForWrite(const Schema& input_ ICEBERG_ARROW_ASSIGN_OR_RETURN(auto arrow_schema, ::arrow::ImportSchema(&c_schema)); ICEBERG_ARROW_ASSIGN_OR_RETURN(auto result, ::arrow::ImportRecordBatch(&aligned, arrow_schema)); + ICEBERG_ARROW_RETURN_NOT_OK(result->ValidateFull()); return result; } +Result> AlignForWrite(const Schema& input_schema, + const Schema& write_schema, + std::string_view input_json) { + return AlignForWrite(input_schema, write_schema, *MakeBatch(input_schema, input_json)); +} + SchemaField DefaultedField(int32_t id, std::string name, std::shared_ptr type, bool optional, Literal initial_default, Literal write_default) { - return SchemaField(id, std::move(name), std::move(type), optional, {}, - std::make_shared(std::move(initial_default)), - std::make_shared(std::move(write_default))); + return {id, + std::move(name), + std::move(type), + optional, + {}, + std::make_shared(std::move(initial_default)), + std::make_shared(std::move(write_default))}; } } // namespace @@ -287,4 +300,152 @@ TEST(AlignBatchForWriteTest, RejectsInputFieldsOutsideWriteSchema) { EXPECT_THAT(result, HasErrorMessage("Input field id 2")); } +TEST(AlignBatchForWriteTest, PreservesCollectionsWhenAddingAndReorderingFields) { + auto list_field = SchemaField::MakeOptional( + 1, "items", list(SchemaField::MakeOptional(2, "element", int32()))); + auto map_field = + SchemaField::MakeOptional(3, "properties", + map(SchemaField::MakeRequired(4, "key", string()), + SchemaField::MakeOptional(5, "value", int32()))); + Schema input_schema({list_field, map_field}); + Schema write_schema({ + map_field, + list_field, + DefaultedField(6, "name", string(), /*optional=*/false, Literal::String("initial"), + Literal::String("written")), + }); + + ICEBERG_UNWRAP_OR_FAIL(auto actual, + AlignForWrite(input_schema, write_schema, + R"([[[1,2],[["a",3]]],[null,[]],[[],null]])")); + auto expected = MakeBatch( + write_schema, + R"([[[["a",3]],[1,2],"written"],[[],null,"written"],[null,[],"written"]])"); + EXPECT_TRUE(actual->Equals(*expected)) << actual->ToString(); +} + +TEST(AlignBatchForWriteTest, AlignsStructsInsideSlicedNestedCollections) { + auto input_item = std::make_shared(std::vector{ + SchemaField::MakeRequired(4, "name", string()), + SchemaField::MakeRequired(5, "score", int32()), + }); + auto write_item = std::make_shared(std::vector{ + SchemaField::MakeRequired(5, "rating", int32()), + SchemaField::MakeRequired(4, "renamed", string()), + DefaultedField(11, "age", int32(), /*optional=*/false, Literal::Int(10), + Literal::Int(18)), + }); + auto input_value = std::make_shared(std::vector{ + SchemaField::MakeRequired(9, "label", string()), + SchemaField::MakeRequired(10, "code", int32()), + }); + auto write_value = std::make_shared(std::vector{ + SchemaField::MakeRequired(10, "code_new", int32()), + SchemaField::MakeRequired(9, "title", string()), + DefaultedField(12, "flag", boolean(), /*optional=*/false, Literal::Boolean(false), + Literal::Boolean(true)), + }); + auto nested_list = [](std::shared_ptr element) { + return list(SchemaField::MakeOptional( + 2, "element", list(SchemaField::MakeOptional(3, "element", std::move(element))))); + }; + auto nested_map = [](std::shared_ptr value) { + return map(SchemaField::MakeRequired(7, "key", string()), + SchemaField::MakeOptional(8, "value", std::move(value))); + }; + Schema input_schema({SchemaField::MakeOptional(1, "items", nested_list(input_item)), + SchemaField::MakeOptional(6, "mapping", nested_map(input_value))}); + Schema write_schema({SchemaField::MakeOptional(1, "items", nested_list(write_item)), + SchemaField::MakeOptional(6, "mapping", nested_map(write_value))}); + auto input = MakeBatch(input_schema, R"([ + [[[]],[]], + [[[{"name":"alice","score":1},null],[]], + [["a",{"label":"bob","code":2}],["b",null]]], + [null,null], + [[],[]] + ])"); + + ICEBERG_UNWRAP_OR_FAIL(auto actual, + AlignForWrite(input_schema, write_schema, *input->Slice(1, 3))); + auto expected = MakeBatch(write_schema, R"([ + [[[{"rating":1,"renamed":"alice","age":18},null],[]], + [["a",{"code_new":2,"title":"bob","flag":true}],["b",null]]], + [null,null], + [[],[]] + ])"); + EXPECT_TRUE(actual->Equals(*expected)) << actual->ToString(); +} + +TEST(AlignBatchForWriteTest, PreservesSlicedStructOffsetsAndParentNulls) { + auto input_info = std::make_shared( + std::vector{SchemaField::MakeRequired(4, "age", int32())}); + auto write_info = std::make_shared(std::vector{ + SchemaField::MakeRequired(4, "age", int32()), + DefaultedField(5, "score", int32(), /*optional=*/false, Literal::Int(0), + Literal::Int(99)), + }); + auto input_person = std::make_shared(std::vector{ + SchemaField::MakeRequired(2, "name", string()), + SchemaField::MakeOptional(3, "info", input_info), + }); + auto write_person = std::make_shared(std::vector{ + SchemaField::MakeOptional(3, "info", write_info), + SchemaField::MakeRequired(2, "username", string()), + SchemaField::MakeOptional(6, "comment", string()), + }); + Schema input_schema({SchemaField::MakeOptional(1, "person", input_person)}); + Schema write_schema({SchemaField::MakeOptional(1, "person", write_person)}); + auto input = MakeBatch(input_schema, R"([ + [{"name":"first","info":{"age":10}}], + [null], + [{"name":"bob","info":null}], + [{"name":"carol","info":{"age":30}}] + ])"); + + ICEBERG_UNWRAP_OR_FAIL(auto actual, + AlignForWrite(input_schema, write_schema, *input->Slice(1, 3))); + auto expected = MakeBatch(write_schema, R"([ + [null], + [{"info":null,"username":"bob","comment":null}], + [{"info":{"age":30,"score":99},"username":"carol","comment":null}] + ])"); + EXPECT_TRUE(actual->Equals(*expected)) << actual->ToString(); +} + +TEST(AlignBatchForWriteTest, RejectsUnknownFieldsInsideCollections) { + auto input_struct = std::make_shared(std::vector{ + SchemaField::MakeRequired(4, "value", int32()), + SchemaField::MakeOptional(5, "extra", string()), + }); + auto write_struct = std::make_shared( + std::vector{SchemaField::MakeRequired(4, "value", int32())}); + for (bool use_map : {false, true}) { + auto wrap = [use_map](std::shared_ptr child) -> std::shared_ptr { + if (use_map) { + return map(SchemaField::MakeRequired(2, "key", string()), + SchemaField::MakeOptional(3, "value", std::move(child))); + } + return list(SchemaField::MakeOptional( + 2, "element", list(SchemaField::MakeOptional(3, "element", std::move(child))))); + }; + Schema input_schema({SchemaField::MakeOptional(1, "items", wrap(input_struct))}); + Schema write_schema({SchemaField::MakeOptional(1, "items", wrap(write_struct))}); + + auto result = AlignForWrite(input_schema, write_schema, R"([[[]]])"); + + EXPECT_THAT(result, IsError(ErrorKind::kInvalidSchema)); + EXPECT_THAT(result, HasErrorMessage("Input field id 5")); + } +} + +TEST(AlignBatchForWriteTest, MaterializesUnknownInputAsTypedNulls) { + Schema input_schema({SchemaField::MakeOptional(1, "value", unknown())}); + Schema write_schema({SchemaField::MakeOptional(1, "value", int32())}); + + ICEBERG_UNWRAP_OR_FAIL(auto actual, + AlignForWrite(input_schema, write_schema, R"([[null],[null]])")); + auto expected = MakeBatch(write_schema, R"([[null],[null]])"); + EXPECT_TRUE(actual->Equals(*expected)) << actual->ToString(); +} + } // namespace iceberg::internal diff --git a/src/iceberg/test/data_writer_test.cc b/src/iceberg/test/data_writer_test.cc index 71832d8b43..bfd99892f2 100644 --- a/src/iceberg/test/data_writer_test.cc +++ b/src/iceberg/test/data_writer_test.cc @@ -71,7 +71,6 @@ class DataWriterTest : public ::testing::Test { return DataWriterOptions{ .path = "test_data.parquet", .schema = schema_, - .input_schema = schema_, .spec = partition_spec_, .partition = std::move(partition), .format = FileFormatType::kParquet, @@ -139,22 +138,11 @@ class DataWriterFormatTest : public DataWriterTest, public ::testing::WithParamInterface> {}; -TEST_F(DataWriterTest, RejectsMissingInputSchema) { - auto options = MakeDefaultOptions(); - options.input_schema.reset(); - - auto result = DataWriter::Make(options); - - EXPECT_THAT(result, IsError(ErrorKind::kInvalidArgument)); - EXPECT_THAT(result, HasErrorMessage("Input schema must not be null")); -} - TEST_P(DataWriterFormatTest, CreateWithFormat) { auto [format, path] = GetParam(); DataWriterOptions options{ .path = path, .schema = schema_, - .input_schema = schema_, .spec = partition_spec_, .partition = PartitionValues{}, .format = format, @@ -174,12 +162,12 @@ TEST_P(DataWriterFormatTest, WriteRowLineage) { DataWriterOptions options{ .path = path, .schema = schema, - .input_schema = schema, .spec = partition_spec_, .partition = PartitionValues{}, .format = format, .io = file_io_, .properties = FormatProperties(format), + .input_schema = schema, }; ICEBERG_UNWRAP_OR_FAIL(auto writer, DataWriter::Make(options)); diff --git a/src/iceberg/test/default_value_test.cc b/src/iceberg/test/default_value_test.cc index ce76f7dfc3..02a612d520 100644 --- a/src/iceberg/test/default_value_test.cc +++ b/src/iceberg/test/default_value_test.cc @@ -81,6 +81,7 @@ struct WriteDefaultEndToEndParam { std::string name; FileFormatType format; std::string path; + bool avro_skip_datum = true; }; class DefaultValueEndToEndTest @@ -118,11 +119,44 @@ class WriteDefaultEndToEndTest } static std::unordered_map FormatProperties( - FileFormatType format) { - if (format == FileFormatType::kParquet) { + const WriteDefaultEndToEndParam& param) { + if (param.format == FileFormatType::kParquet) { return {{"write.parquet.compression-codec", "uncompressed"}}; } - return {}; + return {{"write.avro.skip-datum", param.avro_skip_datum ? "true" : "false"}}; + } + + void WriteAndRead(const std::shared_ptr& input_schema, + const std::shared_ptr& write_schema, + std::string_view input_json, std::string_view expected_json) { + const auto& param = GetParam(); + DataWriterOptions options{ + .path = param.path, + .schema = write_schema, + .spec = PartitionSpec::Unpartitioned(), + .partition = PartitionValues{}, + .format = param.format, + .io = file_io_, + .properties = FormatProperties(param), + .input_schema = input_schema, + }; + ICEBERG_UNWRAP_OR_FAIL(auto writer, DataWriter::Make(options)); + auto input = CreateArray(*input_schema, input_json); + ArrowArray arrow_array; + ASSERT_TRUE(::arrow::ExportArray(*input, &arrow_array).ok()); + ASSERT_THAT(writer->Write(&arrow_array), IsOk()); + ASSERT_THAT(writer->Close(), IsOk()); + + ICEBERG_UNWRAP_OR_FAIL(auto reader, ReaderFactoryRegistry::Open( + param.format, {.path = param.path, + .io = file_io_, + .projection = write_schema})); + ICEBERG_UNWRAP_OR_FAIL(auto batch, reader->Next()); + ASSERT_TRUE(batch.has_value()); + auto expected = CreateArray(*write_schema, expected_json); + auto actual = ::arrow::ImportArray(&batch.value(), expected->type()).ValueOrDie(); + ASSERT_TRUE(actual->Equals(expected)) + << "actual: " << actual->ToString() << "\nexpected: " << expected->ToString(); } std::shared_ptr file_io_; @@ -229,7 +263,6 @@ TEST_P(DefaultValueEndToEndTest, WriteEvolveReadFillsInitialDefault) { } TEST_P(WriteDefaultEndToEndTest, MissingColumnUsesWriteDefault) { - const auto& param = GetParam(); auto input_schema = std::make_shared( std::vector{SchemaField::MakeRequired(1, "id", int32())}); auto write_schema = std::make_shared(std::vector{ @@ -238,38 +271,18 @@ TEST_P(WriteDefaultEndToEndTest, MissingColumnUsesWriteDefault) { std::make_shared(Literal::Int(42)), std::make_shared(Literal::Int(7))), }); - DataWriterOptions options{ - .path = param.path, - .schema = write_schema, - .input_schema = input_schema, - .spec = PartitionSpec::Unpartitioned(), - .partition = PartitionValues{}, - .format = param.format, - .io = file_io_, - .properties = FormatProperties(param.format), - }; - - ICEBERG_UNWRAP_OR_FAIL(auto writer, DataWriter::Make(options)); - auto input = CreateArray(*input_schema, R"([[1], [2]])"); - ArrowArray arrow_array; - ASSERT_TRUE(::arrow::ExportArray(*input, &arrow_array).ok()); - ASSERT_THAT(writer->Write(&arrow_array), IsOk()); - ASSERT_THAT(writer->Close(), IsOk()); - - ICEBERG_UNWRAP_OR_FAIL( - auto reader, ReaderFactoryRegistry::Open( - param.format, - {.path = param.path, .io = file_io_, .projection = write_schema})); - ICEBERG_UNWRAP_OR_FAIL(auto batch, reader->Next()); - ASSERT_TRUE(batch.has_value()); - - auto expected = CreateArray(*write_schema, R"([[1, 7], [2, 7]])"); - auto actual = ::arrow::ImportArray(&batch.value(), expected->type()).ValueOrDie(); - // This file is written after the column exists, so the missing input column uses // write-default (7), not initial-default (42). - ASSERT_TRUE(actual->Equals(expected)) - << "actual: " << actual->ToString() << "\nexpected: " << expected->ToString(); + WriteAndRead(input_schema, write_schema, R"([[1], [2]])", R"([[1, 7], [2, 7]])"); +} + +TEST_P(WriteDefaultEndToEndTest, UnknownInputWritesTypedNulls) { + auto input_schema = std::make_shared( + std::vector{SchemaField::MakeOptional(1, "id", unknown())}); + auto write_schema = std::make_shared( + std::vector{SchemaField::MakeOptional(1, "id", int32())}); + + WriteAndRead(input_schema, write_schema, R"([[null], [null]])", R"([[null], [null]])"); } namespace { @@ -349,10 +362,12 @@ INSTANTIATE_TEST_SUITE_P( INSTANTIATE_TEST_SUITE_P( Formats, WriteDefaultEndToEndTest, - ::testing::Values(WriteDefaultEndToEndParam{"parquet", FileFormatType::kParquet, - "write-default.parquet"}, - WriteDefaultEndToEndParam{"avro", FileFormatType::kAvro, - "write-default.avro"}), + ::testing::Values( + WriteDefaultEndToEndParam{"parquet", FileFormatType::kParquet, + "write-default.parquet"}, + WriteDefaultEndToEndParam{"avro", FileFormatType::kAvro, "write-default.avro"}, + WriteDefaultEndToEndParam{"avro_generic_datum", FileFormatType::kAvro, + "write-default-generic.avro", false}), [](const ::testing::TestParamInfo& info) { return info.param.name; }); diff --git a/src/iceberg/test/schema_util_test.cc b/src/iceberg/test/schema_util_test.cc index 1138704ced..d3e87cce84 100644 --- a/src/iceberg/test/schema_util_test.cc +++ b/src/iceberg/test/schema_util_test.cc @@ -242,10 +242,9 @@ TEST(SchemaUtilTest, ProjectMissingFieldWithWriteDefault) { std::make_shared(Literal::Int(7))), }); - auto projection_result = - Project(expected_schema, source_schema, /*prune_source=*/false, - ProjectionOptions{ - .default_policy = ProjectionOptions::DefaultPolicy::kWrite}); + auto projection_result = Project( + expected_schema, source_schema, /*prune_source=*/false, + ProjectionOptions{.default_policy = ProjectionOptions::DefaultPolicy::kWrite}); ASSERT_THAT(projection_result, IsOk()); const auto& projection = *projection_result; @@ -265,10 +264,9 @@ TEST(SchemaUtilTest, ProjectWritePolicyIgnoresInitialDefault) { /*doc=*/{}, std::make_shared(Literal::String("n/a"))), }); - auto projection_result = - Project(expected_schema, source_schema, /*prune_source=*/false, - ProjectionOptions{ - .default_policy = ProjectionOptions::DefaultPolicy::kWrite}); + auto projection_result = Project( + expected_schema, source_schema, /*prune_source=*/false, + ProjectionOptions{.default_policy = ProjectionOptions::DefaultPolicy::kWrite}); ASSERT_THAT(projection_result, IsOk()); ASSERT_EQ(projection_result->fields[1].kind, FieldProjection::Kind::kNull); } @@ -280,23 +278,45 @@ TEST(SchemaUtilTest, ProjectWritePolicyRejectsMissingRequired) { SchemaField::MakeRequired(/*field_id=*/10, "extra", iceberg::string()), }); - auto projection_result = - Project(expected_schema, source_schema, /*prune_source=*/false, - ProjectionOptions{ - .default_policy = ProjectionOptions::DefaultPolicy::kWrite}); + auto projection_result = Project( + expected_schema, source_schema, /*prune_source=*/false, + ProjectionOptions{.default_policy = ProjectionOptions::DefaultPolicy::kWrite}); ASSERT_THAT(projection_result, IsError(ErrorKind::kInvalidSchema)); ASSERT_THAT(projection_result, HasErrorMessage("without a write-default")); } +TEST(SchemaUtilTest, ProjectWritePolicyUnknownSourceProducesNull) { + Schema source_schema( + {SchemaField::MakeOptional(/*field_id=*/1, "value", iceberg::unknown())}); + Schema expected_schema({ + SchemaField(/*field_id=*/1, "value", iceberg::int32(), /*optional=*/true, + /*doc=*/{}, /*initial_default=*/nullptr, + std::make_shared(Literal::Int(7))), + }); + + auto write_projection = Project( + expected_schema, source_schema, /*prune_source=*/false, + ProjectionOptions{.default_policy = ProjectionOptions::DefaultPolicy::kWrite, + .allow_type_promotion = false}); + ASSERT_THAT(write_projection, IsOk()); + ASSERT_EQ(write_projection->fields.size(), 1); + // A supplied null stays null, even when the field has a write-default. + ASSERT_EQ(write_projection->fields[0].kind, FieldProjection::Kind::kNull); + + auto read_projection = Project(expected_schema, source_schema, /*prune_source=*/false); + ASSERT_THAT(read_projection, IsOk()); + ASSERT_EQ(read_projection->fields.size(), 1); + AssertProjectedField(read_projection->fields[0], 0); +} + TEST(SchemaUtilTest, ProjectWithoutTypePromotion) { Schema source_schema( {SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int32())}); Schema expected_schema( {SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int64())}); - auto projection_result = - Project(expected_schema, source_schema, /*prune_source=*/false, - ProjectionOptions{.allow_type_promotion = false}); + auto projection_result = Project(expected_schema, source_schema, /*prune_source=*/false, + ProjectionOptions{.allow_type_promotion = false}); ASSERT_THAT(projection_result, IsError(ErrorKind::kNotSupported)); ASSERT_THAT(projection_result, HasErrorMessage("type promotion is not allowed")); } From 2e5c33080c73208d0472f34aa65002a11d21288d Mon Sep 17 00:00:00 2001 From: Xin Huang Date: Mon, 12 Oct 2026 09:30:36 -0700 Subject: [PATCH 5/9] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor(data):=20Cent?= =?UTF-8?q?ralize=20schema=20alignment?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Always align batches while retaining the declared or default input schema. - Forward equal-schema batches without importing or exporting Arrow C Data. --- src/iceberg/arrow/arrow_c_data_util.cc | 5 ++++- src/iceberg/avro/avro_writer.cc | 19 ++++++---------- src/iceberg/parquet/parquet_writer.cc | 18 ++++++---------- src/iceberg/test/arrow_c_data_util_test.cc | 25 ++++++++++++++++++++++ 4 files changed, 41 insertions(+), 26 deletions(-) diff --git a/src/iceberg/arrow/arrow_c_data_util.cc b/src/iceberg/arrow/arrow_c_data_util.cc index 90964f2dff..a9c9290664 100644 --- a/src/iceberg/arrow/arrow_c_data_util.cc +++ b/src/iceberg/arrow/arrow_c_data_util.cc @@ -261,8 +261,11 @@ Result AlignBatchForWrite(ArrowArray* input_batch, const Schema& inp const Schema& write_schema, ::arrow::MemoryPool* pool) { ICEBERG_PRECHECK(input_batch != nullptr, "input_batch must not be null"); - ICEBERG_PRECHECK(pool != nullptr, "pool must not be null"); internal::ArrowArrayGuard input_guard(input_batch); + if (input_schema == write_schema) { + return std::exchange(*input_batch, ArrowArray{}); + } + ICEBERG_PRECHECK(pool != nullptr, "pool must not be null"); ICEBERG_ASSIGN_OR_RAISE(auto input_arrow_schema, ImportIcebergSchema(input_schema)); ICEBERG_ARROW_ASSIGN_OR_RETURN( diff --git a/src/iceberg/avro/avro_writer.cc b/src/iceberg/avro/avro_writer.cc index cedb112028..94deed502f 100644 --- a/src/iceberg/avro/avro_writer.cc +++ b/src/iceberg/avro/avro_writer.cc @@ -183,10 +183,7 @@ class AvroWriter::Impl { public: Status Open(const WriterOptions& options) { write_schema_ = options.schema; - input_schema_ = options.input_schema; - if (input_schema_ != nullptr && *input_schema_ == *write_schema_) { - input_schema_.reset(); - } + input_schema_ = options.input_schema ? options.input_schema : write_schema_; ::avro::NodePtr root; ICEBERG_RETURN_UNEXPECTED(ToAvroNodeVisitor{}.Visit(*write_schema_, &root)); @@ -235,16 +232,12 @@ class AvroWriter::Impl { } Status Write(ArrowArray* data) { - ArrowArray aligned_data{}; + ICEBERG_ASSIGN_OR_RAISE( + auto aligned_data, arrow::AlignBatchForWrite(data, *input_schema_, *write_schema_, + ::arrow::default_memory_pool())); internal::ArrowArrayGuard aligned_guard(&aligned_data); - if (input_schema_ != nullptr) { - ICEBERG_ASSIGN_OR_RAISE( - aligned_data, arrow::AlignBatchForWrite(data, *input_schema_, *write_schema_, - ::arrow::default_memory_pool())); - data = &aligned_data; - } - ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batch, - ::arrow::ImportRecordBatch(data, arrow_schema_)); + ICEBERG_ARROW_ASSIGN_OR_RETURN( + auto batch, ::arrow::ImportRecordBatch(&aligned_data, arrow_schema_)); ICEBERG_ARROW_ASSIGN_OR_RETURN(auto struct_array, batch->ToStructArray()); for (int64_t i = 0; i < struct_array->length(); i++) { diff --git a/src/iceberg/parquet/parquet_writer.cc b/src/iceberg/parquet/parquet_writer.cc index 176d8248fb..c9371f2bcd 100644 --- a/src/iceberg/parquet/parquet_writer.cc +++ b/src/iceberg/parquet/parquet_writer.cc @@ -242,10 +242,7 @@ class ParquetWriter::Impl { public: Status Open(const WriterOptions& options) { schema_ = options.schema; - input_schema_ = options.input_schema; - if (input_schema_ != nullptr && *input_schema_ == *schema_) { - input_schema_.reset(); - } + input_schema_ = options.input_schema ? options.input_schema : schema_; ICEBERG_ASSIGN_OR_RAISE(auto compression, ParseCompression(options.properties)); ICEBERG_ASSIGN_OR_RAISE(auto compression_level, ParseCodecLevel(options.properties)); @@ -288,15 +285,12 @@ class ParquetWriter::Impl { } Status Write(ArrowArray* array) { - ArrowArray aligned_array{}; + ICEBERG_ASSIGN_OR_RAISE( + auto aligned_array, + arrow::AlignBatchForWrite(array, *input_schema_, *schema_, pool_)); internal::ArrowArrayGuard aligned_guard(&aligned_array); - if (input_schema_ != nullptr) { - ICEBERG_ASSIGN_OR_RAISE(aligned_array, arrow::AlignBatchForWrite( - array, *input_schema_, *schema_, pool_)); - array = &aligned_array; - } - ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batch, - ::arrow::ImportRecordBatch(array, arrow_schema_)); + ICEBERG_ARROW_ASSIGN_OR_RETURN( + auto batch, ::arrow::ImportRecordBatch(&aligned_array, arrow_schema_)); ICEBERG_ARROW_ASSIGN_OR_RETURN(auto struct_array, batch->ToStructArray()); FieldMetricsCollector field_metrics_collector(field_metrics_, *metrics_config_, diff --git a/src/iceberg/test/arrow_c_data_util_test.cc b/src/iceberg/test/arrow_c_data_util_test.cc index 9c12de052e..c19feb9a02 100644 --- a/src/iceberg/test/arrow_c_data_util_test.cc +++ b/src/iceberg/test/arrow_c_data_util_test.cc @@ -234,6 +234,31 @@ TEST(ProjectBatchTest, ProjectionRejectsNestedPruning) { EXPECT_THAT(projection, IsError(ErrorKind::kInvalidArgument)); } +TEST(AlignBatchForWriteTest, IdenticalSchemasTransferBatchWithoutConversion) { + auto input_schema = MakeFullSchema(); + auto write_schema = MakeFullSchema(); + auto input = + MakeBatch(*input_schema, R"([[1,"alice",1.0],[2,null,2.0],[3,"carol",3.0]])") + ->Slice(1, 2); + ArrowArray c_array; + ASSERT_TRUE(::arrow::ExportRecordBatch(*input, &c_array).ok()); + auto original = c_array; + + ICEBERG_UNWRAP_OR_FAIL(auto aligned, + arrow::AlignBatchForWrite(&c_array, *input_schema, *write_schema, + ::arrow::default_memory_pool())); + ArrowArrayGuard aligned_guard(&aligned); + EXPECT_EQ(c_array.release, nullptr); + EXPECT_EQ(aligned.private_data, original.private_data); + EXPECT_EQ(aligned.release, original.release); + EXPECT_EQ(aligned.buffers, original.buffers); + EXPECT_EQ(aligned.children, original.children); + + auto actual = ::arrow::ImportRecordBatch(&aligned, input->schema()).ValueOrDie(); + ASSERT_TRUE(actual->ValidateFull().ok()); + EXPECT_TRUE(actual->Equals(*input)); +} + TEST(AlignBatchForWriteTest, ReordersAndMaterializesWriteDefaultsAndNulls) { Schema input_schema({SchemaField::MakeOptional(2, "name", string()), SchemaField::MakeRequired(1, "id", int32())}); From 23d3a06f0769351d69a95bb1ab67a4efa5767a5b Mon Sep 17 00:00:00 2001 From: Xin Huang Date: Tue, 13 Oct 2026 22:42:42 -0700 Subject: [PATCH 6/9] =?UTF-8?q?=F0=9F=90=9B=20fix(arrow):=20Correct=20writ?= =?UTF-8?q?e=20alignment=20DLL=20export?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/iceberg/arrow_c_data_util_internal.h | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/iceberg/arrow_c_data_util_internal.h b/src/iceberg/arrow_c_data_util_internal.h index fe2d25d668..114337ea9e 100644 --- a/src/iceberg/arrow_c_data_util_internal.h +++ b/src/iceberg/arrow_c_data_util_internal.h @@ -31,6 +31,7 @@ #include #include "iceberg/arrow_c_data.h" +#include "iceberg/iceberg_bundle_export.h" #include "iceberg/iceberg_export.h" #include "iceberg/result.h" #include "iceberg/type_fwd.h" @@ -256,10 +257,10 @@ namespace arrow { /// preserved. Structs, lists, and maps are aligned recursively. /// /// `input_batch` is consumed even when alignment fails. -ICEBERG_EXPORT Result AlignBatchForWrite(ArrowArray* input_batch, - const Schema& input_schema, - const Schema& write_schema, - ::arrow::MemoryPool* pool); +ICEBERG_BUNDLE_EXPORT Result AlignBatchForWrite(ArrowArray* input_batch, + const Schema& input_schema, + const Schema& write_schema, + ::arrow::MemoryPool* pool); } // namespace arrow From f202b30898aad766ae3e902d7380d7ec6df8dd43 Mon Sep 17 00:00:00 2001 From: Xin Huang Date: Tue, 13 Oct 2026 23:10:31 -0700 Subject: [PATCH 7/9] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor(data):=20Drop?= =?UTF-8?q?=20unrelated=20schema=20check?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/iceberg/data/data_writer.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/src/iceberg/data/data_writer.cc b/src/iceberg/data/data_writer.cc index a61cad4a6b..5288abdf20 100644 --- a/src/iceberg/data/data_writer.cc +++ b/src/iceberg/data/data_writer.cc @@ -33,7 +33,6 @@ namespace iceberg { class DataWriter::Impl { public: static Result> Make(DataWriterOptions options) { - ICEBERG_PRECHECK(options.schema != nullptr, "Write schema must not be null"); WriterOptions writer_options{ .path = options.path, .schema = options.schema, From 02ca050829a2ce9b6cff1957aae05781ce60e154 Mon Sep 17 00:00:00 2001 From: Xin Huang Date: Tue, 13 Oct 2026 23:28:42 -0700 Subject: [PATCH 8/9] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor(data):=20Narr?= =?UTF-8?q?ow=20write=20alignment=20scope?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/iceberg/arrow/arrow_c_data_util.cc | 11 ++---- src/iceberg/schema_util.cc | 55 +++++++++++--------------- src/iceberg/schema_util.h | 32 +++++---------- src/iceberg/test/data_writer_test.cc | 1 - src/iceberg/test/schema_util_test.cc | 25 +++++------- 5 files changed, 47 insertions(+), 77 deletions(-) diff --git a/src/iceberg/arrow/arrow_c_data_util.cc b/src/iceberg/arrow/arrow_c_data_util.cc index a9c9290664..4ca2f09d65 100644 --- a/src/iceberg/arrow/arrow_c_data_util.cc +++ b/src/iceberg/arrow/arrow_c_data_util.cc @@ -149,7 +149,6 @@ Result> ImportIcebergSchema(const Schema& schem ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(schema, &c_schema)); internal::ArrowSchemaGuard c_schema_guard(&c_schema); ICEBERG_ARROW_ASSIGN_OR_RETURN(auto arrow_schema, ::arrow::ImportSchema(&c_schema)); - c_schema_guard.Release(); return arrow_schema; } @@ -271,17 +270,13 @@ Result AlignBatchForWrite(ArrowArray* input_batch, const Schema& inp ICEBERG_ARROW_ASSIGN_OR_RETURN( auto input_record_batch, ::arrow::ImportRecordBatch(input_batch, input_arrow_schema)); - input_guard.Release(); ICEBERG_ASSIGN_OR_RAISE(auto write_arrow_schema, ImportIcebergSchema(write_schema)); ICEBERG_RETURN_UNEXPECTED(RejectUnknownInputFields(input_schema, write_schema)); - ICEBERG_ASSIGN_OR_RAISE( - auto projection, - Project( - write_schema, input_schema, /*prune_source=*/false, - ProjectionOptions{.default_policy = ProjectionOptions::DefaultPolicy::kWrite, - .allow_type_promotion = false})); + ICEBERG_ASSIGN_OR_RAISE(auto projection, + Project(write_schema, input_schema, /*prune_source=*/false, + ProjectionMode::kWrite)); ICEBERG_ASSIGN_OR_RAISE( auto columns, MaterializeProjection(projection.fields, input_record_batch->columns(), diff --git a/src/iceberg/schema_util.cc b/src/iceberg/schema_util.cc index e6894d923e..a24a12fa57 100644 --- a/src/iceberg/schema_util.cc +++ b/src/iceberg/schema_util.cc @@ -35,8 +35,7 @@ namespace iceberg { namespace { -Status ValidateSchemaEvolution(const Type& expected_type, const Type& source_type, - bool allow_type_promotion) { +Status ValidateSchemaEvolution(const Type& expected_type, const Type& source_type) { if (expected_type.is_nested()) { // Nested type requires identical type ids but their sub-fields are checked // recursively and individually. @@ -53,10 +52,6 @@ Status ValidateSchemaEvolution(const Type& expected_type, const Type& source_typ if (source_type.type_id() == TypeId::kUnknown && expected_type.is_primitive()) { return {}; } - if (!allow_type_promotion) { - return NotSupported("Cannot project {} as {}: type promotion is not allowed", - source_type, expected_type); - } switch (expected_type.type_id()) { case TypeId::kLong: { @@ -88,13 +83,11 @@ Status ValidateSchemaEvolution(const Type& expected_type, const Type& source_typ } Result ProjectNested(const Type& expected_type, const Type& source_type, - bool prune_source, - const ProjectionOptions& options); + bool prune_source, ProjectionMode mode); Result ProjectField(const SchemaField& expected_field, const SchemaField& source_field, size_t source_index, - bool prune_source, - const ProjectionOptions& options) { + bool prune_source, ProjectionMode mode) { FieldProjection projection; if (expected_field.type()->type_id() == TypeId::kUnknown) { @@ -111,9 +104,7 @@ Result ProjectField(const SchemaField& expected_field, expected_field.field_id()); } if (source_field.type()->type_id() == TypeId::kUnknown && - (expected_field.type()->is_nested() || - (options.default_policy == ProjectionOptions::DefaultPolicy::kWrite && - expected_field.type()->is_primitive()))) { + (expected_field.type()->is_nested() || mode == ProjectionMode::kWrite)) { // The write path must materialize nulls with the expected type rather than copy // the source's unknown array into a differently typed column. projection.kind = FieldProjection::Kind::kNull; @@ -121,12 +112,17 @@ Result ProjectField(const SchemaField& expected_field, } if (expected_field.type()->is_nested()) { - ICEBERG_ASSIGN_OR_RAISE(projection, - ProjectNested(*expected_field.type(), *source_field.type(), - prune_source, options)); + ICEBERG_ASSIGN_OR_RAISE( + projection, + ProjectNested(*expected_field.type(), *source_field.type(), prune_source, mode)); + } else if (mode == ProjectionMode::kWrite) { + if (*expected_field.type() != *source_field.type()) { + return NotSupported("Cannot project {} as {}: type promotion is not allowed", + *source_field.type(), *expected_field.type()); + } } else { - ICEBERG_RETURN_UNEXPECTED(ValidateSchemaEvolution( - *expected_field.type(), *source_field.type(), options.allow_type_promotion)); + ICEBERG_RETURN_UNEXPECTED( + ValidateSchemaEvolution(*expected_field.type(), *source_field.type())); } // If `prune_source` is false, all fields will be read so the local index is exactly @@ -138,8 +134,7 @@ Result ProjectField(const SchemaField& expected_field, } Result ProjectNested(const Type& expected_type, const Type& source_type, - bool prune_source, - const ProjectionOptions& options) { + bool prune_source, ProjectionMode mode) { if (!expected_type.is_nested()) { return InvalidSchema("Expected a nested type, but got {}", expected_type); } @@ -179,16 +174,14 @@ Result ProjectNested(const Type& expected_type, const Type& sou FieldProjection child_projection; if (auto iter = source_field_map.find(field_id); iter != source_field_map.cend()) { - ICEBERG_ASSIGN_OR_RAISE( - child_projection, - ProjectField(expected_field, *iter->second.field, iter->second.local_index, - prune_source, options)); + ICEBERG_ASSIGN_OR_RAISE(child_projection, + ProjectField(expected_field, *iter->second.field, + iter->second.local_index, prune_source, mode)); } else if (MetadataColumns::IsMetadataColumn(field_id)) { child_projection.kind = FieldProjection::Kind::kMetadata; - } else if (const auto& default_value = - options.default_policy == ProjectionOptions::DefaultPolicy::kWrite - ? expected_field.write_default() - : expected_field.initial_default(); + } else if (const auto& default_value = mode == ProjectionMode::kWrite + ? expected_field.write_default() + : expected_field.initial_default(); default_value != nullptr) { child_projection.kind = FieldProjection::Kind::kDefault; child_projection.from = *default_value; @@ -196,7 +189,7 @@ Result ProjectNested(const Type& expected_type, const Type& sou child_projection.kind = FieldProjection::Kind::kNull; } else { // TODO(gangwu): support constant value - if (options.default_policy == ProjectionOptions::DefaultPolicy::kWrite) { + if (mode == ProjectionMode::kWrite) { return InvalidSchema("Missing required field: {} without a write-default", expected_field.ToString()); } @@ -216,10 +209,10 @@ Result ProjectNested(const Type& expected_type, const Type& sou Result Project(const Schema& expected_schema, const Schema& source_schema, bool prune_source, - const ProjectionOptions& options) { + ProjectionMode mode) { ICEBERG_ASSIGN_OR_RAISE( auto field_projection, - ProjectNested(expected_schema, source_schema, prune_source, options)); + ProjectNested(expected_schema, source_schema, prune_source, mode)); return SchemaProjection{std::move(field_projection.children)}; } diff --git a/src/iceberg/schema_util.h b/src/iceberg/schema_util.h index 7223abee0e..8895e9133f 100644 --- a/src/iceberg/schema_util.h +++ b/src/iceberg/schema_util.h @@ -81,23 +81,12 @@ struct ICEBERG_EXPORT SchemaProjection { std::vector fields; }; -/// \brief Options to control schema projection behavior. -struct ICEBERG_EXPORT ProjectionOptions { - /// \brief Which default value fills a field that is missing in the source schema. - enum class DefaultPolicy { - /// \brief Use the v3 `initial-default`; for reading data written before the field - /// existed. - kInitial, - /// \brief Use the v3 `write-default`; for aligning data to the write schema. - kWrite, - }; - - /// \brief The default value policy for missing fields. - DefaultPolicy default_policy = DefaultPolicy::kInitial; - /// \brief Whether type promotion (e.g. int to long) is allowed for projected fields. - /// Must be disabled when projecting for the write path, where the source values are - /// stored as-is and therefore must exactly match the expected type. - bool allow_type_promotion = true; +/// \brief Whether schema projection is used for reading or writing. +enum class ProjectionMode { + /// \brief Use initial-default and allow type promotion. + kRead, + /// \brief Use write-default and require matching types. + kWrite, }; /// \brief Project the expected schema on top of the source schema. @@ -109,12 +98,11 @@ struct ICEBERG_EXPORT ProjectionOptions { /// column pruning, so `prune_source` is set to true in this case such that the `from` /// field in `FieldProjection` exactly reflects the position (relative to its nesting /// level) to get the column value from the reader. -/// \param options Options controlling default values and type promotion. +/// \param mode Whether to project for reading or writing. /// \return The projection result. -ICEBERG_EXPORT Result Project(const Schema& expected_schema, - const Schema& source_schema, - bool prune_source, - const ProjectionOptions& options = {}); +ICEBERG_EXPORT Result Project( + const Schema& expected_schema, const Schema& source_schema, bool prune_source, + ProjectionMode mode = ProjectionMode::kRead); ICEBERG_EXPORT std::string_view ToString(FieldProjection::Kind kind); ICEBERG_EXPORT std::string ToString(const FieldProjection& projection); diff --git a/src/iceberg/test/data_writer_test.cc b/src/iceberg/test/data_writer_test.cc index bfd99892f2..d95361675e 100644 --- a/src/iceberg/test/data_writer_test.cc +++ b/src/iceberg/test/data_writer_test.cc @@ -167,7 +167,6 @@ TEST_P(DataWriterFormatTest, WriteRowLineage) { .format = format, .io = file_io_, .properties = FormatProperties(format), - .input_schema = schema, }; ICEBERG_UNWRAP_OR_FAIL(auto writer, DataWriter::Make(options)); diff --git a/src/iceberg/test/schema_util_test.cc b/src/iceberg/test/schema_util_test.cc index d3e87cce84..02a4a943c1 100644 --- a/src/iceberg/test/schema_util_test.cc +++ b/src/iceberg/test/schema_util_test.cc @@ -242,9 +242,8 @@ TEST(SchemaUtilTest, ProjectMissingFieldWithWriteDefault) { std::make_shared(Literal::Int(7))), }); - auto projection_result = Project( - expected_schema, source_schema, /*prune_source=*/false, - ProjectionOptions{.default_policy = ProjectionOptions::DefaultPolicy::kWrite}); + auto projection_result = Project(expected_schema, source_schema, /*prune_source=*/false, + ProjectionMode::kWrite); ASSERT_THAT(projection_result, IsOk()); const auto& projection = *projection_result; @@ -264,9 +263,8 @@ TEST(SchemaUtilTest, ProjectWritePolicyIgnoresInitialDefault) { /*doc=*/{}, std::make_shared(Literal::String("n/a"))), }); - auto projection_result = Project( - expected_schema, source_schema, /*prune_source=*/false, - ProjectionOptions{.default_policy = ProjectionOptions::DefaultPolicy::kWrite}); + auto projection_result = Project(expected_schema, source_schema, /*prune_source=*/false, + ProjectionMode::kWrite); ASSERT_THAT(projection_result, IsOk()); ASSERT_EQ(projection_result->fields[1].kind, FieldProjection::Kind::kNull); } @@ -278,9 +276,8 @@ TEST(SchemaUtilTest, ProjectWritePolicyRejectsMissingRequired) { SchemaField::MakeRequired(/*field_id=*/10, "extra", iceberg::string()), }); - auto projection_result = Project( - expected_schema, source_schema, /*prune_source=*/false, - ProjectionOptions{.default_policy = ProjectionOptions::DefaultPolicy::kWrite}); + auto projection_result = Project(expected_schema, source_schema, /*prune_source=*/false, + ProjectionMode::kWrite); ASSERT_THAT(projection_result, IsError(ErrorKind::kInvalidSchema)); ASSERT_THAT(projection_result, HasErrorMessage("without a write-default")); } @@ -294,10 +291,8 @@ TEST(SchemaUtilTest, ProjectWritePolicyUnknownSourceProducesNull) { std::make_shared(Literal::Int(7))), }); - auto write_projection = Project( - expected_schema, source_schema, /*prune_source=*/false, - ProjectionOptions{.default_policy = ProjectionOptions::DefaultPolicy::kWrite, - .allow_type_promotion = false}); + auto write_projection = Project(expected_schema, source_schema, /*prune_source=*/false, + ProjectionMode::kWrite); ASSERT_THAT(write_projection, IsOk()); ASSERT_EQ(write_projection->fields.size(), 1); // A supplied null stays null, even when the field has a write-default. @@ -309,14 +304,14 @@ TEST(SchemaUtilTest, ProjectWritePolicyUnknownSourceProducesNull) { AssertProjectedField(read_projection->fields[0], 0); } -TEST(SchemaUtilTest, ProjectWithoutTypePromotion) { +TEST(SchemaUtilTest, ProjectWriteRejectsTypePromotion) { Schema source_schema( {SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int32())}); Schema expected_schema( {SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int64())}); auto projection_result = Project(expected_schema, source_schema, /*prune_source=*/false, - ProjectionOptions{.allow_type_promotion = false}); + ProjectionMode::kWrite); ASSERT_THAT(projection_result, IsError(ErrorKind::kNotSupported)); ASSERT_THAT(projection_result, HasErrorMessage("type promotion is not allowed")); } From e53ce78a875e3be0bf7bc82a7fd46a646cabe1ee Mon Sep 17 00:00:00 2001 From: Xin Huang Date: Tue, 13 Oct 2026 23:36:24 -0700 Subject: [PATCH 9/9] =?UTF-8?q?=E2=9C=85=20test(data):=20Trim=20duplicate?= =?UTF-8?q?=20alignment=20coverage?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/iceberg/test/arrow_c_data_util_test.cc | 65 ---------------------- src/iceberg/test/default_value_test.cc | 17 +++--- src/iceberg/test/schema_util_test.cc | 34 ----------- 3 files changed, 9 insertions(+), 107 deletions(-) diff --git a/src/iceberg/test/arrow_c_data_util_test.cc b/src/iceberg/test/arrow_c_data_util_test.cc index c19feb9a02..d3a15c587c 100644 --- a/src/iceberg/test/arrow_c_data_util_test.cc +++ b/src/iceberg/test/arrow_c_data_util_test.cc @@ -259,23 +259,6 @@ TEST(AlignBatchForWriteTest, IdenticalSchemasTransferBatchWithoutConversion) { EXPECT_TRUE(actual->Equals(*input)); } -TEST(AlignBatchForWriteTest, ReordersAndMaterializesWriteDefaultsAndNulls) { - Schema input_schema({SchemaField::MakeOptional(2, "name", string()), - SchemaField::MakeRequired(1, "id", int32())}); - Schema write_schema({ - SchemaField::MakeRequired(1, "id", int32()), - SchemaField::MakeOptional(2, "name", string()), - DefaultedField(3, "score", int32(), /*optional=*/false, Literal::Int(42), - Literal::Int(7)), - SchemaField::MakeOptional(4, "comment", string()), - }); - - ICEBERG_UNWRAP_OR_FAIL(auto actual, AlignForWrite(input_schema, write_schema, - R"([["alice",1],[null,2]])")); - auto expected = MakeBatch(write_schema, R"([[1,"alice",7,null],[2,null,7,null]])"); - EXPECT_TRUE(actual->Equals(*expected)) << "actual:\n" << actual->ToString(); -} - TEST(AlignBatchForWriteTest, RejectsMissingRequiredFieldWithoutWriteDefault) { Schema input_schema({SchemaField::MakeRequired(1, "id", int32())}); Schema write_schema({SchemaField::MakeRequired(1, "id", int32()), @@ -287,44 +270,6 @@ TEST(AlignBatchForWriteTest, RejectsMissingRequiredFieldWithoutWriteDefault) { EXPECT_THAT(result, HasErrorMessage("without a write-default")); } -TEST(AlignBatchForWriteTest, RecursivelyAlignsNestedStructs) { - auto input_person = std::make_shared( - std::vector{SchemaField::MakeRequired(2, "name", string())}); - auto write_person = std::make_shared(std::vector{ - SchemaField::MakeRequired(2, "name", string()), - DefaultedField(3, "age", int32(), /*optional=*/false, Literal::Int(10), - Literal::Int(18)), - }); - Schema input_schema({SchemaField::MakeOptional(1, "person", std::move(input_person))}); - Schema write_schema({SchemaField::MakeOptional(1, "person", std::move(write_person))}); - - ICEBERG_UNWRAP_OR_FAIL(auto actual, AlignForWrite(input_schema, write_schema, - R"([[{"name":"alice"}],[null]])")); - auto expected = MakeBatch(write_schema, R"([[{"name":"alice","age":18}],[null]])"); - EXPECT_TRUE(actual->Equals(*expected)) << "actual:\n" << actual->ToString(); -} - -TEST(AlignBatchForWriteTest, RejectsIncompatibleTypes) { - Schema input_schema({SchemaField::MakeRequired(1, "id", int32())}); - Schema write_schema({SchemaField::MakeRequired(1, "id", int64())}); - - auto result = AlignForWrite(input_schema, write_schema, R"([[1]])"); - - EXPECT_THAT(result, IsError(ErrorKind::kNotSupported)); - EXPECT_THAT(result, HasErrorMessage("type promotion is not allowed")); -} - -TEST(AlignBatchForWriteTest, RejectsInputFieldsOutsideWriteSchema) { - Schema input_schema({SchemaField::MakeRequired(1, "id", int32()), - SchemaField::MakeOptional(2, "extra", string())}); - Schema write_schema({SchemaField::MakeRequired(1, "id", int32())}); - - auto result = AlignForWrite(input_schema, write_schema, R"([[1,"extra"]])"); - - EXPECT_THAT(result, IsError(ErrorKind::kInvalidSchema)); - EXPECT_THAT(result, HasErrorMessage("Input field id 2")); -} - TEST(AlignBatchForWriteTest, PreservesCollectionsWhenAddingAndReorderingFields) { auto list_field = SchemaField::MakeOptional( 1, "items", list(SchemaField::MakeOptional(2, "element", int32()))); @@ -463,14 +408,4 @@ TEST(AlignBatchForWriteTest, RejectsUnknownFieldsInsideCollections) { } } -TEST(AlignBatchForWriteTest, MaterializesUnknownInputAsTypedNulls) { - Schema input_schema({SchemaField::MakeOptional(1, "value", unknown())}); - Schema write_schema({SchemaField::MakeOptional(1, "value", int32())}); - - ICEBERG_UNWRAP_OR_FAIL(auto actual, - AlignForWrite(input_schema, write_schema, R"([[null],[null]])")); - auto expected = MakeBatch(write_schema, R"([[null],[null]])"); - EXPECT_TRUE(actual->Equals(*expected)) << actual->ToString(); -} - } // namespace iceberg::internal diff --git a/src/iceberg/test/default_value_test.cc b/src/iceberg/test/default_value_test.cc index 02a612d520..95080910f3 100644 --- a/src/iceberg/test/default_value_test.cc +++ b/src/iceberg/test/default_value_test.cc @@ -46,8 +46,6 @@ #include "iceberg/file_reader.h" #include "iceberg/file_writer.h" #include "iceberg/parquet/parquet_register.h" -#include "iceberg/partition_spec.h" -#include "iceberg/row/partition_values.h" #include "iceberg/schema.h" #include "iceberg/schema_field.h" #include "iceberg/schema_internal.h" @@ -133,8 +131,6 @@ class WriteDefaultEndToEndTest DataWriterOptions options{ .path = param.path, .schema = write_schema, - .spec = PartitionSpec::Unpartitioned(), - .partition = PartitionValues{}, .format = param.format, .io = file_io_, .properties = FormatProperties(param), @@ -263,17 +259,22 @@ TEST_P(DefaultValueEndToEndTest, WriteEvolveReadFillsInitialDefault) { } TEST_P(WriteDefaultEndToEndTest, MissingColumnUsesWriteDefault) { - auto input_schema = std::make_shared( - std::vector{SchemaField::MakeRequired(1, "id", int32())}); + auto input_schema = std::make_shared(std::vector{ + SchemaField::MakeOptional(2, "name", string()), + SchemaField::MakeRequired(1, "id", int32()), + }); auto write_schema = std::make_shared(std::vector{ SchemaField::MakeRequired(1, "id", int32()), - SchemaField(2, "added", int32(), /*optional=*/false, /*doc=*/{}, + SchemaField::MakeOptional(2, "name", string()), + SchemaField(3, "added", int32(), /*optional=*/false, /*doc=*/{}, std::make_shared(Literal::Int(42)), std::make_shared(Literal::Int(7))), + SchemaField::MakeOptional(4, "comment", string()), }); // This file is written after the column exists, so the missing input column uses // write-default (7), not initial-default (42). - WriteAndRead(input_schema, write_schema, R"([[1], [2]])", R"([[1, 7], [2, 7]])"); + WriteAndRead(input_schema, write_schema, R"([["alice", 1], [null, 2]])", + R"([[1, "alice", 7, null], [2, null, 7, null]])"); } TEST_P(WriteDefaultEndToEndTest, UnknownInputWritesTypedNulls) { diff --git a/src/iceberg/test/schema_util_test.cc b/src/iceberg/test/schema_util_test.cc index 02a4a943c1..31cebe3507 100644 --- a/src/iceberg/test/schema_util_test.cc +++ b/src/iceberg/test/schema_util_test.cc @@ -232,27 +232,6 @@ TEST(SchemaUtilTest, ProjectPresentFieldIgnoresInitialDefault) { AssertProjectedField(projection_result->fields[0], 0); } -TEST(SchemaUtilTest, ProjectMissingFieldWithWriteDefault) { - // The write path fills a missing field with `write-default`, not `initial-default`. - Schema source_schema = CreateFlatSchema(); - Schema expected_schema({ - SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int64()), - SchemaField(/*field_id=*/10, "extra", iceberg::int32(), /*optional=*/false, - /*doc=*/{}, std::make_shared(Literal::Int(42)), - std::make_shared(Literal::Int(7))), - }); - - auto projection_result = Project(expected_schema, source_schema, /*prune_source=*/false, - ProjectionMode::kWrite); - ASSERT_THAT(projection_result, IsOk()); - - const auto& projection = *projection_result; - ASSERT_EQ(projection.fields.size(), 2); - AssertProjectedField(projection.fields[0], 0); - ASSERT_EQ(projection.fields[1].kind, FieldProjection::Kind::kDefault); - ASSERT_EQ(std::get(projection.fields[1].from), Literal::Int(7)); -} - TEST(SchemaUtilTest, ProjectWritePolicyIgnoresInitialDefault) { // With the write default source, a field having only `initial-default` behaves as if // it had no default: optional becomes null and required is rejected. @@ -269,19 +248,6 @@ TEST(SchemaUtilTest, ProjectWritePolicyIgnoresInitialDefault) { ASSERT_EQ(projection_result->fields[1].kind, FieldProjection::Kind::kNull); } -TEST(SchemaUtilTest, ProjectWritePolicyRejectsMissingRequired) { - Schema source_schema = CreateFlatSchema(); - Schema expected_schema({ - SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int64()), - SchemaField::MakeRequired(/*field_id=*/10, "extra", iceberg::string()), - }); - - auto projection_result = Project(expected_schema, source_schema, /*prune_source=*/false, - ProjectionMode::kWrite); - ASSERT_THAT(projection_result, IsError(ErrorKind::kInvalidSchema)); - ASSERT_THAT(projection_result, HasErrorMessage("without a write-default")); -} - TEST(SchemaUtilTest, ProjectWritePolicyUnknownSourceProducesNull) { Schema source_schema( {SchemaField::MakeOptional(/*field_id=*/1, "value", iceberg::unknown())});