diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs index 352d690a54b82..a94a38f47f75d 100644 --- a/datafusion/physical-expr/src/equivalence/class.rs +++ b/datafusion/physical-expr/src/equivalence/class.rs @@ -12,6 +12,7 @@ use super::ProjectionMapping; use crate::expressions::{Column, Literal}; use crate::physical_expr::add_offset_to_expr; use crate::projection::ProjectionTargets; +use crate::scalar_subquery::ScalarSubqueryExpr; use crate::{PhysicalExpr, PhysicalExprRef, PhysicalSortExpr, PhysicalSortRequirement}; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; @@ -631,8 +632,8 @@ impl EquivalenceGroup { aug_mapping: &AugmentedMapping, expr: &Arc, ) -> Option> { - // Literals don't need to be projected - if expr.downcast_ref::().is_some() { + // These leaves do not read columns from the projected input schema. + if expr.is::() || expr.is::() { return Some(Arc::clone(expr)); } @@ -742,7 +743,11 @@ impl EquivalenceGroup { pub fn project(&self, mapping: &ProjectionMapping) -> Self { let projected_classes = self.iter().map(|cls| { let new_exprs = self.project_expressions(mapping, cls.iter()); - EquivalenceClass::new(new_exprs.flatten()) + let mut projected = EquivalenceClass::new(new_exprs.flatten()); + // Projection changes column references, not constant values or how + // they are distributed across partitions. + projected.constant = projected.constant.or_else(|| cls.constant.clone()); + projected }); // The key is the source expression, and the value is the equivalence @@ -1334,6 +1339,49 @@ mod tests { Ok(()) } + #[test] + fn test_project_preserves_constant_metadata() -> Result<()> { + let schema = abc_schema(); + let expr = binary(col("b", &schema)?, Operator::Divide, lit(2i32), &schema)?; + let mapping = ProjectionMapping::from_indices(&[1], &schema)?; + let dropped_mapping = ProjectionMapping::from_indices(&[0], &schema)?; + let projected_b: Arc = Arc::new(Column::new("b", 0)); + let projected_expr: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&projected_b), + Operator::Divide, + lit(2i32), + )); + + for across in [ + AcrossPartitions::Uniform(Some(ScalarValue::Int32(Some(2)))), + AcrossPartitions::Uniform(None), + AcrossPartitions::Heterogeneous, + ] { + let mut group = EquivalenceGroup::default(); + // The class has no literal member, so `EquivalenceClass::new` cannot + // restore the constant. + group.add_constant(ConstExpr::new(Arc::clone(&expr), across.clone())); + + let mut projected = group.project(&mapping); + assert_eq!( + projected.is_expr_constant(&projected_expr), + Some(across.clone()) + ); + // A constant expression does not imply that its argument is constant. + assert_eq!(projected.is_expr_constant(&projected_b), None); + assert!(group.project(&dropped_mapping).is_empty()); + + let heterogeneous = across == AcrossPartitions::Heterogeneous; + assert_eq!(projected.clear_per_partition_constants(), heterogeneous); + assert_eq!( + projected.is_expr_constant(&projected_expr), + (!heterogeneous).then_some(across) + ); + } + + Ok(()) + } + #[test] fn test_project_constant_of_equivalent_source() -> Result<()> { // - columns: [a, b, c]. diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 605e54d1543cb..1e12ee4194a49 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -3093,7 +3093,9 @@ SELECT g, percentile_cont(x, m) FROM (SELECT value % 4 AS g, value AS x, 0.25 AS 2 26 3 27 -query error DataFusion error: Execution error: Percentile value for 'PERCENTILE_CONT' must be constant across the aggregation, found differing values 0.5 and 0.9 +# The partial states are merged in whichever order the partitions finish, +# so the two differing values can be reported in either order +query error DataFusion error: Execution error: Percentile value for 'PERCENTILE_CONT' must be constant across the aggregation, found differing values (0\.5 and 0\.9|0\.9 and 0\.5) SELECT g, percentile_cont(x, m) FROM (SELECT value % 4 AS g, value AS x, CASE WHEN value <= 50 THEN 0.5 ELSE 0.9 END AS m FROM generate_series(1, 100)) GROUP BY g query IR diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index b795b050f919f..d2be4f690022b 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -6343,3 +6343,103 @@ DROP TABLE projection_networks; statement ok set datafusion.execution.target_partitions = 4; + +# A computed constant must survive the projection embedded into a sort-merge +# join, even when only the expression's argument is projected. +# https://github.com/apache/datafusion/issues/26624 +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +statement ok +CREATE TABLE constant_projection (a INT, b INT, s VARCHAR) AS VALUES +(1, 1, 'x'), (2, 0, 'X'), (3, 1, 'y'), (4, NULL, NULL), (5, 1, 'X'); + +query II +SELECT t1.b + 1 AS x, t2.a +FROM constant_projection t1 JOIN constant_projection t2 ON t1.a = t2.a +WHERE t1.b + 1 = 2 +ORDER BY x, t2.a; +---- +2 1 +2 3 +2 5 + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT t1.b + 1 AS x, t2.a +FROM constant_projection t1 JOIN constant_projection t2 ON t1.a = t2.a +WHERE t1.b + 1 = 2 +ORDER BY x, t2.a; +---- +physical_plan +01)SortPreservingMergeExec: [x@0 ASC NULLS LAST, a@1 ASC NULLS LAST] +02)--SortExec: expr=[a@1 ASC NULLS LAST], preserve_partitioning=[true] +03)----ProjectionExec: expr=[CAST(b@0 AS Int64) + 1 as x, a@1 as a] +04)------SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)], projection=[b@1, a@2] +05)--------SortExec: expr=[a@0 ASC], preserve_partitioning=[true] +06)----------RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=1 +07)------------FilterExec: CAST(b@1 AS Int64) + 1 = 2 +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] +09)--------SortExec: expr=[a@0 ASC], preserve_partitioning=[true] +10)----------RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=1 +11)------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +reset datafusion.explain.physical_plan_only; + +query II +SELECT t1.b + 1 AS x, t2.a +FROM constant_projection t1 LEFT JOIN constant_projection t2 ON t1.a = t2.a +WHERE t1.b + 1 = 2 +ORDER BY x, t2.a; +---- +2 1 +2 3 +2 5 + +# NULL padding invalidates the right side's constant, so x still needs sorting. +query II +SELECT t2.b + 1 AS x, t1.a +FROM constant_projection t1 +LEFT JOIN (SELECT * FROM constant_projection WHERE b + 1 = 2) t2 ON t1.a = t2.a +ORDER BY x, t1.a; +---- +2 1 +2 3 +2 5 +NULL 2 +NULL 4 + +query TI +SELECT upper(t1.s) AS x, t2.a +FROM constant_projection t1 JOIN constant_projection t2 ON t1.a = t2.a +WHERE upper(t1.s) = 'X' +ORDER BY x, t2.a; +---- +X 1 +X 2 +X 5 + +statement ok +DROP TABLE constant_projection; + +# The expression must also survive when it contains a scalar-subquery leaf. +statement ok +CREATE TABLE subquery_projection (a INT, b INT) AS VALUES (1, 1), (2, 0); + +query III +SELECT * FROM ( + SELECT t1.b + (SELECT max(a) FROM subquery_projection) AS x + FROM subquery_projection t1 JOIN subquery_projection t2 ON t1.a = t2.a +) s JOIN subquery_projection t3 ON s.x = t3.a +WHERE s.x = 2; +---- +2 2 0 + +statement ok +DROP TABLE subquery_projection; + +statement ok +reset datafusion.optimizer.prefer_hash_join;