diff --git a/AGENTS.md b/AGENTS.md index 103ff4c..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 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 c62e4d5..89274d3 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,91 @@ 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(", ")) + } } -// 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. +/// The position of each field from the top level of `schema` down to `source_id`. +fn source_path(schema: &IcebergSchema, source_id: i32) -> Result> { + 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) +} + +/// 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<()> { + if fields != sources { + return plan_err!( + "PartitionExpr needs exactly one source per partition field: expected \ + {fields}, got {sources}" + ); + } + Ok(()) +} + +/// 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 +350,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 +358,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()) } } @@ -247,23 +409,29 @@ impl std::hash::Hash for PartitionExpr { // 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); } } #[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 +439,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 +480,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 +516,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 +562,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 +673,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, 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()), + vec![Arc::new(Column::new("id", 0))], ) .unwrap(); @@ -503,33 +720,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 +790,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"))); @@ -569,6 +804,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 @@ -585,8 +836,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 +849,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 +949,7 @@ mod tests { sort_order, "/test/table".to_string(), FormatVersion::V2, - std::collections::HashMap::new(), + HashMap::new(), ) .unwrap(); @@ -742,10 +978,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 +1010,7 @@ mod tests { sort_order, "/test/table".to_string(), FormatVersion::V2, - std::collections::HashMap::new(), + HashMap::new(), ) .unwrap(); @@ -816,12 +1048,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(()) +}