Skip to content

Commit 75bd63f

Browse files
committed
feat(avro): apply column default values when reading missing fields (3/4)
When a column is present in the read schema but missing from an Avro data file (written before the column existed), fill it with the column's v3 initial-default instead of null. Reuses the shared arrow/literal_util materializer (merged in #792) and adds AppendDefaultToBuilder for the row-by-row Avro decode paths, plus a kDefault projection branch in the Avro schema/data projection. Part 3 of the v3 column-default-values work (POC #731), built on the schema support in #746 and the Parquet read path in #792.
1 parent cd4ca42 commit 75bd63f

8 files changed

Lines changed: 138 additions & 0 deletions

File tree

src/iceberg/arrow/literal_util.cc

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
#include <utility>
2323
#include <vector>
2424

25+
#include <arrow/array/builder_base.h>
2526
#include <arrow/array/util.h>
2627
#include <arrow/buffer.h>
2728
#include <arrow/compute/api.h>
@@ -195,4 +196,27 @@ Result<std::shared_ptr<::arrow::Array>> MakeDefaultArray(
195196
return array;
196197
}
197198

199+
Status AppendDefaultToBuilder(const Literal& literal, ::arrow::ArrayBuilder* builder) {
200+
// The builder's own memory pool is not exposed, so the small scalar buffer uses the
201+
// default pool.
202+
ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<::arrow::Scalar> scalar,
203+
ToArrowScalar(literal, ::arrow::default_memory_pool()));
204+
205+
// For an extension builder (e.g. `arrow.uuid`) target its storage type: ToArrowScalar
206+
// yields the storage scalar (fixed_size_binary(16) for uuid) and Scalar::CastTo has no
207+
// kernel that targets an extension type. This mirrors MakeDefaultArray's extension
208+
// handling.
209+
std::shared_ptr<::arrow::DataType> target_type = builder->type();
210+
if (target_type->id() == ::arrow::Type::EXTENSION) {
211+
target_type = internal::checked_cast<const ::arrow::ExtensionType&>(*target_type)
212+
.storage_type();
213+
}
214+
215+
if (!scalar->type->Equals(*target_type)) {
216+
ICEBERG_ARROW_ASSIGN_OR_RETURN(scalar, scalar->CastTo(target_type));
217+
}
218+
ICEBERG_ARROW_RETURN_NOT_OK(builder->AppendScalar(*scalar));
219+
return {};
220+
}
221+
198222
} // namespace iceberg::arrow

src/iceberg/arrow/literal_util_internal.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,4 +44,8 @@ Result<std::shared_ptr<::arrow::Array>> MakeDefaultArray(
4444
const Literal& literal, const std::shared_ptr<::arrow::DataType>& type,
4545
int64_t num_rows, ::arrow::MemoryPool* pool);
4646

47+
/// \brief Append the literal value once to `builder`, e.g. to materialize a missing
48+
/// field with a default value while building rows one at a time.
49+
Status AppendDefaultToBuilder(const Literal& literal, ::arrow::ArrayBuilder* builder);
50+
4751
} // namespace iceberg::arrow

src/iceberg/avro/avro_data_util.cc

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
#include <avro/Types.hh>
3232

3333
#include "iceberg/arrow/arrow_status_internal.h"
34+
#include "iceberg/arrow/literal_util_internal.h"
3435
#include "iceberg/avro/avro_data_util_internal.h"
3536
#include "iceberg/avro/avro_schema_util_internal.h"
3637
#include "iceberg/metadata_columns.h"
@@ -88,6 +89,9 @@ Status AppendStructToBuilder(const ::avro::NodePtr& avro_node,
8889
metadata_context, field_builder));
8990
} else if (field_projection.kind == FieldProjection::Kind::kNull) {
9091
ICEBERG_ARROW_RETURN_NOT_OK(field_builder->AppendNull());
92+
} else if (field_projection.kind == FieldProjection::Kind::kDefault) {
93+
ICEBERG_RETURN_UNEXPECTED(arrow::AppendDefaultToBuilder(
94+
std::get<Literal>(field_projection.from), field_builder));
9195
} else if (field_projection.kind == FieldProjection::Kind::kMetadata) {
9296
int32_t field_id = expected_field.field_id();
9397
if (field_id == MetadataColumns::kFilePathColumnId) {
@@ -463,6 +467,11 @@ Status AppendFieldToBuilder(const ::avro::NodePtr& avro_node,
463467
return {};
464468
}
465469

470+
if (projection.kind == FieldProjection::Kind::kDefault) {
471+
return arrow::AppendDefaultToBuilder(std::get<Literal>(projection.from),
472+
array_builder);
473+
}
474+
466475
if (avro_node->type() == ::avro::AVRO_UNION) {
467476
size_t branch = avro_datum.unionBranch();
468477
if (avro_node->leafAt(branch)->type() == ::avro::AVRO_NULL) {

src/iceberg/avro/avro_direct_decoder.cc

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
#include <avro/Types.hh>
3030

3131
#include "iceberg/arrow/arrow_status_internal.h"
32+
#include "iceberg/arrow/literal_util_internal.h"
3233
#include "iceberg/avro/avro_direct_decoder_internal.h"
3334
#include "iceberg/avro/avro_schema_util_internal.h"
3435
#include "iceberg/metadata_columns.h"
@@ -209,6 +210,9 @@ Status DecodeStructToBuilder(const ::avro::NodePtr& avro_node, ::avro::Decoder&
209210
auto* field_builder = struct_builder->field_builder(static_cast<int>(proj_idx));
210211
if (field_projection.kind == FieldProjection::Kind::kNull) {
211212
ICEBERG_ARROW_RETURN_NOT_OK(field_builder->AppendNull());
213+
} else if (field_projection.kind == FieldProjection::Kind::kDefault) {
214+
ICEBERG_RETURN_UNEXPECTED(arrow::AppendDefaultToBuilder(
215+
std::get<Literal>(field_projection.from), field_builder));
212216
} else if (field_projection.kind == FieldProjection::Kind::kMetadata) {
213217
int32_t field_id = expected_field.field_id();
214218
if (field_id == MetadataColumns::kFilePathColumnId) {

src/iceberg/avro/avro_schema_util.cc

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -752,6 +752,10 @@ Result<FieldProjection> ProjectStruct(const StructType& struct_type,
752752
iter->second.local_index, prune_source));
753753
} else if (MetadataColumns::IsMetadataColumn(field_id)) {
754754
child_projection.kind = FieldProjection::Kind::kMetadata;
755+
} else if (expected_field.initial_default() != nullptr) {
756+
// Rows written before the field existed assume its `initial-default` value.
757+
child_projection.kind = FieldProjection::Kind::kDefault;
758+
child_projection.from = *expected_field.initial_default();
755759
} else if (expected_field.optional()) {
756760
child_projection.kind = FieldProjection::Kind::kNull;
757761
} else {

src/iceberg/test/avro_data_test.cc

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131

3232
#include "iceberg/avro/avro_data_util_internal.h"
3333
#include "iceberg/avro/avro_schema_util_internal.h"
34+
#include "iceberg/expression/literal.h"
3435
#include "iceberg/schema.h"
3536
#include "iceberg/schema_internal.h"
3637
#include "iceberg/schema_util.h"
@@ -662,6 +663,43 @@ TEST(AppendDatumToBuilderTest, StructWithMissingOptionalField) {
662663
avro_data, expected_json));
663664
}
664665

666+
TEST(AppendDatumToBuilderTest, StructWithMissingDefaultFields) {
667+
Schema iceberg_schema({
668+
SchemaField::MakeRequired(1, "id", iceberg::int32()),
669+
// Missing required field with an initial-default: filled with the default.
670+
SchemaField(2, "score", iceberg::int64(), /*optional=*/false, /*doc=*/{},
671+
std::make_shared<const Literal>(Literal::Long(100))),
672+
// Missing optional field with an initial-default: also filled, not null.
673+
SchemaField(3, "grade", iceberg::string(), /*optional=*/true, /*doc=*/{},
674+
std::make_shared<const Literal>(Literal::String("A"))),
675+
});
676+
677+
// Create Avro schema that only has the id field (missing score and grade).
678+
std::string avro_schema_json = R"({
679+
"type": "record",
680+
"name": "person",
681+
"fields": [
682+
{"name": "id", "type": "int", "field-id": 1}
683+
]
684+
})";
685+
auto avro_schema = ::avro::compileJsonSchemaFromString(avro_schema_json);
686+
687+
std::vector<::avro::GenericDatum> avro_data;
688+
for (int i = 0; i < 2; ++i) {
689+
::avro::GenericDatum avro_datum(avro_schema.root());
690+
auto& record = avro_datum.value<::avro::GenericRecord>();
691+
record.fieldAt(0).value<int32_t>() = i + 1;
692+
avro_data.push_back(avro_datum);
693+
}
694+
695+
const std::string expected_json = R"([
696+
{"id": 1, "score": 100, "grade": "A"},
697+
{"id": 2, "score": 100, "grade": "A"}
698+
])";
699+
ASSERT_NO_FATAL_FAILURE(VerifyAppendDatumToBuilder(iceberg_schema, avro_schema.root(),
700+
avro_data, expected_json));
701+
}
702+
665703
TEST(AppendDatumToBuilderTest, NestedStructWithMissingOptionalFields) {
666704
Schema iceberg_schema({
667705
SchemaField::MakeRequired(1, "id", iceberg::int32()),

src/iceberg/test/avro_test.cc

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
#include "iceberg/avro/avro_register.h"
4343
#include "iceberg/avro/avro_stream_internal.h"
4444
#include "iceberg/avro/avro_writer.h"
45+
#include "iceberg/expression/literal.h"
4546
#include "iceberg/file_reader.h"
4647
#include "iceberg/metadata_columns.h"
4748
#include "iceberg/schema.h"
@@ -265,6 +266,34 @@ TEST_F(AvroReaderTest, ReadTwoFields) {
265266
ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader));
266267
}
267268

269+
TEST_F(AvroReaderTest, ReadMissingFieldsWithDefaults) {
270+
// The file contains only fields 1 and 2; the projected schema adds fields 3 and 4
271+
// with initial-defaults, which are filled for all rows written before the columns
272+
// existed.
273+
CreateSimpleAvroFile();
274+
275+
auto schema = std::make_shared<Schema>(std::vector<SchemaField>{
276+
SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()),
277+
SchemaField::MakeOptional(2, "name", std::make_shared<StringType>()),
278+
SchemaField(3, "score", std::make_shared<LongType>(), /*optional=*/false,
279+
/*doc=*/{}, std::make_shared<const Literal>(Literal::Long(100))),
280+
SchemaField(4, "status", std::make_shared<StringType>(), /*optional=*/true,
281+
/*doc=*/{}, std::make_shared<const Literal>(Literal::String("active"))),
282+
});
283+
284+
auto reader_result = ReaderFactoryRegistry::Open(
285+
FileFormatType::kAvro,
286+
{.path = temp_avro_file_, .io = file_io_, .projection = schema});
287+
ASSERT_THAT(reader_result, IsOk());
288+
auto reader = std::move(reader_result.value());
289+
290+
ASSERT_NO_FATAL_FAILURE(VerifyNextBatch(*reader,
291+
R"([[1, "Alice", 100, "active"],
292+
[2, "Bob", 100, "active"],
293+
[3, "Charlie", 100, "active"]])"));
294+
ASSERT_NO_FATAL_FAILURE(VerifyExhausted(*reader));
295+
}
296+
268297
TEST_F(AvroReaderTest, RoundTripWithGenericFileIO) {
269298
file_io_ = std::make_shared<iceberg::test::StdFileIO>();
270299
temp_avro_file_ = CreateNewTempFilePathWithSuffix(".avro");

src/iceberg/test/literal_util_test.cc

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
#include <vector>
2626

2727
#include <arrow/array.h>
28+
#include <arrow/array/builder_primitive.h>
2829
#include <arrow/buffer.h>
2930
#include <arrow/extension/uuid.h>
3031
#include <arrow/scalar.h>
@@ -186,6 +187,31 @@ TEST(LiteralUtilTest, MakeDefaultArrayWrapsExtensionType) {
186187
}
187188
}
188189

190+
TEST(LiteralUtilTest, AppendDefaultToBuilderAppendsValue) {
191+
::arrow::Int64Builder builder;
192+
ASSERT_THAT(AppendDefaultToBuilder(Literal::Long(42), &builder), IsOk());
193+
ASSERT_THAT(AppendDefaultToBuilder(Literal::Long(42), &builder), IsOk());
194+
195+
std::shared_ptr<::arrow::Array> array;
196+
ASSERT_TRUE(builder.Finish(&array).ok());
197+
ASSERT_EQ(array->length(), 2);
198+
const auto& long_array = static_cast<const ::arrow::Int64Array&>(*array);
199+
ASSERT_EQ(long_array.Value(0), 42);
200+
ASSERT_EQ(long_array.Value(1), 42);
201+
}
202+
203+
TEST(LiteralUtilTest, AppendDefaultToBuilderCastsToBuilderType) {
204+
// The literal's natural type (int32) differs from the builder type (int64); the value
205+
// is cast to the builder type.
206+
::arrow::Int64Builder builder;
207+
ASSERT_THAT(AppendDefaultToBuilder(Literal::Int(7), &builder), IsOk());
208+
209+
std::shared_ptr<::arrow::Array> array;
210+
ASSERT_TRUE(builder.Finish(&array).ok());
211+
ASSERT_EQ(array->length(), 1);
212+
ASSERT_EQ(static_cast<const ::arrow::Int64Array&>(*array).Value(0), 7);
213+
}
214+
189215
} // namespace
190216

191217
} // namespace iceberg::arrow

0 commit comments

Comments
 (0)