c-thiel commented on code in PR #2880:
URL: https://github.com/apache/iceberg-rust/pull/2880#discussion_r4446344375
##########
crates/iceberg/src/arrow/reader/projection.rs:
##########
@@ -237,6 +269,184 @@ impl ArrowReader {
}
}
+ /// Walks an Arrow `Fields` tree, recording `leaf_idx → variant_field_id`
for every
+ /// leaf that sits inside a variant column (top-level or nested in a
struct/list/map).
+ ///
+ /// The leaf numbering must match [`arrow_schema::Fields::filter_leaves`],
since the
+ /// resulting map is consulted by index in the `filter_leaves` passes
above. In
+ /// particular it mirrors `filter_leaves`' handling of
`Dictionary`/`RunEndEncoded`
+ /// (unwrapped to their value type) and `Union` (each member counted) —
otherwise a
+ /// mismatch would silently shift every variant leaf after it onto the
wrong column.
+ fn collect_variant_leaves(
+ fields: &Fields,
+ leaf_idx: &mut usize,
+ variant_parent: Option<i32>,
+ iceberg_schema: &Schema,
+ leaf_field_id_set: &HashSet<i32>,
+ out: &mut HashMap<usize, i32>,
+ ) -> Result<()> {
+ for field in fields {
+ Self::collect_variant_leaves_in_field(
+ field,
+ leaf_idx,
+ variant_parent,
+ iceberg_schema,
+ leaf_field_id_set,
+ out,
+ )?;
+ }
+ Ok(())
+ }
+
+ fn collect_variant_leaves_in_field(
+ field: &FieldRef,
+ leaf_idx: &mut usize,
+ variant_parent: Option<i32>,
+ iceberg_schema: &Schema,
+ leaf_field_id_set: &HashSet<i32>,
+ out: &mut HashMap<usize, i32>,
+ ) -> Result<()> {
+ // Once inside a variant, stay inside; otherwise check whether this
field is itself
+ // a variant column (its embedded field id resolves to
`Type::Variant`).
+ let entering_variant = variant_parent.is_none();
+ let effective_variant = variant_parent.or_else(|| {
+ Self::field_variant_id(field, iceberg_schema)
+ .filter(|fid| leaf_field_id_set.contains(fid))
+ });
+
+ // Reject shredded variants: a `typed_value` sub-field means the
payload is shredded,
+ // which we can't reconstruct yet. Projecting only metadata/value
would silently drop
+ // it, so fail loudly instead.
+ if entering_variant
+ && effective_variant.is_some()
+ && let DataType::Struct(sub) = field.data_type()
+ && sub.iter().any(|f| f.name() == "typed_value")
+ {
+ return Err(Error::new(
+ ErrorKind::FeatureUnsupported,
+ "Reading shredded variant columns is not supported yet: found
a `typed_value` \
+ sub-field. Only unshredded variants (metadata + value) can be
read.",
+ ));
+ }
+
+ // Mirror `Fields::filter_leaves`: unwrap `Dictionary`/`RunEndEncoded`
to their value
+ // type before deciding whether this is a leaf or a nested type to
descend into.
+ let data_type = match field.data_type() {
+ DataType::Dictionary(_, value) => value.as_ref(),
+ DataType::RunEndEncoded(_, value) => value.data_type(),
+ other => other,
+ };
+
+ match data_type {
+ DataType::Struct(sub) => {
+ Self::collect_variant_leaves(
+ sub,
+ leaf_idx,
+ effective_variant,
+ iceberg_schema,
+ leaf_field_id_set,
+ out,
+ )?;
+ }
+ DataType::List(inner)
+ | DataType::LargeList(inner)
+ | DataType::FixedSizeList(inner, _)
+ | DataType::Map(inner, _) => {
+ Self::collect_variant_leaves_in_field(
+ inner,
+ leaf_idx,
+ effective_variant,
+ iceberg_schema,
+ leaf_field_id_set,
+ out,
+ )?;
+ }
+ DataType::Union(union_fields, _) => {
+ for (_, inner) in union_fields.iter() {
+ Self::collect_variant_leaves_in_field(
+ inner,
+ leaf_idx,
+ effective_variant,
+ iceberg_schema,
+ leaf_field_id_set,
+ out,
+ )?;
+ }
+ }
+ _ => {
+ if let Some(vid) = effective_variant {
+ out.insert(*leaf_idx, vid);
+ }
+ *leaf_idx += 1;
+ }
+ }
+ Ok(())
+ }
+
+ /// If `field`'s embedded Parquet field id resolves to `Type::Variant` in
the Iceberg
+ /// schema, returns that id.
+ fn field_variant_id(field: &FieldRef, iceberg_schema: &Schema) ->
Option<i32> {
+ let fid = field
+ .metadata()
+ .get(PARQUET_FIELD_ID_META_KEY)
+ .and_then(|s| i32::from_str(s).ok())?;
+ let iceberg_field = iceberg_schema.field_by_id(fid)?;
+ matches!(iceberg_field.field_type.as_ref(),
Type::Variant(_)).then_some(fid)
+ }
+
+ /// Returns `fields` with the canonical `arrow.parquet.variant` extension
attached to every
+ /// field the Iceberg schema declares as a variant (recursing into
struct/list/map to reach
+ /// nested variants).
+ ///
+ /// This is how the read path recovers variant identity: the extension lets
+ /// `arrow_schema_to_schema` fold the storage struct back into
`Type::Variant` instead of
+ /// descending into its id-less `metadata`/`value` sub-fields. Keying on
the Iceberg schema
+ /// (rather than requiring the Parquet `variant` annotation to be present
on the Arrow field)
+ /// mirrors iceberg-java's `TypeWithSchemaVisitor`, which recognizes a
variant by the
+ /// annotation OR the Iceberg type.
+ fn attach_variant_extensions(fields: &Fields, iceberg_schema: &Schema) ->
Fields {
Review Comment:
Agreed, done in c2835f9c3. Before anything else, projection now collects the
projected field ids that resolve to `Type::Variant`. If there are none, both
the collect pass and the tag pass are skipped, and the fields are used as they
are, so non-variant reads don't pay for any of this. `type_contains_variant` is
back as well. It guards the path for files without field ids, where variant
reads are still rejected as before this PR.
--
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]
