diff --git a/src/iceberg/arrow/arrow_c_data_util.cc b/src/iceberg/arrow/arrow_c_data_util.cc index 5c8ea5bbf3..4ca2f09d65 100644 --- a/src/iceberg/arrow/arrow_c_data_util.cc +++ b/src/iceberg/arrow/arrow_c_data_util.cc @@ -22,20 +22,33 @@ #include #include #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 +141,156 @@ 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)); + return arrow_schema; +} + +// 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_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()->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()))); + } + } + 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 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_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()); + } + 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)); + 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"); + 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( + auto input_record_batch, + ::arrow::ImportRecordBatch(input_batch, input_arrow_schema)); + + 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, + ProjectionMode::kWrite)); + 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..114337ea9e 100644 --- a/src/iceberg/arrow_c_data_util_internal.h +++ b/src/iceberg/arrow_c_data_util_internal.h @@ -31,10 +31,15 @@ #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" +namespace arrow { +class MemoryPool; +} + namespace iceberg { /// \brief Cached state for ProjectBatch over one input/output schema pair. @@ -242,4 +247,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. Structs, lists, and maps are aligned recursively. +/// +/// `input_batch` is consumed even when alignment fails. +ICEBERG_BUNDLE_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/avro/avro_writer.cc b/src/iceberg/avro/avro_writer.cc index 01d8cb3b14..94deed502f 100644 --- a/src/iceberg/avro/avro_writer.cc +++ b/src/iceberg/avro/avro_writer.cc @@ -32,6 +32,8 @@ #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" #include "iceberg/avro/avro_metrics.h" @@ -181,6 +183,7 @@ class AvroWriter::Impl { public: Status Open(const WriterOptions& options) { write_schema_ = options.schema; + input_schema_ = options.input_schema ? options.input_schema : write_schema_; ::avro::NodePtr root; ICEBERG_RETURN_UNEXPECTED(ToAvroNodeVisitor{}.Visit(*write_schema_, &root)); @@ -229,8 +232,12 @@ class AvroWriter::Impl { } Status Write(ArrowArray* data) { - ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batch, - ::arrow::ImportRecordBatch(data, arrow_schema_)); + ICEBERG_ASSIGN_OR_RAISE( + auto aligned_data, arrow::AlignBatchForWrite(data, *input_schema_, *write_schema_, + ::arrow::default_memory_pool())); + internal::ArrowArrayGuard aligned_guard(&aligned_data); + 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++) { @@ -272,6 +279,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 diff --git a/src/iceberg/data/data_writer.cc b/src/iceberg/data/data_writer.cc index 512f77cc96..5288abdf20 100644 --- a/src/iceberg/data/data_writer.cc +++ b/src/iceberg/data/data_writer.cc @@ -38,6 +38,7 @@ class DataWriter::Impl { .schema = options.schema, .io = options.io, .properties = WriterProperties::FromMap(options.properties), + .input_schema = options.input_schema, }; ICEBERG_ASSIGN_OR_RAISE(auto writer, diff --git a/src/iceberg/data/data_writer.h b/src/iceberg/data/data_writer.h index 9a2002f0db..39d9f50873 100644 --- a/src/iceberg/data/data_writer.h +++ b/src/iceberg/data/data_writer.h @@ -48,6 +48,8 @@ struct ICEBERG_DATA_EXPORT DataWriterOptions { 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 4b1045a409..d4aea3431a 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; @@ -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. diff --git a/src/iceberg/parquet/parquet_writer.cc b/src/iceberg/parquet/parquet_writer.cc index 08c3941383..c9371f2bcd 100644 --- a/src/iceberg/parquet/parquet_writer.cc +++ b/src/iceberg/parquet/parquet_writer.cc @@ -39,8 +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" @@ -239,6 +242,7 @@ class ParquetWriter::Impl { public: Status Open(const WriterOptions& options) { schema_ = options.schema; + 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)); @@ -281,8 +285,12 @@ class ParquetWriter::Impl { } Status Write(ArrowArray* array) { - ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batch, - ::arrow::ImportRecordBatch(array, arrow_schema_)); + ICEBERG_ASSIGN_OR_RAISE( + auto aligned_array, + arrow::AlignBatchForWrite(array, *input_schema_, *schema_, pool_)); + internal::ArrowArrayGuard aligned_guard(&aligned_array); + 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_, @@ -342,6 +350,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. diff --git a/src/iceberg/schema_util.cc b/src/iceberg/schema_util.cc index b9f32346a9..a24a12fa57 100644 --- a/src/iceberg/schema_util.cc +++ b/src/iceberg/schema_util.cc @@ -83,11 +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); + bool prune_source, ProjectionMode mode); Result ProjectField(const SchemaField& expected_field, const SchemaField& source_field, size_t source_index, - bool prune_source) { + bool prune_source, ProjectionMode mode) { FieldProjection projection; if (expected_field.type()->type_id() == TypeId::kUnknown) { @@ -104,7 +104,9 @@ 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() || 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; return projection; } @@ -112,7 +114,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, 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())); @@ -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, ProjectionMode mode) { if (!expected_type.is_nested()) { return InvalidSchema("Expected a nested type, but got {}", expected_type); } @@ -169,17 +176,23 @@ 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, mode)); } 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 = mode == ProjectionMode::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 (mode == ProjectionMode::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 +208,11 @@ Result ProjectNested(const Type& expected_type, const Type& sou } // namespace Result Project(const Schema& expected_schema, - const Schema& source_schema, bool prune_source) { - ICEBERG_ASSIGN_OR_RAISE(auto field_projection, - ProjectNested(expected_schema, source_schema, prune_source)); + const Schema& source_schema, bool prune_source, + ProjectionMode mode) { + ICEBERG_ASSIGN_OR_RAISE( + auto field_projection, + 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 421d00e0ad..8895e9133f 100644 --- a/src/iceberg/schema_util.h +++ b/src/iceberg/schema_util.h @@ -81,6 +81,14 @@ struct ICEBERG_EXPORT SchemaProjection { std::vector fields; }; +/// \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. /// /// \param expected_schema The expected schema. @@ -90,10 +98,11 @@ 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 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); +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/arrow_c_data_util_test.cc b/src/iceberg/test/arrow_c_data_util_test.cc index c93a928a3d..d3a15c587c 100644 --- a/src/iceberg/test/arrow_c_data_util_test.cc +++ b/src/iceberg/test/arrow_c_data_util_test.cc @@ -21,24 +21,31 @@ /// Verifies ProjectBatch behavior across registered implementations. #include +#include +#include #include #include #include #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/result.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 +126,42 @@ std::shared_ptr MakeFullSchema() { SchemaField::MakeOptional(3, "score", float64())}); } +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_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)); + 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 {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 +234,178 @@ 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, 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, 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")); + } +} + } // namespace iceberg::internal diff --git a/src/iceberg/test/default_value_test.cc b/src/iceberg/test/default_value_test.cc index f85f2e0d22..95080910f3 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,6 +40,7 @@ #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" @@ -72,6 +75,13 @@ struct DefaultValueEndToEndParam { bool avro_skip_datum = true; }; +struct WriteDefaultEndToEndParam { + std::string name; + FileFormatType format; + std::string path; + bool avro_skip_datum = true; +}; + class DefaultValueEndToEndTest : public UpdateTestBase, public ::testing::WithParamInterface { @@ -86,6 +96,68 @@ 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( + const WriteDefaultEndToEndParam& param) { + if (param.format == FileFormatType::kParquet) { + return {{"write.parquet.compression-codec", "uncompressed"}}; + } + 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, + .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_; +}; + TEST_P(DefaultValueEndToEndTest, WriteEvolveReadFillsInitialDefault) { const auto& param = GetParam(); ICEBERG_UNWRAP_OR_FAIL(auto original_schema, table_->schema()); @@ -186,6 +258,34 @@ TEST_P(DefaultValueEndToEndTest, WriteEvolveReadFillsInitialDefault) { ASSERT_FALSE(next_batch.has_value()); } +TEST_P(WriteDefaultEndToEndTest, MissingColumnUsesWriteDefault) { + 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::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"([["alice", 1], [null, 2]])", + R"([[1, "alice", 7, null], [2, null, 7, null]])"); +} + +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 { // Two-row column of `json` values at `type`, for the simple cases. @@ -261,6 +361,18 @@ 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"}, + WriteDefaultEndToEndParam{"avro_generic_datum", FileFormatType::kAvro, + "write-default-generic.avro", false}), + [](const ::testing::TestParamInfo& info) { + return info.param.name; + }); + } // namespace } // namespace iceberg diff --git a/src/iceberg/test/schema_util_test.cc b/src/iceberg/test/schema_util_test.cc index 9a3eff8870..31cebe3507 100644 --- a/src/iceberg/test/schema_util_test.cc +++ b/src/iceberg/test/schema_util_test.cc @@ -232,6 +232,56 @@ TEST(SchemaUtilTest, ProjectPresentFieldIgnoresInitialDefault) { AssertProjectedField(projection_result->fields[0], 0); } +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, + ProjectionMode::kWrite); + ASSERT_THAT(projection_result, IsOk()); + ASSERT_EQ(projection_result->fields[1].kind, FieldProjection::Kind::kNull); +} + +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, + 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. + 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, 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, + ProjectionMode::kWrite); + 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({