From a9d853af6849f623b4ea5597aa9484f8458d7dab Mon Sep 17 00:00:00 2001 From: osipovartem Date: Thu, 15 Oct 2026 13:49:31 +0300 Subject: [PATCH] test: cover nested field IDs in Parquet-backed lists and maps --- .../tests/nested_schema_evolution.rs | 313 ++++++++++++++++++ 1 file changed, 313 insertions(+) create mode 100644 datafusion_iceberg/tests/nested_schema_evolution.rs diff --git a/datafusion_iceberg/tests/nested_schema_evolution.rs b/datafusion_iceberg/tests/nested_schema_evolution.rs new file mode 100644 index 00000000..f506e7bf --- /dev/null +++ b/datafusion_iceberg/tests/nested_schema_evolution.rs @@ -0,0 +1,313 @@ +use std::sync::Arc; + +use datafusion::arrow::array::{ + ArrayRef, Int32Array, ListArray, MapArray, StringArray, StructArray, +}; +use datafusion::arrow::buffer::OffsetBuffer; +use datafusion::arrow::datatypes::{DataType, Field, Fields, Schema as ArrowSchema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::dataframe::DataFrameWriteOptions; +use datafusion::prelude::SessionContext; +use datafusion_iceberg::catalog::catalog::IcebergCatalog; +use iceberg_rust::catalog::commit::{CommitTable, TableRequirement, TableUpdate}; +use iceberg_rust::catalog::identifier::Identifier; +use iceberg_rust::catalog::tabular::Tabular; +use iceberg_rust::catalog::Catalog; +use iceberg_rust::object_store::ObjectStoreBuilder; +use iceberg_rust::spec::namespace::Namespace; +use iceberg_rust::spec::schema::Schema; +use iceberg_rust::spec::types::{ListType, MapType, PrimitiveType, StructField, StructType, Type}; +use iceberg_rust::table::Table; +use iceberg_sql_catalog::SqlCatalog; + +fn nested_fields(old_names: bool, extra_id: i32) -> StructType { + let mut fields = vec![ + StructField::new( + extra_id - 2, + if old_names { "old" } else { "renamed" }, + true, + Type::Primitive(PrimitiveType::Int), + None, + ), + StructField::new( + extra_id - 1, + if old_names { "other" } else { "other_renamed" }, + true, + Type::Primitive(PrimitiveType::Int), + None, + ), + ]; + if !old_names { + fields.reverse(); + fields.push(StructField::new( + extra_id, + "added", + false, + Type::Primitive(PrimitiveType::Int), + None, + )); + } + StructType::new(fields) +} + +fn iceberg_schema(old_names: bool) -> Schema { + let fields = StructType::new(vec![ + StructField::new(1, "id", true, Type::Primitive(PrimitiveType::Int), None), + StructField::new( + 2, + "items", + true, + Type::List(ListType { + element_id: 3, + element_required: true, + element: Box::new(Type::Struct(nested_fields(old_names, 6))), + }), + None, + ), + StructField::new( + 7, + "attrs", + true, + Type::Map(MapType { + key_id: 8, + key: Box::new(Type::Primitive(PrimitiveType::String)), + value_id: 9, + value_required: true, + value: Box::new(Type::Struct(nested_fields(old_names, 12))), + }), + None, + ), + ]); + Schema::from_struct_type(fields, if old_names { 0 } else { 1 }, None) +} + +fn old_batch() -> RecordBatch { + let list_fields = Fields::from(vec![ + Field::new("old", DataType::Int32, false), + Field::new("other", DataType::Int32, false), + ]); + let list_values = StructArray::try_new( + list_fields.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![10, 20])), + ], + None, + ) + .unwrap(); + let element = Arc::new(Field::new("element", DataType::Struct(list_fields), false)); + let items = Arc::new( + ListArray::try_new( + Arc::clone(&element), + OffsetBuffer::from_lengths([2]), + Arc::new(list_values), + None, + ) + .unwrap(), + ) as ArrayRef; + + let map_value_fields = Fields::from(vec![ + Field::new("old", DataType::Int32, false), + Field::new("other", DataType::Int32, false), + ]); + let map_values = StructArray::try_new( + map_value_fields.clone(), + vec![ + Arc::new(Int32Array::from(vec![3])), + Arc::new(Int32Array::from(vec![30])), + ], + None, + ) + .unwrap(); + let entries_fields = Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Struct(map_value_fields), false), + ]); + let entries = StructArray::try_new( + entries_fields.clone(), + vec![Arc::new(StringArray::from(vec!["k"])), Arc::new(map_values)], + None, + ) + .unwrap(); + let entry = Arc::new(Field::new( + "entries", + DataType::Struct(entries_fields), + false, + )); + let attrs = Arc::new( + MapArray::try_new( + Arc::clone(&entry), + OffsetBuffer::from_lengths([1]), + entries, + None, + false, + ) + .unwrap(), + ) as ArrayRef; + + RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("items", DataType::List(element), false), + Field::new("attrs", DataType::Map(entry, false), false), + ])), + vec![ + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + items, + attrs, + ], + ) + .unwrap() +} + +#[tokio::test] +async fn parquet_backed_list_and_map_follow_nested_field_ids() { + let catalog: Arc = Arc::new( + SqlCatalog::new("sqlite://", "test", ObjectStoreBuilder::memory()) + .await + .unwrap(), + ); + catalog + .create_namespace(&Namespace::try_new(&["public".to_owned()]).unwrap(), None) + .await + .unwrap(); + let identifier = Identifier::new(&["public".to_owned()], "nested_evolution"); + Table::builder() + .with_name("nested_evolution") + .with_location("/nested_evolution") + .with_schema(iceberg_schema(true)) + .build(identifier.namespace(), Arc::clone(&catalog)) + .await + .unwrap(); + + let ctx = SessionContext::new(); + ctx.register_catalog( + "warehouse", + Arc::new( + IcebergCatalog::new(Arc::clone(&catalog), None) + .await + .unwrap(), + ), + ); + ctx.read_batch(old_batch()) + .unwrap() + .write_table( + "warehouse.public.nested_evolution", + DataFrameWriteOptions::default(), + ) + .await + .unwrap(); + + let Tabular::Table(table) = Arc::clone(&catalog) + .load_tabular(&identifier) + .await + .unwrap() + else { + panic!("expected table"); + }; + catalog + .update_table(CommitTable { + identifier, + requirements: vec![TableRequirement::AssertCurrentSchemaId { + current_schema_id: *table.current_schema().unwrap().schema_id(), + }], + updates: vec![ + TableUpdate::AddSchema { + schema: iceberg_schema(false), + last_column_id: Some(12), + }, + TableUpdate::SetCurrentSchema { schema_id: -1 }, + ], + }) + .await + .unwrap(); + + let batches = ctx + .sql("SELECT items, attrs FROM warehouse.public.nested_evolution") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 1); + let items = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(items.value_offsets(), &[0, 2]); + let list_values = items.values(); + let list_values = list_values.as_any().downcast_ref::().unwrap(); + assert_eq!( + list_values.column_names(), + vec!["other_renamed", "renamed", "added"] + ); + assert_eq!( + list_values + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ref(), + &[10, 20] + ); + assert_eq!( + list_values + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ref(), + &[1, 2] + ); + assert_eq!(list_values.column(2).null_count(), 2); + + let attrs = batches[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(attrs.value_offsets(), &[0, 1]); + assert_eq!( + attrs + .entries() + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + "k" + ); + let map_values = attrs + .entries() + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + map_values.column_names(), + vec!["other_renamed", "renamed", "added"] + ); + assert_eq!( + map_values + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 30 + ); + assert_eq!( + map_values + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 3 + ); + assert_eq!(map_values.column(2).null_count(), 1); +}