c-thiel commented on code in PR #2880:
URL: https://github.com/apache/iceberg-rust/pull/2880#discussion_r4446307005
##########
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:
I looked into it, and the guard would unfortunately reject valid columns.
Arrow's `DataType::is_primitive()` only covers numeric and temporal types, so
it's false for `Utf8`, `Binary`, `Boolean` and `FixedSizeBinary`. Those are
ordinary leaf columns here, and they should pass through unchanged.
The catch-all also covers Dictionary, RunEndEncoded and Union. Schemas read
from Parquet never put a variant under those, and I added a comment on that arm
saying so. The defensive part now happens earlier: when projection enters a
projected variant, `validate_unshredded_variant` checks its shape and returns
an error naming the field id if it's wrong.
--
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]
