NoahKusaba commented on code in PR #14:
URL: https://github.com/apache/datafusion-iceberg/pull/14#discussion_r4093962795
##########
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:
Done in `fff3f51` — took your suggested line verbatim. I wrote it by hand
rather than through the suggestion UI, so GitHub may still show it as
unapplied; the file matches it character for character.
Added `test_try_new_error_paths` covering all three cases. Your predicted
errors were exact: `DataInvalid => Cannot create partition calculator for
unpartitioned table`, `Unexpected, context: { field_id: 2 } => Field not
found`, and `DataInvalid => boolean is not a valid input type of bucket
transform`. I folded them into one table-driven test since all three exercise
the same forwarding call; each row asserts on the distinguishing phrase so an
upstream rewording won't break them.
##########
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:
You're right, and the wording wasn't just loose, it was wrong. Two tables'
specs with identical fields but different `spec_id`s do compute the same values
and compare unequal, so "equal when they compute the same partition values"
described something stricter than the code actually does. Took your suggestion
as-is in `fff3f51`.
##########
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:
Added as `test_partition_expr_differs_on_schema_with_shared_id` in
`fff3f51`, with `part` as `Int` vs `Long` under a shared `schema_id`. It
asserts the two `schema_id`s are equal and the expressions are not, so the ids
`Hash` uses cannot be what separates them.
--
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]
