diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index bd5736bc17..e6de7ab4d8 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -122,6 +122,7 @@ pub fn iceberg::arrow::ArrowSchemaVisitor::map(&mut self, map: &arrow_schema::da pub fn iceberg::arrow::ArrowSchemaVisitor::primitive(&mut self, p: &arrow_schema::datatype::DataType) -> iceberg::Result pub fn iceberg::arrow::ArrowSchemaVisitor::schema(&mut self, schema: &arrow_schema::schema::Schema, values: alloc::vec::Vec) -> iceberg::Result pub fn iceberg::arrow::ArrowSchemaVisitor::struct(&mut self, fields: &arrow_schema::fields::Fields, results: alloc::vec::Vec) -> iceberg::Result +pub fn iceberg::arrow::ArrowSchemaVisitor::variant(&mut self, field: &arrow_schema::field::FieldRef) -> iceberg::Result where Self: core::marker::Sized pub fn iceberg::arrow::arrow_primitive_to_literal(primitive_array: &arrow_array::array::ArrayRef, ty: &iceberg::spec::Type) -> iceberg::Result>> pub fn iceberg::arrow::arrow_schema_to_schema(schema: &arrow_schema::schema::Schema) -> iceberg::Result pub fn iceberg::arrow::arrow_schema_to_schema_auto_assign_ids(schema: &arrow_schema::schema::Schema) -> iceberg::Result diff --git a/crates/iceberg/src/arrow/schema.rs b/crates/iceberg/src/arrow/schema.rs index f47176bc95..9cb6ac23bc 100644 --- a/crates/iceberg/src/arrow/schema.rs +++ b/crates/iceberg/src/arrow/schema.rs @@ -164,6 +164,21 @@ pub trait ArrowSchemaVisitor { /// Called when see a primitive type. fn primitive(&mut self, p: &DataType) -> Result; + + /// Called when a field carries the `arrow.parquet.variant` extension type. + /// + /// The default treats the variant as its underlying Arrow struct storage: it + /// re-enters normal traversal, preserving the behavior of visitors that don't + /// special-case variants. A visitor that produces Iceberg types (or otherwise + /// needs variant identity) should override this to fold the struct into a + /// single logical variant. + /// + /// Takes the `&FieldRef` rather than a `&DataType` because the variant signal + /// lives on the field's metadata, not on its data type. + fn variant(&mut self, field: &FieldRef) -> Result + where Self: Sized { + visit_type(field.data_type(), self) + } } /// Visiting a type in post order. @@ -201,14 +216,14 @@ fn visit_type(r#type: &DataType, visitor: &mut V) -> Resu let key_result = { visitor.before_map_key(key_field)?; - let ret = visit_type(key_field.data_type(), visitor)?; + let ret = visit_field(key_field, visitor)?; visitor.after_map_key(key_field)?; ret }; let value_result = { visitor.before_map_value(value_field)?; - let ret = visit_type(value_field.data_type(), visitor)?; + let ret = visit_field(value_field, visitor)?; visitor.after_map_value(value_field)?; ret }; @@ -229,6 +244,16 @@ fn visit_type(r#type: &DataType, visitor: &mut V) -> Resu } } +/// Dispatch a field: fold it into a variant when it carries the +/// `arrow.parquet.variant` extension type, otherwise visit its data type. +fn visit_field(field: &FieldRef, visitor: &mut V) -> Result { + if field.extension_type_name() == Some(VariantExtensionType::NAME) { + visitor.variant(field) + } else { + visit_type(field.data_type(), visitor) + } +} + /// Visit list types in post order. fn visit_list( data_type: &DataType, @@ -236,7 +261,7 @@ fn visit_list( visitor: &mut V, ) -> Result { visitor.before_list_element(element_field)?; - let value = visit_type(element_field.data_type(), visitor)?; + let value = visit_field(element_field, visitor)?; visitor.after_list_element(element_field)?; visitor.list(data_type, value) } @@ -246,7 +271,7 @@ fn visit_struct(fields: &Fields, visitor: &mut V) -> Resu let mut results = Vec::with_capacity(fields.len()); for field in fields { visitor.before_field(field)?; - let result = visit_type(field.data_type(), visitor)?; + let result = visit_field(field, visitor)?; visitor.after_field(field)?; results.push(result); } @@ -262,7 +287,7 @@ pub(crate) fn visit_schema( let mut results = Vec::with_capacity(schema.fields().len()); for field in schema.fields() { visitor.before_field(field)?; - let result = visit_type(field.data_type(), visitor)?; + let result = visit_field(field, visitor)?; visitor.after_field(field)?; results.push(result); } @@ -544,6 +569,21 @@ impl ArrowSchemaVisitor for ArrowSchemaConverter { )), } } + + fn variant(&mut self, field: &FieldRef) -> Result { + // The extension may only sit on struct storage (mirrors + // `VariantExtensionType::supports_data_type`). + if !matches!(field.data_type(), DataType::Struct(_)) { + return Err(Error::new( + ErrorKind::DataInvalid, + "arrow.parquet.variant extension requires Struct storage", + )); + } + // Fold the whole struct into a single logical variant without descending: + // the storage sub-fields carry no Iceberg field id, so visiting them would + // fail. The enclosing field's own id is read by the caller. + Ok(Type::Variant(VariantType)) + } } struct ToArrowSchemaConverter; @@ -2080,6 +2120,114 @@ mod tests { assert_eq!(value.extension_type_name(), Some("arrow.parquet.variant")); } + /// The unshredded Arrow storage of a variant: `metadata` (required) + `value` + /// (nullable) binary, with no field ids on the sub-fields. + fn variant_storage() -> DataType { + DataType::Struct(Fields::from(vec![ + Field::new("metadata", DataType::Binary, false), + Field::new("value", DataType::Binary, true), + ])) + } + + #[test] + fn test_variant_arrow_field_folds_to_iceberg_variant() { + // A field tagged with the arrow.parquet.variant extension is folded into an + // atomic Type::Variant; its storage sub-fields (which carry no field id) are + // never descended into. + let field = simple_field("v", variant_storage(), true, "1") + .with_extension_type(VariantExtensionType); + let arrow_schema = ArrowSchema::new(vec![field]); + + let converted = arrow_schema_to_schema(&arrow_schema).unwrap(); + let expected = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "v", Type::Variant(VariantType)).into(), + ]) + .build() + .unwrap(); + pretty_assertions::assert_eq!(converted, expected); + } + + #[test] + fn test_variant_schema_round_trips() { + // Iceberg -> Arrow -> Iceberg is the identity for variants at every position: + // top-level, nested in a struct, as a list element, and as a map value. + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "v", Type::Variant(VariantType)).into(), + NestedField::optional( + 2, + "s", + Type::Struct(StructType::new(vec![ + NestedField::optional(3, "sv", Type::Variant(VariantType)).into(), + ])), + ) + .into(), + NestedField::optional( + 4, + "l", + Type::List(ListType::new( + NestedField::optional(5, "element", Type::Variant(VariantType)).into(), + )), + ) + .into(), + NestedField::optional( + 6, + "m", + Type::Map(MapType::new( + NestedField::map_key_element(7, Type::Primitive(PrimitiveType::String)) + .into(), + NestedField::map_value_element(8, Type::Variant(VariantType), false).into(), + )), + ) + .into(), + ]) + .build() + .unwrap(); + + let arrow_schema = schema_to_arrow_schema(&schema).unwrap(); + let round_tripped = arrow_schema_to_schema(&arrow_schema).unwrap(); + pretty_assertions::assert_eq!(round_tripped, schema); + } + + #[test] + fn test_variant_recognized_with_auto_assigned_ids() { + // Recognition also works when the Arrow schema has no field ids: the variant + // field gets an auto-assigned id and its storage is still not descended into. + let field = + Field::new("v", variant_storage(), true).with_extension_type(VariantExtensionType); + let arrow_schema = ArrowSchema::new(vec![field]); + + let converted = arrow_schema_to_schema_auto_assign_ids(&arrow_schema).unwrap(); + let expected = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "v", Type::Variant(VariantType)).into(), + ]) + .build() + .unwrap(); + pretty_assertions::assert_eq!(converted, expected); + } + + #[test] + fn test_variant_extension_on_non_struct_storage_is_rejected() { + // The extension may only sit on struct storage. A hand-injected tag on a + // non-struct field is rejected rather than silently reinterpreted as a variant. + let field = Field::new("v", DataType::Int32, true).with_metadata(HashMap::from([ + (PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string()), + ( + arrow_schema::extension::EXTENSION_TYPE_NAME_KEY.to_string(), + VariantExtensionType::NAME.to_string(), + ), + ])); + let arrow_schema = ArrowSchema::new(vec![field]); + + let err = arrow_schema_to_schema(&arrow_schema).unwrap_err(); + assert!( + err.to_string().contains("requires Struct storage"), + "unexpected error: {err}" + ); + } + #[test] fn test_type_conversion() { // test primitive type