[branch-56] fix: preserve constant equivalence properties through projection (backport #27123)

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<dyn PhysicalExpr>,
     ) -> Option<Arc<dyn PhysicalExpr>> {
-        // Literals don't need to be projected
-        if expr.downcast_ref::<Literal>().is_some() {
+        // These leaves do not read columns from the projected input schema.
+        if expr.is::<Literal>() || expr.is::<ScalarSubqueryExpr>() {
             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<dyn PhysicalExpr> = Arc::new(Column::new("b", 0));
+        let projected_expr: Arc<dyn PhysicalExpr> = 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].

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

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;