dannycjones commented on code in PR #2880:
URL: https://github.com/apache/iceberg-rust/pull/2880#discussion_r3915957271
##########
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,
+ )?;
+ }
Review Comment:
nitpick: I'd prefer to see more explicit names here so it's clear how we're
recursing.
```suggestion
DataType::Struct(subfields) => {
Self::collect_variant_leaves(
subfields,
leaf_idx,
effective_variant,
iceberg_schema,
leaf_field_id_set,
out,
)?;
}
DataType::List(inner_field)
| DataType::LargeList(inner_field)
| DataType::FixedSizeList(inner_field, _)
| DataType::Map(inner_field, _) => {
Self::collect_variant_leaves_in_field(
inner_field,
leaf_idx,
effective_variant,
iceberg_schema,
leaf_field_id_set,
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,
+ )?;
+ }
+ }
Review Comment:
Why do we care about union?
##########
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)
+ }
Review Comment:
nit: `variant_field_id` as the function name would read better to me
##########
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:
I have a similar concern I think, although do let me know if its unrelated.
I don't understand why we're attaching the extension at this point. The
Arrow schema is passed in here, determined within the FileScanTaskReader.
Would it be better to handle this when we determine the Arrow schema? I see
we do a bunch of corrections to the Arrow schema in that module to coerce
timestamps and project meta columns.
##########
crates/iceberg/src/arrow/reader/projection.rs:
##########
@@ -599,31 +792,257 @@ message schema {
assert_eq!(mask, ProjectionMask::leaves(&parquet_schema, vec![0]));
}
+ fn field_id_meta(id: i32) -> HashMap<String, String> {
+ HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(),
id.to_string())])
+ }
+
+ /// Arrow `Struct(metadata: Binary, value: Binary)` — the unshredded
variant layout.
+ fn variant_arrow_fields() -> Fields {
+ Fields::from(vec![
+ Field::new("metadata", DataType::Binary, false),
+ Field::new("value", DataType::Binary, false),
+ ])
+ }
+
+ /// A variant storage group as it arrives from Parquet: `Struct(metadata,
value)` with the
+ /// field id on the group and no field ids on the sub-fields. The
`arrow.parquet.variant`
+ /// extension is intentionally ABSENT — the read path recovers variant
identity from the
+ /// Iceberg schema and self-attaches it, so this exercises files that lack
the annotation.
+ fn variant_arrow_field(name: &str, id: i32) -> Field {
+ Field::new(name, DataType::Struct(variant_arrow_fields()), false)
+ .with_metadata(field_id_meta(id))
+ }
+
+ /// A variant is a Parquet group whose leaves carry no field id, so it is
projected via
+ /// its enclosing group's field id and all of its leaves are read together.
#[test]
- fn test_arrow_projection_mask_variant_is_unsupported() {
- // Reading variant columns is not supported yet: projecting one
(top-level or
- // nested) must fail loudly rather than return a partial/incorrect
batch.
+ fn test_arrow_projection_mask_variant() {
+ // c1 (String, id 1) + v (Variant, id 2).
+ let schema = Arc::new(
+ Schema::builder()
+ .with_schema_id(1)
+ .with_fields(vec![
+ NestedField::required(1, "c1",
Type::Primitive(PrimitiveType::String)).into(),
+ NestedField::required(2, "v",
Type::Variant(VariantType)).into(),
+ ])
+ .build()
+ .unwrap(),
+ );
+ let arrow_schema = Arc::new(ArrowSchema::new(vec![
+ Field::new("c1", DataType::Utf8,
false).with_metadata(field_id_meta(1)),
+ variant_arrow_field("v", 2),
+ ]));
+ let message_type = "
+message schema {
+ required binary c1 (STRING) = 1;
+ required group v = 2 {
+ required binary metadata;
+ required binary value;
+ }
+}
+";
+ let parquet_schema =
+
SchemaDescriptor::new(Arc::new(parse_message_type(message_type).unwrap()));
+
+ // Both fields: all three leaves.
+ let mask = ArrowReader::get_arrow_projection_mask(
+ &[1, 2],
+ &schema,
+ &parquet_schema,
+ &arrow_schema,
+ false,
+ )
+ .expect("projection mask for c1 + v");
Review Comment:
Inclined to either change the expect message or use `unwrap`.
```suggestion
// Both fields: all three leaves.
let mask = ArrowReader::get_arrow_projection_mask(
&[1, 2],
&schema,
&parquet_schema,
&arrow_schema,
false,
)
.expect("all leaf types should be accepted by ArrowReader");
```
Applies to later uses of `expect` in this file too.
##########
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);
+ }
Review Comment:
```suggestion
if let Some(variant_field_id) = effective_variant {
out.insert(*leaf_idx, variant_field_id);
}
```
##########
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<()> {
Review Comment:
If we keep passing in the mutable hashmap, I think we should use something
more descriptive:
```suggestion
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>,
leaf_idx_to_field_id: &mut HashMap<usize, i32>,
) -> Result<()> {
```
##########
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
+ };
Review Comment:
I don't follow this logic super well.
In particular, it's really unclear to me why we want to pass in mutable leaf
index and hash map versus just returning a new hashmap at the end.
##########
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 assume the intention here is to recurse into complex types, or simply
return primitive types.
I think it would be good to be defensive here and guard on the type being
primitive. What do you think?
```suggestion
// 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,
),
data_type if data_type.is_primitive() => return field.clone(),
data_type => return iceberg::Error(
ErrorKind::Unexpected,
format!("Unexpected Arrow data type {data_type} when walking
Arrow record batch schema"),
),
};
```
--
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]
