From 7400eddef39bdd03e32724cfec96d9f6ba673f2d Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 5 Nov 2026 12:26:52 -0500 Subject: [PATCH] fix: skip a leading per-partition constant when finding ordering heads A normal ordering keeps a constant whose value differs between partitions (`AcrossPartitions::Heterogeneous`, since #23528), so such a constant can lead a normal ordering. `SortKeyWalk` assumed every ordering's head was at index 0 until the first key was pinned, so it treated the constant as the head and did not see the expression after it. A window ordered by `(total DESC, k)` over an input ordered by `[k, total]`, with `k` constant within each partition, then failed `SanityCheckPlan`. Compute the initial heads past leading fixed expressions, as the walk already does after each pin. Closes #26762 Co-Authored-By: Claude Fable 5.1 --- .../src/equivalence/properties/walk.rs | 56 ++++++++++++++++++- datafusion/sqllogictest/test_files/window.slt | 51 +++++++++++++++++ 2 files changed, 106 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-expr/src/equivalence/properties/walk.rs b/datafusion/physical-expr/src/equivalence/properties/walk.rs index e0487cdd23dfa..826e6cb222418 100644 --- a/datafusion/physical-expr/src/equivalence/properties/walk.rs +++ b/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 . + #[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`, diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 8508ca6ff66bb..54dee349d8be8 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/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;