diff --git a/src/iceberg/avro/avro_data_util.cc b/src/iceberg/avro/avro_data_util.cc index 7931b0719..d81bea977 100644 --- a/src/iceberg/avro/avro_data_util.cc +++ b/src/iceberg/avro/avro_data_util.cc @@ -17,12 +17,15 @@ * under the License. */ +#include + #include #include #include #include #include #include +#include #include #include #include @@ -31,6 +34,7 @@ #include #include "iceberg/arrow/arrow_status_internal.h" +#include "iceberg/arrow/literal_util_internal.h" #include "iceberg/avro/avro_data_util_internal.h" #include "iceberg/avro/avro_schema_util_internal.h" #include "iceberg/metadata_columns.h" @@ -88,6 +92,8 @@ Status AppendStructToBuilder(const ::avro::NodePtr& avro_node, metadata_context, field_builder)); } else if (field_projection.kind == FieldProjection::Kind::kNull) { ICEBERG_ARROW_RETURN_NOT_OK(field_builder->AppendNull()); + } else if (field_projection.kind == FieldProjection::Kind::kDefault) { + ICEBERG_RETURN_UNEXPECTED(AppendDefaultToBuilder(field_projection, field_builder)); } else if (field_projection.kind == FieldProjection::Kind::kMetadata) { int32_t field_id = expected_field.field_id(); if (field_id == MetadataColumns::kFilePathColumnId) { @@ -466,6 +472,10 @@ Status AppendFieldToBuilder(const ::avro::NodePtr& avro_node, return {}; } + if (projection.kind == FieldProjection::Kind::kDefault) { + return AppendDefaultToBuilder(projection, array_builder); + } + const bool is_row_lineage = MetadataColumns::IsRowLineageColumn(projected_field.field_id()); @@ -497,6 +507,128 @@ Status AppendFieldToBuilder(const ::avro::NodePtr& avro_node, } // namespace +namespace { + +Result> MakeDefaultScalar( + const Literal& literal, const std::shared_ptr<::arrow::DataType>& builder_type) { + // The builder's own memory pool is not exposed, so the small scalar buffer uses the + // default pool. + ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<::arrow::Scalar> scalar, + arrow::ToArrowScalar(literal, ::arrow::default_memory_pool())); + + // For an extension builder (e.g. `arrow.uuid`) target its storage type: ToArrowScalar + // yields the storage scalar (fixed_size_binary(16) for uuid) and Scalar::CastTo has no + // kernel that targets an extension type. This mirrors MakeDefaultArray's extension + // handling. + std::shared_ptr<::arrow::DataType> target_type = builder_type; + if (target_type->id() == ::arrow::Type::EXTENSION) { + target_type = internal::checked_cast(*target_type) + .storage_type(); + } + + if (!scalar->type->Equals(*target_type)) { + ICEBERG_ARROW_ASSIGN_OR_RETURN(scalar, scalar->CastTo(target_type)); + } + return scalar; +} + +Status PrepareDefaultScalars(std::span projections, + ::arrow::ArrayBuilder* builder) { + auto* struct_builder = internal::checked_cast<::arrow::StructBuilder*>(builder); + if (static_cast(struct_builder->num_fields()) != projections.size()) { + return InvalidArgument( + "Inconsistent number of struct builder fields ({}) and projections ({})", + struct_builder->num_fields(), projections.size()); + } + + for (size_t i = 0; i < projections.size(); ++i) { + auto& field_projection = projections[i]; + auto* field_builder = struct_builder->field_builder(static_cast(i)); + + if (field_projection.kind == FieldProjection::Kind::kDefault) { + if (dynamic_cast(field_projection.attributes.get()) != + nullptr) { + continue; + } + auto attrs = std::make_shared(); + ICEBERG_ASSIGN_OR_RAISE(attrs->scalar, + MakeDefaultScalar(std::get(field_projection.from), + field_builder->type())); + field_projection.attributes = std::move(attrs); + continue; + } + + if (field_projection.kind != FieldProjection::Kind::kProjected || + field_projection.children.empty()) { + continue; + } + + switch (field_builder->type()->id()) { + case ::arrow::Type::STRUCT: + ICEBERG_RETURN_UNEXPECTED( + PrepareDefaultScalars(field_projection.children, field_builder)); + break; + case ::arrow::Type::LIST: { + // List projections store a single child for the element. Defaults only appear + // on nested struct fields of that element. + auto* list_builder = internal::checked_cast<::arrow::ListBuilder*>(field_builder); + auto& element_projection = field_projection.children[0]; + if (element_projection.kind == FieldProjection::Kind::kProjected && + !element_projection.children.empty() && + list_builder->value_builder()->type()->id() == ::arrow::Type::STRUCT) { + ICEBERG_RETURN_UNEXPECTED(PrepareDefaultScalars(element_projection.children, + list_builder->value_builder())); + } + break; + } + case ::arrow::Type::MAP: { + auto* map_builder = internal::checked_cast<::arrow::MapBuilder*>(field_builder); + if (field_projection.children.size() >= 1 && + !field_projection.children[0].children.empty() && + map_builder->key_builder()->type()->id() == ::arrow::Type::STRUCT) { + ICEBERG_RETURN_UNEXPECTED(PrepareDefaultScalars( + field_projection.children[0].children, map_builder->key_builder())); + } + if (field_projection.children.size() >= 2 && + !field_projection.children[1].children.empty() && + map_builder->item_builder()->type()->id() == ::arrow::Type::STRUCT) { + ICEBERG_RETURN_UNEXPECTED(PrepareDefaultScalars( + field_projection.children[1].children, map_builder->item_builder())); + } + break; + } + default: + break; + } + } + return {}; +} + +} // namespace + +Status PrepareDefaultScalars(SchemaProjection& projection, + ::arrow::ArrayBuilder* root_builder) { + return PrepareDefaultScalars(projection.fields, root_builder); +} + +Status AppendDefaultToBuilder(const Literal& literal, ::arrow::ArrayBuilder* builder) { + ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<::arrow::Scalar> scalar, + MakeDefaultScalar(literal, builder->type())); + ICEBERG_ARROW_RETURN_NOT_OK(builder->AppendScalar(*scalar)); + return {}; +} + +Status AppendDefaultToBuilder(const FieldProjection& projection, + ::arrow::ArrayBuilder* builder) { + if (const auto* attrs = + dynamic_cast(projection.attributes.get()); + attrs != nullptr && attrs->scalar != nullptr) { + ICEBERG_ARROW_RETURN_NOT_OK(builder->AppendScalar(*attrs->scalar)); + return {}; + } + return AppendDefaultToBuilder(std::get(projection.from), builder); +} + Status AppendDatumToBuilder(const ::avro::NodePtr& avro_node, const ::avro::GenericDatum& avro_datum, const SchemaProjection& projection, diff --git a/src/iceberg/avro/avro_data_util_internal.h b/src/iceberg/avro/avro_data_util_internal.h index 281187f12..f894f204f 100644 --- a/src/iceberg/avro/avro_data_util_internal.h +++ b/src/iceberg/avro/avro_data_util_internal.h @@ -19,14 +19,48 @@ #pragma once +#include + #include +#include #include #include "iceberg/arrow/metadata_column_util_internal.h" +#include "iceberg/expression/literal.h" #include "iceberg/schema_util.h" namespace iceberg::avro { +/// \brief Cached Arrow scalar for a `kDefault` field projection. +/// +/// Materialized once (see `PrepareDefaultScalars`) so row-by-row Avro decode only +/// needs `AppendScalar` instead of repeating `ToArrowScalar` / `CastTo` per row. +struct AvroDefaultAttributes : FieldProjection::ExtraAttributes { + std::shared_ptr<::arrow::Scalar> scalar; +}; + +/// \brief Precompute cast Arrow scalars for every `kDefault` field under `projection`. +/// +/// Walks `root_builder` in lockstep with the projection so each default is cast to the +/// builder's Arrow type once per scan. Safe to call repeatedly; existing +/// `AvroDefaultAttributes` entries are left unchanged. +Status PrepareDefaultScalars(SchemaProjection& projection, + ::arrow::ArrayBuilder* root_builder); + +/// \brief Append a literal once to `builder` while decoding Avro row-by-row. +/// +/// Used to materialize `FieldProjection::Kind::kDefault`. Shares `ToArrowScalar` with +/// Parquet's batch path (`MakeDefaultArray`); the append shape stays Avro-local because +/// Avro builds Arrow arrays via per-row `ArrayBuilder`s rather than whole-column arrays. +/// Prefer the `FieldProjection` overload after `PrepareDefaultScalars` so the scalar is +/// reused across rows. +Status AppendDefaultToBuilder(const Literal& literal, ::arrow::ArrayBuilder* builder); + +/// \brief Append a `kDefault` projection, reusing a scalar cached on +/// `projection.attributes` when present. +Status AppendDefaultToBuilder(const FieldProjection& projection, + ::arrow::ArrayBuilder* builder); + /// \brief Append an Avro datum to an Arrow array builder. /// /// This function handles schema evolution by using the provided projection to map diff --git a/src/iceberg/avro/avro_direct_decoder.cc b/src/iceberg/avro/avro_direct_decoder.cc index b66a93109..8f563b3ba 100644 --- a/src/iceberg/avro/avro_direct_decoder.cc +++ b/src/iceberg/avro/avro_direct_decoder.cc @@ -29,6 +29,7 @@ #include #include "iceberg/arrow/arrow_status_internal.h" +#include "iceberg/avro/avro_data_util_internal.h" #include "iceberg/avro/avro_direct_decoder_internal.h" #include "iceberg/avro/avro_schema_util_internal.h" #include "iceberg/metadata_columns.h" @@ -209,6 +210,8 @@ Status DecodeStructToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder& auto* field_builder = struct_builder->field_builder(static_cast(proj_idx)); if (field_projection.kind == FieldProjection::Kind::kNull) { ICEBERG_ARROW_RETURN_NOT_OK(field_builder->AppendNull()); + } else if (field_projection.kind == FieldProjection::Kind::kDefault) { + ICEBERG_RETURN_UNEXPECTED(AppendDefaultToBuilder(field_projection, field_builder)); } else if (field_projection.kind == FieldProjection::Kind::kMetadata) { int32_t field_id = expected_field.field_id(); if (field_id == MetadataColumns::kFilePathColumnId) { diff --git a/src/iceberg/avro/avro_reader.cc b/src/iceberg/avro/avro_reader.cc index b0977bcdb..f333a38cb 100644 --- a/src/iceberg/avro/avro_reader.cc +++ b/src/iceberg/avro/avro_reader.cc @@ -434,6 +434,8 @@ class AvroReader::Impl { builder_result.status().message()); } context_->builder_ = builder_result.MoveValueUnsafe(); + ICEBERG_RETURN_UNEXPECTED( + PrepareDefaultScalars(projection_, context_->builder_.get())); backend_->InitReadContext(backend_->GetReaderSchema()); return {}; diff --git a/src/iceberg/avro/avro_schema_util.cc b/src/iceberg/avro/avro_schema_util.cc index 14b464cee..5e6bee957 100644 --- a/src/iceberg/avro/avro_schema_util.cc +++ b/src/iceberg/avro/avro_schema_util.cc @@ -752,6 +752,10 @@ Result ProjectStruct(const StructType& struct_type, iter->second.local_index, prune_source)); } 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. + child_projection.kind = FieldProjection::Kind::kDefault; + child_projection.from = *expected_field.initial_default(); } else if (expected_field.optional()) { child_projection.kind = FieldProjection::Kind::kNull; } else { diff --git a/src/iceberg/test/avro_data_test.cc b/src/iceberg/test/avro_data_test.cc index 7731f58d3..46d940bb1 100644 --- a/src/iceberg/test/avro_data_test.cc +++ b/src/iceberg/test/avro_data_test.cc @@ -19,8 +19,12 @@ #include +#include +#include #include #include +#include +#include #include #include #include @@ -31,6 +35,7 @@ #include "iceberg/avro/avro_data_util_internal.h" #include "iceberg/avro/avro_schema_util_internal.h" +#include "iceberg/expression/literal.h" #include "iceberg/schema.h" #include "iceberg/schema_internal.h" #include "iceberg/schema_util.h" @@ -662,6 +667,97 @@ TEST(AppendDatumToBuilderTest, StructWithMissingOptionalField) { avro_data, expected_json)); } +TEST(AppendDatumToBuilderTest, StructWithMissingDefaultFields) { + Schema iceberg_schema({ + SchemaField::MakeRequired(1, "id", iceberg::int32()), + // Missing required field with an initial-default: filled with the default. + SchemaField(2, "score", iceberg::int64(), /*optional=*/false, /*doc=*/{}, + std::make_shared(Literal::Long(100))), + // Missing optional field with an initial-default: also filled, not null. + SchemaField(3, "grade", iceberg::string(), /*optional=*/true, /*doc=*/{}, + std::make_shared(Literal::String("A"))), + }); + + // Create Avro schema that only has the id field (missing score and grade). + std::string avro_schema_json = R"({ + "type": "record", + "name": "person", + "fields": [ + {"name": "id", "type": "int", "field-id": 1} + ] + })"; + auto avro_schema = ::avro::compileJsonSchemaFromString(avro_schema_json); + + std::vector<::avro::GenericDatum> avro_data; + for (int i = 0; i < 2; ++i) { + ::avro::GenericDatum avro_datum(avro_schema.root()); + auto& record = avro_datum.value<::avro::GenericRecord>(); + record.fieldAt(0).value() = i + 1; + avro_data.push_back(avro_datum); + } + + const std::string expected_json = R"([ + {"id": 1, "score": 100, "grade": "A"}, + {"id": 2, "score": 100, "grade": "A"} + ])"; + ASSERT_NO_FATAL_FAILURE(VerifyAppendDatumToBuilder(iceberg_schema, avro_schema.root(), + avro_data, expected_json)); +} + +TEST(AppendDefaultToBuilderTest, AppendsValue) { + ::arrow::Int64Builder builder; + ASSERT_THAT(AppendDefaultToBuilder(Literal::Long(42), &builder), IsOk()); + ASSERT_THAT(AppendDefaultToBuilder(Literal::Long(42), &builder), IsOk()); + + std::shared_ptr<::arrow::Array> array; + ASSERT_TRUE(builder.Finish(&array).ok()); + ASSERT_EQ(array->length(), 2); + const auto& long_array = static_cast(*array); + ASSERT_EQ(long_array.Value(0), 42); + ASSERT_EQ(long_array.Value(1), 42); +} + +TEST(AppendDefaultToBuilderTest, CastsToBuilderType) { + // The literal's natural type (int32) differs from the builder type (int64); the value + // is cast to the builder type. + ::arrow::Int64Builder builder; + ASSERT_THAT(AppendDefaultToBuilder(Literal::Int(7), &builder), IsOk()); + + std::shared_ptr<::arrow::Array> array; + ASSERT_TRUE(builder.Finish(&array).ok()); + ASSERT_EQ(array->length(), 1); + ASSERT_EQ(static_cast(*array).Value(0), 7); +} + +TEST(AppendDefaultToBuilderTest, ReusesPreparedScalar) { + auto pool = ::arrow::default_memory_pool(); + auto child = std::make_shared<::arrow::Int64Builder>(pool); + ::arrow::StructBuilder struct_builder( + ::arrow::struct_({::arrow::field("d", ::arrow::int64())}), pool, {child}); + + FieldProjection projection; + projection.kind = FieldProjection::Kind::kDefault; + projection.from = Literal::Long(42); + + SchemaProjection schema_projection; + schema_projection.fields.push_back(projection); + ASSERT_THAT(PrepareDefaultScalars(schema_projection, &struct_builder), IsOk()); + ASSERT_NE(dynamic_cast( + schema_projection.fields[0].attributes.get()), + nullptr); + + auto* field_builder = struct_builder.field_builder(0); + ASSERT_THAT(AppendDefaultToBuilder(schema_projection.fields[0], field_builder), IsOk()); + ASSERT_THAT(AppendDefaultToBuilder(schema_projection.fields[0], field_builder), IsOk()); + + std::shared_ptr<::arrow::Array> array; + ASSERT_TRUE(field_builder->Finish(&array).ok()); + ASSERT_EQ(array->length(), 2); + const auto& long_array = static_cast(*array); + ASSERT_EQ(long_array.Value(0), 42); + ASSERT_EQ(long_array.Value(1), 42); +} + TEST(AppendDatumToBuilderTest, NestedStructWithMissingOptionalFields) { Schema iceberg_schema({ SchemaField::MakeRequired(1, "id", iceberg::int32()), diff --git a/src/iceberg/test/avro_test.cc b/src/iceberg/test/avro_test.cc index 604cf5b0a..20156a57c 100644 --- a/src/iceberg/test/avro_test.cc +++ b/src/iceberg/test/avro_test.cc @@ -42,6 +42,7 @@ #include "iceberg/avro/avro_register.h" #include "iceberg/avro/avro_stream_internal.h" #include "iceberg/avro/avro_writer.h" +#include "iceberg/expression/literal.h" #include "iceberg/file_reader.h" #include "iceberg/metadata_columns.h" #include "iceberg/schema.h" @@ -318,6 +319,34 @@ TEST_F(AvroReaderTest, ReadTwoFields) { ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); } +TEST_F(AvroReaderTest, ReadMissingFieldsWithDefaults) { + // The file contains only fields 1 and 2; the projected schema adds fields 3 and 4 + // with initial-defaults, which are filled for all rows written before the columns + // existed. + CreateSimpleAvroFile(); + + auto schema = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", std::make_shared()), + SchemaField::MakeOptional(2, "name", std::make_shared()), + SchemaField(3, "score", std::make_shared(), /*optional=*/false, + /*doc=*/{}, std::make_shared(Literal::Long(100))), + SchemaField(4, "status", std::make_shared(), /*optional=*/true, + /*doc=*/{}, std::make_shared(Literal::String("active"))), + }); + + auto reader_result = ReaderFactoryRegistry::Open( + FileFormatType::kAvro, + {.path = temp_avro_file_, .io = file_io_, .projection = schema}); + ASSERT_THAT(reader_result, IsOk()); + auto reader = std::move(reader_result.value()); + + ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader, + R"([[1, "Alice", 100, "active"], + [2, "Bob", 100, "active"], + [3, "Charlie", 100, "active"]])")); + ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader)); +} + TEST_F(AvroReaderTest, RoundTripWithGenericFileIO) { file_io_ = std::make_shared(); temp_avro_file_ = CreateNewTempFilePathWithSuffix(".avro");