diff --git a/crates/iceberg/Cargo.toml b/crates/iceberg/Cargo.toml index 66eadfb7ec..4b8e41f859 100644 --- a/crates/iceberg/Cargo.toml +++ b/crates/iceberg/Cargo.toml @@ -42,7 +42,7 @@ arrow-array = { workspace = true } arrow-buffer = { workspace = true } arrow-cast = { workspace = true } arrow-ord = { workspace = true } -arrow-schema = { workspace = true } +arrow-schema = { workspace = true, features = ["canonical_extension_types"] } arrow-select = { workspace = true } arrow-string = { workspace = true } as-any = { workspace = true } diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 530aa4cbf3..65c68d46e8 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::uuid(&mut self, field: &arrow_schema::field::FieldRef) -> iceberg::Result where Self: core::marker::Sized 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 diff --git a/crates/iceberg/src/arrow/schema.rs b/crates/iceberg/src/arrow/schema.rs index 923c043c74..6999f7d400 100644 --- a/crates/iceberg/src/arrow/schema.rs +++ b/crates/iceberg/src/arrow/schema.rs @@ -26,7 +26,7 @@ use arrow_array::{ FixedSizeBinaryArray, Float32Array, Float64Array, Int32Array, Int64Array, Scalar, StringArray, TimestampMicrosecondArray, TimestampNanosecondArray, }; -use arrow_schema::extension::ExtensionType; +use arrow_schema::extension::{ExtensionType, Uuid as UuidExtensionType}; use arrow_schema::{ ArrowError, DataType, Field, FieldRef, Fields, Schema as ArrowSchema, TimeUnit, }; @@ -179,6 +179,20 @@ pub trait ArrowSchemaVisitor { where Self: Sized { visit_type(field.data_type(), self) } + + /// Called when a field carries the `arrow.uuid` extension type. + /// + /// The default treats the field as its underlying Arrow storage, preserving + /// the behavior of visitors that don't special-case UUIDs. A visitor that + /// produces Iceberg types should override this to fold the storage into a + /// single logical UUID. + /// + /// Takes the `&FieldRef` rather than a `&DataType` because the UUID signal + /// lives on the field's metadata, not on its data type. + fn uuid(&mut self, field: &FieldRef) -> Result + where Self: Sized { + visit_type(field.data_type(), self) + } } /// Visiting a type in post order. @@ -249,6 +263,8 @@ fn visit_type(r#type: &DataType, visitor: &mut V) -> Resu fn visit_field(field: &FieldRef, visitor: &mut V) -> Result { if field.extension_type_name() == Some(VariantExtensionType::NAME) { visitor.variant(field) + } else if field.extension_type_name() == Some(UuidExtensionType::NAME) { + visitor.uuid(field) } else { visit_type(field.data_type(), visitor) } @@ -584,6 +600,23 @@ impl ArrowSchemaVisitor for ArrowSchemaConverter { // fail. The enclosing field's own id is read by the caller. Ok(Type::Variant(VariantType)) } + + fn uuid(&mut self, field: &FieldRef) -> Result { + // The extension expects a fixed size binary. + // See: Apache Arrow spec https://arrow.apache.org/docs/format/CanonicalExtensions.html#uuid + let field_data_type = field.data_type(); + let expected_data_type = DataType::FixedSizeBinary(16); + if *field_data_type != expected_data_type { + let error_message = format!( + "{} extension requires {expected_data_type} storage, found {field_data_type}", + UuidExtensionType::NAME, + ); + + return Err(Error::new(ErrorKind::DataInvalid, error_message)); + } + + Ok(Type::Primitive(PrimitiveType::Uuid)) + } } struct ToArrowSchemaConverter; @@ -634,9 +667,17 @@ impl SchemaVisitor for ToArrowSchemaConverter { // `arrow.parquet.variant` extension type so consumers read it as a Variant, not a struct. let arrow_field = if field.field_type.is_variant() { arrow_field.with_extension_type(VariantExtensionType) + } else if matches!( + field.field_type.as_ref(), + Type::Primitive(PrimitiveType::Uuid) + ) { + // A uuid column's storage is FixedSizeBinary(16); tag the field with the canonical + // `arrow.uuid` extension type so consumers read it as a UUID, not opaque bytes. + arrow_field.with_extension_type(UuidExtensionType) } else { arrow_field }; + Ok(ArrowSchemaOrFieldOrType::Field(arrow_field)) } @@ -1393,6 +1434,7 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use arrow_schema::extension::Uuid as DataTypeUuidExt; use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit}; use super::*; @@ -1699,6 +1741,74 @@ mod tests { pretty_assertions::assert_eq!(converted_schema, schema); } + #[test] + fn test_converting_uuid_from_arrow_to_iceberg_to_arrow_should_give_uuid() { + let arrow_schema = ArrowSchema::new(vec![ + simple_field("uuid_field", DataType::FixedSizeBinary(16), false, "1") + .with_extension_type(DataTypeUuidExt), + ]); + let output = + schema_to_arrow_schema(&arrow_schema_to_schema(&arrow_schema).unwrap()).unwrap(); + + assert_eq!(arrow_schema, output); + } + + #[test] + fn test_converting_uuid_from_iceberg_to_arrow_to_iceberg_should_give_uuid() { + let iceberg_schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "uuid_field", Type::Primitive(PrimitiveType::Uuid)).into(), + ]) + .build() + .unwrap(); + let output = + arrow_schema_to_schema(&schema_to_arrow_schema(&iceberg_schema).unwrap()).unwrap(); + + assert_eq!(iceberg_schema, output); + } + + #[test] + fn test_arrow_schema_to_schema_should_convert_uuid_when_fixed_size_binary() { + let converted_schema = arrow_schema_to_schema(&ArrowSchema::new(vec![ + simple_field("uuid_field", DataType::FixedSizeBinary(16), false, "1") + .with_extension_type(DataTypeUuidExt), + ])) + .unwrap(); + + let expected = Schema::builder() + .with_fields([NestedField::required( + 1, + "uuid_field", + Type::Primitive(PrimitiveType::Uuid), + ) + .into()]) + .build() + .unwrap(); + + pretty_assertions::assert_eq!(expected, converted_schema); + } + + #[test] + fn test_arrow_schema_to_schema_should_reject_uuid_when_not_a_fixed_size_binary() { + // The field must be built correctly, and then changed to the wrong type, + // to avoid Arrow's own validation and panic. + let mut field = simple_field( + "incorrect_uuid_field", + DataType::FixedSizeBinary(16), + false, + "1", + ) + .with_extension_type(DataTypeUuidExt); + field = field.with_data_type(DataType::Utf8); + + let error = arrow_schema_to_schema(&ArrowSchema::new(vec![field])).unwrap_err(); + + pretty_assertions::assert_eq!( + "arrow.uuid extension requires FixedSizeBinary(16) storage, found Utf8", + error.message() + ); + } + fn arrow_schema_for_schema_to_arrow_schema_test() -> ArrowSchema { let fields = Fields::from(vec![ simple_field("key", DataType::Int32, false, "28"), @@ -1791,7 +1901,8 @@ mod tests { ), simple_field("map", map, false, "16"), simple_field("struct", r#struct, false, "17"), - simple_field("uuid", DataType::FixedSizeBinary(16), false, "30"), + simple_field("uuid", DataType::FixedSizeBinary(16), false, "30") + .with_extension_type(DataTypeUuidExt), Field::new( "v", DataType::Struct(Fields::from(vec![ @@ -2007,10 +2118,10 @@ mod tests { #[test] fn test_schema_to_arrow_schema() { - let arrow_schema = arrow_schema_for_schema_to_arrow_schema_test(); + let expected_arrow_schema = arrow_schema_for_schema_to_arrow_schema_test(); let schema = iceberg_schema_for_schema_to_arrow_schema(); let converted_arrow_schema = schema_to_arrow_schema(&schema).unwrap(); - assert_eq!(converted_arrow_schema, arrow_schema); + assert_eq!(converted_arrow_schema, expected_arrow_schema); } #[test] @@ -2056,6 +2167,27 @@ mod tests { ); } + #[test] + fn test_schema_to_arrow_schema_should_convert_uuid_to_arrow_uuid() { + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "uuid_field", Type::Primitive(PrimitiveType::Uuid)).into(), + ]) + .build() + .unwrap(); + + let arrow_schema = schema_to_arrow_schema(&schema).unwrap(); + let field = arrow_schema.field_with_name("uuid_field").unwrap(); + + assert_eq!(field.extension_type_name(), Some(UuidExtensionType::NAME)); + // Attaching the extension type must not clobber the Iceberg field id. + assert_eq!( + field.metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&"1".to_string()) + ); + assert_eq!(*field.data_type(), DataType::FixedSizeBinary(16)); + } + #[test] fn test_variant_nested_in_list_and_map_carries_arrow_extension_type() { // A variant nested in a list element or map value keeps the arrow.parquet.variant @@ -2487,6 +2619,8 @@ mod tests { Field::new("id", DataType::Int64, false), Field::new("name", DataType::Utf8, true), Field::new("price", DataType::Decimal128(10, 2), false), + Field::new("uuid", DataType::FixedSizeBinary(16), false) + .with_extension_type(DataTypeUuidExt), Field::new( "created_at", DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())), @@ -2538,9 +2672,9 @@ mod tests { let schema = arrow_schema_to_schema_auto_assign_ids(&arrow_schema).unwrap(); // Build expected schema with exact field IDs following level-order assignment: - // Level 0: id=1, name=2, price=3, created_at=4, tags=5, address=6, attributes=7, orders=8 - // Level 1: tags.element=9, address.{street=10,city=11,zip=12}, attributes.{key=13,value=14}, orders.element=15 - // Level 2: orders.element.{order_id=16,amount=17} + // Level 0: id=1, name=2, price=3, uuid=4, created_at=5, tags=6, address=7, attributes=8, orders=9 + // Level 1: tags.element=10, address.{street=11,city=12,zip=13}, attributes.{key=14,value=15}, orders.element=16 + // Level 2: orders.element.{order_id=17,amount=18} let expected = Schema::builder() .with_fields(vec![ NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), @@ -2554,14 +2688,15 @@ mod tests { }), ) .into(), - NestedField::optional(4, "created_at", Type::Primitive(PrimitiveType::Timestamptz)) + NestedField::required(4, "uuid", Type::Primitive(PrimitiveType::Uuid)).into(), + NestedField::optional(5, "created_at", Type::Primitive(PrimitiveType::Timestamptz)) .into(), NestedField::optional( - 5, + 6, "tags", Type::List(ListType { element_field: NestedField::list_element( - 9, + 10, Type::Primitive(PrimitiveType::String), false, ) @@ -2570,29 +2705,29 @@ mod tests { ) .into(), NestedField::optional( - 6, + 7, "address", Type::Struct(StructType::new(vec![ - NestedField::optional(10, "street", Type::Primitive(PrimitiveType::String)) + NestedField::optional(11, "street", Type::Primitive(PrimitiveType::String)) .into(), - NestedField::required(11, "city", Type::Primitive(PrimitiveType::String)) + NestedField::required(12, "city", Type::Primitive(PrimitiveType::String)) .into(), - NestedField::optional(12, "zip", Type::Primitive(PrimitiveType::Int)) + NestedField::optional(13, "zip", Type::Primitive(PrimitiveType::Int)) .into(), ])), ) .into(), NestedField::optional( - 7, + 8, "attributes", Type::Map(MapType { key_field: NestedField::map_key_element( - 13, + 14, Type::Primitive(PrimitiveType::String), ) .into(), value_field: NestedField::map_value_element( - 14, + 15, Type::Primitive(PrimitiveType::String), false, ) @@ -2601,20 +2736,20 @@ mod tests { ) .into(), NestedField::optional( - 8, + 9, "orders", Type::List(ListType { element_field: NestedField::list_element( - 15, + 16, Type::Struct(StructType::new(vec![ NestedField::required( - 16, + 17, "order_id", Type::Primitive(PrimitiveType::Long), ) .into(), NestedField::required( - 17, + 18, "amount", Type::Primitive(PrimitiveType::Double), ) @@ -2631,6 +2766,6 @@ mod tests { .unwrap(); pretty_assertions::assert_eq!(schema, expected); - assert_eq!(schema.highest_field_id(), 17); + assert_eq!(schema.highest_field_id(), 18); } }