From b1917d0026427ff83448d9c93b9ba560b791fb55 Mon Sep 17 00:00:00 2001 From: Noah Date: Wed, 28 Oct 2026 19:39:23 -0400 Subject: [PATCH 1/3] fix: read partition source columns through PartitionExpr children PartitionExpr read its source columns by position in the input batch but reported no children. When ProjectionPushdown merged the partition projection into the projection below it, nothing rewrote those positions, so partition values were computed from the wrong columns. Files landed in the wrong partitions and partition pruning skipped them. PartitionExpr now has one child per partition field, the top-level column holding its source, and evaluates the transforms itself, descending into structs for nested sources and nulling values under a null parent. The optimizer rewrites the children like any other column reference. The public `try_new` now requires these sources, so a codec must pass the decoded children; only `project_with_partition`, which checks that its input is laid out like the table, builds them from the table layout. Spec and schema errors are reported as iceberg DataInvalid errors. Closes #35 --- AGENTS.md | 2 +- crates/datafusion/src/write/project.rs | 630 ++++++++++++------ .../tests/integration_datafusion_test.rs | 128 +++- 3 files changed, 545 insertions(+), 215 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 103ff4c..a52aaa8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -89,7 +89,7 @@ input - **`TaskWriter`** picks one of three writers: `UnpartitionedWriter`; `FanoutWriter` (accepts unsorted input, keeps many files open); or `ClusteredWriter` (requires sorted input, which is why the sort step exists). - **Writer limits:** only Parquet is supported. The writer honors `write.parquet.*` table properties and table encryption. It matches columns by name, because DataFusion batches carry no field IDs. -- **`PartitionExpr`** is public and retains its partition spec and table schema so that distributed engines can serialize them and rebuild the expression on workers with `try_new`. Its equality and hash are by value; don't go back to pointer equality. +- **`PartitionExpr`** is public and retains its partition spec and table schema so that distributed engines can serialize them and rebuild the expression on workers with `try_new`. Its children are the top-level columns holding each partition field's source (it descends into structs itself, nulling values under a null parent), and it never reads the batch by position: `ProjectionPushdown` may merge it into a lower projection and only rewrites columns it can see. That is why the public `try_new` requires the sources: a codec passes the decoded children, and only `project_with_partition` may assume a table-shaped input (crate-private `try_new_over_table`). Its equality and hash are by value; don't go back to pointer equality. ### Errors diff --git a/crates/datafusion/src/write/project.rs b/crates/datafusion/src/write/project.rs index c62e4d5..c1d1eed 100644 --- a/crates/datafusion/src/write/project.rs +++ b/crates/datafusion/src/write/project.rs @@ -19,20 +19,28 @@ use std::sync::Arc; -use datafusion::arrow::array::RecordBatch; -use datafusion::arrow::datatypes::{DataType, Schema as ArrowSchema}; -use datafusion::common::plan_err; +use datafusion::arrow::array::{Array, ArrayRef, RecordBatch, StructArray, make_array}; +use datafusion::arrow::buffer::NullBuffer; +use datafusion::arrow::datatypes::{DataType, Fields, Schema as ArrowSchema}; +use datafusion::common::cast::as_struct_array; +use datafusion::common::{internal_err, plan_err}; use datafusion::error::Result; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr_common::physical_expr::fmt_sql; use datafusion::physical_plan::projection::ProjectionExec; use datafusion::physical_plan::{ColumnarValue, ExecutionPlan}; use iceberg::arrow::{ - PROJECTED_PARTITION_VALUE_COLUMN, PartitionValueCalculator, schema_to_arrow_schema, - strip_metadata_from_schema, + PROJECTED_PARTITION_VALUE_COLUMN, schema_to_arrow_schema, strip_metadata_from_schema, + type_to_arrow_type, +}; +use iceberg::spec::{ + PartitionSpec, Schema as IcebergSchema, SchemaRef as IcebergSchemaRef, StructType, + Type, }; -use iceberg::spec::{PartitionSpec, SchemaRef as IcebergSchemaRef}; use iceberg::table::Table; +use iceberg::transform::{BoxedTransformFunction, create_transform_function}; +use iceberg::{Error, ErrorKind}; use crate::to_datafusion_error; @@ -88,7 +96,7 @@ pub fn project_with_partition( projection_exprs.push((column_expr, field.name().clone())); } - let partition_expr = Arc::new(PartitionExpr::try_new( + let partition_expr = Arc::new(PartitionExpr::try_new_over_table( partition_spec.clone(), table_schema.clone(), )?); @@ -100,30 +108,56 @@ pub fn project_with_partition( /// PhysicalExpr implementation for partition value calculation /// -/// The [`PartitionValueCalculator`] cannot be serialized, so the spec and schema it -/// was built from are retained for [`Self::try_new`] to rebuild from. +/// Each partition field is computed from one child, the top-level column holding its +/// source, so [`Self::children`] are exactly the columns the expression reads. +/// Optimizer rules that move the expression onto a different input, such as merging +/// it into the projection below, rewrite those children like any other column +/// reference. +/// +/// The transform functions cannot be serialized, so the spec and schema they were +/// built from are retained. A codec rebuilds the expression by passing them to +/// [`Self::try_new`] together with the decoded children. #[derive(Debug, Clone)] pub struct PartitionExpr { - calculator: Arc, partition_spec: Arc, table_schema: IcebergSchemaRef, + transforms: Arc<[FieldTransform]>, + partition_fields: Fields, + sources: Vec>, +} + +/// How one partition field is computed from its child, derived from the spec and +/// schema. +#[derive(Debug)] +struct FieldTransform { + /// Struct field positions from the child down to a nested source; empty for a + /// top-level source. + nested_path: Vec, + transform: BoxedTransformFunction, } impl PartitionExpr { - /// Builds the expression from the spec and schema that define it. + /// Builds the expression from the spec and schema that define it, reading each + /// partition field's source from the matching expression in `sources`. /// - /// The calculator is built here rather than passed in, so it cannot drift from - /// the retained inputs. + /// `sources` holds one expression per partition field: the top-level input + /// column that holds that field's source column. A codec passes the decoded + /// children of the expression it serialized, which already point at the right + /// input columns. The transforms and partition type are derived here rather than + /// passed in, so they cannot drift from the retained spec and schema. /// /// # Errors /// - /// Returns an error if the spec is unpartitioned or cannot be bound to the schema. + /// Returns an error if the spec is unpartitioned, cannot be bound to the schema, + /// or does not have exactly one source per partition field. /// /// # Example /// /// ``` /// use std::sync::Arc; /// + /// use datafusion::physical_expr::PhysicalExpr; + /// use datafusion::physical_expr::expressions::Column; /// use iceberg_datafusion::write::PartitionExpr; /// use iceberg::spec::{NestedField, PartitionSpec, PrimitiveType, Schema, Transform, Type}; /// @@ -139,12 +173,14 @@ impl PartitionExpr { /// .add_partition_field("id", "id_bucket", Transform::Bucket(16))? /// .build()?, /// ); - /// let expr = PartitionExpr::try_new(spec, schema)?; + /// let expr = PartitionExpr::try_new(spec, schema, vec![Arc::new(Column::new("id", 0))])?; /// - /// // A worker rebuilds an equal expression from the two retained inputs. + /// // A worker rebuilds an equal expression from the retained spec and schema and + /// // the decoded children. /// let rebuilt = PartitionExpr::try_new( /// Arc::new(expr.partition_spec().as_ref().clone()), /// Arc::new(expr.table_schema().as_ref().clone()), + /// expr.children().into_iter().cloned().collect(), /// )?; /// assert_eq!(expr, rebuilt); /// # Ok::<(), Box>(()) @@ -152,21 +188,67 @@ impl PartitionExpr { pub fn try_new( partition_spec: Arc, table_schema: IcebergSchemaRef, + sources: Vec>, ) -> Result { - let calculator = PartitionValueCalculator::try_new( - partition_spec.as_ref(), - table_schema.as_ref(), - ) - .map_err(to_datafusion_error)?; + if partition_spec.is_unpartitioned() { + return Err(to_datafusion_error(Error::new( + ErrorKind::DataInvalid, + "Cannot compute partition values for an unpartitioned spec", + ))); + } + check_source_count(partition_spec.fields().len(), sources.len())?; + let partition_type = partition_spec + .partition_type(&table_schema) + .map_err(to_datafusion_error)?; + let DataType::Struct(partition_fields) = + type_to_arrow_type(&Type::Struct(partition_type)) + .map_err(to_datafusion_error)? + else { + return internal_err!( + "Iceberg partition type must convert to an Arrow struct" + ); + }; + let transforms = partition_spec + .fields() + .iter() + .map(|pf| { + Ok(FieldTransform { + nested_path: source_path(&table_schema, pf.source_id)?[1..].to_vec(), + transform: create_transform_function(&pf.transform) + .map_err(to_datafusion_error)?, + }) + }) + .collect::>>()?; Ok(Self { - calculator: Arc::new(calculator), partition_spec, table_schema, + transforms: transforms.into(), + partition_fields, + sources, }) } + /// Builds the expression over an input laid out like the table, which + /// [`project_with_partition`] checks before using it. + pub(crate) fn try_new_over_table( + partition_spec: Arc, + table_schema: IcebergSchemaRef, + ) -> Result { + let sources = partition_spec + .fields() + .iter() + .map(|pf| { + let index = source_path(&table_schema, pf.source_id)?[0]; + let column = &table_schema.as_struct().fields()[index]; + Ok(Arc::new(Column::new(&column.name, index)) as Arc) + }) + .collect::>>()?; + Self::try_new(partition_spec, table_schema, sources) + } + /// The partition spec this expression computes values for. With - /// [`Self::table_schema`], all [`Self::try_new`] needs to rebuild an equal one. + /// [`Self::table_schema`] and the children, all [`Self::try_new`] needs to + /// rebuild an equal one. pub fn partition_spec(&self) -> &Arc { &self.partition_spec } @@ -176,15 +258,90 @@ impl PartitionExpr { pub fn table_schema(&self) -> &IcebergSchemaRef { &self.table_schema } + + fn fmt_with( + &self, + f: &mut std::fmt::Formatter<'_>, + fmt_source: impl Fn(&dyn PhysicalExpr) -> String, + ) -> std::fmt::Result { + let fields: Vec = self + .partition_spec + .fields() + .iter() + .zip(&self.sources) + .map(|(pf, source)| { + format!( + "{}: {}({})", + pf.name, + pf.transform, + fmt_source(source.as_ref()) + ) + }) + .collect(); + write!(f, "iceberg_partition_values({})", fields.join(", ")) + } +} + +/// The position of each field from the top level of `schema` down to `source_id`. +fn source_path(schema: &IcebergSchema, source_id: i32) -> Result> { + field_path(schema.as_struct(), source_id).ok_or_else(|| { + to_datafusion_error(Error::new( + ErrorKind::DataInvalid, + format!("Partition source field {source_id} not found in table schema"), + )) + }) +} + +fn field_path(fields: &StructType, field_id: i32) -> Option> { + fields + .fields() + .iter() + .enumerate() + .find_map(|(index, field)| { + let mut path = if field.id == field_id { + vec![] + } else if let Type::Struct(nested) = field.field_type.as_ref() { + field_path(nested, field_id)? + } else { + return None; + }; + path.insert(0, index); + Some(path) + }) +} + +fn check_source_count(fields: usize, sources: usize) -> Result<()> { + if fields != sources { + return plan_err!( + "PartitionExpr needs exactly one source per partition field: expected \ + {fields}, got {sources}" + ); + } + Ok(()) } -// Equal when the spec and schema are equal. The calculator is derived from them, so -// it takes no part. Both comparisons include `spec_id` and `schema_id`, which `Hash` -// relies on. +/// Child `index` of a struct array, null wherever the struct is null. Arrow leaves +/// a child's values under a null parent unspecified, so they must not be read. +fn struct_field(array: &ArrayRef, index: usize) -> Result { + let parent = as_struct_array(array)?; + let child = parent.column(index); + let Some(parent_nulls) = parent.nulls() else { + return Ok(Arc::clone(child)); + }; + let nulls = NullBuffer::union(Some(parent_nulls), child.logical_nulls().as_ref()); + Ok(make_array( + child.to_data().into_builder().nulls(nulls).build()?, + )) +} + +// Equal when the spec, schema and sources are equal. The transforms and partition +// type are derived from the first two, so they take no part. Both comparisons include +// `spec_id` and `schema_id`, which `Hash` relies on. impl PartialEq for PartitionExpr { fn eq(&self, other: &Self) -> bool { self.partition_spec == other.partition_spec && self.table_schema == other.table_schema + && self.sources == other.sources } } @@ -192,7 +349,7 @@ impl Eq for PartitionExpr {} impl PhysicalExpr for PartitionExpr { fn data_type(&self, _input_schema: &ArrowSchema) -> Result { - Ok(self.calculator.partition_arrow_type().clone()) + Ok(DataType::Struct(self.partition_fields.clone())) } fn nullable(&self, _input_schema: &ArrowSchema) -> Result { @@ -200,44 +357,48 @@ impl PhysicalExpr for PartitionExpr { } fn evaluate(&self, batch: &RecordBatch) -> Result { - let array = self - .calculator - .calculate(batch) - .map_err(to_datafusion_error)?; - Ok(ColumnarValue::Array(array)) + let values = self + .sources + .iter() + .zip(self.transforms.iter()) + .map(|(source, field)| { + let mut source = source.evaluate(batch)?.into_array(batch.num_rows())?; + for &index in &field.nested_path { + source = struct_field(&source, index)?; + } + field + .transform + .transform(source) + .map_err(to_datafusion_error) + }) + .collect::>>()?; + let array = StructArray::try_new(self.partition_fields.clone(), values, None)?; + Ok(ColumnarValue::Array(Arc::new(array))) } fn children(&self) -> Vec<&Arc> { - vec![] + self.sources.iter().collect() } fn with_new_children( self: Arc, - _children: Vec>, + children: Vec>, ) -> Result> { - Ok(self) + check_source_count(self.sources.len(), children.len())?; + Ok(Arc::new(Self { + sources: children, + ..self.as_ref().clone() + })) } fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let field_names: Vec = self - .partition_spec - .fields() - .iter() - .map(|pf| format!("{}({})", pf.transform, pf.name)) - .collect(); - write!(f, "iceberg_partition_values[{}]", field_names.join(", ")) + self.fmt_with(f, |source| fmt_sql(source).to_string()) } } impl std::fmt::Display for PartitionExpr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let field_names: Vec<&str> = self - .partition_spec - .fields() - .iter() - .map(|pf| pf.name.as_str()) - .collect(); - write!(f, "iceberg_partition_values({})", field_names.join(", ")) + self.fmt_with(f, |source| source.to_string()) } } @@ -252,18 +413,23 @@ impl std::hash::Hash for PartitionExpr { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; - use datafusion::arrow::array::{ArrayRef, Int32Array, StructArray}; + use datafusion::arrow::array::{ArrayRef, Int32Array, StringArray, StructArray}; use datafusion::arrow::datatypes::{DataType, Field, Fields}; + use datafusion::arrow::util::pretty::pretty_format_columns; + use datafusion::parquet::arrow::PARQUET_FIELD_ID_META_KEY; use datafusion::physical_plan::empty::EmptyExec; - use iceberg::Runtime; + use iceberg::io::FileIO; use iceberg::spec::{ - NestedField, PrimitiveType, Schema, StructType, Transform, Type, + FormatVersion, NestedField, PrimitiveType, Schema, StructType, Transform, Type, }; + use iceberg::{Runtime, TableIdent}; use super::*; + use crate::from_datafusion_error; fn hash_of(expr: &PartitionExpr) -> u64 { let mut hasher = DefaultHasher::new(); @@ -271,6 +437,31 @@ mod tests { hasher.finish() } + /// Evaluates `expr` over `batch` and renders the partition values as a table. + fn partition_values(expr: &dyn PhysicalExpr, batch: &RecordBatch) -> String { + let values = expr + .evaluate(batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + pretty_format_columns(PROJECTED_PARTITION_VALUE_COLUMN, &[values]) + .unwrap() + .to_string() + } + + /// Partitions [`schema_with_part`] by `identity(id)`, leaving `part` unused. + fn partitioned_on_id() -> PartitionExpr { + let table_schema = schema_with_part(PrimitiveType::Int); + let partition_spec = Arc::new( + PartitionSpec::builder(table_schema.clone()) + .add_partition_field("id", "id_partition", Transform::Identity) + .unwrap() + .build() + .unwrap(), + ); + PartitionExpr::try_new_over_table(partition_spec, table_schema).unwrap() + } + // `part` typed per caller, so a spec built against one schema can be paired with // a schema that no longer matches it, as happens when a codec rebuilds. fn schema_with_part(part_type: PrimitiveType) -> IcebergSchemaRef { @@ -287,33 +478,6 @@ mod tests { ) } - #[test] - fn test_partition_calculator_basic() { - let table_schema = Schema::builder() - .with_schema_id(0) - .with_fields(vec![ - NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)) - .into(), - NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)) - .into(), - ]) - .build() - .unwrap(); - - let partition_spec = PartitionSpec::builder(Arc::new(table_schema.clone())) - .add_partition_field("id", "id_partition", Transform::Identity) - .unwrap() - .build() - .unwrap(); - - let calculator = - PartitionValueCalculator::try_new(&partition_spec, &table_schema).unwrap(); - - // Verify partition type - assert_eq!(calculator.partition_type().fields().len(), 1); - assert_eq!(calculator.partition_type().fields()[0].name, "id_partition"); - } - #[test] fn test_partition_expr_with_projection() { let table_schema = Schema::builder() @@ -350,8 +514,11 @@ mod tests { } let partition_expr = Arc::new( - PartitionExpr::try_new(partition_spec, Arc::new(table_schema.clone())) - .unwrap(), + PartitionExpr::try_new_over_table( + partition_spec, + Arc::new(table_schema.clone()), + ) + .unwrap(), ); projection_exprs .push((partition_expr, PROJECTED_PARTITION_VALUE_COLUMN.to_string())); @@ -393,39 +560,85 @@ mod tests { arrow_schema.clone(), vec![ Arc::new(Int32Array::from(vec![10, 20, 30])), - Arc::new(datafusion::arrow::array::StringArray::from(vec![ - "a", "b", "c", - ])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), ], ) .unwrap(); - let partition_spec = Arc::new(partition_spec); - let calculator = - PartitionValueCalculator::try_new(&partition_spec, &table_schema).unwrap(); - let partition_type = calculator.partition_arrow_type().clone(); - let expr = PartitionExpr::try_new(partition_spec, Arc::new(table_schema.clone())) - .unwrap(); + let expr = PartitionExpr::try_new_over_table( + Arc::new(partition_spec), + Arc::new(table_schema), + ) + .unwrap(); + // Partition fields are optional and carry their partition field id. + let partition_type = DataType::Struct(Fields::from(vec![ + Field::new("id_partition", DataType::Int32, true).with_metadata( + HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1000".to_string(), + )]), + ), + ])); assert_eq!(expr.data_type(&arrow_schema).unwrap(), partition_type); assert!(!expr.nullable(&arrow_schema).unwrap()); + insta::assert_snapshot!(partition_values(&expr, &batch), @r" + +--------------------+ + | _partition | + +--------------------+ + | {id_partition: 10} | + | {id_partition: 20} | + | {id_partition: 30} | + +--------------------+ + "); + } - let result = expr.evaluate(&batch).unwrap(); - match result { - ColumnarValue::Array(array) => { - let struct_array = array.as_any().downcast_ref::().unwrap(); - let id_partition = struct_array - .column_by_name("id_partition") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(id_partition.value(0), 10); - assert_eq!(id_partition.value(1), 20); - assert_eq!(id_partition.value(2), 30); - } - _ => panic!("Expected array result"), - } + #[test] + fn test_partition_expr_children_are_its_source_columns() { + let expr = Arc::new(partitioned_on_id()); + // `part` is not a source, so it is not a child. + assert_eq!( + expr.to_string(), + "iceberg_partition_values(id_partition: identity(id@0))" + ); + + // An input laid out unlike the table, as when the optimizer merges the + // partition projection into one that reorders columns. Rewriting the child + // is all it takes to read `id` from its new position. + let remapped = expr + .with_new_children(vec![Arc::new(Column::new("id", 2))]) + .unwrap(); + assert_eq!( + remapped.to_string(), + "iceberg_partition_values(id_partition: identity(id@2))" + ); + let batch = RecordBatch::try_from_iter(vec![ + ("x", Arc::new(Int32Array::from(vec![7, 8])) as ArrayRef), + ("part", Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef), + ("id", Arc::new(Int32Array::from(vec![10, 20])) as ArrayRef), + ]) + .unwrap(); + insta::assert_snapshot!(partition_values(remapped.as_ref(), &batch), @r" + +--------------------+ + | _partition | + +--------------------+ + | {id_partition: 10} | + | {id_partition: 20} | + +--------------------+ + "); + } + + #[test] + fn test_with_new_children_requires_one_child_per_partition_field() { + let expr = Arc::new(partitioned_on_id()); + + let err = expr + .with_new_children(vec![ + Arc::new(Column::new("id", 0)), + Arc::new(Column::new("part", 1)), + ]) + .unwrap_err(); + insta::assert_snapshot!(err.strip_backtrace(), @"Error during planning: PartitionExpr needs exactly one source per partition field: expected 1, got 2"); } #[test] @@ -458,13 +671,15 @@ mod tests { ) .unwrap(); - let expr = PartitionExpr::try_new(partition_spec, table_schema).unwrap(); + let expr = + PartitionExpr::try_new_over_table(partition_spec, table_schema).unwrap(); - // Rebuild from deep copies, as a codec would; cloning the Arcs instead - // would let a pointer-based impl pass. + // Rebuild from deep copies and the children, as a codec would; cloning the + // Arcs instead would let a pointer-based impl pass. let rebuilt = PartitionExpr::try_new( Arc::new(expr.partition_spec().as_ref().clone()), Arc::new(expr.table_schema().as_ref().clone()), + expr.children().into_iter().cloned().collect(), ) .unwrap(); @@ -503,33 +718,50 @@ mod tests { .unwrap(), ); - // try_new takes the spec and schema separately, so a codec can pair two that - // do not fit. project_with_partition reaches none of these: it returns early - // when unpartitioned and reads both from the same table metadata. + let part = || vec![Arc::new(Column::new("part", 1)) as Arc]; + + // try_new takes the spec, schema and sources separately, so a codec can pass + // ones that do not fit. project_with_partition reaches none of these: it + // returns early when unpartitioned and derives all three from the table. let cases = [ ( Arc::new(PartitionSpec::builder(int_schema.clone()).build().unwrap()), int_schema.clone(), - "unpartitioned", - ), - ( - on_part(Transform::Identity), - without_part, - "Field not found", + vec![], ), + (on_part(Transform::Identity), int_schema.clone(), vec![]), + (on_part(Transform::Identity), without_part, part()), ( on_part(Transform::Bucket(4)), schema_with_part(PrimitiveType::Boolean), - "not a valid input type of bucket transform", + part(), ), ]; - for (spec, schema, expected) in cases { - let err = PartitionExpr::try_new(spec, schema) - .unwrap_err() - .to_string(); - assert!(err.contains(expected), "expected {expected:?}, got {err}"); - } + // Each error with the iceberg kind that survives `from_datafusion_error`. + let errors: Vec = cases + .into_iter() + .map(|(spec, schema, sources)| { + let err = PartitionExpr::try_new(spec, schema, sources).unwrap_err(); + let message = err.strip_backtrace(); + format!("{}: {message}", from_datafusion_error(err).kind()) + }) + .collect(); + insta::assert_snapshot!(errors.join("\n\n"), @r" + DataInvalid: IcebergError(DataInvalid) + caused by + Execution error: DataInvalid => Cannot compute partition values for an unpartitioned spec + + Unexpected: Error during planning: PartitionExpr needs exactly one source per partition field: expected 1, got 0 + + DataInvalid: IcebergError(DataInvalid) + caused by + Execution error: DataInvalid => Partition source field 2 not found in table schema + + DataInvalid: IcebergError(DataInvalid) + caused by + Execution error: DataInvalid => boolean is not a valid input type of bucket transform + "); } #[test] @@ -556,7 +788,8 @@ mod tests { .unwrap(), ) }; - let expr = |spec| PartitionExpr::try_new(spec, table_schema.clone()).unwrap(); + let expr = + |spec| PartitionExpr::try_new_over_table(spec, table_schema.clone()).unwrap(); // The same inputs describe the same expression. assert_eq!(expr(spec(1, "id")), expr(spec(1, "id"))); @@ -585,8 +818,9 @@ mod tests { .unwrap(), ); - let on_int = PartitionExpr::try_new(shared.clone(), int_schema).unwrap(); - let on_long = PartitionExpr::try_new(shared, long_schema).unwrap(); + let on_int = + PartitionExpr::try_new_over_table(shared.clone(), int_schema).unwrap(); + let on_long = PartitionExpr::try_new_over_table(shared, long_schema).unwrap(); assert_eq!( on_int.table_schema().schema_id(), @@ -597,90 +831,74 @@ mod tests { #[test] fn test_nested_partition() { - let address_struct = StructType::new(vec![ + let address_type = StructType::new(vec![ NestedField::required(3, "street", Type::Primitive(PrimitiveType::String)) .into(), NestedField::required(4, "city", Type::Primitive(PrimitiveType::String)) .into(), ]); + let table_schema = Arc::new( + Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)) + .into(), + NestedField::optional(2, "address", Type::Struct(address_type)) + .into(), + ]) + .build() + .unwrap(), + ); + let partition_spec = Arc::new( + PartitionSpec::builder(table_schema.clone()) + .add_partition_field( + "address.city", + "city_partition", + Transform::Identity, + ) + .unwrap() + .build() + .unwrap(), + ); + let expr = + PartitionExpr::try_new_over_table(partition_spec, table_schema).unwrap(); + // The child is the top-level column that holds the nested source. + assert_eq!( + expr.to_string(), + "iceberg_partition_values(city_partition: identity(address@1))" + ); - let table_schema = Schema::builder() - .with_schema_id(0) - .with_fields(vec![ - NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)) - .into(), - NestedField::required(2, "address", Type::Struct(address_struct)).into(), - ]) - .build() - .unwrap(); - - let partition_spec = PartitionSpec::builder(Arc::new(table_schema.clone())) - .add_partition_field("address.city", "city_partition", Transform::Identity) - .unwrap() - .build() - .unwrap(); - - let struct_fields = Fields::from(vec![ - Field::new("street", DataType::Utf8, false), - Field::new("city", DataType::Utf8, false), - ]); - - let arrow_schema = Arc::new(ArrowSchema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("address", DataType::Struct(struct_fields), false), - ])); - - let street_array = Arc::new(datafusion::arrow::array::StringArray::from(vec![ - "123 Main St", - "456 Oak Ave", - ])); - let city_array = Arc::new(datafusion::arrow::array::StringArray::from(vec![ - "New York", - "Los Angeles", - ])); - - let struct_array = StructArray::from(vec![ - ( - Arc::new(Field::new("street", DataType::Utf8, false)), - street_array as ArrayRef, - ), - ( - Arc::new(Field::new("city", DataType::Utf8, false)), - city_array as ArrayRef, - ), - ]); - - let batch = RecordBatch::try_new( - arrow_schema.clone(), + // The third address is null, but its `city` slot still holds a value, + // which must not become a partition value. + let address = StructArray::new( + Fields::from(vec![ + Field::new("street", DataType::Utf8, false), + Field::new("city", DataType::Utf8, false), + ]), vec![ - Arc::new(Int32Array::from(vec![1, 2])), - Arc::new(struct_array), + Arc::new(StringArray::from(vec!["123 Main St", "456 Oak Ave", ""])), + Arc::new(StringArray::from(vec!["New York", "Los Angeles", "stale"])), ], - ) + Some(NullBuffer::from(vec![true, true, false])), + ); + let batch = RecordBatch::try_from_iter(vec![ + ("id", Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef), + ("address", Arc::new(address) as ArrayRef), + ]) .unwrap(); - - let calculator = - PartitionValueCalculator::try_new(&partition_spec, &table_schema).unwrap(); - let array = calculator.calculate(&batch).unwrap(); - - let struct_array = array.as_any().downcast_ref::().unwrap(); - let city_partition = struct_array - .column_by_name("city_partition") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - - assert_eq!(city_partition.value(0), "New York"); - assert_eq!(city_partition.value(1), "Los Angeles"); + insta::assert_snapshot!(partition_values(&expr, &batch), @r" + +-------------------------------+ + | _partition | + +-------------------------------+ + | {city_partition: New York} | + | {city_partition: Los Angeles} | + | {city_partition: } | + +-------------------------------+ + "); } #[tokio::test] async fn test_schema_validation_matching_schemas() { - use iceberg::TableIdent; - use iceberg::io::FileIO; - use iceberg::spec::{FormatVersion, NestedField, PrimitiveType, Schema, Type}; - let table_schema = Arc::new( Schema::builder() .with_fields(vec![ @@ -713,7 +931,7 @@ mod tests { sort_order, "/test/table".to_string(), FormatVersion::V2, - std::collections::HashMap::new(), + HashMap::new(), ) .unwrap(); @@ -742,10 +960,6 @@ mod tests { #[tokio::test] async fn test_schema_validation_mismatched_schemas() { - use iceberg::TableIdent; - use iceberg::io::FileIO; - use iceberg::spec::{FormatVersion, NestedField, PrimitiveType, Schema, Type}; - let table_schema = Arc::new( Schema::builder() .with_fields(vec![ @@ -778,7 +992,7 @@ mod tests { sort_order, "/test/table".to_string(), FormatVersion::V2, - std::collections::HashMap::new(), + HashMap::new(), ) .unwrap(); @@ -816,12 +1030,6 @@ mod tests { #[tokio::test] async fn test_schema_validation_with_metadata_differences() { - use std::collections::HashMap; - - use iceberg::TableIdent; - use iceberg::io::FileIO; - use iceberg::spec::{FormatVersion, NestedField, PrimitiveType, Schema, Type}; - let table_schema = Arc::new( Schema::builder() .with_fields(vec![ diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 42d350a..e60d49b 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -22,19 +22,21 @@ use std::error::Error; use std::sync::Arc; use std::vec; -use datafusion::arrow::array::{Array, AsArray, RecordBatch, StringArray, UInt64Array}; +use datafusion::arrow::array::{ + Array, AsArray, Int32Array, RecordBatch, StringArray, UInt64Array, +}; use datafusion::arrow::compute::{ cast, concat_batches, sort_to_indices, take_record_batch, }; use datafusion::arrow::datatypes::{DataType, Field, Int64Type, Schema as ArrowSchema}; use datafusion::arrow::util::pretty::pretty_format_batches; -use datafusion::datasource::TableProvider; use datafusion::datasource::source::DataSourceExec; +use datafusion::datasource::{MemTable, TableProvider}; 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; use datafusion::physical_plan::common::collect; +use datafusion::physical_plan::{ExecutionPlan, displayable}; use iceberg::io::LocalFsStorageFactory; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{ @@ -1280,3 +1282,123 @@ async fn test_scan_after_schema_evolution_reads_provider_columns() Ok(()) } + +/// An INSERT whose input ends in a projection that renames and reorders columns gets +/// the partition projection merged into it by the optimizer. The partition values +/// must still come from the columns being written (apache/datafusion-iceberg#35). +#[tokio::test] +async fn test_insert_into_partitioned_through_reordering_projection() +-> Result<(), Box> { + let iceberg_catalog = get_iceberg_catalog().await; + let namespace = NamespaceIdent::new("test_reordered_insert".to_string()); + set_test_namespace(&iceberg_catalog, &namespace).await?; + + let schema = Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "val", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build()?; + let partition_spec = UnboundPartitionSpec::builder() + .with_spec_id(0) + .add_partition_field( + UnboundPartitionField::builder() + .source_ids(vec![1]) + .name("id") + .transform(Transform::Identity) + .build()?, + )? + .build(); + let creation = TableCreation::builder() + .name("partitioned_by_id".to_string()) + .location(temp_path()) + .schema(schema) + .partition_spec(partition_spec) + .properties(HashMap::new()) + .build(); + iceberg_catalog.create_table(&namespace, creation).await?; + + let client = Arc::new(iceberg_catalog); + let catalog = Arc::new(IcebergCatalogProvider::try_new(client.clone()).await?); + // One partition keeps repartitioning out of the plan. + let ctx = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + ctx.register_catalog("catalog", catalog); + + let source_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let source = RecordBatch::try_new( + source_schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![100, 200])), + ], + )?; + ctx.register_table( + "source", + Arc::new(MemTable::try_new(source_schema, vec![vec![source]])?), + )?; + + // `id` is written from `b`, the second source column, so the partition values + // must read `b@1` and not the first column of the merged projection's input. + let plan = ctx + .sql( + "INSERT INTO catalog.test_reordered_insert.partitioned_by_id + SELECT b AS id, a AS val FROM source", + ) + .await? + .create_physical_plan() + .await?; + insta::assert_snapshot!(displayable(plan.as_ref()).indent(false), @r" + IcebergCommitExec: table=test_reordered_insert.partitioned_by_id + IcebergWriteExec: table=test_reordered_insert.partitioned_by_id + ProjectionExec: expr=[b@1 as id, a@0 as val, iceberg_partition_values(id: identity(b@1)) as _partition] + DataSourceExec: partitions=1, partition_sizes=[1] + "); + insta::assert_snapshot!(run(plan.as_ref(), &ctx).await?, @r" + +-------+ + | count | + +-------+ + | 2 | + +-------+ + "); + + let table_ident = TableIdent::new(namespace, "partitioned_by_id".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 partitions = HashSet::new(); + for entry in manifest_list.entries() { + let manifest = table.manifest_reader().read(entry).await?; + for entry in manifest.entries() { + partitions.insert(entry.data_file().partition().clone()); + } + } + assert_eq!( + partitions, + HashSet::from([ + Struct::from_iter([Some(Literal::int(100))]), + Struct::from_iter([Some(Literal::int(200))]), + ]) + ); + + // A filter on the partition column prunes by these values, so it only finds + // the row if the values are right. + let batches = ctx + .sql("SELECT * FROM catalog.test_reordered_insert.partitioned_by_id WHERE id = 100") + .await? + .collect() + .await?; + insta::assert_snapshot!(pretty_format_batches(&batches)?.to_string(), @r" + +-----+-----+ + | id | val | + +-----+-----+ + | 100 | 1 | + +-----+-----+ + "); + + Ok(()) +} From 25523909d4bb545cdc0864ba517e4c17f3d88f63 Mon Sep 17 00:00:00 2001 From: Noah Date: Thu, 29 Oct 2026 15:31:21 -0400 Subject: [PATCH 2/3] test: check that PartitionExpr equality depends on its sources Also build source paths by pushing positions on the way down rather than prepending them, and say why Hash leaves the sources out. --- crates/datafusion/src/write/project.rs | 62 +++++++++++++++++--------- 1 file changed, 41 insertions(+), 21 deletions(-) diff --git a/crates/datafusion/src/write/project.rs b/crates/datafusion/src/write/project.rs index c1d1eed..79e7b6b 100644 --- a/crates/datafusion/src/write/project.rs +++ b/crates/datafusion/src/write/project.rs @@ -284,30 +284,31 @@ impl PartitionExpr { /// The position of each field from the top level of `schema` down to `source_id`. fn source_path(schema: &IcebergSchema, source_id: i32) -> Result> { - field_path(schema.as_struct(), source_id).ok_or_else(|| { - to_datafusion_error(Error::new( + let mut path = Vec::new(); + if !push_field_path(schema.as_struct(), source_id, &mut path) { + return Err(to_datafusion_error(Error::new( ErrorKind::DataInvalid, format!("Partition source field {source_id} not found in table schema"), - )) - }) + ))); + } + Ok(path) } -fn field_path(fields: &StructType, field_id: i32) -> Option> { - fields - .fields() - .iter() - .enumerate() - .find_map(|(index, field)| { - let mut path = if field.id == field_id { - vec![] - } else if let Type::Struct(nested) = field.field_type.as_ref() { - field_path(nested, field_id)? - } else { - return None; - }; - path.insert(0, index); - Some(path) - }) +/// Appends the positions from `fields` down to `field_id` to `path`, returning +/// whether the field was found. `path` is left as it was if not. +fn push_field_path(fields: &StructType, field_id: i32, path: &mut Vec) -> bool { + fields.fields().iter().enumerate().any(|(index, field)| { + path.push(index); + let found = field.id == field_id + || matches!( + field.field_type.as_ref(), + Type::Struct(nested) if push_field_path(nested, field_id, path) + ); + if !found { + path.pop(); + } + found + }) } fn check_source_count(fields: usize, sources: usize) -> Result<()> { @@ -405,7 +406,10 @@ impl std::fmt::Display for PartitionExpr { impl std::hash::Hash for PartitionExpr { fn hash(&self, state: &mut H) { // Neither PartitionSpec nor Schema implements Hash, so hash their ids: - // equal expressions agree on both, and collisions fall through to eq. + // equal expressions agree on both, and collisions fall through to eq. The + // sources are left out on purpose: equal expressions also have equal + // sources, so hashing the ids alone is enough, and expressions that differ + // only in their sources are told apart by eq. self.partition_spec.spec_id().hash(state); self.table_schema.schema_id().hash(state); } @@ -802,6 +806,22 @@ mod tests { assert_ne!(expr(spec(7, "id")), expr(spec(7, "part"))); } + #[test] + fn test_partition_expr_differs_on_sources() { + // The same spec and schema reading `id` from a different input column, as + // after the optimizer moves the expression onto a reordered input. Equating + // the two would let DataFusion reuse one in place of the other. + let expr = partitioned_on_id(); + let moved = PartitionExpr::try_new( + expr.partition_spec().clone(), + expr.table_schema().clone(), + vec![Arc::new(Column::new("id", 2))], + ) + .unwrap(); + + assert_ne!(expr, moved); + } + #[test] fn test_partition_expr_differs_on_schema_with_shared_id() { // Two tables can both be on schema 0 with `part` typed differently, which From 0eae2ff4ff3990d05c7b0d6050ef50229c3f7c7f Mon Sep 17 00:00:00 2001 From: Noah Date: Thu, 29 Oct 2026 19:37:16 -0400 Subject: [PATCH 3/3] Hash PartitionExpr sources and rebuild from decoded children in tests Hash the sources along with the spec and schema ids, so expressions that differ only in their inputs no longer collide. The rebuild test now passes freshly built children, as a codec would, so it fails if sources are compared by pointer. Clarify in AGENTS.md that PartitionExpr reads the batch only through its children. --- AGENTS.md | 2 +- crates/datafusion/src/write/project.rs | 10 ++++------ 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index a52aaa8..ab894f3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -89,7 +89,7 @@ input - **`TaskWriter`** picks one of three writers: `UnpartitionedWriter`; `FanoutWriter` (accepts unsorted input, keeps many files open); or `ClusteredWriter` (requires sorted input, which is why the sort step exists). - **Writer limits:** only Parquet is supported. The writer honors `write.parquet.*` table properties and table encryption. It matches columns by name, because DataFusion batches carry no field IDs. -- **`PartitionExpr`** is public and retains its partition spec and table schema so that distributed engines can serialize them and rebuild the expression on workers with `try_new`. Its children are the top-level columns holding each partition field's source (it descends into structs itself, nulling values under a null parent), and it never reads the batch by position: `ProjectionPushdown` may merge it into a lower projection and only rewrites columns it can see. That is why the public `try_new` requires the sources: a codec passes the decoded children, and only `project_with_partition` may assume a table-shaped input (crate-private `try_new_over_table`). Its equality and hash are by value; don't go back to pointer equality. +- **`PartitionExpr`** is public and retains its partition spec and table schema so that distributed engines can serialize them and rebuild the expression on workers with `try_new`, passing the decoded children. Its children are the top-level columns holding each partition field's source, and it reads the batch only through them, so optimizer rules like `ProjectionPushdown` that move it onto a different input rewrite them like any other column reference. Only `project_with_partition`, which checks that its input matches the table, derives the children from the table layout (crate-private `try_new_over_table`). Its equality and hash are by value; don't go back to pointer equality. ### Errors diff --git a/crates/datafusion/src/write/project.rs b/crates/datafusion/src/write/project.rs index 79e7b6b..89274d3 100644 --- a/crates/datafusion/src/write/project.rs +++ b/crates/datafusion/src/write/project.rs @@ -406,12 +406,10 @@ impl std::fmt::Display for PartitionExpr { impl std::hash::Hash for PartitionExpr { fn hash(&self, state: &mut H) { // Neither PartitionSpec nor Schema implements Hash, so hash their ids: - // equal expressions agree on both, and collisions fall through to eq. The - // sources are left out on purpose: equal expressions also have equal - // sources, so hashing the ids alone is enough, and expressions that differ - // only in their sources are told apart by eq. + // equal expressions agree on both, and collisions fall through to eq. self.partition_spec.spec_id().hash(state); self.table_schema.schema_id().hash(state); + self.sources.hash(state); } } @@ -678,12 +676,12 @@ mod tests { let expr = PartitionExpr::try_new_over_table(partition_spec, table_schema).unwrap(); - // Rebuild from deep copies and the children, as a codec would; cloning the + // Rebuild from deep copies, as a codec would after decoding; cloning the // Arcs instead would let a pointer-based impl pass. let rebuilt = PartitionExpr::try_new( Arc::new(expr.partition_spec().as_ref().clone()), Arc::new(expr.table_schema().as_ref().clone()), - expr.children().into_iter().cloned().collect(), + vec![Arc::new(Column::new("id", 0))], ) .unwrap();