fix: treat an uncorrelated scalar subquery as a constant in ordering checks

datafusion/physical-expr/src/equivalence/class.rs

@@ -29,6 +29,7 @@ use crate::{PhysicalExpr, PhysicalExprRef, PhysicalSortExpr, PhysicalSortRequire
 
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::{JoinType, Result, ScalarValue};
+use datafusion_expr::sort_properties::SortProperties;
 use datafusion_physical_expr_common::physical_expr::format_physical_expr_list;
 
 use indexmap::{IndexMap, IndexSet};
@@ -361,8 +362,9 @@ impl EquivalenceGroup {
             return;
         }
         // If the expression is not in any equivalence class, but has the same
-        // constant value with some class, add it to that class:
-        if let AcrossPartitions::Uniform(_) = &const_expr.across_partitions {
+        // constant value with some class, add it to that class. An unknown
+        // value is never known to match, so it gets a class of its own:
+        if let AcrossPartitions::Uniform(Some(_)) = &const_expr.across_partitions {
             for (idx, cls) in self.classes.iter_mut().enumerate() {
                 if cls
                     .constant
@@ -792,9 +794,21 @@ impl EquivalenceGroup {
         // TODO: This function should be able to return values of non-literal
         //       complex constants as well; e.g. it should return `8` for the
         //       expression `3 + 5`, not an unknown `heterogenous` value.
+        // A volatile expression can change from row to row, even when all of
+        // its inputs are constant
+        if expr.is_volatile_node() {
+            return None;
+        }
         let children = expr.children();
         if children.is_empty() {
-            return None;
+            // A leaf that reads no column can still be constant, e.g. an
+            // uncorrelated scalar subquery, which reports itself as a
+            // singleton. It does not depend on the row or the partition, so
+            // its value is the same everywhere, even if it is not known yet.
+            let singleton = expr
+                .get_properties(&[])
+                .is_ok_and(|props| props.sort_properties == SortProperties::Singleton);
+            return singleton.then_some(AcrossPartitions::Uniform(None));
         }
         for child in children {
             self.is_expr_constant(child)?;
@@ -822,21 +836,35 @@ impl EquivalenceGroup {
         let group = match join_type {
             JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::Right => {
                 let mut result = Self::new(
-                    self.iter().cloned().chain(
-                        right_equivalences
-                            .iter()
-                            .map(|cls| {
-                                let mut cls = cls.try_with_offset(left_size as _)?;
-                                // Rows without a match on the right side are
-                                // padded with NULLs, so its constants no
-                                // longer hold.
-                                if matches!(join_type, JoinType::Left | JoinType::Full) {
-                                    cls.constant = None;
-                                }
-                                Ok(cls)
-                            })
-                            .collect::<Result<Vec<_>>>()?,
-                    ),
+                    self.iter()
+                        .cloned()
+                        .map(|mut cls| {
+                            // Rows without a match on the left side are
+                            // padded with NULLs, so its constants no
+                            // longer hold.
+                            if matches!(join_type, JoinType::Right | JoinType::Full) {
+                                cls.constant = None;
+                            }
+                            cls
+                        })
+                        .chain(
+                            right_equivalences
+                                .iter()
+                                .map(|cls| {
+                                    let mut cls = cls.try_with_offset(left_size as _)?;
+                                    // Rows without a match on the right side are
+                                    // padded with NULLs, so its constants no
+                                    // longer hold.
+                                    if matches!(
+                                        join_type,
+                                        JoinType::Left | JoinType::Full
+                                    ) {
+                                        cls.constant = None;
+                                    }
+                                    Ok(cls)
+                                })
+                                .collect::<Result<Vec<_>>>()?,
+                        ),
                 );
                 // In we have an inner join, expressions in the "on" condition
                 // are equal in the resulting table.
@@ -1362,6 +1390,117 @@ mod tests {
         Ok(())
     }
 
+    /// An uncorrelated scalar subquery is a leaf with no columns, so the only
+    /// way to know it is constant is to ask the expression itself.
+    /// <https://github.com/apache/datafusion/issues/26572>
+    #[test]
+    fn test_scalar_subquery_is_uniform_constant() -> Result<()> {
+        use crate::scalar_subquery::ScalarSubqueryExpr;
+        use datafusion_expr::physical_planning_context::{
+            ScalarSubqueryResults, SubqueryIndex,
+        };
+
+        let subquery: Arc<dyn PhysicalExpr> = Arc::new(ScalarSubqueryExpr::new(
+            DataType::Int32,
+            true,
+            SubqueryIndex::new(0),
+            ScalarSubqueryResults::new(1),
+        ));
+        let group = EquivalenceGroup::default();
+        assert_eq!(
+            group.is_expr_constant(&subquery),
+            Some(AcrossPartitions::Uniform(None))
+        );
+        // A column is a leaf too, but it is not constant
+        let schema = abc_schema();
+        assert_eq!(group.is_expr_constant(&col("a", &schema)?), None);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_volatile_function_of_constants_is_not_constant() -> Result<()> {
+        use crate::ScalarFunctionExpr;
+        use crate::scalar_subquery::ScalarSubqueryExpr;
+        use datafusion_common::config::ConfigOptions;
+        use datafusion_expr::physical_planning_context::{
+            ScalarSubqueryResults, SubqueryIndex,
+        };
+        use datafusion_expr::{ColumnarValue, Volatility, create_udf};
+
+        let schema = abc_schema();
+        let group = EquivalenceGroup::default();
+        let subquery: Arc<dyn PhysicalExpr> = Arc::new(ScalarSubqueryExpr::new(
+            DataType::Int32,
+            true,
+            SubqueryIndex::new(0),
+            ScalarSubqueryResults::new(1),
+        ));
+        let call = |volatility, arg: Arc<dyn PhysicalExpr>| -> Result<_> {
+            let udf = Arc::new(create_udf(
+                "f",
+                vec![DataType::Int32],
+                DataType::Int32,
+                volatility,
+                Arc::new(|args: &[ColumnarValue]| Ok(args[0].clone())),
+            ));
+            let expr = ScalarFunctionExpr::try_new(
+                udf,
+                vec![arg],
+                &schema,
+                Arc::new(ConfigOptions::new()),
+            )?;
+            Ok(Arc::new(expr) as Arc<dyn PhysicalExpr>)
+        };
+
+        // A volatile function can return a new value for every row, even
+        // when its argument is constant
+        for arg in [lit(1i32), Arc::clone(&subquery)] {
+            let volatile = call(Volatility::Volatile, arg)?;
+            assert_eq!(group.is_expr_constant(&volatile), None);
+        }
+        // A stable function of a constant is constant
+        let stable = call(Volatility::Stable, subquery)?;
+        assert_eq!(
+            group.is_expr_constant(&stable),
+            Some(AcrossPartitions::Heterogeneous)
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_unknown_constants_stay_apart() -> Result<()> {
+        // Two constants whose values are unknown are not known to be equal,
+        // so they must not be merged into one class
+        let schema = abc_schema();
+        let mut group = EquivalenceGroup::default();
+        group.add_constant(ConstExpr::new(
+            col("a", &schema)?,
+            AcrossPartitions::Uniform(None),
+        ));
+        group.add_constant(ConstExpr::new(
+            col("b", &schema)?,
+            AcrossPartitions::Uniform(None),
+        ));
+        assert_eq!(group.len(), 2);
+
+        // Two constants with the same known value are merged
+        let mut group = EquivalenceGroup::default();
+        let one = Some(ScalarValue::Int32(Some(1)));
+        group.add_constant(ConstExpr::new(
+            col("a", &schema)?,
+            AcrossPartitions::Uniform(one.clone()),
+        ));
+        group.add_constant(ConstExpr::new(
+            col("b", &schema)?,
+            AcrossPartitions::Uniform(one),
+        ));
+        assert_eq!(group.len(), 1);
+
+        Ok(())
+    }
+
     /// A constant on the right side of a join stays constant in the join
     /// output, unless unmatched rows pad the right side with NULLs.
     /// <https://github.com/apache/datafusion/issues/26283>
@@ -1390,6 +1529,31 @@ mod tests {
         Ok(())
     }
 
+    /// The same for a constant on the left side, whose rows are padded with
+    /// NULLs by a right or full join.
+    #[test]
+    fn test_join_keeps_left_constants() -> Result<()> {
+        let schema = abc_schema();
+        let mut left = EquivalenceGroup::default();
+        left.add_constant(ConstExpr::from(col("a", &schema)?));
+        let left_a = col("a", &schema)?;
+        for (join_type, expected) in [
+            (JoinType::Inner, true),
+            (JoinType::Left, true),
+            (JoinType::Right, false),
+            (JoinType::Full, false),
+        ] {
+            let joined = left.join(&EquivalenceGroup::default(), &join_type, 4, &[])?;
+            assert_eq!(
+                joined.is_expr_constant(&left_a).is_some(),
+                expected,
+                "{join_type}"
+            );
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn test_has_same_classes_ignores_the_map() -> Result<()> {
         // `map` indexes into `classes`, so equal classes must imply equal

datafusion/physical-expr/src/equivalence/properties/union.rs

@@ -58,8 +58,12 @@ fn calculate_union_binary(
                 .map(|rhs_const| {
                     let mut const_expr = lhs_const.clone();
                     // If both sides have matching constant values, preserve it.
-                    // Otherwise, set fall back to heterogeneous values.
-                    if lhs_const.across_partitions != rhs_const.across_partitions {
+                    // Otherwise, set fall back to heterogeneous values. Two
+                    // unknown values are not known to match, so they fall back
+                    // too.
+                    if lhs_const.across_partitions != rhs_const.across_partitions
+                        || lhs_const.across_partitions == AcrossPartitions::Uniform(None)
+                    {
                         const_expr.across_partitions = AcrossPartitions::Heterogeneous;
                     }
                     const_expr

datafusion/sqllogictest/test_files/subquery.slt

@@ -4245,3 +4245,133 @@ SELECT * FROM (VALUES (1)) AS x(d)
 WHERE x.d IN (SELECT t.c FROM (SELECT * FROM (VALUES (1)) AS y(c) WHERE y.c = x.d) AS t WHERE x.d > 0);
 ----
 1
+
+# Regression test for #26572: an uncorrelated scalar subquery is constant, so a
+# requirement on it is satisfied by any ordering, even when it is not projected
+# from a literal.
+statement ok
+SET datafusion.execution.batch_size = 1;
+
+statement ok
+CREATE TABLE t26572 (a INT) AS VALUES (1), (2);
+
+query II
+SELECT t.a, (SELECT max(a) FROM t26572) AS m
+FROM t26572 AS u RIGHT JOIN t26572 AS t ON u.a = t.a
+ORDER BY t.a, m LIMIT 5;
+----
+1 2
+2 2
+
+query II
+SELECT t.a, (SELECT max(a) FROM t26572) AS m
+FROM t26572 AS u RIGHT JOIN t26572 AS t ON u.a = t.a
+ORDER BY m, t.a DESC LIMIT 5;
+----
+2 2
+1 2
+
+# Control, passes without the fix: a volatile leaf such as random() is not
+# constant, so it stays in the sort.
+query TT
+EXPLAIN SELECT a FROM t26572 ORDER BY a, random();
+----
+logical_plan
+01)Sort: t26572.a ASC NULLS LAST, random() ASC NULLS LAST
+02)--TableScan: t26572 projection=[a]
+physical_plan
+01)SortExec: expr=[a@0 ASC NULLS LAST, random() ASC NULLS LAST], preserve_partitioning=[false]
+02)--DataSourceExec: partitions=1, partition_sizes=[2]
+
+# Two different subqueries are both constant, but not known to be equal, so a
+# UNION of them is not constant and the sort stays.
+query TT
+EXPLAIN SELECT m FROM (
+  SELECT (SELECT max(a) FROM t26572) AS m FROM t26572
+  UNION ALL
+  SELECT (SELECT min(a) FROM t26572) AS m FROM t26572
+) ORDER BY m;
+----
+logical_plan
+01)Sort: m ASC NULLS LAST
+02)--Union
+03)----Projection: (<subquery>) AS m
+04)------Subquery:
+05)--------Aggregate: groupBy=[[]], aggr=[[max(t26572.a)]]
+06)----------TableScan: t26572 projection=[a]
+07)------TableScan: t26572 projection=[]
+08)----Projection: (<subquery>) AS m
+09)------Subquery:
+10)--------Aggregate: groupBy=[[]], aggr=[[min(t26572.a)]]
+11)----------TableScan: t26572 projection=[a]
+12)------TableScan: t26572 projection=[]
+physical_plan
+01)ScalarSubqueryExec: subqueries=2
+02)--SortPreservingMergeExec: [m@0 ASC NULLS LAST]
+03)----UnionExec
+04)------ProjectionExec: expr=[scalar_subquery(<pending>) as m]
+05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+06)----------DataSourceExec: partitions=1, partition_sizes=[2]
+07)------ProjectionExec: expr=[scalar_subquery(<pending>) as m]
+08)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+09)----------DataSourceExec: partitions=1, partition_sizes=[2]
+10)--AggregateExec: mode=Final, gby=[], aggr=[max(t26572.a)]
+11)----CoalescePartitionsExec
+12)------AggregateExec: mode=Partial, gby=[], aggr=[max(t26572.a)]
+13)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+14)----------DataSourceExec: partitions=1, partition_sizes=[2]
+15)--AggregateExec: mode=Final, gby=[], aggr=[min(t26572.a)]
+16)----CoalescePartitionsExec
+17)------AggregateExec: mode=Partial, gby=[], aggr=[min(t26572.a)]
+18)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+19)----------DataSourceExec: partitions=1, partition_sizes=[2]
+
+query I
+SELECT m FROM (
+  SELECT (SELECT max(a) FROM t26572) AS m FROM t26572
+  UNION ALL
+  SELECT (SELECT min(a) FROM t26572) AS m FROM t26572
+) ORDER BY m;
+----
+1
+1
+2
+2
+
+# Nor are two different subqueries in one projection the same expression.
+query III
+SELECT (SELECT max(a) FROM t26572) AS x, (SELECT min(a) FROM t26572) AS y, a
+FROM t26572 ORDER BY y, x, a DESC;
+----
+2 1 2
+2 1 1
+
+statement ok
+DROP TABLE t26572;
+
+# A LEFT JOIN can be swapped into a physical RIGHT join with the subquery side
+# on the left. Its rows without a match are padded with NULLs, so `m` is not
+# constant there and the sort on it must stay.
+statement ok
+CREATE TABLE p26572 (a INT, b INT) AS VALUES (1, 10), (2, 20), (3, 30), (2, 5), (NULL, 7);
+
+query III
+SELECT x.a, s.a, s.m
+FROM p26572 x LEFT JOIN (
+  SELECT a, (SELECT max(a) FROM p26572) AS m FROM p26572 WHERE a > 1
+) s ON x.a = s.a
+ORDER BY s.m, x.a, s.a;
+----
+2 2 3
+2 2 3
+2 2 3
+2 2 3
+3 3 3
+1 NULL NULL
+NULL NULL NULL
+
+statement ok
+DROP TABLE p26572;
+
+statement ok
+RESET datafusion.execution.batch_size;