This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-2930-5dc649bec6aee9cc25f33e9b6bb2462378646e18 in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git
commit f83bec42298a1cf51a42ea2aee26c041a15db5b4 Author: Joseph Lenton <[email protected]> AuthorDate: Thu Oct 15 20:04:13 2026 +0000 fix: arrow schema to schema supports uuids (#2930) * chore: add a test capturing the change * fix: add uuid support to arrow_schema_to_schema * chore: add uuid to the large arrow_schema_to_schema_auto_assign_ids test * chore: update the public api * feat: add a type check on the uuid conversion * chore: add a test for schema_to_arrow_schema for uuid * feat: add UUID support for schema_to_arrow_schema * docs: add uuid extension to visit_field documentation * refactor: remove duplicate UUID extension import from test * chore: add tests for fixed size binary of wrong size * feat: enable parquet writing uuid with uuid extension * chore: fmt Cargo.toml --- Cargo.lock | 2 + crates/iceberg/Cargo.toml | 8 +- crates/iceberg/public-api.txt | 1 + crates/iceberg/src/arrow/schema.rs | 215 ++++++++++++++++++--- .../src/writer/file_writer/parquet_writer.rs | 73 +++++++ 5 files changed, 273 insertions(+), 26 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3bb94e339..9e92c31b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -423,6 +423,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"; checksum = "10fab8d4563491417ba801fab29d205104d20d4bdf37bda6cd1cf425cff598cd" dependencies = [ "bitflags 2.13.0", + "serde_core", + "serde_json", ] [[package]] diff --git a/crates/iceberg/Cargo.toml b/crates/iceberg/Cargo.toml index f22ed59b9..99ec5f433 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 } @@ -65,7 +65,11 @@ moka = { version = "0.12.10", features = ["future"] } murmur3 = { workspace = true } once_cell = { workspace = true } ordered-float = { workspace = true } -parquet = { workspace = true, features = ["async", "encryption"] } +parquet = { workspace = true, features = [ + "async", + "encryption", + "arrow_canonical_extension_types", +] } rand = { workspace = true } reqwest = { workspace = true } roaring = { workspace = true } diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 00040053a..ad1af4389 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -123,6 +123,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<Self::T> pub fn iceberg::arrow::ArrowSchemaVisitor::schema(&mut self, schema: &arrow_schema::schema::Schema, values: alloc::vec::Vec<Self::T>) -> iceberg::Result<Self::U> pub fn iceberg::arrow::ArrowSchemaVisitor::struct(&mut self, fields: &arrow_schema::fields::Fields, results: alloc::vec::Vec<Self::T>) -> iceberg::Result<Self::T> +pub fn iceberg::arrow::ArrowSchemaVisitor::uuid(&mut self, field: &arrow_schema::field::FieldRef) -> iceberg::Result<Self::T> where Self: core::marker::Sized pub fn iceberg::arrow::ArrowSchemaVisitor::variant(&mut self, field: &arrow_schema::field::FieldRef) -> iceberg::Result<Self::T> 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<alloc::vec::Vec<core::option::Option<iceberg::spec::Literal>>> pub fn iceberg::arrow::arrow_schema_to_schema(schema: &arrow_schema::schema::Schema) -> iceberg::Result<iceberg::spec::Schema> diff --git a/crates/iceberg/src/arrow/schema.rs b/crates/iceberg/src/arrow/schema.rs index 0c673c6bc..748b97b3a 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, }; @@ -184,6 +184,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<Self::T> + where Self: Sized { + visit_type(field.data_type(), self) + } } /// Visiting a type in post order. @@ -241,11 +255,15 @@ fn visit_type<V: ArrowSchemaVisitor>(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. +/// Dispatch a field, and fold it into following extension types when carried: +/// - `arrow.parquet.variant` -- fold into a variant, +/// - `arrow.uuid` -- fold it into a UUID type, +/// - and in all other cases we visit it's data. fn visit_field<V: ArrowSchemaVisitor>(field: &FieldRef, visitor: &mut V) -> Result<V::T> { 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) } @@ -562,6 +580,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<Self::T> { + // 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; @@ -612,9 +647,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)) } @@ -1695,6 +1738,105 @@ 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(UuidExtensionType), + ]); + 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(UuidExtensionType), + ])) + .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 field = simple_field_with_uuid_ext(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() + ); + } + + #[test] + fn test_arrow_schema_to_schema_should_reject_uuid_when_fixed_size_binary_that_is_less_than_16() + { + let field = simple_field_with_uuid_ext(DataType::FixedSizeBinary(8)); + let error = arrow_schema_to_schema(&ArrowSchema::new(vec![field])).unwrap_err(); + + pretty_assertions::assert_eq!( + "arrow.uuid extension requires FixedSizeBinary(16) storage, found FixedSizeBinary(8)", + error.message() + ); + } + + #[test] + fn test_arrow_schema_to_schema_should_reject_uuid_when_fixed_size_binary_that_is_more_than_16() + { + let field = simple_field_with_uuid_ext(DataType::FixedSizeBinary(17)); + let error = arrow_schema_to_schema(&ArrowSchema::new(vec![field])).unwrap_err(); + + pretty_assertions::assert_eq!( + "arrow.uuid extension requires FixedSizeBinary(16) storage, found FixedSizeBinary(17)", + error.message() + ); + } + + fn simple_field_with_uuid_ext(field_data_type: DataType) -> Field { + let mut field = simple_field("uuid_field", DataType::FixedSizeBinary(16), false, "1") + .with_extension_type(UuidExtensionType); + + // We want to allow setting a UUID extension to the wrong data type. + // This is to test our logic against bad schemas. + // + // To avoid Arrow's own internal validation and panic, we need to build + // the field correctly above, and then manually override the type here. + field = field.with_data_type(field_data_type); + + field + } + fn arrow_schema_for_schema_to_arrow_schema_test() -> ArrowSchema { let fields = Fields::from(vec![ simple_field("key", DataType::Int32, false, "28"), @@ -1787,7 +1929,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(UuidExtensionType), Field::new( "v", DataType::Struct(Fields::from(vec![ @@ -2003,10 +2146,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] @@ -2052,6 +2195,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 @@ -2490,6 +2654,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(UuidExtensionType), Field::new( "created_at", DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())), @@ -2541,9 +2707,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(), @@ -2557,14 +2723,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, ) @@ -2573,29 +2740,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, ) @@ -2604,20 +2771,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), ) @@ -2634,7 +2801,7 @@ mod tests { .unwrap(); pretty_assertions::assert_eq!(schema, expected); - assert_eq!(schema.highest_field_id(), 17); + assert_eq!(schema.highest_field_id(), 18); } fn decimal_byte_array_stats(min: Vec<u8>, max: Vec<u8>) -> Statistics { diff --git a/crates/iceberg/src/writer/file_writer/parquet_writer.rs b/crates/iceberg/src/writer/file_writer/parquet_writer.rs index 7cdc48888..a471aea5c 100644 --- a/crates/iceberg/src/writer/file_writer/parquet_writer.rs +++ b/crates/iceberg/src/writer/file_writer/parquet_writer.rs @@ -2497,6 +2497,79 @@ mod tests { assert_eq!(std::fs::read_dir(temp_dir.path()).unwrap().count(), 0); } + #[tokio::test] + async fn test_parquet_writer_uuid_logical_type() { + let temp_dir = TempDir::new().unwrap(); + let file_io = FileIO::new_with_fs(); + let location_gen = DefaultLocationGenerator::with_data_location( + temp_dir.path().to_str().unwrap().to_string(), + ); + let file_name_gen = + DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet); + let output_file = file_io + .new_output(location_gen.generate_location(None, &file_name_gen.generate_file_name())) + .unwrap(); + + let schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required(0, "uuid", Type::Primitive(PrimitiveType::Uuid)).into(), + ]) + .build() + .unwrap(), + ); + let arrow_schema: ArrowSchemaRef = Arc::new(schema_to_arrow_schema(&schema).unwrap()); + let col = Arc::new( + arrow_array::FixedSizeBinaryArray::try_from_iter( + [Uuid::from_u128(0), Uuid::from_u128(1)] + .into_iter() + .map(|uuid| uuid.into_bytes()), + ) + .unwrap(), + ) as ArrayRef; + let to_write = RecordBatch::try_new(arrow_schema, vec![col]).unwrap(); + + let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema) + .build(output_file) + .await + .unwrap(); + pw.write(&to_write).await.unwrap(); + let data_file = pw + .close() + .await + .unwrap() + .into_iter() + .next() + .unwrap() + .content(DataContentType::Data) + .partition(Struct::empty()) + .partition_spec_id(0) + .build() + .unwrap(); + + let raw = file_io + .new_input(data_file.file_path.clone()) + .unwrap() + .read() + .await + .unwrap(); + let reader = ParquetRecordBatchReaderBuilder::try_new(raw).unwrap(); + let column = reader.metadata().file_metadata().schema_descr().column(0); + assert_eq!( + ( + column.physical_type(), + column.type_length(), + column.logical_type_ref() + ), + ( + parquet::basic::Type::FIXED_LEN_BYTE_ARRAY, + 16, + Some(&parquet::basic::LogicalType::Uuid) + ) + ); + } + #[test] fn test_min_max_aggregator() { let schema = Arc::new(
