c-thiel commented on code in PR #2880:
URL: https://github.com/apache/iceberg-rust/pull/2880#discussion_r4446472161
##########
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 {
+ fields
+ .iter()
+ .map(|field| Self::attach_variant_extension(field, iceberg_schema))
+ .collect()
+ }
+
+ fn attach_variant_extension(field: &FieldRef, iceberg_schema: &Schema) ->
FieldRef {
+ // A variant's storage is a struct; tag it and don't descend — its
metadata/value
+ // children are not themselves variants.
+ if Self::field_variant_id(field, iceberg_schema).is_some()
+ && matches!(field.data_type(), DataType::Struct(_))
+ {
+ return Arc::new(
+ field
+ .as_ref()
+ .clone()
+ .with_extension_type(VariantExtensionType),
+ );
+ }
+ // Otherwise recurse into containers to reach nested variants.
+ let data_type = match field.data_type() {
+ DataType::Struct(children) => {
+ DataType::Struct(Self::attach_variant_extensions(children,
iceberg_schema))
+ }
+ DataType::List(child) => {
+ DataType::List(Self::attach_variant_extension(child,
iceberg_schema))
+ }
+ DataType::LargeList(child) => {
+ DataType::LargeList(Self::attach_variant_extension(child,
iceberg_schema))
+ }
+ DataType::FixedSizeList(child, len) => {
+ DataType::FixedSizeList(Self::attach_variant_extension(child,
iceberg_schema), *len)
+ }
+ DataType::Map(child, sorted) => DataType::Map(
+ Self::attach_variant_extension(child, iceberg_schema),
+ *sorted,
+ ),
+ _ => return field.clone(),
Review Comment:
Agreed, I added that comment in c2835f9c3, on the `_` arm. It explains that
the numbering walk enters Dictionary, RunEndEncoded and Union only to stay in
step with `filter_leaves`, and that schemas read from Parquet never put a
variant under them.
It also can't go wrong silently. A variant stored under a Dictionary isn't a
group, so validation rejects it with `DataInvalid`. And parquet-rs never
produces RunEndEncoded or Union columns. Longer term, #3609 would replace both
walks with a single one based on the Parquet schema, which removes the
asymmetry entirely.
--
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]
