Skip to content

Commit 1b8045b

Browse files
huan233uschuan233usc
committed
recurse through nested collection builders when preparing default scalars
1 parent 289a7c9 commit 1b8045b

2 files changed

Lines changed: 145 additions & 45 deletions

File tree

‎src/iceberg/avro/avro_data_util.cc‎

Lines changed: 49 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -532,8 +532,52 @@ Result<std::shared_ptr<::arrow::Scalar>> MakeDefaultScalar(
532532
return scalar;
533533
}
534534

535-
Status PrepareDefaultScalars(std::span<FieldProjection> projections,
536-
::arrow::ArrayBuilder* builder) {
535+
Status PrepareStructDefaultScalars(std::span<FieldProjection> projections,
536+
::arrow::ArrayBuilder* builder);
537+
538+
// Recurse into whatever nested builder this projection describes, so a default cached for
539+
// a struct field is found at any nesting depth (e.g. `list<list<struct<...>>>`) instead
540+
// of only when the collection's child is an immediate struct.
541+
Status PrepareNestedDefaultScalars(FieldProjection& projection,
542+
::arrow::ArrayBuilder* builder) {
543+
if (projection.kind != FieldProjection::Kind::kProjected ||
544+
projection.children.empty()) {
545+
return {};
546+
}
547+
548+
switch (builder->type()->id()) {
549+
case ::arrow::Type::STRUCT:
550+
return PrepareStructDefaultScalars(projection.children, builder);
551+
case ::arrow::Type::LIST: {
552+
// List projections store a single child for the element.
553+
auto* list_builder = internal::checked_cast<::arrow::ListBuilder*>(builder);
554+
return PrepareNestedDefaultScalars(projection.children[0],
555+
list_builder->value_builder());
556+
}
557+
case ::arrow::Type::LARGE_LIST: {
558+
auto* list_builder = internal::checked_cast<::arrow::LargeListBuilder*>(builder);
559+
return PrepareNestedDefaultScalars(projection.children[0],
560+
list_builder->value_builder());
561+
}
562+
case ::arrow::Type::MAP: {
563+
auto* map_builder = internal::checked_cast<::arrow::MapBuilder*>(builder);
564+
if (projection.children.size() >= 1) {
565+
ICEBERG_RETURN_UNEXPECTED(PrepareNestedDefaultScalars(
566+
projection.children[0], map_builder->key_builder()));
567+
}
568+
if (projection.children.size() >= 2) {
569+
ICEBERG_RETURN_UNEXPECTED(PrepareNestedDefaultScalars(
570+
projection.children[1], map_builder->item_builder()));
571+
}
572+
return {};
573+
}
574+
default:
575+
return {};
576+
}
577+
}
578+
579+
Status PrepareStructDefaultScalars(std::span<FieldProjection> projections,
580+
::arrow::ArrayBuilder* builder) {
537581
auto* struct_builder = internal::checked_cast<::arrow::StructBuilder*>(builder);
538582
if (static_cast<size_t>(struct_builder->num_fields()) != projections.size()) {
539583
return InvalidArgument(
@@ -558,48 +602,8 @@ Status PrepareDefaultScalars(std::span<FieldProjection> projections,
558602
continue;
559603
}
560604

561-
if (field_projection.kind != FieldProjection::Kind::kProjected ||
562-
field_projection.children.empty()) {
563-
continue;
564-
}
565-
566-
switch (field_builder->type()->id()) {
567-
case ::arrow::Type::STRUCT:
568-
ICEBERG_RETURN_UNEXPECTED(
569-
PrepareDefaultScalars(field_projection.children, field_builder));
570-
break;
571-
case ::arrow::Type::LIST: {
572-
// List projections store a single child for the element. Defaults only appear
573-
// on nested struct fields of that element.
574-
auto* list_builder = internal::checked_cast<::arrow::ListBuilder*>(field_builder);
575-
auto& element_projection = field_projection.children[0];
576-
if (element_projection.kind == FieldProjection::Kind::kProjected &&
577-
!element_projection.children.empty() &&
578-
list_builder->value_builder()->type()->id() == ::arrow::Type::STRUCT) {
579-
ICEBERG_RETURN_UNEXPECTED(PrepareDefaultScalars(element_projection.children,
580-
list_builder->value_builder()));
581-
}
582-
break;
583-
}
584-
case ::arrow::Type::MAP: {
585-
auto* map_builder = internal::checked_cast<::arrow::MapBuilder*>(field_builder);
586-
if (field_projection.children.size() >= 1 &&
587-
!field_projection.children[0].children.empty() &&
588-
map_builder->key_builder()->type()->id() == ::arrow::Type::STRUCT) {
589-
ICEBERG_RETURN_UNEXPECTED(PrepareDefaultScalars(
590-
field_projection.children[0].children, map_builder->key_builder()));
591-
}
592-
if (field_projection.children.size() >= 2 &&
593-
!field_projection.children[1].children.empty() &&
594-
map_builder->item_builder()->type()->id() == ::arrow::Type::STRUCT) {
595-
ICEBERG_RETURN_UNEXPECTED(PrepareDefaultScalars(
596-
field_projection.children[1].children, map_builder->item_builder()));
597-
}
598-
break;
599-
}
600-
default:
601-
break;
602-
}
605+
ICEBERG_RETURN_UNEXPECTED(
606+
PrepareNestedDefaultScalars(field_projection, field_builder));
603607
}
604608
return {};
605609
}
@@ -608,7 +612,7 @@ Status PrepareDefaultScalars(std::span<FieldProjection> projections,
608612

609613
Status PrepareDefaultScalars(SchemaProjection& projection,
610614
::arrow::ArrayBuilder* root_builder) {
611-
return PrepareDefaultScalars(projection.fields, root_builder);
615+
return PrepareStructDefaultScalars(projection.fields, root_builder);
612616
}
613617

614618
Status AppendDefaultToBuilder(const Literal& literal, ::arrow::ArrayBuilder* builder) {

‎src/iceberg/test/avro_data_test.cc‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -704,6 +704,63 @@ TEST(AppendDatumToBuilderTest, StructWithMissingDefaultFields) {
704704
avro_data, expected_json));
705705
}
706706

707+
TEST(AppendDatumToBuilderTest, NestedListOfListStructWithMissingDefaultField) {
708+
// A default on a struct nested two collection levels deep must still be filled, which
709+
// requires recursing through both list builders rather than only an immediate struct.
710+
auto inner_struct = std::make_shared<StructType>(std::vector<SchemaField>{
711+
SchemaField::MakeRequired(5, "x", iceberg::int32()),
712+
SchemaField(6, "y", iceberg::int64(), /*optional=*/false, /*doc=*/{},
713+
std::make_shared<const Literal>(Literal::Long(7))),
714+
});
715+
auto inner_list =
716+
std::make_shared<ListType>(SchemaField::MakeRequired(4, "element", inner_struct));
717+
auto outer_list =
718+
std::make_shared<ListType>(SchemaField::MakeRequired(3, "element", inner_list));
719+
Schema iceberg_schema({
720+
SchemaField::MakeRequired(1, "id", iceberg::int32()),
721+
SchemaField::MakeRequired(2, "matrix", outer_list),
722+
});
723+
724+
// The Avro schema's innermost record only has `x`, so `y` comes from its default.
725+
std::string avro_schema_json = R"({
726+
"type": "record",
727+
"name": "outer",
728+
"fields": [
729+
{"name": "id", "type": "int", "field-id": 1},
730+
{"name": "matrix", "field-id": 2, "type": {
731+
"type": "array", "element-id": 3, "items": {
732+
"type": "array", "element-id": 4, "items": {
733+
"type": "record", "name": "point",
734+
"fields": [{"name": "x", "type": "int", "field-id": 5}]
735+
}
736+
}
737+
}}
738+
]
739+
})";
740+
auto avro_schema = ::avro::compileJsonSchemaFromString(avro_schema_json);
741+
742+
std::vector<::avro::GenericDatum> avro_data;
743+
::avro::GenericDatum avro_datum(avro_schema.root());
744+
auto& record = avro_datum.value<::avro::GenericRecord>();
745+
record.fieldAt(0).value<int32_t>() = 1;
746+
auto& outer_array = record.fieldAt(1).value<::avro::GenericArray>();
747+
::avro::GenericDatum inner_datum(avro_schema.root()->leafAt(1)->leafAt(0));
748+
auto& inner_array = inner_datum.value<::avro::GenericArray>();
749+
for (int32_t x : {10, 20}) {
750+
::avro::GenericDatum point_datum(avro_schema.root()->leafAt(1)->leafAt(0)->leafAt(0));
751+
point_datum.value<::avro::GenericRecord>().fieldAt(0).value<int32_t>() = x;
752+
inner_array.value().push_back(point_datum);
753+
}
754+
outer_array.value().push_back(inner_datum);
755+
avro_data.push_back(avro_datum);
756+
757+
const std::string expected_json = R"([
758+
{"id": 1, "matrix": [[{"x": 10, "y": 7}, {"x": 20, "y": 7}]]}
759+
])";
760+
ASSERT_NO_FATAL_FAILURE(VerifyAppendDatumToBuilder(iceberg_schema, avro_schema.root(),
761+
avro_data, expected_json));
762+
}
763+
707764
TEST(AppendDefaultToBuilderTest, AppendsValue) {
708765
::arrow::Int64Builder builder;
709766
ASSERT_THAT(AppendDefaultToBuilder(Literal::Long(42), &builder), IsOk());
@@ -758,6 +815,45 @@ TEST(AppendDefaultToBuilderTest, ReusesPreparedScalar) {
758815
ASSERT_EQ(long_array.Value(1), 42);
759816
}
760817

818+
TEST(AppendDefaultToBuilderTest, PreparesScalarUnderNestedCollections) {
819+
// A default under `list<list<struct<...>>>` must be prepared too, so decoding reuses
820+
// the cached scalar instead of rebuilding it for every element.
821+
auto pool = ::arrow::default_memory_pool();
822+
auto leaf = std::make_shared<::arrow::Int64Builder>(pool);
823+
auto point_builder = std::make_shared<::arrow::StructBuilder>(
824+
::arrow::struct_({::arrow::field("y", ::arrow::int64())}), pool,
825+
std::vector<std::shared_ptr<::arrow::ArrayBuilder>>{leaf});
826+
auto inner_list = std::make_shared<::arrow::ListBuilder>(pool, point_builder);
827+
auto outer_list = std::make_shared<::arrow::ListBuilder>(pool, inner_list);
828+
::arrow::StructBuilder root_builder(
829+
::arrow::struct_({::arrow::field("matrix", outer_list->type())}), pool,
830+
std::vector<std::shared_ptr<::arrow::ArrayBuilder>>{outer_list});
831+
832+
FieldProjection leaf_default;
833+
leaf_default.kind = FieldProjection::Kind::kDefault;
834+
leaf_default.from = Literal::Long(7);
835+
836+
FieldProjection point_projection;
837+
point_projection.kind = FieldProjection::Kind::kProjected;
838+
point_projection.children.push_back(leaf_default);
839+
840+
FieldProjection inner_element;
841+
inner_element.kind = FieldProjection::Kind::kProjected;
842+
inner_element.children.push_back(point_projection);
843+
844+
FieldProjection outer_element;
845+
outer_element.kind = FieldProjection::Kind::kProjected;
846+
outer_element.children.push_back(inner_element);
847+
848+
SchemaProjection schema_projection;
849+
schema_projection.fields.push_back(outer_element);
850+
ASSERT_THAT(PrepareDefaultScalars(schema_projection, &root_builder), IsOk());
851+
852+
const auto& prepared = schema_projection.fields[0].children[0].children[0].children[0];
853+
ASSERT_NE(dynamic_cast<const AvroDefaultAttributes*>(prepared.attributes.get()),
854+
nullptr);
855+
}
856+
761857
TEST(AppendDatumToBuilderTest, NestedStructWithMissingOptionalFields) {
762858
Schema iceberg_schema({
763859
SchemaField::MakeRequired(1, "id", iceberg::int32()),

0 commit comments

Comments
 (0)