From 2642687767727d81b4523ce3fdb248148189c770 Mon Sep 17 00:00:00 2001 From: Christian Thiel Date: Thu, 23 Jul 2026 07:53:55 +0200 Subject: [PATCH] feat(arrow): read unshredded variant columns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Project variant columns on read by mapping their metadata/value leaves (which carry no field id — only the enclosing group does) to the variant's field id, mirroring Java's PruneColumns: the whole group is projected with no type-promotion check. Shredded variants (typed_value present) are rejected until we can reconstruct them. collect_variant_leaves numbers leaves exactly like Fields::filter_leaves, including its Dictionary/RunEndEncoded unwrapping and Union descent, so a sibling of those kinds can't shift a variant's leaves onto the wrong column. The pruned schema used for the type-promotion check includes variant leaves too, keeping map<_, variant> a well-formed two-field map. Variant identity is recovered from the Iceberg schema: fields the table schema declares as variants are tagged with the arrow.parquet.variant extension so arrow_schema_to_schema folds the storage struct back into Type::Variant. This mirrors iceberg-java's TypeWithSchemaVisitor, which recognizes a variant by the Parquet annotation OR the Iceberg type, and avoids depending on the Parquet variant annotation being present. --- crates/iceberg/src/arrow/reader/projection.rs | 1064 +++++++++++++++-- 1 file changed, 991 insertions(+), 73 deletions(-) diff --git a/crates/iceberg/src/arrow/reader/projection.rs b/crates/iceberg/src/arrow/reader/projection.rs index 9fdc9fed65..5d88fcde6c 100644 --- a/crates/iceberg/src/arrow/reader/projection.rs +++ b/crates/iceberg/src/arrow/reader/projection.rs @@ -23,12 +23,14 @@ use std::collections::{HashMap, HashSet}; use std::str::FromStr; use std::sync::Arc; -use arrow_schema::{Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef}; +use arrow_schema::{ + DataType, Field, FieldRef, Fields, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, +}; use parquet::arrow::{PARQUET_FIELD_ID_META_KEY, ProjectionMask}; use parquet::schema::types::{SchemaDescriptor, Type as ParquetType}; use super::{ArrowReader, CollectFieldIdVisitor}; -use crate::arrow::arrow_schema_to_schema; +use crate::arrow::{VariantExtensionType, arrow_schema_to_schema}; use crate::error::Result; use crate::expr::BoundPredicate; use crate::expr::visitors::bound_predicate_visitor::visit; @@ -81,8 +83,9 @@ impl ArrowReader { Self::include_leaf_field_id(&map_type.key_field, field_ids); Self::include_leaf_field_id(&map_type.value_field, field_ids); } - // Variant projection is rejected earlier (in `get_arrow_projection_mask`); this - // arm only keeps the match exhaustive. Treat it as a leaf, like a primitive. + // A variant is an Iceberg leaf type but a Parquet group: only the parent + // group carries the field id, so we project it like a primitive. Its + // metadata/value leaves are resolved later via `collect_variant_leaves`. Type::Variant(_) => { field_ids.push(field.id); } @@ -124,19 +127,6 @@ impl ArrowReader { return Ok(ProjectionMask::all()); } - // Reading variant columns is not supported yet (see #2188 follow-ups): reject any - // projection that touches a variant, rather than returning a partial/incorrect batch. - for field_id in field_ids { - if let Some(field) = iceberg_schema_of_task.field_by_id(*field_id) - && type_contains_variant(&field.field_type) - { - return Err(Error::new( - ErrorKind::FeatureUnsupported, - "Reading variant columns is not supported yet", - )); - } - } - if use_fallback { // Position-based projection necessary because file lacks embedded field IDs Self::get_arrow_projection_mask_fallback(field_ids, parquet_schema) @@ -171,16 +161,50 @@ impl ArrowReader { arrow_schema: &ArrowSchemaRef, type_promotion_is_valid: fn(Option<&PrimitiveType>, Option<&PrimitiveType>) -> bool, ) -> Result { - 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> = 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 = 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); + // Pre-project only the fields that have been selected, possibly avoiding converting - // some Arrow types that are not yet supported. - let mut projected_fields: HashMap = HashMap::new(); + // some Arrow types that are not yet supported. Variant leaves are included here too + // (they carry no field id) so nested containers — e.g. a map's key_value struct — + // stay well-formed for `arrow_schema_to_schema` below. + let mut projected_fields: HashMap = HashMap::new(); let projected_arrow_schema = ArrowSchema::new_with_metadata( - fields.filter_leaves(|_, f| { + tagged_fields.filter_leaves(|idx, f| { + if variant_leaves.contains_key(&idx) { + return true; + } f.metadata() .get(PARQUET_FIELD_ID_META_KEY) .and_then(|field_id| i32::from_str(field_id).ok()) @@ -193,7 +217,15 @@ impl ArrowReader { ); let iceberg_schema = arrow_schema_to_schema(&projected_arrow_schema)?; - fields.filter_leaves(|idx, field| { + tagged_fields.filter_leaves(|idx, field| { + // Variant sub-fields: the parent group carries the field id, not the leaf, and + // `Type::Variant` is not a primitive, so skip the type-promotion check (matching + // iceberg-java, which projects the variant group whole with no type check). + if let Some(&variant_field_id) = variant_leaves.get(&idx) { + column_map.entry(variant_field_id).or_default().push(idx); + return true; + } + let Some(field_id) = projected_fields.get(field).cloned() else { return false; }; @@ -215,7 +247,7 @@ impl ArrowReader { return false; } - column_map.insert(field_id, idx); + column_map.entry(field_id).or_default().push(idx); true }); @@ -223,8 +255,8 @@ impl ArrowReader { // We only project existing columns; RecordBatchTransformer adds default/NULL values. let mut indices = vec![]; for field_id in leaf_field_ids { - if let Some(col_idx) = column_map.get(field_id) { - indices.push(*col_idx); + if let Some(col_indices) = column_map.get(field_id) { + indices.extend_from_slice(col_indices); } } @@ -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, + iceberg_schema: &Schema, + leaf_field_id_set: &HashSet, + out: &mut HashMap, + ) -> 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, + iceberg_schema: &Schema, + leaf_field_id_set: &HashSet, + out: &mut HashMap, + ) -> 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 { + 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(), + }; + Arc::new(field.as_ref().clone().with_data_type(data_type)) + } + /// Fallback projection for Parquet files without field IDs. /// Uses position-based matching: field ID N → column position N-1. /// Projects entire top-level columns (including nested content) for iceberg-java compatibility. @@ -265,23 +475,6 @@ impl ArrowReader { } } -/// Whether `field_type` is, or transitively contains, a variant type. -fn type_contains_variant(field_type: &Type) -> bool { - match field_type { - Type::Variant(_) => true, - Type::Struct(s) => s - .fields() - .iter() - .any(|f| type_contains_variant(&f.field_type)), - Type::List(l) => type_contains_variant(&l.element_field.field_type), - Type::Map(m) => { - type_contains_variant(&m.key_field.field_type) - || type_contains_variant(&m.value_field.field_type) - } - Type::Primitive(_) => false, - } -} - /// Build the map of parquet field id to Parquet column index in the schema. /// Returns None if the Parquet file doesn't have field IDs embedded (e.g., migrated tables). pub(super) fn build_field_id_map( @@ -482,13 +675,13 @@ pub(super) fn add_fallback_field_ids_to_arrow_schema( #[cfg(test)] mod tests { - use std::collections::HashMap; + use std::collections::{HashMap, HashSet}; use std::fs::File; use std::sync::Arc; use arrow_array::cast::AsArray; use arrow_array::{Array, ArrayRef, RecordBatch, StringArray}; - use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit}; + use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema, TimeUnit}; use futures::TryStreamExt; use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY, ProjectionMask}; use parquet::basic::Compression; @@ -599,31 +792,257 @@ message schema { assert_eq!(mask, ProjectionMask::leaves(&parquet_schema, vec![0])); } + fn field_id_meta(id: i32) -> HashMap { + 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"); + assert_eq!(mask, ProjectionMask::leaves(&parquet_schema, vec![0, 1, 2])); + + // Variant only: its two leaves (metadata, value). + let mask_variant = ArrowReader::get_arrow_projection_mask( + &[2], + &schema, + &parquet_schema, + &arrow_schema, + false, + ) + .expect("projection mask for v only"); + assert_eq!( + mask_variant, + ProjectionMask::leaves(&parquet_schema, vec![1, 2]) + ); + + // Primitive only: variant leaves must not leak in. + let mask_primitive = ArrowReader::get_arrow_projection_mask( + &[1], + &schema, + &parquet_schema, + &arrow_schema, + false, + ) + .expect("projection mask for c1 only"); + assert_eq!( + mask_primitive, + ProjectionMask::leaves(&parquet_schema, vec![0]) + ); + } + + /// A shredded variant (with a `typed_value` sub-field) must be rejected, not silently + /// projected as metadata/value only. + #[test] + fn test_arrow_projection_mask_variant_shredded_is_rejected() { + 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 shredded = Field::new( + "v", + DataType::Struct(Fields::from(vec![ + Field::new("metadata", DataType::Binary, false), + Field::new("value", DataType::Binary, true), + Field::new("typed_value", DataType::Int64, true), + ])), + false, + ) + .with_metadata(field_id_meta(2)); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("c1", DataType::Utf8, false).with_metadata(field_id_meta(1)), + shredded, + ])); + let message_type = " +message schema { + required binary c1 (STRING) = 1; + required group v = 2 { + required binary metadata; + optional binary value; + optional int64 typed_value; + } +} +"; + let parquet_schema = + SchemaDescriptor::new(Arc::new(parse_message_type(message_type).unwrap())); + + let err = ArrowReader::get_arrow_projection_mask( + &[1, 2], + &schema, + &parquet_schema, + &arrow_schema, + false, + ) + .expect_err("shredded variant must be rejected"); + assert_eq!(err.kind(), ErrorKind::FeatureUnsupported); + assert!( + err.message().contains("shredded variant"), + "unexpected error message: {err}" + ); + } + + /// A variant nested inside a struct must have both of its sub-leaves projected. + #[test] + fn test_arrow_projection_mask_variant_nested_in_struct() { let schema = Arc::new( Schema::builder() .with_schema_id(1) .with_fields(vec![ - NestedField::optional(1, "id", Type::Primitive(PrimitiveType::Int)).into(), - NestedField::optional(2, "v", Type::Variant(VariantType)).into(), NestedField::required( - 3, - "s", + 1, + "parent", Type::Struct(StructType::new(vec![ - NestedField::optional(4, "vv", Type::Variant(VariantType)).into(), + NestedField::required(2, "c2", Type::Primitive(PrimitiveType::String)) + .into(), + NestedField::required(3, "v", Type::Variant(VariantType)).into(), ])), ) .into(), + ]) + .build() + .unwrap(), + ); + let parent_fields = Fields::from(vec![ + Field::new("c2", DataType::Utf8, false).with_metadata(field_id_meta(2)), + variant_arrow_field("v", 3), + ]); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("parent", DataType::Struct(parent_fields), false) + .with_metadata(field_id_meta(1)), + ])); + let message_type = " +message schema { + required group parent = 1 { + required binary c2 (STRING) = 2; + required group v = 3 { + required binary metadata; + required binary value; + } + } +} +"; + let parquet_schema = + SchemaDescriptor::new(Arc::new(parse_message_type(message_type).unwrap())); + + // Nested variant: both of its leaves. + let mask_variant = ArrowReader::get_arrow_projection_mask( + &[3], + &schema, + &parquet_schema, + &arrow_schema, + false, + ) + .expect("projection mask for nested variant"); + assert_eq!( + mask_variant, + ProjectionMask::leaves(&parquet_schema, vec![1, 2]), + "variant nested in a struct was dropped from projection" + ); + + // Sibling primitive: no variant leaves. + let mask_primitive = ArrowReader::get_arrow_projection_mask( + &[2], + &schema, + &parquet_schema, + &arrow_schema, + false, + ) + .expect("projection mask for nested primitive"); + assert_eq!( + mask_primitive, + ProjectionMask::leaves(&parquet_schema, vec![0]) + ); + + // Both. + let mask_both = ArrowReader::get_arrow_projection_mask( + &[2, 3], + &schema, + &parquet_schema, + &arrow_schema, + false, + ) + .expect("projection mask for nested primitive + variant"); + assert_eq!( + mask_both, + ProjectionMask::leaves(&parquet_schema, vec![0, 1, 2]) + ); + } + + /// `map`: the map's key leaf and the variant value's metadata/value + /// leaves must all be projected, and the pruned schema must stay a well-formed 2-field + /// map (regression test for the projected-schema path). + #[test] + fn test_arrow_projection_mask_map_with_variant_value() { + let schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ NestedField::required( - 5, + 1, "m", Type::Map(crate::spec::MapType::required( - 6, + 2, Type::Primitive(PrimitiveType::String), - 7, + 3, Type::Variant(VariantType), )), ) @@ -632,26 +1051,176 @@ message schema { .build() .unwrap(), ); - // The parquet/arrow schemas are irrelevant: the variant is rejected before they - // are consulted, so an empty descriptor is enough to drive the code path. - let parquet_schema = SchemaDescriptor::new(Arc::new( - parse_message_type("message schema { optional int32 id = 1; }").unwrap(), - )); - let arrow_schema = Arc::new(ArrowSchema::empty()); - - // 2 = top-level variant, 3 = struct containing a variant, 4 = the nested variant, - // 5 = map, plus a mix with a non-variant sibling. - for projected in [vec![2], vec![3], vec![4], vec![5], vec![1, 2]] { - let err = ArrowReader::get_arrow_projection_mask( - &projected, - &schema, - &parquet_schema, - &arrow_schema, + let entries = Field::new( + "key_value", + DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false).with_metadata(field_id_meta(2)), + variant_arrow_field("value", 3), + ])), + false, + ); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("m", DataType::Map(Arc::new(entries), false), false) + .with_metadata(field_id_meta(1)), + ])); + let message_type = " +message schema { + required group m (MAP) = 1 { + repeated group key_value { + required binary key (STRING) = 2; + required group value = 3 { + required binary metadata; + required binary value; + } + } + } +} +"; + let parquet_schema = + SchemaDescriptor::new(Arc::new(parse_message_type(message_type).unwrap())); + + let mask = ArrowReader::get_arrow_projection_mask( + &[1], + &schema, + &parquet_schema, + &arrow_schema, + false, + ) + .expect("projection mask for map"); + assert_eq!(mask, ProjectionMask::leaves(&parquet_schema, vec![0, 1, 2])); + } + + /// `list`: the variant element's metadata/value leaves must both be projected. + #[test] + fn test_arrow_projection_mask_list_with_variant() { + let schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required( + 1, + "l", + Type::List(crate::spec::ListType::new( + NestedField::list_element(2, Type::Variant(VariantType), true).into(), + )), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "l", + DataType::List(Arc::new(variant_arrow_field("element", 2))), false, ) - .expect_err("variant projection must be rejected"); - assert_eq!(err.kind(), ErrorKind::FeatureUnsupported, "{err}"); - } + .with_metadata(field_id_meta(1)), + ])); + let message_type = " +message schema { + required group l (LIST) = 1 { + repeated group list { + required group element = 2 { + required binary metadata; + required binary value; + } + } + } +} +"; + let parquet_schema = + SchemaDescriptor::new(Arc::new(parse_message_type(message_type).unwrap())); + + let mask = ArrowReader::get_arrow_projection_mask( + &[1], + &schema, + &parquet_schema, + &arrow_schema, + false, + ) + .expect("projection mask for list"); + assert_eq!(mask, ProjectionMask::leaves(&parquet_schema, vec![0, 1])); + } + + /// `collect_variant_leaves` must number leaves exactly like `Fields::filter_leaves`, + /// including `filter_leaves`' `Dictionary`/`RunEndEncoded` unwrapping and `Union` descent. + /// Otherwise a sibling of those kinds shifts every following variant leaf onto the wrong + /// column. + #[test] + fn test_collect_variant_leaves_numbering_matches_filter_leaves() { + use arrow_schema::{UnionFields, UnionMode}; + + let union_fields: UnionFields = [ + (0i8, Arc::new(Field::new("i", DataType::Int32, false))), + (1i8, Arc::new(Field::new("l", DataType::Int64, false))), + ] + .into_iter() + .collect(); + + // Siblings that each expand to multiple leaves under `filter_leaves`, placed before the + // variant: a Dictionary(Struct{a,b}), a Union{i32,i64}, and a RunEndEncoded(Struct{x,y}) + // — 2 leaves each. `filter_leaves` unwraps Dictionary/RunEndEncoded to their value type + // and descends Union members; our walk must match. + let dict_struct = DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]))), + ); + let union = DataType::Union(union_fields, UnionMode::Sparse); + let run_end_encoded = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new( + "values", + DataType::Struct(Fields::from(vec![ + Field::new("x", DataType::Int32, false), + Field::new("y", DataType::Int32, false), + ])), + true, + )), + ); + let fields = Fields::from(vec![ + Field::new("d", dict_struct, false), + Field::new("u", union, false), + Field::new("r", run_end_encoded, false), + variant_arrow_field("v", 2), + ]); + // Only the variant (id 2) is relevant to the iceberg side. + let iceberg_schema = Schema::builder() + .with_fields(vec![ + NestedField::required(2, "v", Type::Variant(VariantType)).into(), + ]) + .build() + .unwrap(); + let leaf_field_id_set = HashSet::from([2]); + + let mut out = HashMap::new(); + let mut leaf_idx = 0usize; + ArrowReader::collect_variant_leaves( + &fields, + &mut leaf_idx, + None, + &iceberg_schema, + &leaf_field_id_set, + &mut out, + ) + .unwrap(); + + // d → leaves 0,1; u → leaves 2,3; r → leaves 4,5; v → metadata 6, value 7. + assert_eq!(out, HashMap::from([(6, 2), (7, 2)])); + + // The total leaf count must equal arrow's own `filter_leaves` count. + let mut arrow_leaf_count = 0usize; + let _ = fields.filter_leaves(|idx, _| { + arrow_leaf_count = idx + 1; + false + }); + assert_eq!( + leaf_idx, arrow_leaf_count, + "leaf numbering diverged from filter_leaves" + ); } /// Test schema evolution: reading old Parquet file (with only column 'a') @@ -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 + // and the reader re-attaches the extension), then read it back through the full reader + // path and assert the raw metadata/value bytes round-trip. + // --------------------------------------------------------------------- + + // Known variant bytes the reader must round-trip unchanged. `value` is distinct per row, + // so a dropped or reordered row is caught, not just the struct shape. + const VARIANT_METADATA: [&[u8]; 3] = [b"\x01", b"\x01", b"\x01"]; + const VARIANT_VALUE: [&[u8]; 3] = [b"\x0a", b"\x0b", b"\x0c"]; + + /// The 3-row variant `StructArray` for [`VARIANT_METADATA`]/[`VARIANT_VALUE`]. + fn variant_struct_array() -> ArrayRef { + use arrow_array::{BinaryArray, StructArray}; + let metadata = Arc::new(BinaryArray::from(VARIANT_METADATA.to_vec())) as ArrayRef; + let value = Arc::new(BinaryArray::from(VARIANT_VALUE.to_vec())) as ArrayRef; + Arc::new(StructArray::new( + variant_arrow_fields(), + vec![metadata, value], + None, + )) as ArrayRef + } + + /// Asserts `col` is `Struct(metadata, value)` and its bytes round-tripped exactly. + fn assert_variant_data(col: &ArrayRef) { + let s = col.as_struct(); + assert_eq!( + s.fields().len(), + 2, + "variant must have exactly metadata + value, got {:?}", + s.fields() + ); + let metadata = s + .column_by_name("metadata") + .expect("missing 'metadata'") + .as_binary::(); + let value = s + .column_by_name("value") + .expect("missing 'value'") + .as_binary::(); + assert_eq!( + metadata.iter().collect::>(), + VARIANT_METADATA + .iter() + .copied() + .map(Some) + .collect::>() + ); + assert_eq!( + value.iter().collect::>(), + VARIANT_VALUE.iter().copied().map(Some).collect::>() + ); + } + + /// Asserts the `id` column holds exactly `[1, 2, 3]`. + fn assert_id_values(col: &ArrayRef) { + let ids = col.as_primitive::(); + assert_eq!(ids.values(), &[1, 2, 3]); + } + + fn write_variant_parquet( + arrow_schema: Arc, + columns: Vec, + dir: &str, + ) -> String { + let batch = RecordBatch::try_new(arrow_schema.clone(), columns).unwrap(); + let path = format!("{dir}/variant.parquet"); + let props = WriterProperties::builder() + .set_compression(Compression::SNAPPY) + .build(); + let mut writer = + ArrowWriter::try_new(File::create(&path).unwrap(), arrow_schema, Some(props)).unwrap(); + writer.write(&batch).expect("writing batch"); + writer.close().unwrap(); + path + } + + async fn read_variant( + schema: Arc, + path: &str, + project_field_ids: Vec, + ) -> Vec { + let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build(); + let tasks = Box::pin(futures::stream::iter(vec![Ok(FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(path).unwrap().len()) + .with_start(0) + .with_length(0) + .with_data_file_path(path.to_string()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema) + .with_project_field_ids(project_field_ids) + .with_case_sensitive(false) + .build())])) as FileScanTaskStream; + reader + .read(tasks) + .unwrap() + .stream() + .try_collect::>() + .await + .unwrap() + } + + /// `id` (int, id 1) + `v` (variant, id 2), 3 rows. + fn write_top_level_variant(dir: &str) -> Arc { + use arrow_array::Int32Array; + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(2, "v", Type::Variant(VariantType)).into(), + ]) + .build() + .unwrap(), + ); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(field_id_meta(1)), + variant_arrow_field("v", 2), + ])); + let id = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + write_variant_parquet(arrow_schema, vec![id, variant_struct_array()], dir); + schema + } + + #[tokio::test] + async fn test_read_variant_full_scan() { + let tmp = TempDir::new().unwrap(); + let dir = tmp.path().to_str().unwrap(); + let schema = write_top_level_variant(dir); + + let batches = read_variant(schema, &format!("{dir}/variant.parquet"), vec![1, 2]).await; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 3); + assert_id_values( + batches[0] + .column_by_name("id") + .expect("'id' column dropped"), + ); + assert_variant_data( + batches[0] + .column_by_name("v") + .expect("variant column dropped"), + ); + } + + #[tokio::test] + async fn test_read_variant_only() { + let tmp = TempDir::new().unwrap(); + let dir = tmp.path().to_str().unwrap(); + let schema = write_top_level_variant(dir); + + let batches = read_variant(schema, &format!("{dir}/variant.parquet"), vec![2]).await; + assert_eq!(batches[0].num_rows(), 3); + assert_variant_data( + batches[0] + .column_by_name("v") + .expect("variant-only projection dropped 'v'"), + ); + assert!(batches[0].column_by_name("id").is_none()); + } + + #[tokio::test] + async fn test_read_variant_sibling_only_excludes_variant() { + let tmp = TempDir::new().unwrap(); + let dir = tmp.path().to_str().unwrap(); + let schema = write_top_level_variant(dir); + + let batches = read_variant(schema, &format!("{dir}/variant.parquet"), vec![1]).await; + assert_eq!(batches[0].num_rows(), 3); + assert_id_values( + batches[0] + .column_by_name("id") + .expect("'id' column dropped"), + ); + assert!( + batches[0].column_by_name("v").is_none(), + "variant leaked into an id-only projection" + ); + } + + #[tokio::test] + async fn test_read_variant_nested_in_struct() { + use arrow_array::{Int32Array, StructArray}; + let tmp = TempDir::new().unwrap(); + let dir = tmp.path().to_str().unwrap(); + + // id (int, 1) + nested struct (2) { payload variant (3) }. + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional( + 2, + "nested", + Type::Struct(StructType::new(vec![ + NestedField::optional(3, "payload", Type::Variant(VariantType)).into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let nested_fields = Fields::from(vec![variant_arrow_field("payload", 3)]); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(field_id_meta(1)), + Field::new("nested", DataType::Struct(nested_fields.clone()), false) + .with_metadata(field_id_meta(2)), + ])); + let id = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let nested = Arc::new(StructArray::new( + nested_fields, + vec![variant_struct_array()], + None, + )) as ArrayRef; + let path = write_variant_parquet(arrow_schema, vec![id, nested], dir); + + let batches = read_variant(schema, &path, vec![2]).await; + assert_eq!(batches[0].num_rows(), 3); + let payload = batches[0] + .column_by_name("nested") + .expect("nested struct dropped") + .as_struct() + .column_by_name("payload") + .expect("nested.payload variant dropped"); + assert_variant_data(payload); + } + + #[tokio::test] + async fn test_read_variant_in_list() { + use arrow_array::{Int32Array, ListArray}; + use arrow_buffer::OffsetBuffer; + let tmp = TempDir::new().unwrap(); + let dir = tmp.path().to_str().unwrap(); + + // id (int, 1) + l list (2, element 3); 3 rows, each list holds one variant. + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional( + 2, + "l", + Type::List(crate::spec::ListType::new( + NestedField::list_element(3, Type::Variant(VariantType), true).into(), + )), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let element_field = Arc::new(variant_arrow_field("element", 3)); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(field_id_meta(1)), + Field::new("l", DataType::List(element_field.clone()), false) + .with_metadata(field_id_meta(2)), + ])); + let id = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let list = Arc::new(ListArray::new( + element_field, + OffsetBuffer::new(vec![0, 1, 2, 3].into()), + variant_struct_array(), + None, + )) as ArrayRef; + let path = write_variant_parquet(arrow_schema, vec![id, list], dir); + + let batches = read_variant(schema, &path, vec![2]).await; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 3); + let list_col = batches[0] + .column_by_name("l") + .expect("list column dropped") + .as_list::(); + // Three single-element lists → the flattened element values are the 3 variants. + assert_variant_data(list_col.values()); + } + + #[tokio::test] + async fn test_read_variant_in_map() { + use arrow_array::{Int32Array, MapArray, StringArray, StructArray}; + use arrow_buffer::OffsetBuffer; + let tmp = TempDir::new().unwrap(); + let dir = tmp.path().to_str().unwrap(); + + // id (int, 1) + m map (2; key 3, value 4); 3 rows, each map one entry. + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional( + 2, + "m", + Type::Map(crate::spec::MapType::new( + NestedField::map_key_element(3, Type::Primitive(PrimitiveType::String)) + .into(), + NestedField::map_value_element(4, Type::Variant(VariantType), true) + .into(), + )), + ) + .into(), + ]) + .build() + .unwrap(), + ); + let key_field = Field::new("key", DataType::Utf8, false).with_metadata(field_id_meta(3)); + let value_field = variant_arrow_field("value", 4); + let entry_fields = Fields::from(vec![key_field, value_field]); + let entries_field = Arc::new(Field::new( + "key_value", + DataType::Struct(entry_fields.clone()), + false, + )); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(field_id_meta(1)), + Field::new("m", DataType::Map(entries_field.clone(), false), false) + .with_metadata(field_id_meta(2)), + ])); + let id = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let entries = StructArray::new( + entry_fields, + vec![ + Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef, + variant_struct_array(), + ], + None, + ); + let map = Arc::new(MapArray::new( + entries_field, + OffsetBuffer::new(vec![0, 1, 2, 3].into()), + entries, + None, + false, + )) as ArrayRef; + let path = write_variant_parquet(arrow_schema, vec![id, map], dir); + + let batches = read_variant(schema, &path, vec![2]).await; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 3); + let map_col = batches[0] + .column_by_name("m") + .expect("map column dropped") + .as_map(); + // The map's value column holds the 3 variants. + assert_variant_data(map_col.values()); + } }