laskoviymishka commented on code in PR #2880:
URL: https://github.com/apache/iceberg-rust/pull/2880#discussion_r4251086029
##########
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")
Review Comment:
The shredded check keys only on the literal name `typed_value`, and a
variant is identified purely from its id resolving to `Type::Variant` — nothing
validates the storage is actually `{metadata, value}`. So `{metadata, value,
foo}`, a missing `metadata`, or a variant-typed id whose Arrow column is a
plain `Binary` (which falls through to the `_` leaf arm and skips the
type-promotion check) all get projected as a variant. On untrusted Parquet
that's silently wrong data instead of a clear error. I'd require the
entering-variant branch to be exactly `DataType::Struct` with children
`{metadata, value}`, both binary, and return `DataInvalid` (with the field id)
otherwise.
##########
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:
This is blackmwk's `:196` point, and I think it's firmer than an observation
— it's a blocker. `tagged_fields` never leaves this function: it feeds the
throwaway `projected_arrow_schema` and the mask, and that's it. The stream
builder still emits a plain `Struct(metadata, value)`, and because
`RecordBatchTransformer::compare_schemas` matches on name/type/nullability and
ignores top-level field metadata, it returns `PassThrough` and drops the
extension. Consumers get an anonymous struct, not a `Variant` — the one thing
the PR set out to fix.
I'd either propagate the tagged schema into the returned batch (or pass a
target schema into the stream builder), or enable the parquet variant feature
and let its reader attach the annotation. Whichever way, an assertion on
`extension_type_name()` on the output is what pins it.
##########
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`,
Review Comment:
The "mirrors iceberg-java" framing here (and on the
`attach_variant_extensions` / `field_variant_id` docs) is a bit strong — Java
projects the variant group wholesale and reads `typed_value` in the shredded
case, whereas we reject shredded entirely. That's a defensible interim
semantic, but rejecting a shredded table is a hard failure for anything written
with `write.parquet.variant.shred`, so I'd call it out as a known limitation
and link a tracking issue in the error. Either cite the specific Java
classes/methods or soften the claim.
##########
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:
The numbering walk (`collect_variant_leaves`) descends
Dictionary/RunEndEncoded/Union, but this tagging walk drops them into the `_`
arm untouched — so a variant under one of those gets counted as a leaf but
never tagged. Mostly theoretical for Parquet-derived schemas, but the two walks
handling different type sets is the kind of asymmetry that bites later; at
least worth a comment calling it out.
##########
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()
Review Comment:
`attach_variant_extensions` tags every variant in the schema, not just
projected ones — it never consults `leaf_field_id_set`. So a shredded variant
we never asked for still gets run through `with_extension_type` on a `SELECT
id`, and the shredded rejection in `collect_variant_leaves` won't catch it
because that only fires for projected variants. Depending on what
`VariantExtensionType::supports_data_type` accepts, that either panics or
mis-tags a shredded struct — and before this PR that same scan was fine, since
the old guard only rejected projections that actually touched a variant. I'd
gate the tag on the projected set and add a test with an unprojected shredded
sibling.
##########
crates/iceberg/src/arrow/reader/projection.rs:
##########
@@ -2017,4 +2586,353 @@ message schema {
.collect();
assert_eq!(ids, vec![2, 3]);
}
+
+ // ---------------------------------------------------------------------
+ // End-to-end variant read tests (synthetic Parquet).
+ //
+ // Write the unshredded variant layout (a `group { metadata; value }`
tagged with the
+ // `arrow.parquet.variant` extension, so the writer emits the Parquet
variant annotation
Review Comment:
This banner reads backwards: it says the group is tagged with
`arrow.parquet.variant` so the writer emits the annotation and the reader
re-attaches it — but `variant_arrow_field`'s doc says the extension is
intentionally absent, so these tests actually exercise the no-annotation,
recover-from-Iceberg-schema path. Worth fixing the comment to match.
More importantly, every assertion here goes through `as_struct()` + raw
bytes, so these would pass even if the tagging were deleted from the output
path (which, per the `:196` thread, it effectively is). I'd add a case where
the input field carries the real annotation/extension and assert
`extension_type_name()` on the output (or a round-trip through
`arrow_schema_to_schema` to `Type::Variant`) — that's what actually pins the
PR's goal.
##########
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(
Review Comment:
This hand-rolls arrow-rs's private leaf-numbering (Dictionary/RunEndEncoded
unwrap, Union members, list/map inner), and if a future arrow release adds a
child-bearing type the indices shift silently and the mask selects the wrong
Parquet columns — no error, just wrong data. The numbering test only pins the
current arrow version and only exercises Dictionary/Union/RunEndEncoded
siblings, not List/Map/LargeList/FixedSizeList, which are the common production
shapes.
If there's a way not to re-derive it, I'd lean on the Parquet
`SchemaDescriptor` column paths as the authoritative leaf order, or track
variant ancestry through a single `tagged_fields.filter_leaves` closure instead
of a parallel walk. If the manual walk stays, at minimum extend the test to
collect `(idx, name)` from `filter_leaves` and compare per-leaf against the
variant keys.
##########
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:
This rebuilds the entire field tree (every struct/list/map re-`Arc`'d) on
every file read, even when the schema has no variant at all. The cheap guard is
exactly the `type_contains_variant` helper this PR removes — I'd keep a version
of it and skip both the collect and tag passes when the projection holds no
variant ids.
##########
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> {
Review Comment:
`field_variant_id` resolves variants only through the embedded Parquet field
id, so on files that rely on name-mapping to assign ids (no embedded ids) it
finds nothing and the variant is silently dropped from the projection rather
than erroring. Probably fine to defer, but worth a comment noting the
limitation so it's not mistaken for a guarantee.
--
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]
