From 5576ec3b901d08bc258bb105e0ddc1879aa2055e Mon Sep 17 00:00:00 2001 From: Haresh Khanna Date: Thu, 8 Oct 2026 19:52:10 +0100 Subject: [PATCH 1/5] Use avro for datafiles collection from write-exec --- AGENTS.md | 2 +- crates/datafusion/src/physical_plan/commit.rs | 98 +++++++------ crates/datafusion/src/physical_plan/write.rs | 134 +++++++++++++----- 3 files changed, 156 insertions(+), 78 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index bc1fa61..7bf173e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -82,7 +82,7 @@ input → project_with_partition partitioned tables only: appends a `_partition` struct column computed by `PartitionExpr` → repartition hash on `_partition` for identity/bucket transforms, otherwise round-robin; uses session target_partitions → sort_by_partition only when table property `write.datafusion.fanout.enabled` is false (default: true) - → IcebergWriteExec per partition: `TaskWriter` writes Parquet files and emits JSON-serialized DataFiles in a `data_files` column + → IcebergWriteExec per partition: `TaskWriter` writes Parquet files and emits Avro-serialized DataFiles in a `data_files: LargeBinary` column → CoalescePartitionsExec → IcebergCommitExec single partition: one `fast_append` transaction; outputs a `count` row ``` diff --git a/crates/datafusion/src/physical_plan/commit.rs b/crates/datafusion/src/physical_plan/commit.rs index 468bf5a..78bcb1f 100644 --- a/crates/datafusion/src/physical_plan/commit.rs +++ b/crates/datafusion/src/physical_plan/commit.rs @@ -18,7 +18,9 @@ use std::fmt::{Debug, Formatter}; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, RecordBatch, StringArray, UInt64Array}; +use datafusion::arrow::array::{ + Array, ArrayRef, LargeBinaryArray, RecordBatch, UInt64Array, +}; use datafusion::arrow::datatypes::{ DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, }; @@ -35,7 +37,7 @@ use datafusion::physical_plan::{ }; use futures::StreamExt; use iceberg::Catalog; -use iceberg::spec::{DataFile, deserialize_data_file_from_json}; +use iceberg::spec::{DataFile, read_data_files_from_avro}; use iceberg::table::Table; use iceberg::transaction::{ApplyTransactionAction, Transaction}; @@ -237,6 +239,7 @@ impl ExecutionPlan for IcebergCommitExec { let spec_id = self.table.metadata().default_partition_spec_id(); let partition_type = self.table.metadata().default_partition_type().clone(); let current_schema = self.table.metadata().current_schema().clone(); + let format_version = self.table.metadata().format_version(); let catalog = Arc::clone(&self.catalog); @@ -259,35 +262,39 @@ impl ExecutionPlan for IcebergCommitExec { ) })? .as_any() - .downcast_ref::() + .downcast_ref::() .ok_or_else(|| { internal_datafusion_err!( - "Expected 'data_files' column to be StringArray" + "Expected 'data_files' column to be LargeBinaryArray" ) })?; - // Deserialize all data files from the StringArray - let batch_files: Vec = files_array - .into_iter() - .flatten() - .map(|f| -> Result { - // Parse JSON to DataFileSerde and convert to DataFile - deserialize_data_file_from_json( - f, - spec_id, - &partition_type, - ¤t_schema, + // Each value is a separate Avro container with one or more data files. + files_array.iter().try_for_each(|payload| { + let mut payload = payload.ok_or_else(|| { + DataFusionError::Execution( + "Null Avro data file payload".to_string(), ) - .map_err(to_datafusion_error) - }) - .collect::>()?; - - // add record_counts from the current batch to total record count - total_record_count += - batch_files.iter().map(|f| f.record_count()).sum::(); - - // Add all deserialized files to our collection - data_files.extend(batch_files); + })?; + if payload.is_empty() { + return Err(DataFusionError::Execution( + "Empty Avro data file payload".to_string(), + )); + } + let files = read_data_files_from_avro( + &mut payload, + ¤t_schema, + spec_id, + &partition_type, + format_version, + ) + .map_err(to_datafusion_error)?; + + total_record_count += + files.iter().map(|f| f.record_count()).sum::(); + data_files.extend(files); + Ok(()) + })?; } // If no data files were collected, return a single row with count = 0 @@ -343,7 +350,7 @@ mod tests { use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{ DataContentType, DataFileBuilder, DataFileFormat, NestedField, PrimitiveType, - Schema, Struct, Type, + Schema, Struct, Type, write_data_files_to_avro, }; use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation, TableIdent}; @@ -355,15 +362,15 @@ mod tests { #[derive(Debug)] struct MockWriteExec { schema: Arc, - data_files_json: Vec, + data_files_avro: Vec>, plan_properties: Arc, } impl MockWriteExec { - fn new(data_files_json: Vec) -> Self { + fn new(data_files_avro: Vec>) -> Self { let schema = Arc::new(ArrowSchema::new(vec![Field::new( DATA_FILES_COL_NAME, - DataType::Utf8, + DataType::LargeBinary, false, )])); @@ -376,7 +383,7 @@ mod tests { Self { schema, - data_files_json, + data_files_avro, plan_properties, } } @@ -419,8 +426,9 @@ mod tests { _context: Arc, ) -> Result { // Create a record batch with the serialized data files - let array = - Arc::new(StringArray::from(self.data_files_json.clone())) as ArrayRef; + let array = Arc::new(LargeBinaryArray::from_iter_values( + self.data_files_avro.iter(), + )) as ArrayRef; let batch = RecordBatch::try_new(self.schema.clone(), vec![array])?; // Create a stream that returns this batch @@ -439,7 +447,7 @@ mod tests { DisplayFormatType::Default | DisplayFormatType::Verbose | DisplayFormatType::TreeRender => { - write!(f, "MockDataFilesExec: files={}", self.data_files_json.len()) + write!(f, "MockDataFilesExec: files={}", self.data_files_avro.len()) } } } @@ -507,23 +515,27 @@ mod tests { .partition(Struct::empty()) .build()?; - // Serialize data files to JSON + // Serialize each task's files into a separate Avro container let partition_type = table.metadata().default_partition_type().clone(); - let data_file1_json = iceberg::spec::serialize_data_file_to_json( - data_file1.clone(), + let mut data_file1_avro = Vec::new(); + write_data_files_to_avro( + &mut data_file1_avro, + [data_file1.clone()], &partition_type, table.metadata().format_version(), )?; - let data_file2_json = iceberg::spec::serialize_data_file_to_json( - data_file2.clone(), + let mut data_file2_avro = Vec::new(); + write_data_files_to_avro( + &mut data_file2_avro, + [data_file2.clone()], &partition_type, table.metadata().format_version(), )?; // Create a mock execution plan that returns the serialized data files let input_exec = - Arc::new(MockWriteExec::new(vec![data_file1_json, data_file2_json])); + Arc::new(MockWriteExec::new(vec![data_file1_avro, data_file2_avro])); let commit_exec = IcebergCommitExec::new(table.clone(), catalog.clone(), input_exec); @@ -630,12 +642,14 @@ mod tests { .partition_spec_id(table.metadata().default_partition_spec_id()) .partition(Struct::empty()) .build()?; - let json = iceberg::spec::serialize_data_file_to_json( - data_file, + let mut payload = Vec::new(); + write_data_files_to_avro( + &mut payload, + [data_file], &partition_type, table.metadata().format_version(), )?; - partitions.push(Arc::new(MockWriteExec::new(vec![json]))); + partitions.push(Arc::new(MockWriteExec::new(vec![payload]))); } let input = UnionExec::try_new(partitions)?; assert_eq!(input.properties().partitioning.partition_count(), 2); diff --git a/crates/datafusion/src/physical_plan/write.rs b/crates/datafusion/src/physical_plan/write.rs index 86f364c..243f056 100644 --- a/crates/datafusion/src/physical_plan/write.rs +++ b/crates/datafusion/src/physical_plan/write.rs @@ -19,7 +19,7 @@ use std::fmt::{Debug, Formatter}; use std::str::FromStr; use std::sync::Arc; -use datafusion::arrow::array::{ArrayRef, RecordBatch, StringArray}; +use datafusion::arrow::array::{ArrayRef, LargeBinaryArray, RecordBatch}; use datafusion::arrow::datatypes::{ DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, }; @@ -36,7 +36,9 @@ use datafusion::physical_plan::{ }; use futures::StreamExt; use iceberg::arrow::FieldMatchMode; -use iceberg::spec::{DataFileFormat, serialize_data_file_to_json}; +use iceberg::spec::{ + DataFile, DataFileFormat, FormatVersion, StructType, write_data_files_to_avro, +}; use iceberg::table::Table; use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder; use iceberg::writer::file_writer::ParquetWriterBuilder; @@ -103,9 +105,21 @@ impl IcebergWriteExec { )) } - // Create a record batch with serialized data files - fn make_result_batch(data_files: Vec) -> Result { - let files_array = Arc::new(StringArray::from(data_files)) as ArrayRef; + // Each row holds one Avro container, amortizing its schema over all files from a task. + fn make_result_batch( + data_files: Vec, + partition_type: &StructType, + format_version: FormatVersion, + ) -> Result { + if data_files.is_empty() { + return Ok(RecordBatch::new_empty(Self::make_result_schema())); + } + + let mut buffer = Vec::new(); + write_data_files_to_avro(&mut buffer, data_files, partition_type, format_version) + .map_err(to_datafusion_error)?; + let files_array = + Arc::new(LargeBinaryArray::from_vec(vec![buffer.as_slice()])) as ArrayRef; RecordBatch::try_new(Self::make_result_schema(), vec![files_array]).map_err(|e| { DataFusionError::ArrowError( @@ -119,7 +133,7 @@ impl IcebergWriteExec { // Define a schema. Arc::new(ArrowSchema::new(vec![Field::new( DATA_FILES_COL_NAME, - DataType::Utf8, + DataType::LargeBinary, false, )])) } @@ -207,11 +221,12 @@ impl ExecutionPlan for IcebergWriteExec { /// +------------------+ /// | data_files | /// +------------------+ - /// | "{"file_path":.. | <- JSON string representing a data file + /// | Avro bytes | <- Avro container holding the task's data files /// +------------------+ /// ``` /// - /// Each row in the output contains a JSON string representing a data file that was written. + /// Each row contains an Avro container with the data files written by this task. + /// Tasks that write no files return an empty batch. /// /// This output can be used by a subsequent operation to commit the added files to the table. fn execute( @@ -304,20 +319,7 @@ impl ExecutionPlan for IcebergWriteExec { let data_files = task_writer.close().await?; - // Convert builders to data files and then to JSON strings - let data_files_strs: Vec = data_files - .into_iter() - .map(|data_file| { - serialize_data_file_to_json( - data_file, - &partition_type, - format_version, - ) - .map_err(to_datafusion_error) - }) - .collect::>>()?; - - Self::make_result_batch(data_files_strs) + Self::make_result_batch(data_files, &partition_type, format_version) }) .boxed(); @@ -349,8 +351,8 @@ mod tests { use futures::{StreamExt, stream}; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{ - DataFileFormat, NestedField, PrimitiveType, Schema, Type, - deserialize_data_file_from_json, + DataContentType, DataFileBuilder, DataFileFormat, Literal, NestedField, + PrimitiveType, Schema, Struct, Type, read_data_files_from_avro, }; use iceberg::{ Catalog, CatalogBuilder, Error, ErrorKind, MemoryCatalog, NamespaceIdent, @@ -488,6 +490,60 @@ mod tests { .build() } + #[test] + fn test_float_partition_metadata_roundtrip() { + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "f", Type::Primitive(PrimitiveType::Float)) + .into(), + ]) + .build() + .unwrap(); + let files: Vec<_> = [ + Some(123.456_f32), + Some(f32::NAN), + Some(f32::INFINITY), + Some(f32::NEG_INFINITY), + None, + ] + .into_iter() + .enumerate() + .map(|(i, value)| { + DataFileBuilder::default() + .content(DataContentType::Data) + .file_path(format!("memory://data/{i}.parquet")) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(1024) + .record_count(1) + .partition(Struct::from_iter([value.map(Literal::float)])) + .build() + .unwrap() + }) + .collect(); + let batch = IcebergWriteExec::make_result_batch( + files.clone(), + schema.as_struct(), + FormatVersion::V2, + ) + .unwrap(); + assert_eq!(batch.num_rows(), 1); + let mut payload = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + let decoded = read_data_files_from_avro( + &mut payload, + &schema, + 0, + schema.as_struct(), + FormatVersion::V2, + ) + .unwrap(); + assert_eq!(decoded, files); + } + #[tokio::test] async fn test_iceberg_write_exec() -> Result<(), Box> { // 1. Set up test environment @@ -568,7 +624,7 @@ mod tests { result_batch.schema().as_ref(), &ArrowSchema::new(vec![Field::new( DATA_FILES_COL_NAME, - DataType::Utf8, + DataType::LargeBinary, false )]) ); @@ -576,26 +632,29 @@ mod tests { // Check data assert_eq!(result_batch.num_rows(), 1, "Expected one data file"); - // Get the data file JSON - let data_file_json = result_batch + // Get the Avro container + let mut data_files_avro = result_batch .column(0) .as_any() - .downcast_ref::() - .expect("Expected StringArray") + .downcast_ref::() + .expect("Expected LargeBinaryArray") .value(0); - // Deserialize the data file JSON + // Deserialize the Avro container let partition_type = table.metadata().default_partition_type(); let spec_id = table.metadata().default_partition_spec_id(); let schema = table.metadata().current_schema(); - let data_file = deserialize_data_file_from_json( - data_file_json, + let data_files = read_data_files_from_avro( + &mut data_files_avro, + schema, spec_id, partition_type, - schema, + table.metadata().format_version(), ) - .expect("Failed to deserialize data file JSON"); + .expect("Failed to deserialize Avro data files"); + assert_eq!(data_files.len(), 1); + let data_file = &data_files[0]; // Verify data file properties assert_eq!( @@ -680,12 +739,17 @@ mod tests { write_exec.schema().as_ref(), &ArrowSchema::new(vec![Field::new( DATA_FILES_COL_NAME, - DataType::Utf8, + DataType::LargeBinary, false )]), "IcebergWriteExec should advertise the data_files schema, not the table schema" ); + let mut stream = write_exec.execute(0, Arc::new(TaskContext::default()))?; + let batch = stream.next().await.unwrap()?; + assert_eq!(batch.num_rows(), 0); + assert_eq!(batch.schema(), write_exec.schema()); + Ok(()) } } From 9aa7f53e4862e56dad4d3033861f20a38ac6f114 Mon Sep 17 00:00:00 2001 From: Haresh Khanna Date: Fri, 9 Oct 2026 10:44:15 +0100 Subject: [PATCH 2/5] Remove avro buffer re-alloc --- crates/datafusion/src/physical_plan/write.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/crates/datafusion/src/physical_plan/write.rs b/crates/datafusion/src/physical_plan/write.rs index 243f056..ad84e37 100644 --- a/crates/datafusion/src/physical_plan/write.rs +++ b/crates/datafusion/src/physical_plan/write.rs @@ -20,6 +20,7 @@ use std::str::FromStr; use std::sync::Arc; use datafusion::arrow::array::{ArrayRef, LargeBinaryArray, RecordBatch}; +use datafusion::arrow::buffer::{Buffer, OffsetBuffer}; use datafusion::arrow::datatypes::{ DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, }; @@ -118,8 +119,11 @@ impl IcebergWriteExec { let mut buffer = Vec::new(); write_data_files_to_avro(&mut buffer, data_files, partition_type, format_version) .map_err(to_datafusion_error)?; - let files_array = - Arc::new(LargeBinaryArray::from_vec(vec![buffer.as_slice()])) as ArrayRef; + let files_array = Arc::new(LargeBinaryArray::new( + OffsetBuffer::from_lengths([buffer.len()]), + Buffer::from_vec(buffer), + None, + )) as ArrayRef; RecordBatch::try_new(Self::make_result_schema(), vec![files_array]).map_err(|e| { DataFusionError::ArrowError( From e9dc3f4a1054f9e1b7c053c215fb71af04f24d25 Mon Sep 17 00:00:00 2001 From: Haresh Khanna Date: Fri, 9 Oct 2026 11:03:06 +0100 Subject: [PATCH 3/5] Adds an integration test --- .../tests/integration_datafusion_test.rs | 66 +++++++++++++++---- 1 file changed, 53 insertions(+), 13 deletions(-) diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 76e379c..6edd38f 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -29,7 +29,7 @@ use datafusion::arrow::compute::{ use datafusion::arrow::datatypes::{DataType, Field, Int64Type, Schema as ArrowSchema}; use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::datasource::TableProvider; -use datafusion::execution::context::SessionContext; +use datafusion::execution::context::{SessionConfig, SessionContext}; use datafusion::parquet::arrow::PARQUET_FIELD_ID_META_KEY; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; @@ -860,7 +860,7 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { let namespace = NamespaceIdent::new("test_partitioned_write".to_string()); set_test_namespace(&iceberg_catalog, &namespace).await?; - // Create a schema with a partition column + // Create a schema with string and nullable float partition columns let schema = Schema::builder() .with_schema_id(0) .with_fields(vec![ @@ -869,13 +869,20 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { .into(), NestedField::required(3, "value", Type::Primitive(PrimitiveType::String)) .into(), + NestedField::optional( + 4, + "float_partition", + Type::Primitive(PrimitiveType::Float), + ) + .into(), ]) .build()?; - // Create partition spec with identity transform on category + // Identity-partition on both category and the float value let partition_spec = UnboundPartitionSpec::builder() .with_spec_id(0) .add_partition_field(2, "category", Transform::Identity)? + .add_partition_field(4, "float_partition", Transform::Identity)? .build(); // Create the partitioned table @@ -892,7 +899,9 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { let client = Arc::new(iceberg_catalog); let catalog = Arc::new(IcebergCatalogProvider::try_new(client.clone()).await?); - let ctx = SessionContext::new(); + // One task must return all five partition files in the same Avro payload. + let ctx = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); ctx.register_catalog("catalog", catalog); // Insert data with multiple partition values in a single batch @@ -901,11 +910,12 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { r#" INSERT INTO catalog.test_partitioned_write.partitioned_table VALUES - (1, 'electronics', 'laptop'), - (2, 'electronics', 'phone'), - (3, 'books', 'novel'), - (4, 'books', 'textbook'), - (5, 'clothing', 'shirt') + (1, 'electronics', 'laptop', CAST(123.456 AS REAL)), + (2, 'electronics', 'phone', CAST(123.456 AS REAL)), + (3, 'books', 'novel', CAST('NaN' AS REAL)), + (4, 'books', 'textbook', CAST('Infinity' AS REAL)), + (5, 'clothing', 'shirt', CAST('-Infinity' AS REAL)), + (6, 'clothing', 'coat', CAST(NULL AS REAL)) "#, ) .await @@ -919,9 +929,9 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { .as_any() .downcast_ref::() .unwrap(); - assert_eq!(rows_inserted.value(0), 5); + assert_eq!(rows_inserted.value(0), 6); - // Query the table to verify data + // Read all rows without a filter, including NaN, infinities, and null. let df = ctx .sql("SELECT * FROM catalog.test_partitioned_write.partitioned_table ORDER BY id") .await @@ -935,7 +945,8 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { expect![[r#" Field { "id": Int32, metadata: {"PARQUET:field_id": "1"} }, Field { "category": Utf8, metadata: {"PARQUET:field_id": "2"} }, - Field { "value": Utf8, metadata: {"PARQUET:field_id": "3"} }"#]], + Field { "value": Utf8, metadata: {"PARQUET:field_id": "3"} }, + Field { "float_partition": nullable Float32, metadata: {"PARQUET:field_id": "4"} }"#]], expect![[r#" id: PrimitiveArray [ @@ -944,6 +955,7 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { 3, 4, 5, + 6, ], category: StringArray [ @@ -952,6 +964,7 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { "books", "books", "clothing", + "clothing", ], value: StringArray [ @@ -960,14 +973,41 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { "novel", "textbook", "shirt", + "coat", + ], + float_partition: PrimitiveArray + [ + 123.456, + 123.456, + NaN, + inf, + -inf, + null, ]"#]], &[], Some("id"), ); - // Verify that data files exist under correct partition paths + // Verify all five partition files and six rows were committed. let table_ident = TableIdent::new(namespace.clone(), "partitioned_table".to_string()); let table = client.load_table(&table_ident).await?; + + let snapshot = table.metadata().current_snapshot().unwrap(); + let manifest_list = table.manifest_list_reader(snapshot).load().await?; + let mut file_count = 0; + let mut row_count = 0; + for entry in manifest_list.entries() { + let manifest = table.manifest_reader().read(entry).await?; + file_count += manifest.entries().len(); + row_count += manifest + .entries() + .iter() + .map(|entry| entry.data_file().record_count()) + .sum::(); + } + assert_eq!(file_count, 5); + assert_eq!(row_count, 6); + let table_location = table.metadata().location(); let file_io = table.file_io(); From ee3f7728a83b31f0e22cc50b2145ceaa1b47a536 Mon Sep 17 00:00:00 2001 From: Haresh Khanna Date: Mon, 12 Oct 2026 13:07:58 +0100 Subject: [PATCH 4/5] Bump iceberg-rust --- Cargo.lock | 110 +++++++++++++++--- Cargo.toml | 4 +- .../datafusion/src/physical_plan/project.rs | 20 ++-- .../src/physical_plan/repartition.rs | 19 ++- crates/datafusion/src/table/mod.rs | 11 +- .../tests/integration_datafusion_test.rs | 65 ++++++++++- crates/sqllogictest/src/engine/datafusion.rs | 10 +- 7 files changed, 192 insertions(+), 47 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index afb81d2..a02472c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -66,6 +66,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "aliasable" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "250f629c0161ad8107cf89319e990051fae62832fd343083bea452d93e2205fd" + [[package]] name = "alloc-no-stdlib" version = "2.0.4" @@ -154,26 +160,26 @@ checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" [[package]] name = "apache-avro" -version = "0.21.0" +version = "0.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36fa98bc79671c7981272d91a8753a928ff6a1cd8e4f20a44c45bd5d313840bf" +checksum = "312c1ea69e5fe9966e0029fb95aca8790100b85aff4f0d3b00a9337c74069a9c" dependencies = [ "bigdecimal", "bon", "crc32fast", - "digest 0.10.7", + "digest 0.11.3", "log", - "miniz_oxide", + "miniz_oxide 0.9.1", "num-bigint 0.4.8", + "ouroboros", "quad-rand", - "rand 0.9.5", + "rand 0.10.2", "regex-lite", "serde", "serde_bytes", "serde_json", "snap", - "strum", - "strum_macros 0.27.2", + "strum 0.28.0", "thiserror", "uuid", "zstd", @@ -1232,7 +1238,7 @@ version = "4.6.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f9c751b79415d4e559e3d1fcf128e09e720eb673a06d26cf6f392d37d75b66e0" dependencies = [ - "heck", + "heck 0.5.0", "proc-macro2", "quote", "syn 3.0.3", @@ -2712,7 +2718,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ "crc32fast", - "miniz_oxide", + "miniz_oxide 0.8.9", "zlib-rs", ] @@ -2986,6 +2992,12 @@ dependencies = [ "foldhash 0.2.0", ] +[[package]] +name = "heck" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" + [[package]] name = "heck" version = "0.5.0" @@ -3179,7 +3191,7 @@ dependencies = [ [[package]] name = "iceberg" version = "0.10.1" -source = "git+https://github.com/apache/iceberg-rust?rev=8cb2adeddb4f8da1ca7bd86ca303337f011c12e2#8cb2adeddb4f8da1ca7bd86ca303337f011c12e2" +source = "git+https://github.com/apache/iceberg-rust?rev=5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3#5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3" dependencies = [ "aes-gcm", "anyhow", @@ -3202,7 +3214,6 @@ dependencies = [ "chrono", "crc32fast", "derive_builder", - "expect-test", "fastnum", "flate2", "fnv", @@ -3224,7 +3235,7 @@ dependencies = [ "serde_json", "serde_repr", "serde_with", - "strum", + "strum 0.27.2", "tokio", "tracing", "typed-builder", @@ -3238,7 +3249,7 @@ dependencies = [ [[package]] name = "iceberg-catalog-rest" version = "0.10.1" -source = "git+https://github.com/apache/iceberg-rust?rev=8cb2adeddb4f8da1ca7bd86ca303337f011c12e2#8cb2adeddb4f8da1ca7bd86ca303337f011c12e2" +source = "git+https://github.com/apache/iceberg-rust?rev=5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3#5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3" dependencies = [ "async-trait", "chrono", @@ -3280,7 +3291,7 @@ dependencies = [ [[package]] name = "iceberg-property-macro" version = "0.10.1" -source = "git+https://github.com/apache/iceberg-rust?rev=8cb2adeddb4f8da1ca7bd86ca303337f011c12e2#8cb2adeddb4f8da1ca7bd86ca303337f011c12e2" +source = "git+https://github.com/apache/iceberg-rust?rev=5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3#5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3" dependencies = [ "proc-macro2", "quote", @@ -3810,6 +3821,15 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "miniz_oxide" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" +dependencies = [ + "adler2", +] + [[package]] name = "mio" version = "1.2.2" @@ -4019,6 +4039,30 @@ dependencies = [ "num-traits", ] +[[package]] +name = "ouroboros" +version = "0.18.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e0f050db9c44b97a94723127e6be766ac5c340c48f2c4bb3ffa11713744be59" +dependencies = [ + "aliasable", + "ouroboros_macro", + "static_assertions", +] + +[[package]] +name = "ouroboros_macro" +version = "0.18.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c7028bdd3d43083f6d8d4d5187680d0d3560d54df4cc9d752005268b41e64d0" +dependencies = [ + "heck 0.4.1", + "proc-macro2", + "proc-macro2-diagnostics", + "quote", + "syn 2.0.119", +] + [[package]] name = "outref" version = "0.5.2" @@ -4239,6 +4283,19 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "proc-macro2-diagnostics" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af066a9c399a26e020ada66a034357a868728e72cd426f3adcd35f80d88d88c8" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "version_check", + "yansi", +] + [[package]] name = "psm" version = "0.1.32" @@ -5095,6 +5152,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "strsim" version = "0.11.1" @@ -5110,13 +5173,22 @@ dependencies = [ "strum_macros 0.27.2", ] +[[package]] +name = "strum" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9628de9b8791db39ceda2b119bbe13134770b56c138ec1d3af810d045c04f9bd" +dependencies = [ + "strum_macros 0.28.0", +] + [[package]] name = "strum_macros" version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7695ce3845ea4b33927c055a39dc438a45b059f7c1b3d91d38d10355fb8cbca7" dependencies = [ - "heck", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.119", @@ -5128,7 +5200,7 @@ version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab85eea0270ee17587ed4156089e10b9e6880ee688791d45a905f5b1ca36f664" dependencies = [ - "heck", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.119", @@ -5982,6 +6054,12 @@ version = "0.13.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" +[[package]] +name = "yansi" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049" + [[package]] name = "yoke" version = "0.8.3" diff --git a/Cargo.toml b/Cargo.toml index f1933d3..12767c1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,8 +37,8 @@ datafusion = "55.0.0" datafusion-cli = "55.0.0" datafusion-iceberg = { path = "crates/datafusion" } datafusion-sqllogictest = "55.0.0" -iceberg = { git = "https://github.com/apache/iceberg-rust", rev = "8cb2adeddb4f8da1ca7bd86ca303337f011c12e2" } -iceberg-catalog-rest = { git = "https://github.com/apache/iceberg-rust", rev = "8cb2adeddb4f8da1ca7bd86ca303337f011c12e2" } +iceberg = { git = "https://github.com/apache/iceberg-rust", rev = "5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3" } +iceberg-catalog-rest = { git = "https://github.com/apache/iceberg-rust", rev = "5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3" } tokio = { version = "1.47", default-features = false, features = ["sync", "rt-multi-thread"] } toml = "1.1" tracing = "0.1.41" diff --git a/crates/datafusion/src/physical_plan/project.rs b/crates/datafusion/src/physical_plan/project.rs index 9e335c8..c8db587 100644 --- a/crates/datafusion/src/physical_plan/project.rs +++ b/crates/datafusion/src/physical_plan/project.rs @@ -258,10 +258,10 @@ mod tests { use datafusion::arrow::array::{ArrayRef, Int32Array, StructArray}; use datafusion::arrow::datatypes::{DataType, Field, Fields}; use datafusion::physical_plan::empty::EmptyExec; + use iceberg::Runtime; use iceberg::spec::{ NestedField, PrimitiveType, Schema, StructType, Transform, Type, }; - use iceberg::test_utils::test_runtime; use super::*; @@ -675,8 +675,8 @@ mod tests { assert_eq!(city_partition.value(1), "Los Angeles"); } - #[test] - fn test_schema_validation_matching_schemas() { + #[tokio::test] + async fn test_schema_validation_matching_schemas() { use iceberg::TableIdent; use iceberg::io::FileIO; use iceberg::spec::{FormatVersion, NestedField, PrimitiveType, Schema, Type}; @@ -732,7 +732,7 @@ mod tests { .identifier(TableIdent::from_strs(["test", "table"]).unwrap()) .file_io(FileIO::new_with_fs()) .metadata_location("/test/metadata.json") - .runtime(test_runtime()) + .runtime(Runtime::current()) .build() .unwrap(); @@ -740,8 +740,8 @@ mod tests { assert!(result.is_ok(), "Schema validation should pass"); } - #[test] - fn test_schema_validation_mismatched_schemas() { + #[tokio::test] + async fn test_schema_validation_mismatched_schemas() { use iceberg::TableIdent; use iceberg::io::FileIO; use iceberg::spec::{FormatVersion, NestedField, PrimitiveType, Schema, Type}; @@ -797,7 +797,7 @@ mod tests { .identifier(TableIdent::from_strs(["test", "table"]).unwrap()) .file_io(FileIO::new_with_fs()) .metadata_location("/test/metadata.json") - .runtime(test_runtime()) + .runtime(Runtime::current()) .build() .unwrap(); @@ -814,8 +814,8 @@ mod tests { ); } - #[test] - fn test_schema_validation_with_metadata_differences() { + #[tokio::test] + async fn test_schema_validation_with_metadata_differences() { use std::collections::HashMap; use iceberg::TableIdent; @@ -876,7 +876,7 @@ mod tests { .identifier(TableIdent::from_strs(["test", "table"]).unwrap()) .file_io(FileIO::new_with_fs()) .metadata_location("/test/metadata.json") - .runtime(test_runtime()) + .runtime(Runtime::current()) .build() .unwrap(); diff --git a/crates/datafusion/src/physical_plan/repartition.rs b/crates/datafusion/src/physical_plan/repartition.rs index 0dbda0c..e6fccff 100644 --- a/crates/datafusion/src/physical_plan/repartition.rs +++ b/crates/datafusion/src/physical_plan/repartition.rs @@ -177,14 +177,13 @@ mod tests { }; use datafusion::execution::TaskContext; use datafusion::physical_plan::empty::EmptyExec; - use iceberg::TableIdent; use iceberg::io::FileIO; use iceberg::spec::{ NestedField, NullOrder, PrimitiveType, Schema, SortDirection, SortField, SortOrder, Transform, Type, }; use iceberg::table::Table; - use iceberg::test_utils::test_runtime; + use iceberg::{Runtime, TableIdent}; use super::*; @@ -226,7 +225,7 @@ mod tests { .identifier(TableIdent::from_strs(["test", "table"]).unwrap()) .file_io(FileIO::new_with_fs()) .metadata_location("/test/metadata.json") - .runtime(test_runtime()) + .runtime(Runtime::current()) .build() .unwrap() } @@ -378,7 +377,7 @@ mod tests { .identifier(TableIdent::from_strs(["test", "bucketed_table"]).unwrap()) .file_io(FileIO::new_with_fs()) .metadata_location("/test/bucketed_metadata.json") - .runtime(test_runtime()) + .runtime(Runtime::current()) .build() .unwrap(); @@ -463,7 +462,7 @@ mod tests { ) .file_io(FileIO::new_with_fs()) .metadata_location("/test/partitioned_bucketed_metadata.json") - .runtime(test_runtime()) + .runtime(Runtime::current()) .build() .unwrap(); @@ -544,7 +543,7 @@ mod tests { .identifier(TableIdent::from_strs(["test", "none_table"]).unwrap()) .file_io(FileIO::new_with_fs()) .metadata_location("/test/none_metadata.json") - .runtime(test_runtime()) + .runtime(Runtime::current()) .build() .unwrap(); @@ -614,7 +613,7 @@ mod tests { .identifier(TableIdent::from_strs(["test", "range_only_table"]).unwrap()) .file_io(FileIO::new_with_fs()) .metadata_location("/test/range_only_metadata.json") - .runtime(test_runtime()) + .runtime(Runtime::current()) .build() .unwrap(); @@ -689,7 +688,7 @@ mod tests { ) .file_io(FileIO::new_with_fs()) .metadata_location("/test/mixed_transforms_metadata.json") - .runtime(test_runtime()) + .runtime(Runtime::current()) .build() .unwrap(); @@ -769,7 +768,7 @@ mod tests { .identifier(TableIdent::from_strs(["test", "temporal_partition"]).unwrap()) .file_io(FileIO::new_with_fs()) .metadata_location("/test/temporal_metadata.json") - .runtime(test_runtime()) + .runtime(Runtime::current()) .build() .unwrap(); @@ -840,7 +839,7 @@ mod tests { .identifier(TableIdent::from_strs(["test", "identity_partition"]).unwrap()) .file_io(FileIO::new_with_fs()) .metadata_location("/test/identity_metadata.json") - .runtime(test_runtime()) + .runtime(Runtime::current()) .build() .unwrap(); diff --git a/crates/datafusion/src/table/mod.rs b/crates/datafusion/src/table/mod.rs index dc20a4a..9c2296c 100644 --- a/crates/datafusion/src/table/mod.rs +++ b/crates/datafusion/src/table/mod.rs @@ -683,7 +683,7 @@ mod tests { async fn get_partitioned_test_catalog_and_table( fanout_enabled: Option, ) -> (Arc, NamespaceIdent, String, TempDir) { - use iceberg::spec::{Transform, UnboundPartitionSpec}; + use iceberg::spec::{Transform, UnboundPartitionField, UnboundPartitionSpec}; let temp_dir = TempDir::new().unwrap(); let warehouse_path = temp_dir.path().to_str().unwrap().to_string(); @@ -722,7 +722,14 @@ mod tests { let partition_spec = UnboundPartitionSpec::builder() .with_spec_id(0) - .add_partition_field(2, "category", Transform::Identity) + .add_partition_field( + UnboundPartitionField::builder() + .source_ids(vec![2]) + .name("category") + .transform(Transform::Identity) + .build() + .unwrap(), + ) .unwrap() .build(); diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 6edd38f..602fa1d 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -41,13 +41,13 @@ use datafusion_iceberg::{ IcebergCatalogProvider, IcebergMetadataTableProvider, IcebergStaticTableProvider, IcebergTableProvider, }; -use expect_test::expect; +use expect_test::{Expect, expect}; use iceberg::io::LocalFsStorageFactory; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{ - NestedField, PrimitiveType, Schema, StructType, Transform, Type, UnboundPartitionSpec, + NestedField, PrimitiveType, Schema, StructType, Transform, Type, + UnboundPartitionField, UnboundPartitionSpec, }; -use iceberg::test_utils::check_record_batches; use iceberg::transaction::{AddColumn, ApplyTransactionAction, Transaction}; use iceberg::{ Catalog, CatalogBuilder, MemoryCatalog, NamespaceIdent, Result as IcebergResult, @@ -55,6 +55,49 @@ use iceberg::{ }; use tempfile::TempDir; +/// Snapshot assertions over record batches, optionally sorted by one column. +fn check_record_batches( + record_batches: Vec, + expected_schema: Expect, + expected_data: Expect, + ignore_check_columns: &[&str], + sort_column: Option<&str>, +) { + let first_batch = record_batches.first().expect("Empty record batches"); + let mut batch = concat_batches(&first_batch.schema(), &record_batches).unwrap(); + if let Some(sort_column) = sort_column { + let column = batch.column_by_name(sort_column).unwrap(); + let indices = sort_to_indices(column, None, None).unwrap(); + batch = take_record_batch(&batch, &indices).unwrap(); + } + + expected_schema.assert_eq( + &batch + .schema() + .fields() + .iter() + .map(ToString::to_string) + .collect::>() + .join(",\n"), + ); + expected_data.assert_eq( + &batch + .schema() + .fields() + .iter() + .zip(batch.columns()) + .map(|(field, column)| { + if ignore_check_columns.contains(&field.name().as_str()) { + format!("{}: (skipped)", field.name()) + } else { + format!("{}: {column:?}", field.name()) + } + }) + .collect::>() + .join(",\n"), + ); +} + fn temp_path() -> String { let temp_dir = TempDir::new().unwrap(); temp_dir.path().to_str().unwrap().to_string() @@ -881,8 +924,20 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { // Identity-partition on both category and the float value let partition_spec = UnboundPartitionSpec::builder() .with_spec_id(0) - .add_partition_field(2, "category", Transform::Identity)? - .add_partition_field(4, "float_partition", Transform::Identity)? + .add_partition_field( + UnboundPartitionField::builder() + .source_ids(vec![2]) + .name("category") + .transform(Transform::Identity) + .build()?, + )? + .add_partition_field( + UnboundPartitionField::builder() + .source_ids(vec![4]) + .name("float_partition") + .transform(Transform::Identity) + .build()?, + )? .build(); // Create the partitioned table diff --git a/crates/sqllogictest/src/engine/datafusion.rs b/crates/sqllogictest/src/engine/datafusion.rs index e5fc045..0f51907 100644 --- a/crates/sqllogictest/src/engine/datafusion.rs +++ b/crates/sqllogictest/src/engine/datafusion.rs @@ -27,7 +27,7 @@ use iceberg::encryption::kms::MemoryKmsClientFactory; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{ NestedField, PrimitiveType, Schema, TableProperties, Transform, Type, - UnboundPartitionSpec, + UnboundPartitionField, UnboundPartitionSpec, }; use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation}; use indicatif::ProgressBar; @@ -137,7 +137,13 @@ impl DataFusionEngine { let partition_spec = UnboundPartitionSpec::builder() .with_spec_id(0) - .add_partition_field(2, "category", Transform::Identity)? + .add_partition_field( + UnboundPartitionField::builder() + .source_ids(vec![2]) + .name("category") + .transform(Transform::Identity) + .build()?, + )? .build(); catalog From fee7f0972971c5fe589fd70abd1076f7b46889ca Mon Sep 17 00:00:00 2001 From: Haresh Khanna Date: Mon, 12 Oct 2026 13:10:19 +0100 Subject: [PATCH 5/5] Use internal error and enhance tests --- crates/datafusion/src/physical_plan/commit.rs | 51 +++++++++++-- .../tests/integration_datafusion_test.rs | 76 ++++++++++++++----- 2 files changed, 103 insertions(+), 24 deletions(-) diff --git a/crates/datafusion/src/physical_plan/commit.rs b/crates/datafusion/src/physical_plan/commit.rs index 78bcb1f..1d30935 100644 --- a/crates/datafusion/src/physical_plan/commit.rs +++ b/crates/datafusion/src/physical_plan/commit.rs @@ -272,14 +272,10 @@ impl ExecutionPlan for IcebergCommitExec { // Each value is a separate Avro container with one or more data files. files_array.iter().try_for_each(|payload| { let mut payload = payload.ok_or_else(|| { - DataFusionError::Execution( - "Null Avro data file payload".to_string(), - ) + internal_datafusion_err!("Null Avro data file payload") })?; if payload.is_empty() { - return Err(DataFusionError::Execution( - "Empty Avro data file payload".to_string(), - )); + return internal_err!("Empty Avro data file payload"); } let files = read_data_files_from_avro( &mut payload, @@ -670,6 +666,49 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_iceberg_commit_exec_rejects_empty_payload() + -> Result<(), Box> { + let catalog = Arc::new( + MemoryCatalogBuilder::default() + .load( + "memory", + HashMap::from([( + MEMORY_CATALOG_WAREHOUSE.to_string(), + "memory://root".to_string(), + )]), + ) + .await?, + ); + let namespace = NamespaceIdent::new("test_empty_payload".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)) + .into(), + ]) + .build()?; + let creation = TableCreation::builder() + .name("test_table".to_string()) + .schema(schema) + .location("memory://root/test_table".to_string()) + .build(); + let table = catalog.create_table(&namespace, creation).await?; + let input = Arc::new(MockWriteExec::new(vec![vec![]])); + let commit_exec = IcebergCommitExec::new(table.clone(), catalog.clone(), input); + + let stream = commit_exec.execute(0, Arc::new(TaskContext::default()))?; + let err = collect(stream).await.unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(message) if message == "Empty Avro data file payload" + )); + let table = catalog.load_table(table.identifier()).await?; + assert!(table.metadata().current_snapshot().is_none()); + + Ok(()) + } + #[tokio::test] async fn test_iceberg_commit_exec_empty_insert() -> Result<(), Box> { diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 602fa1d..b1e4d89 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -17,7 +17,7 @@ //! Integration tests for Iceberg Datafusion with Hive Metastore. -use std::collections::{BTreeSet, HashMap}; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::error::Error; use std::sync::Arc; use std::vec; @@ -45,7 +45,7 @@ use expect_test::{Expect, expect}; use iceberg::io::LocalFsStorageFactory; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{ - NestedField, PrimitiveType, Schema, StructType, Transform, Type, + Literal, NestedField, PrimitiveType, Schema, Struct, StructType, Transform, Type, UnboundPartitionField, UnboundPartitionSpec, }; use iceberg::transaction::{AddColumn, ApplyTransactionAction, Transaction}; @@ -903,7 +903,7 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { let namespace = NamespaceIdent::new("test_partitioned_write".to_string()); set_test_namespace(&iceberg_catalog, &namespace).await?; - // Create a schema with string and nullable float partition columns + // Create a schema with string and nullable float/double partition columns let schema = Schema::builder() .with_schema_id(0) .with_fields(vec![ @@ -918,10 +918,16 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { Type::Primitive(PrimitiveType::Float), ) .into(), + NestedField::optional( + 5, + "double_partition", + Type::Primitive(PrimitiveType::Double), + ) + .into(), ]) .build()?; - // Identity-partition on both category and the float value + // Identity-partition on category and both floating-point values let partition_spec = UnboundPartitionSpec::builder() .with_spec_id(0) .add_partition_field( @@ -938,6 +944,13 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { .transform(Transform::Identity) .build()?, )? + .add_partition_field( + UnboundPartitionField::builder() + .source_ids(vec![5]) + .name("double_partition") + .transform(Transform::Identity) + .build()?, + )? .build(); // Create the partitioned table @@ -965,12 +978,12 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { r#" INSERT INTO catalog.test_partitioned_write.partitioned_table VALUES - (1, 'electronics', 'laptop', CAST(123.456 AS REAL)), - (2, 'electronics', 'phone', CAST(123.456 AS REAL)), - (3, 'books', 'novel', CAST('NaN' AS REAL)), - (4, 'books', 'textbook', CAST('Infinity' AS REAL)), - (5, 'clothing', 'shirt', CAST('-Infinity' AS REAL)), - (6, 'clothing', 'coat', CAST(NULL AS REAL)) + (1, 'electronics', 'laptop', CAST(123.456 AS REAL), CAST(123.456 AS DOUBLE)), + (2, 'electronics', 'phone', CAST(123.456 AS REAL), CAST(123.456 AS DOUBLE)), + (3, 'books', 'novel', CAST('NaN' AS REAL), CAST('NaN' AS DOUBLE)), + (4, 'books', 'textbook', CAST('Infinity' AS REAL), CAST('Infinity' AS DOUBLE)), + (5, 'clothing', 'shirt', CAST('-Infinity' AS REAL), CAST('-Infinity' AS DOUBLE)), + (6, 'clothing', 'coat', CAST(NULL AS REAL), CAST(NULL AS DOUBLE)) "#, ) .await @@ -1001,7 +1014,8 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { Field { "id": Int32, metadata: {"PARQUET:field_id": "1"} }, Field { "category": Utf8, metadata: {"PARQUET:field_id": "2"} }, Field { "value": Utf8, metadata: {"PARQUET:field_id": "3"} }, - Field { "float_partition": nullable Float32, metadata: {"PARQUET:field_id": "4"} }"#]], + Field { "float_partition": nullable Float32, metadata: {"PARQUET:field_id": "4"} }, + Field { "double_partition": nullable Float64, metadata: {"PARQUET:field_id": "5"} }"#]], expect![[r#" id: PrimitiveArray [ @@ -1031,6 +1045,15 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { "coat", ], float_partition: PrimitiveArray + [ + 123.456, + 123.456, + NaN, + inf, + -inf, + null, + ], + double_partition: PrimitiveArray [ 123.456, 123.456, @@ -1043,7 +1066,7 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { Some("id"), ); - // Verify all five partition files and six rows were committed. + // Verify the committed files, rows, and partition values. let table_ident = TableIdent::new(namespace.clone(), "partitioned_table".to_string()); let table = client.load_table(&table_ident).await?; @@ -1051,17 +1074,34 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { let manifest_list = table.manifest_list_reader(snapshot).load().await?; let mut file_count = 0; let mut row_count = 0; + let mut partitions = HashSet::new(); for entry in manifest_list.entries() { let manifest = table.manifest_reader().read(entry).await?; - file_count += manifest.entries().len(); - row_count += manifest - .entries() - .iter() - .map(|entry| entry.data_file().record_count()) - .sum::(); + for entry in manifest.entries() { + file_count += 1; + row_count += entry.data_file().record_count(); + partitions.insert(entry.data_file().partition().clone()); + } } assert_eq!(file_count, 5); assert_eq!(row_count, 6); + let partition = |category: &str, float: Option, double: Option| { + Struct::from_iter([ + Some(Literal::string(category)), + float.map(Literal::float), + double.map(Literal::double), + ]) + }; + assert_eq!( + partitions, + HashSet::from([ + partition("electronics", Some(123.456), Some(123.456)), + partition("books", Some(f32::NAN), Some(f64::NAN)), + partition("books", Some(f32::INFINITY), Some(f64::INFINITY)), + partition("clothing", Some(f32::NEG_INFINITY), Some(f64::NEG_INFINITY)), + partition("clothing", None, None), + ]) + ); let table_location = table.metadata().location(); let file_io = table.file_io();