mbutrovich commented on code in PR #14:
URL: https://github.com/apache/datafusion-iceberg/pull/14#discussion_r4087575482
##########
crates/datafusion/src/physical_plan/project.rs:
##########
@@ -91,38 +87,72 @@ pub fn project_with_partition(
projection_exprs.push((column_expr, field.name().clone()));
}
- let partition_expr = Arc::new(PartitionExpr::new(calculator,
partition_spec.clone()));
+ let partition_expr = Arc::new(PartitionExpr::try_new(
+ partition_spec.clone(),
+ table_schema.clone(),
+ )?);
projection_exprs.push((partition_expr,
PROJECTED_PARTITION_VALUE_COLUMN.to_string()));
let projection = ProjectionExec::try_new(projection_exprs, input)?;
Ok(Arc::new(projection))
}
/// 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.
#[derive(Debug, Clone)]
-struct PartitionExpr {
+pub struct PartitionExpr {
calculator: Arc<PartitionValueCalculator>,
partition_spec: Arc<PartitionSpec>,
+ table_schema: IcebergSchemaRef,
}
impl PartitionExpr {
- fn new(
- calculator: PartitionValueCalculator,
+ /// Builds the expression from the spec and schema that define it.
+ ///
+ /// The calculator is built here rather than passed in, so it cannot drift
from
+ /// the retained inputs.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the spec cannot be bound to the schema.
+ pub fn try_new(
partition_spec: Arc<PartitionSpec>,
- ) -> Self {
- Self {
+ table_schema: IcebergSchemaRef,
+ ) -> DFResult<Self> {
+ let calculator = PartitionValueCalculator::try_new(
+ partition_spec.as_ref(),
+ table_schema.as_ref(),
+ )
+ .map_err(to_datafusion_error)?;
+ Ok(Self {
calculator: Arc::new(calculator),
partition_spec,
- }
+ table_schema,
+ })
+ }
+
+ /// The partition spec this expression computes values for. With
+ /// [`Self::table_schema`], all [`Self::try_new`] needs to rebuild an
equal one.
+ pub fn partition_spec(&self) -> &Arc<PartitionSpec> {
+ &self.partition_spec
+ }
+
+ /// The table schema the spec is bound to. Needed to rebuild: the spec
refers to
+ /// columns by `source_id`, and only the schema resolves those.
+ pub fn table_schema(&self) -> &IcebergSchemaRef {
+ &self.table_schema
}
}
-// Manual PartialEq/Eq implementations for pointer-based equality
-// (two PartitionExpr are equal if they share the same calculator and
partition_spec instances)
+// Equal when they compute the same partition values, which the spec and schema
+// decide entirely. The calculator is derived from those two, so it takes no
part:
+// comparing it by pointer would make an expression unequal to its own rebuild.
Review Comment:
This comment says two expressions are equal when they compute the same
partition values, but `eq` is stricter than that. The Iceberg spec rules out
two equivalent specs with different ids within one table
([spec](https://github.com/apache/iceberg/blob/81045840187d172df9a9abb32a7c39420a4f10fc/format/spec.md#L564-L566)),
but specs from two tables can have the same fields and different `spec_id`s.
Those compute the same values and compare unequal, and the same goes for
`schema_id`. That's the right behavior. The partition tuple takes its field ids
and names from the spec
([spec](https://github.com/apache/iceberg/blob/81045840187d172df9a9abb32a7c39420a4f10fc/format/spec.md#L745)),
so equal values don't imply equal output. `Hash` also only hashes the ids and
relies on `eq` comparing them. Could the comment state that dependency, so
nobody later loosens `eq` to match the current wording and breaks `Hash`?
```suggestion
// 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.
```
##########
crates/datafusion/src/physical_plan/project.rs:
##########
@@ -91,38 +87,72 @@ pub fn project_with_partition(
projection_exprs.push((column_expr, field.name().clone()));
}
- let partition_expr = Arc::new(PartitionExpr::new(calculator,
partition_spec.clone()));
+ let partition_expr = Arc::new(PartitionExpr::try_new(
+ partition_spec.clone(),
+ table_schema.clone(),
+ )?);
projection_exprs.push((partition_expr,
PROJECTED_PARTITION_VALUE_COLUMN.to_string()));
let projection = ProjectionExec::try_new(projection_exprs, input)?;
Ok(Arc::new(projection))
}
/// 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.
#[derive(Debug, Clone)]
-struct PartitionExpr {
+pub struct PartitionExpr {
calculator: Arc<PartitionValueCalculator>,
partition_spec: Arc<PartitionSpec>,
+ table_schema: IcebergSchemaRef,
}
impl PartitionExpr {
- fn new(
- calculator: PartitionValueCalculator,
+ /// Builds the expression from the spec and schema that define it.
+ ///
+ /// The calculator is built here rather than passed in, so it cannot drift
from
+ /// the retained inputs.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the spec cannot be bound to the schema.
Review Comment:
Now that `try_new` is public, could the doc cover both failure cases?
`PartitionValueCalculator::try_new` also rejects an unpartitioned spec
([partition_value_calculator.rs](https://github.com/apache/iceberg-rust/blob/665c64e48e8d33797ecb1a421f327edd9b024879/crates/iceberg/src/arrow/partition_value_calculator.rs#L64-L70)),
and that spec binds to any schema. `project_with_partition` returns early for
unpartitioned tables, but a codec calling `try_new` directly has no such guard.
Could you also add tests for the error paths? I tried three on the head
commit. An unpartitioned spec returns `DataInvalid => Cannot create partition
calculator for unpartitioned table`. A spec whose `source_id` is missing from
the schema returns `Unexpected => Field not found`. The Iceberg spec also
requires rejecting a transform that doesn't accept the source type ([transforms
table](https://github.com/apache/iceberg/blob/81045840187d172df9a9abb32a7c39420a4f10fc/format/spec.md#L570-L579)),
and a `bucket[4]` spec paired with a schema where that column is `boolean`
returns `DataInvalid => boolean is not a valid input type of bucket transform`.
A codec receives the spec and schema separately, so this mismatch can reach
`try_new`. All three return errors as expected, but nothing pins that behavior.
```suggestion
/// # Errors
///
/// Returns an error if the spec is unpartitioned or cannot be bound to
the schema.
```
##########
crates/datafusion/src/physical_plan/project.rs:
##########
@@ -339,6 +380,95 @@ mod tests {
}
}
+ #[test]
+ fn test_partition_expr_rebuilds_from_its_retained_parts() {
+ let table_schema = Arc::new(
+ Schema::builder()
+ .with_schema_id(1)
+ .with_fields(vec![
+ NestedField::required(1, "id",
Type::Primitive(PrimitiveType::Int))
+ .into(),
+ ])
+ .build()
+ .unwrap(),
+ );
+ let partition_spec = Arc::new(
+ PartitionSpec::builder(table_schema.clone())
+ .add_partition_field("id", "id_partition", Transform::Identity)
+ .unwrap()
+ .build()
+ .unwrap(),
+ );
+ let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
+ "id",
+ DataType::Int32,
+ false,
+ )]));
+ let batch = RecordBatch::try_new(
+ arrow_schema,
+ vec![Arc::new(Int32Array::from(vec![10, 20, 30]))],
+ )
+ .unwrap();
+
+ let expr = PartitionExpr::try_new(partition_spec,
table_schema).unwrap();
+
+ // Rebuild from deep copies, 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()),
+ )
+ .unwrap();
+
+ let eval = |e: &PartitionExpr| match e.evaluate(&batch).unwrap() {
+ ColumnarValue::Array(array) => array,
+ _ => panic!("Expected array result"),
+ };
+ assert_eq!(&eval(&rebuilt), &eval(&expr));
+
+ // Same values is not enough: it must also compare and hash as the same
+ // expression, or plan-level equality and dedup treat the two as
unrelated.
+ assert_eq!(rebuilt, expr);
+ assert_eq!(hash_of(&rebuilt), hash_of(&expr));
+ }
+
+ #[test]
+ fn test_partition_expr_equality_is_value_based() {
+ let table_schema = Arc::new(
+ Schema::builder()
+ .with_schema_id(1)
+ .with_fields(vec![
+ NestedField::required(1, "id",
Type::Primitive(PrimitiveType::Int))
+ .into(),
+ NestedField::required(2, "part",
Type::Primitive(PrimitiveType::Int))
+ .into(),
+ ])
+ .build()
+ .unwrap(),
+ );
+ let spec = |spec_id: i32, source: &str| {
+ Arc::new(
+ PartitionSpec::builder(table_schema.clone())
+ .with_spec_id(spec_id)
+ .add_partition_field(source, "p", Transform::Identity)
+ .unwrap()
+ .build()
+ .unwrap(),
+ )
+ };
+ let expr = |spec| PartitionExpr::try_new(spec,
table_schema.clone()).unwrap();
+
+ // The same inputs describe the same expression.
+ assert_eq!(expr(spec(1, "id")), expr(spec(1, "id")));
+
+ // Genuinely different specs stay apart.
+ assert_ne!(expr(spec(1, "id")), expr(spec(2, "part")));
+
+ // Same spec_id, different partition column: comparing only the ids
that
+ // Hash uses would wrongly call these equal.
+ assert_ne!(expr(spec(7, "id")), expr(spec(7, "part")));
+ }
Review Comment:
This test varies only the spec. Could you add a case with the same spec and
two schemas that share a `schema_id` but differ in a field type (for example
`part` as `Int` vs `Long`)? Schema ids are unique only within a table
([spec](https://github.com/apache/iceberg/blob/81045840187d172df9a9abb32a7c39420a4f10fc/format/spec.md#L354)),
and promoting `part` within one table produces a new id. So this case comes
from a plan over two tables that both have schema 0. The partition type
differs, and only the struct comparison in `eq` tells the two apart. I checked
it on the head commit and it compares unequal, but no test covers it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
