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 12f844b362..169f48de1e 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..d7f46e062b 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,10 @@ 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 { + Ok(Type::Primitive(PrimitiveType::Uuid)) + } } struct ToArrowSchemaConverter; @@ -1393,6 +1413,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 +1720,27 @@ mod tests { pretty_assertions::assert_eq!(converted_schema, schema); } + #[test] + fn test_arrow_schema_to_schema_should_convert_uuids() { + 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); + } + fn arrow_schema_for_schema_to_arrow_schema_test() -> ArrowSchema { let fields = Fields::from(vec![ simple_field("key", DataType::Int32, false, "28"), @@ -2487,6 +2529,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 +2582,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 +2598,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 +2615,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 +2646,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 +2676,6 @@ mod tests { .unwrap(); pretty_assertions::assert_eq!(schema, expected); - assert_eq!(schema.highest_field_id(), 17); + assert_eq!(schema.highest_field_id(), 18); } }