|
17 | 17 | * under the License. |
18 | 18 | */ |
19 | 19 |
|
| 20 | +#include <span> |
| 21 | + |
20 | 22 | #include <arrow/array/builder_binary.h> |
21 | 23 | #include <arrow/array/builder_decimal.h> |
22 | 24 | #include <arrow/array/builder_nested.h> |
23 | 25 | #include <arrow/array/builder_primitive.h> |
24 | 26 | #include <arrow/extension_type.h> |
25 | 27 | #include <arrow/json/from_string.h> |
| 28 | +#include <arrow/scalar.h> |
26 | 29 | #include <arrow/type.h> |
27 | 30 | #include <arrow/util/decimal.h> |
28 | 31 | #include <avro/Generic.hh> |
|
31 | 34 | #include <avro/Types.hh> |
32 | 35 |
|
33 | 36 | #include "iceberg/arrow/arrow_status_internal.h" |
| 37 | +#include "iceberg/arrow/literal_util_internal.h" |
34 | 38 | #include "iceberg/avro/avro_data_util_internal.h" |
35 | 39 | #include "iceberg/avro/avro_schema_util_internal.h" |
36 | 40 | #include "iceberg/metadata_columns.h" |
@@ -88,6 +92,8 @@ Status AppendStructToBuilder(const ::avro::NodePtr& avro_node, |
88 | 92 | metadata_context, field_builder)); |
89 | 93 | } else if (field_projection.kind == FieldProjection::Kind::kNull) { |
90 | 94 | ICEBERG_ARROW_RETURN_NOT_OK(field_builder->AppendNull()); |
| 95 | + } else if (field_projection.kind == FieldProjection::Kind::kDefault) { |
| 96 | + ICEBERG_RETURN_UNEXPECTED(AppendDefaultToBuilder(field_projection, field_builder)); |
91 | 97 | } else if (field_projection.kind == FieldProjection::Kind::kMetadata) { |
92 | 98 | int32_t field_id = expected_field.field_id(); |
93 | 99 | if (field_id == MetadataColumns::kFilePathColumnId) { |
@@ -466,6 +472,10 @@ Status AppendFieldToBuilder(const ::avro::NodePtr& avro_node, |
466 | 472 | return {}; |
467 | 473 | } |
468 | 474 |
|
| 475 | + if (projection.kind == FieldProjection::Kind::kDefault) { |
| 476 | + return AppendDefaultToBuilder(projection, array_builder); |
| 477 | + } |
| 478 | + |
469 | 479 | const bool is_row_lineage = |
470 | 480 | MetadataColumns::IsRowLineageColumn(projected_field.field_id()); |
471 | 481 |
|
@@ -497,6 +507,148 @@ Status AppendFieldToBuilder(const ::avro::NodePtr& avro_node, |
497 | 507 |
|
498 | 508 | } // namespace |
499 | 509 |
|
| 510 | +namespace { |
| 511 | + |
| 512 | +Result<std::shared_ptr<::arrow::Scalar>> MakeDefaultScalar( |
| 513 | + const Literal& literal, const std::shared_ptr<::arrow::DataType>& builder_type) { |
| 514 | + // The builder's own memory pool is not exposed, so the small scalar buffer uses the |
| 515 | + // default pool. |
| 516 | + ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<::arrow::Scalar> scalar, |
| 517 | + arrow::ToArrowScalar(literal, ::arrow::default_memory_pool())); |
| 518 | + |
| 519 | + // For an extension builder (e.g. `arrow.uuid`) target its storage type: ToArrowScalar |
| 520 | + // yields the storage scalar (fixed_size_binary(16) for uuid) and Scalar::CastTo has no |
| 521 | + // kernel that targets an extension type. This mirrors MakeDefaultArray's extension |
| 522 | + // handling. |
| 523 | + std::shared_ptr<::arrow::DataType> target_type = builder_type; |
| 524 | + if (target_type->id() == ::arrow::Type::EXTENSION) { |
| 525 | + target_type = internal::checked_cast<const ::arrow::ExtensionType&>(*target_type) |
| 526 | + .storage_type(); |
| 527 | + } |
| 528 | + |
| 529 | + if (!scalar->type->Equals(*target_type)) { |
| 530 | + ICEBERG_ARROW_ASSIGN_OR_RETURN(scalar, scalar->CastTo(target_type)); |
| 531 | + } |
| 532 | + return scalar; |
| 533 | +} |
| 534 | + |
| 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) { |
| 581 | + auto* struct_builder = internal::checked_cast<::arrow::StructBuilder*>(builder); |
| 582 | + if (static_cast<size_t>(struct_builder->num_fields()) != projections.size()) { |
| 583 | + return InvalidArgument( |
| 584 | + "Inconsistent number of struct builder fields ({}) and projections ({})", |
| 585 | + struct_builder->num_fields(), projections.size()); |
| 586 | + } |
| 587 | + |
| 588 | + for (size_t i = 0; i < projections.size(); ++i) { |
| 589 | + auto& field_projection = projections[i]; |
| 590 | + auto* field_builder = struct_builder->field_builder(static_cast<int>(i)); |
| 591 | + |
| 592 | + if (field_projection.kind == FieldProjection::Kind::kDefault) { |
| 593 | + // Get-or-create the single Avro attributes container: another Avro attribute may |
| 594 | + // have created it already, so guard on `default_scalar` being unset rather than on |
| 595 | + // the container's presence (otherwise a pre-existing container would skip |
| 596 | + // preparation and silently fall back to per-row scalar rebuilding). |
| 597 | + std::shared_ptr<AvroExtraAttributes> attrs; |
| 598 | + if (field_projection.attributes == nullptr) { |
| 599 | + attrs = std::make_shared<AvroExtraAttributes>(); |
| 600 | + field_projection.attributes = attrs; |
| 601 | + } else { |
| 602 | + // Avro attaches only AvroExtraAttributes; checked_pointer_cast asserts that |
| 603 | + // invariant in debug rather than silently reinterpreting a foreign attributes |
| 604 | + // type. |
| 605 | + attrs = internal::checked_pointer_cast<AvroExtraAttributes>( |
| 606 | + field_projection.attributes); |
| 607 | + } |
| 608 | + if (attrs->default_scalar == nullptr) { |
| 609 | + ICEBERG_ASSIGN_OR_RAISE( |
| 610 | + attrs->default_scalar, |
| 611 | + MakeDefaultScalar(std::get<Literal>(field_projection.from), |
| 612 | + field_builder->type())); |
| 613 | + } |
| 614 | + continue; |
| 615 | + } |
| 616 | + |
| 617 | + ICEBERG_RETURN_UNEXPECTED( |
| 618 | + PrepareNestedDefaultScalars(field_projection, field_builder)); |
| 619 | + } |
| 620 | + return {}; |
| 621 | +} |
| 622 | + |
| 623 | +} // namespace |
| 624 | + |
| 625 | +Status PrepareDefaultScalars(SchemaProjection& projection, |
| 626 | + ::arrow::ArrayBuilder* root_builder) { |
| 627 | + return PrepareStructDefaultScalars(projection.fields, root_builder); |
| 628 | +} |
| 629 | + |
| 630 | +Status AppendDefaultToBuilder(const Literal& literal, ::arrow::ArrayBuilder* builder) { |
| 631 | + ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<::arrow::Scalar> scalar, |
| 632 | + MakeDefaultScalar(literal, builder->type())); |
| 633 | + ICEBERG_ARROW_RETURN_NOT_OK(builder->AppendScalar(*scalar)); |
| 634 | + return {}; |
| 635 | +} |
| 636 | + |
| 637 | +Status AppendDefaultToBuilder(const FieldProjection& projection, |
| 638 | + ::arrow::ArrayBuilder* builder) { |
| 639 | + // Avro projections carry a single attributes type, so once one is attached it is an |
| 640 | + // AvroExtraAttributes; use checked_cast instead of a per-row dynamic_cast. |
| 641 | + if (projection.attributes != nullptr) { |
| 642 | + const auto& attrs = |
| 643 | + internal::checked_cast<const AvroExtraAttributes&>(*projection.attributes); |
| 644 | + if (attrs.default_scalar != nullptr) { |
| 645 | + ICEBERG_ARROW_RETURN_NOT_OK(builder->AppendScalar(*attrs.default_scalar)); |
| 646 | + return {}; |
| 647 | + } |
| 648 | + } |
| 649 | + return AppendDefaultToBuilder(std::get<Literal>(projection.from), builder); |
| 650 | +} |
| 651 | + |
500 | 652 | Status AppendDatumToBuilder(const ::avro::NodePtr& avro_node, |
501 | 653 | const ::avro::GenericDatum& avro_datum, |
502 | 654 | const SchemaProjection& projection, |
|
0 commit comments