fix: skip a leading per-partition constant when finding ordering heads

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

@@ -50,12 +50,32 @@ pub(super) struct SortKeyWalk<'a> {
 
 impl<'a> SortKeyWalk<'a> {
     pub(super) fn new(props: &'a EquivalenceProperties) -> Self {
-        Self {
+        let mut walk = Self {
             props,
             pinned: HashSet::new(),
             implied: vec![],
             head_idxs: vec![],
+        };
+        // A normal ordering keeps a constant whose value differs between
+        // partitions (see `EquivalenceGroup::normalize_sort_exprs`), and such
+        // a constant may lead it. Within a partition it has one value, so it
+        // is fixed and the head is the first expression after it.
+        let normal = &props.oeq_cache.normal_cls;
+        if normal
+            .iter()
+            .any(|ordering| walk.is_fixed(&ordering[0].expr))
+        {
+            walk.head_idxs = normal
+                .iter()
+                .map(|ordering| {
+                    ordering
+                        .iter()
+                        .position(|sort_expr| !walk.is_fixed(&sort_expr.expr))
+                        .unwrap_or(ordering.len())
+                })
+                .collect();
         }
+        walk
     }
 
     /// Returns how many leading keys (in normal form) the existing orderings
@@ -328,6 +348,40 @@ mod tests {
     };
     use datafusion_expr::{ColumnarValue, Operator, Volatility, create_udf};
 
+    /// A normal ordering keeps a constant whose value differs between
+    /// partitions (`AcrossPartitions::Heterogeneous`), and such a constant
+    /// can lead it. Within a partition the constant has one value, so it is
+    /// fixed and the head is the expression after it. With `k` such a
+    /// constant and ordering `[k, b]`, the requirements `[b]` and `[b, k]`
+    /// hold before anything is pinned.
+    /// See <https://github.com/apache/datafusion/issues/26762>.
+    #[test]
+    fn leading_per_partition_constant_is_not_a_head() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("k", DataType::Int32, true),
+            Field::new("b", DataType::Int32, true),
+        ]));
+        let [k, b] = ["k", "b"].map(|name| col(name, &schema).unwrap());
+
+        let mut properties = EquivalenceProperties::new(Arc::clone(&schema));
+        properties.add_constants([ConstExpr::new(
+            Arc::clone(&k),
+            AcrossPartitions::Heterogeneous,
+        )])?;
+        // Added after the constant, so that `k` stays in the normal ordering:
+        properties.add_ordering(
+            [&k, &b].map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))),
+        );
+
+        for required in [vec![&b], vec![&b, &k], vec![&k, &b]] {
+            let required = required
+                .into_iter()
+                .map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr)));
+            assert!(properties.ordering_satisfy(required)?);
+        }
+        Ok(())
+    }
+
     /// Pinning a key can move the head of an ordering to an expression that
     /// is equal to a function of the expressions after it, which implies a
     /// new ordering. With ordering `[k, a, c]` and `CAST(c AS DATE) = a`,

datafusion/sqllogictest/test_files/window.slt

@@ -7834,3 +7834,54 @@ DROP TABLE range_null_peers;
 
 statement ok
 DROP TABLE range_nan_peers;
+
+# A window ordered by `(total DESC, k)` over an input sorted by `[k, total]`,
+# where `k` is constant within each partition (`lower(g)` under `WHERE g = 'a'`).
+# Within a partition `k` has one value, so the input ordering satisfies the
+# window's requirement and no sort is added.
+# https://github.com/apache/datafusion/issues/26762
+statement ok
+set datafusion.execution.target_partitions = 1;
+
+statement ok
+CREATE TABLE per_partition_constant(g VARCHAR, v BIGINT) AS VALUES ('a', 1), ('b', 2);
+
+query TI
+SELECT k, DENSE_RANK() OVER (ORDER BY total DESC, k) AS r
+FROM (
+  SELECT k, SUM(s) OVER (PARTITION BY k) AS total
+  FROM (SELECT lower(g) AS k, SUM(v) AS s FROM per_partition_constant WHERE g = 'a' GROUP BY lower(g))
+);
+----
+a 1
+
+statement ok
+set datafusion.explain.physical_plan_only = true;
+
+query TT
+EXPLAIN SELECT k, DENSE_RANK() OVER (ORDER BY total DESC, k) AS r
+FROM (
+  SELECT k, SUM(s) OVER (PARTITION BY k) AS total
+  FROM (SELECT lower(g) AS k, SUM(v) AS s FROM per_partition_constant WHERE g = 'a' GROUP BY lower(g))
+);
+----
+physical_plan
+01)ProjectionExec: expr=[k@0 as k, dense_rank() ORDER BY [total DESC NULLS FIRST, k ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as r]
+02)--BoundedWindowAggExec: wdw=[dense_rank() ORDER BY [total DESC NULLS FIRST, k ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "dense_rank() ORDER BY [total DESC NULLS FIRST, k ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+03)----ProjectionExec: expr=[k@0 as k, sum(s) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as total]
+04)------WindowAggExec: wdw=[sum(s) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(s) PARTITION BY [k] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
+05)--------ProjectionExec: expr=[lower(per_partition_constant.g)@0 as k, sum(per_partition_constant.v)@1 as s]
+06)----------AggregateExec: mode=Single, gby=[lower(g@0) as lower(per_partition_constant.g)], aggr=[sum(per_partition_constant.v)]
+07)------------FilterExec: g@0 = a
+08)--------------DataSourceExec: partitions=1, partition_sizes=[1]
+
+statement ok
+reset datafusion.explain.physical_plan_only;
+
+statement ok
+DROP TABLE per_partition_constant;
+
+# The SLT runner sets `target_partitions` to 4 instead of using the default, so
+# restore that value rather than resetting.
+statement ok
+set datafusion.execution.target_partitions = 4;