c-thiel commented on code in PR #2880:
URL: https://github.com/apache/iceberg-rust/pull/2880#discussion_r4446219218
##########
crates/iceberg/src/arrow/reader/projection.rs:
##########
@@ -171,16 +161,50 @@ impl ArrowReader {
arrow_schema: &ArrowSchemaRef,
type_promotion_is_valid: fn(Option<&PrimitiveType>,
Option<&PrimitiveType>) -> bool,
) -> Result<ProjectionMask> {
- let mut column_map = HashMap::new();
+ // Maps field_id → leaf column indices. `Vec` because a variant
contributes two
+ // leaves (metadata + value) under a single field id.
+ let mut column_map: HashMap<i32, Vec<usize>> = HashMap::new();
let fields = arrow_schema.fields();
// HashSet for O(1) membership checks instead of O(n) slice scans.
let leaf_field_id_set: HashSet<i32> =
leaf_field_ids.iter().copied().collect();
+ // A variant is an Iceberg leaf type but a Parquet group: its
metadata/value
+ // sub-fields carry no embedded field id, so the field-id scan below
never finds
+ // them. Iceberg-java's `PruneColumns` projects the whole variant
group unchanged
+ // (the enclosing struct/list/map re-adds the original group); we
replicate that by
+ // pre-computing, for every Arrow leaf sitting inside a variant
column, the enclosing
+ // variant's field id (numbering matches `filter_leaves`).
+ let variant_leaves = {
+ let mut out = HashMap::new();
+ let mut leaf_idx = 0usize;
+ Self::collect_variant_leaves(
+ fields,
+ &mut leaf_idx,
+ None,
+ iceberg_schema_of_task,
+ &leaf_field_id_set,
+ &mut out,
+ )?;
+ out
+ };
+
+ // Recover variant identity from the Iceberg schema rather than the
Parquet `variant`
+ // annotation: tag every variant storage struct with the
`arrow.parquet.variant`
+ // extension so `arrow_schema_to_schema` folds it back into
`Type::Variant` instead of
+ // descending into its id-less sub-fields. Mirrors iceberg-java's
`TypeWithSchemaVisitor`,
+ // which keys on the annotation OR the Iceberg type.
+ let tagged_fields = Self::attach_variant_extensions(fields,
iceberg_schema_of_task);
Review Comment:
Thanks to both of you, and sorry this took so long!
@blackmwk You were right. `compare_schemas` ignores field metadata, so a
file stored exactly in the table's Arrow layout passed through unchanged and
lost the extension. After rebasing onto main I found a second problem:
`RecordBatchTransformer` now matches struct children by field id, and a
variant's `metadata` and `value` don't have one, so most variant reads failed
with `DataInvalid`.
Both are fixed in `RecordBatchTransformer`, which now handles variants on
their own. It matches their children by name, casts binary storage to the
table's type, and always sets the extension. Every end-to-end test that returns
a variant now checks `extension_type_name()`: at top level, in a struct, in a
list and in a map. `test_read_variant_stored_as_table_type_gets_extension`
covers the exact case you described.
@dannycjones Good question. The extension on the returned batch is now set
in the transformer, which builds the output from the table schema. The
transformer had to learn about variants anyway, because their children have no
field ids. The tag added here is only used inside this function. It lets
`arrow_schema_to_schema` turn the storage struct into `Type::Variant` instead
of failing on its id-less children, while it builds the schema for the
type-promotion check. If you'd rather have that step next to the other schema
fixes in the pipeline, I'm happy to look into moving it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
