You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ur...@apache.org on 2026/07/22 15:43:18 UTC
(spark) branch master updated: [SPARK-58233][SQL] Push down local sort to reduce redundant sorts within a stage
This is an automated email from the ASF dual-hosted git repository.
uros-b pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 589adcb9b0b4 [SPARK-58233][SQL] Push down local sort to reduce redundant sorts within a stage
589adcb9b0b4 is described below
commit 589adcb9b0b4ee247d37acadd851baac776a5a91
Author: Xiduo You <ul...@gmail.com>
AuthorDate: Wed Jul 22 15:43:04 2026 +0000
[SPARK-58233][SQL] Push down local sort to reduce redundant sorts within a stage
### What changes were proposed in this pull request?
Adds a new physical rule `PushDownLocalSort`, wired into both the AQE `queryStagePreparationRules` and the non-AQE `QueryExecution` preparations (right after `EnsureRequirements`, before `CombineAdjacentAggregation`), gated by a new internal config `spark.sql.execution.pushDownLocalSort` (default `true`).
`EnsureRequirements` inserts one local `SortExec(global = false)` above every operator whose `requiredChildOrdering` is not satisfied. When two such requirements are in a prefix-cover relationship, a stage computes multiple local sorts that only differ in width. The rule pushes the wider local sort down through order-preserving operators onto the narrower one below, widening it so a single sort serves both operators and the redundant upper sort is dropped.
For example, a sort aggregate stacked on a window over the same clustering keys:
```
SortAggregate(key = [a, b, c])
Sort([a, b, c], global = false) <- upper, wider
Window([a], [b])
Sort([a, b], global = false) <- lower, narrower
Exchange(hashpartitioning([a]))
```
becomes:
```
SortAggregate(key = [a, b, c])
Window([a], [b]) requiredChildOrdering [a, b] satisfied by [a, b, c]
Sort([a, b, c], global = false) <- single sort now serves both operators
Exchange(hashpartitioning([a]))
```
Order-preserving operators traversed: `ProjectExec`, `WindowExecBase`, and — only when `spark.sql.execution.pushDownLocalSort.throughCardinalityReducer` is enabled (default `false`) — the cardinality reducers `FilterExec` and `WindowGroupLimitExec`. Crossing a reducer can move the wider sort from the surviving rows to the full input, which may outweigh the saved sort, so it is opt-in. A `ProjectExec` or `FilterExec` is crossed only when its expressions are deterministic (mirroring `El [...]
`CollectMetricsExec` is intentionally excluded: it observes rows in input order, so widening a sort under it would silently change order-sensitive observed metrics (`first`/`last`/`collect_list`).
This is complementary to `RemoveRedundantSorts`, which only removes a sort already satisfied by its child; it never moves a wider sort down to absorb a narrower one.
### Why are the changes needed?
Stacked windows, or a sort aggregate over a window with prefix-compatible ordering requirements, compute several local sorts that a single wider sort could satisfy. Reusing the widest sort saves the redundant sorting work in the stage.
### Does this PR introduce _any_ user-facing change?
No. It only removes redundant local sorts from physical plans; the result of a deterministic query is unchanged. As with any plan change that alters within-partition row order, an order-sensitive window function over a non-unique `ORDER BY` (e.g. `first_value`/`last_value`/`collect_list`) may return a different — but still valid — non-deterministic result. Guarded by `spark.sql.execution.pushDownLocalSort`.
### How was this patch tested?
New `PushDownLocalSortSuite` (AQE on/off variants) with:
- positive SQL cases: stacked windows, filter above the sorts, window to sort aggregate, renaming project;
- negative SQL cases: disjoint orderings, a shuffle between the sorts, sort-direction mismatch, a project that computes the ordering column;
- plan-level cases: multi-operator chain, renaming-project rewrite, three stacked sorts in one pass, non-deterministic project/filter are not crossed, and a cardinality-reducing filter is crossed only when the reducer config is enabled.
Verified the TPCDS/TPCH `PlanStability` golden plans are unchanged.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #57396 from ulysses-you/spark-pushdown-merge-local-sort.
Authored-by: Xiduo You <ul...@gmail.com>
Signed-off-by: Uros Bojanic <22...@users.noreply.github.com>
---
.../org/apache/spark/sql/internal/SQLConf.scala | 26 ++
.../spark/sql/execution/PushDownLocalSort.scala | 156 ++++++++
.../spark/sql/execution/QueryExecution.scala | 4 +
.../execution/adaptive/AdaptiveSparkPlanExec.scala | 4 +
.../sql/execution/PushDownLocalSortSuite.scala | 401 +++++++++++++++++++++
5 files changed, 591 insertions(+)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 5cb317a52b33..ea4c2ffc1559 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -3001,6 +3001,32 @@ object SQLConf {
.booleanConf
.createWithDefault(true)
+ val PUSH_DOWN_LOCAL_SORT_ENABLED =
+ buildConf("spark.sql.execution.pushDownLocalSort")
+ .internal()
+ .doc("When true, pushes a wider local sort down through order-preserving " +
+ "operators to replace a narrower local sort below it, so that a single sort can satisfy " +
+ "multiple operators' ordering requirements. This reduces the total number of local sorts " +
+ "computed in a stage, for example when a sort aggregate is stacked on a window over the " +
+ "same clustering keys.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .booleanConf
+ .createWithDefault(true)
+
+ val PUSH_DOWN_LOCAL_SORT_THROUGH_CARDINALITY_REDUCER_ENABLED =
+ buildConf("spark.sql.execution.pushDownLocalSort.throughCardinalityReducer")
+ .internal()
+ .doc("When true, `spark.sql.execution.pushDownLocalSort` may also push a wider local sort " +
+ "down through cardinality-reducing operators (`FilterExec` and `WindowGroupLimitExec`). " +
+ "This is disabled by default because moving the wider sort below a selective reducer can " +
+ "sort the full input instead of only the surviving rows, which may outweigh the saved " +
+ "sort. Has no effect when `spark.sql.execution.pushDownLocalSort` is false.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .booleanConf
+ .createWithDefault(false)
+
val REPLACE_HASH_WITH_SORT_AGG_ENABLED = buildConf("spark.sql.execution.replaceHashWithSortAgg")
.internal()
.doc("Whether to replace hash aggregate node with sort aggregate based on children's ordering")
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/PushDownLocalSort.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/PushDownLocalSort.scala
new file mode 100644
index 000000000000..161e779a0c9a
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/PushDownLocalSort.scala
@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, AttributeReference, AttributeSet, SortOrder}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.window.{WindowExecBase, WindowGroupLimitExec}
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Pushes a wider local sort down through order-preserving operators onto a narrower local sort
+ * below, widening it so that a single sort satisfies several operators' ordering requirements
+ * instead of re-sorting once per operator.
+ *
+ * `EnsureRequirements` adds one local `SortExec` (`global = false`) above every operator whose
+ * `requiredChildOrdering` is not already satisfied. When such requirements are in a prefix-cover
+ * relationship, this produces multiple local sorts that only differ in width. A canonical case is a
+ * sort aggregate stacked on a window over the same clustering keys, where the aggregate needs a
+ * wider ordering than the window:
+ *
+ * {{{
+ * SortAggregate(key = [a, b, c])
+ * Sort([a, b, c], global = false) <- upper, wider
+ * Window([a], [b])
+ * Sort([a, b], global = false) <- lower, narrower
+ * Exchange(hashpartitioning([a]))
+ * }}}
+ *
+ * Because every operator between the two sorts is order-preserving and the upper ordering
+ * prefix-covers everything required along the way, the wider ordering can be pushed down to widen
+ * the lower sort, and the upper sort then dropped entirely:
+ *
+ * {{{
+ * SortAggregate(key = [a, b, c])
+ * Window([a], [b]) requiredChildOrdering [a, b] is satisfied by [a, b, c]
+ * Sort([a, b, c], global = false) <- single sort now serves both operators
+ * Exchange(hashpartitioning([a]))
+ * }}}
+ *
+ * When a `ProjectExec` on the path renames an ordering column in its output (`b AS x`), the
+ * ordering is rewritten from the project's output space back to its child's space (`x` -> `b`) as
+ * it is pushed through, so a sort over the renamed column is still matched below. Only plain
+ * renames are followed, and the rule never crosses a shuffle or a non-order-preserving operator.
+ */
+object PushDownLocalSort extends Rule[SparkPlan] {
+
+ def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.getConf(SQLConf.PUSH_DOWN_LOCAL_SORT_ENABLED)) {
+ return plan
+ }
+ val throughCardinalityReducer =
+ conf.getConf(SQLConf.PUSH_DOWN_LOCAL_SORT_THROUGH_CARDINALITY_REDUCER_ENABLED)
+
+ plan.transform {
+ case upper @ SortExec(upperOrder, false, child, _) =>
+ pushDown(child, upperOrder, throughCardinalityReducer).getOrElse(upper)
+ }
+ }
+
+ /**
+ * Walks down from `plan` through a chain of order-preserving unary operators, looking for a
+ * lower local `SortExec` that `upperOrder` strictly covers. When found, widens that lower sort
+ * to `upperOrder` and returns the rebuilt subtree (which re-exposes `upperOrder` at its top);
+ * returns `None` if no safe widening applies, leaving the plan untouched. As it crosses an
+ * operator that renames ordering columns, `upperOrder` is rewritten into that operator's child
+ * space so the search continues against the child's own attributes. `throughCardinalityReducer`
+ * controls whether cardinality-reducing operators (`FilterExec`, `WindowGroupLimitExec`) may be
+ * crossed.
+ */
+ private def pushDown(
+ plan: SparkPlan,
+ upperOrder: Seq[SortOrder],
+ throughCardinalityReducer: Boolean): Option[SparkPlan] = plan match {
+ case lower @ SortExec(lowerOrder, false, _, _)
+ // Only widen when the upper ordering strictly covers the lower one. When they are
+ // equivalent the upper sort is plainly redundant and is left to `RemoveRedundantSorts`; a
+ // non-covering ordering cannot serve the lower requirement. The column check keeps the
+ // widened sort well-formed (every key of `upperOrder` is available below the lower sort).
+ if SortOrder.orderingSatisfies(upperOrder, lowerOrder) &&
+ !SortOrder.orderingSatisfies(lowerOrder, upperOrder) &&
+ AttributeSet(upperOrder.flatMap(_.references)).subsetOf(lower.child.outputSet) =>
+ Some(SortExec(upperOrder, global = false, child = lower.child))
+
+ case op: UnaryExecNode if isOrderPreserving(op, throughCardinalityReducer) =>
+ // A `ProjectExec` may rename ordering columns in its output (`b AS x`). Rewrite `upperOrder`
+ // from the operator's output space back to its child's space before pushing further down, so
+ // a sort over the renamed column is still matched below. Only plain renames are followed; an
+ // expression alias leaves the sort key referencing an output attribute the child does not
+ // produce, so the `subsetOf(op.child.outputSet)` check below rejects it.
+ val outputExprs = plan match {
+ case p: ProjectExec => p.projectList
+ case _ => Nil
+ }
+ val rewrittenUpperOrder = if (outputExprs.isEmpty) {
+ upperOrder
+ } else {
+ val aliasToAttributeMap = AttributeMap(outputExprs.collect {
+ case a @ Alias(child: AttributeReference, _) => (a.toAttribute, child: Attribute)
+ })
+ upperOrder.map { _.transformUp {
+ case a: Attribute => aliasToAttributeMap.getOrElse(a, a)
+ }.asInstanceOf[SortOrder]
+ }
+ }
+ if (SortOrder.orderingSatisfies(rewrittenUpperOrder, op.requiredChildOrdering.head) &&
+ AttributeSet(rewrittenUpperOrder.flatMap(_.references)).subsetOf(op.child.outputSet)) {
+ pushDown(op.child, rewrittenUpperOrder, throughCardinalityReducer)
+ .map(newChild => op.withNewChildren(Seq(newChild)))
+ } else {
+ None
+ }
+
+ case _ => None
+ }
+
+ private def isOrderPreserving(
+ plan: UnaryExecNode,
+ throughCardinalityReducer: Boolean): Boolean = plan match {
+ // A non-deterministic project/filter must not be crossed: moving the sort below it changes
+ // which rows a seeded non-deterministic expression is evaluated over (a different row-value
+ // association for a project, a different surviving set for a filter). This mirrors the
+ // determinism guard in the logical `EliminateSorts.canEliminateSort`.
+ case p: ProjectExec => p.projectList.forall(_.deterministic)
+ // `FilterExec` and `WindowGroupLimitExec` are cardinality reducers: pushing the wider sort
+ // below them sorts the full input instead of only the surviving rows, which can outweigh the
+ // saved sort. Only cross them when explicitly enabled. `FilterExec` additionally needs the
+ // determinism guard, for the same reason as the project above.
+ case f: FilterExec => throughCardinalityReducer && f.condition.deterministic
+ case _: WindowGroupLimitExec => throughCardinalityReducer
+ // Pushing a wider sort under a window refines the input order the window sees within ties of
+ // its own `ORDER BY`. For order-sensitive window functions (`first_value`/`last_value`/
+ // `collect_list`, `lead`/`lag`, `ROWS`-frame aggregates) that can change the result -- but only
+ // for a non-unique window `ORDER BY`, where the result is already non-deterministic by Spark's
+ // contract, so this does not change any deterministic result. This differs from
+ // `CollectMetricsExec`, which is deliberately excluded: its observed metric is a documented,
+ // expected-stable side output, not a non-deterministic query result, so refining the order it
+ // sees would be a real surprise.
+ case _: WindowExecBase => true
+ case _ => false
+ }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
index f70eafc70af3..0a1185aa4ea5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
@@ -764,6 +764,10 @@ object QueryExecution {
EnsureRequirements(),
// This rule must be run after `EnsureRequirements`.
InsertSortForLimitAndOffset,
+ // `PushDownLocalSort` pushes a wider local sort down onto a narrower one below it, so a
+ // single sort serves several operators' ordering requirements. It must run after
+ // `EnsureRequirements`, which is what inserts the local sorts it pushes down.
+ PushDownLocalSort,
// `CombineAdjacentAggregation` must run before `ReplaceHashWithSortAgg`: it combines a pair
// of adjacent partial and final aggregate into a single `Complete` mode aggregate, which
// `ReplaceHashWithSortAgg` can then replace with a sort aggregate when the ordering allows.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
index 1905b31b49f9..bfe6a9a3f633 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
@@ -131,6 +131,10 @@ case class AdaptiveSparkPlanExec(
// join drops its child ordering, which `ReplaceHashWithSortAgg` would otherwise rely on to
// turn a hash aggregate into a sort aggregate.
ConvertSortMergeJoinToShuffledHashJoin(ensureRequirements),
+ // `PushDownLocalSort` pushes a wider local sort down onto a narrower one below it, so a
+ // single sort serves several operators' ordering requirements. It must run after
+ // `EnsureRequirements`, which is what inserts the local sorts it pushes down.
+ PushDownLocalSort,
// `CombineAdjacentAggregation` must run before `ReplaceHashWithSortAgg`: it combines a pair
// of adjacent partial and final aggregate into a single `Complete` mode aggregate, which
// `ReplaceHashWithSortAgg` can then replace with a sort aggregate when the ordering allows.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PushDownLocalSortSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PushDownLocalSortSuite.scala
new file mode 100644
index 000000000000..42ea16048d66
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PushDownLocalSortSuite.scala
@@ -0,0 +1,401 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, AttributeReference, IsNotNull, LessThan, Literal, Rand, SortOrder}
+import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
+import org.apache.spark.sql.execution.exchange.ValidateRequirements
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.IntegerType
+
+
+abstract class PushDownLocalSortSuiteBase
+ extends SharedSparkSession
+ with AdaptiveSparkPlanHelper {
+
+ private def checkNumSorts(df: DataFrame, count: Int): Unit = {
+ // Execute first so that, under AQE, the final adaptive plan (after query-stage materialization
+ // and any replanning) is inspected rather than the initial plan.
+ df.collect()
+ val plan = df.queryExecution.executedPlan
+ assert(collectWithSubqueries(plan) { case s: SortExec => s }.length == count)
+ }
+
+ private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = {
+ withSQLConf(SQLConf.PUSH_DOWN_LOCAL_SORT_ENABLED.key -> "true") {
+ val df = sql(query)
+ checkNumSorts(df, enabledCount)
+ val result = df.collect()
+ withSQLConf(SQLConf.PUSH_DOWN_LOCAL_SORT_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkNumSorts(df, disabledCount)
+ checkAnswer(df, result)
+ }
+ }
+ }
+
+ test("Push a wider local sort down across stacked windows with prefix-compatible order specs") {
+ withTempView("t") {
+ spark.range(100).selectExpr("id % 10 as a", "id % 7 as b", "id as c")
+ .createOrReplaceTempView("t")
+ // The narrower window (order by b) is listed first, so it is planned closest to the leaf and
+ // the wider window (order by b, c) ends up above it. Without the rule two local sorts
+ // ([a, b] below the inner window, [a, b, c] above it) are computed. The rule widens the
+ // lower sort to [a, b, c] so it serves both windows and drops the upper sort, leaving a
+ // single [a, b, c] sort.
+ val query =
+ """
+ |SELECT a, b, c,
+ | RANK() OVER (PARTITION BY a ORDER BY b) AS rk,
+ | ROW_NUMBER() OVER (PARTITION BY a ORDER BY b, c) AS rn
+ |FROM t
+ |""".stripMargin
+ checkSorts(query, 1, 2)
+ }
+ }
+
+ test("No-op when the wider sort is already below the narrower one") {
+ withTempView("t") {
+ spark.range(100).selectExpr("id % 10 as a", "id % 7 as b", "id as c")
+ .createOrReplaceTempView("t")
+ // The wider window (order by b, c) is listed first, so it is planned closest to the leaf and
+ // the narrower window (order by b) ends up above it. `EnsureRequirements` inserts only one
+ // sort here: the wider window's [a, b, c] sort already satisfies the narrower window's
+ // [a, b] requirement, so no second sort is added. This rule only pushes a wider sort down,
+ // so it does not fire; the single sort is unchanged whether it is on or off.
+ val query =
+ """
+ |SELECT a, b, c,
+ | ROW_NUMBER() OVER (PARTITION BY a ORDER BY b, c) AS rn,
+ | RANK() OVER (PARTITION BY a ORDER BY b) AS rk
+ |FROM t
+ |""".stripMargin
+ checkSorts(query, 1, 1)
+ }
+ }
+
+ test("Push-down still applies and stays correct with a filter on a window output") {
+ withTempView("t") {
+ spark.range(200).selectExpr("id % 10 as a", "id % 7 as b", "id as c")
+ .createOrReplaceTempView("t")
+ val query =
+ """
+ |SELECT * FROM (
+ | SELECT a, b, c,
+ | RANK() OVER (PARTITION BY a ORDER BY b) AS rk,
+ | ROW_NUMBER() OVER (PARTITION BY a ORDER BY b, c) AS rn
+ | FROM t
+ |) WHERE rn > 1
+ |""".stripMargin
+ // The filter on rn sits above both windows and does not affect the two sorts that feed them,
+ // so the push-down still reduces 2 sorts to 1 and the results are unchanged.
+ checkSorts(query, 1, 2)
+ }
+ }
+
+ test("Push a wider sort down through a window to feed a sort aggregate above it") {
+ withTempView("t") {
+ // `c` varies within each (a, b) group (70 % 13 != 0) and `b` repeats within each partition
+ // `a`, so widening the window's sort from [a, b] to [a, b, c] genuinely reorders rows within
+ // the window's `ORDER BY b` ties. `RANK()` is tie-safe -- its value does not depend on the
+ // order of tied rows -- so the result stays stable while the sort count still drops.
+ spark.range(500).selectExpr("id % 10 as a", "id % 7 as b", "id % 13 as c")
+ .createOrReplaceTempView("t")
+ // Plan shape within one stage: shuffle -> Sort([a,b,c]) -> Window([a],[b]) -> Sort([a,b,c])
+ // -> SortAggregate(group by a,b,c). The window needs [a,b] and the sort aggregate needs the
+ // wider [a,b,c]; the aggregate's grouping keys are clustered-compatible with the window's
+ // partitioning, so no shuffle separates them. The wider [a,b,c] sort feeding the aggregate is
+ // pushed down through the window, replacing the window's [a,b] sort and serving both.
+ // `collect_list` with object-hash aggregation off forces a sort aggregate.
+ withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "false") {
+ val query =
+ """
+ |SELECT a, b, c, collect_list(rk) AS cl
+ |FROM (SELECT a, b, c, RANK() OVER (PARTITION BY a ORDER BY b) AS rk FROM t)
+ |GROUP BY a, b, c
+ |""".stripMargin
+ checkSorts(query, 1, 2)
+ }
+ }
+ }
+
+ test("Push a wider sort down through a renaming project by rewriting the ordering") {
+ withTempView("t") {
+ spark.range(200).selectExpr("id % 10 as a", "id % 7 as b", "id % 5 as c")
+ .createOrReplaceTempView("t")
+ // The inner window sorts by [a, b]; a project then renames b to bb; the outer window sorts by
+ // [a, bb, c] (wider). The wider sort's key `bb` is the renamed `b`, so pushing it down
+ // through the renaming project requires rewriting the ordering from bb back to b. After the
+ // rewrite the [a, b, c] sort is pushed through the inner window, replacing its [a, b] sort.
+ val query =
+ """
+ |SELECT a, bb, c, rn1,
+ | ROW_NUMBER() OVER (PARTITION BY a ORDER BY bb, c) AS rn2
+ |FROM (
+ | SELECT a, b AS bb, c, rn1 FROM (
+ | SELECT a, b, c, RANK() OVER (PARTITION BY a ORDER BY b) AS rn1 FROM t
+ | )
+ |)
+ |""".stripMargin
+ checkSorts(query, 1, 2)
+ }
+ }
+
+ test("Negative: no push-down when the two orderings are disjoint") {
+ withTempView("t") {
+ spark.range(300).selectExpr("id % 10 as a", "id % 7 as b", "id % 5 as c")
+ .createOrReplaceTempView("t")
+ // Two windows over the same partition but disjoint order columns (b vs c). Neither ordering
+ // covers the other, so neither local sort can be pushed onto the other: both survive.
+ val query =
+ """
+ |SELECT a, b, c,
+ | RANK() OVER (PARTITION BY a ORDER BY b) AS rk,
+ | RANK() OVER (PARTITION BY a ORDER BY c) AS rc
+ |FROM t
+ |""".stripMargin
+ checkSorts(query, 2, 2)
+ }
+ }
+
+ test("Negative: no push-down when a shuffle separates the two sorts") {
+ withTempView("t") {
+ spark.range(300).selectExpr("id % 10 as a", "id % 7 as b", "id % 5 as c")
+ .createOrReplaceTempView("t")
+ // Two windows partitioned by different keys (a vs b) force a shuffle between their local
+ // sorts. A local sort is never pushed across a shuffle, so both sorts survive.
+ val query =
+ """
+ |SELECT a, b,
+ | RANK() OVER (PARTITION BY a ORDER BY b) AS rk,
+ | RANK() OVER (PARTITION BY b ORDER BY a) AS rb
+ |FROM t
+ |""".stripMargin
+ checkSorts(query, 2, 2)
+ }
+ }
+
+ test("Negative: no push-down when the sort directions differ") {
+ withTempView("t") {
+ spark.range(300).selectExpr("id % 10 as a", "id % 7 as b", "id % 5 as c")
+ .createOrReplaceTempView("t")
+ // The inner window orders by b ascending; the outer orders by b descending then c. The outer
+ // ordering [a, b DESC, c] does not cover the inner [a, b ASC] because the directions differ,
+ // so the wider sort cannot be pushed down and both sorts survive.
+ val query =
+ """
+ |SELECT a, b, c,
+ | RANK() OVER (PARTITION BY a ORDER BY b) AS rk,
+ | ROW_NUMBER() OVER (PARTITION BY a ORDER BY b DESC, c) AS rn
+ |FROM t
+ |""".stripMargin
+ checkSorts(query, 2, 2)
+ }
+ }
+
+ test("Negative: no push-down when an intermediate project computes the ordering column") {
+ withTempView("t") {
+ spark.range(300).selectExpr("id % 10 as a", "id % 7 as b", "id % 5 as c")
+ .createOrReplaceTempView("t")
+ // Between the two window sorts a project derives `bx = b + 1` (an expression alias, not a
+ // plain rename). The outer window orders by [a, bx, c]; `bx` is only an expression alias, so
+ // it is not rewritten back to `b` and is absent from the project's input. The push-down is
+ // rejected because the rewritten ordering still references a column the child does not
+ // produce, and both sorts survive.
+ val query =
+ """
+ |SELECT a, bx, c, rn1,
+ | ROW_NUMBER() OVER (PARTITION BY a ORDER BY bx, c) AS rn2
+ |FROM (
+ | SELECT a, b + 1 AS bx, c, rn1 FROM (
+ | SELECT a, b, c, RANK() OVER (PARTITION BY a ORDER BY b) AS rn1 FROM t
+ | )
+ |)
+ |""".stripMargin
+ checkSorts(query, 2, 2)
+ }
+ }
+
+ test("Negative: do not push through a non-deterministic project") {
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val scan = LocalTableScanExec(Seq(a, b), Nil, None)
+ val orderA = SortOrder(a, Ascending) :: Nil
+ val orderAB = SortOrder(a, Ascending) :: SortOrder(b, Ascending) :: Nil
+
+ // SortExec([a, b]) <- Project[a, b, rand(0)] <- SortExec([a]). Pushing the sort below the
+ // project would change which rows the seeded random stream is evaluated over, so the rule must
+ // not cross a non-deterministic project (mirrors `EliminateSorts.canEliminateSort`).
+ val lower = SortExec(orderA, global = false, scan)
+ val project = ProjectExec(Seq(a, b, Alias(Rand(Literal(0L)), "r")()), lower)
+ val upper = SortExec(orderAB, global = false, project)
+
+ withSQLConf(SQLConf.PUSH_DOWN_LOCAL_SORT_ENABLED.key -> "true") {
+ assert(PushDownLocalSort(upper).fastEquals(upper), "the plan must be left unchanged")
+ }
+ }
+
+ test("Negative: do not push through a non-deterministic filter") {
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val scan = LocalTableScanExec(Seq(a, b), Nil, None)
+ val orderA = SortOrder(a, Ascending) :: Nil
+ val orderAB = SortOrder(a, Ascending) :: SortOrder(b, Ascending) :: Nil
+
+ // SortExec([a, b]) <- Filter(rand(0) < 0.5) <- SortExec([a]). Even with cardinality-reducer
+ // traversal enabled, a non-deterministic filter must not be crossed: the surviving row set
+ // would change if the sort moved below it.
+ val lower = SortExec(orderA, global = false, scan)
+ val filter = FilterExec(LessThan(Rand(Literal(0L)), Literal(0.5)), lower)
+ val upper = SortExec(orderAB, global = false, filter)
+
+ withSQLConf(
+ SQLConf.PUSH_DOWN_LOCAL_SORT_ENABLED.key -> "true",
+ SQLConf.PUSH_DOWN_LOCAL_SORT_THROUGH_CARDINALITY_REDUCER_ENABLED.key -> "true") {
+ assert(PushDownLocalSort(upper).fastEquals(upper), "the plan must be left unchanged")
+ }
+ }
+
+ test("Cardinality reducer: do not cross a filter unless explicitly enabled") {
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val scan = LocalTableScanExec(Seq(a, b), Nil, None)
+ val orderA = SortOrder(a, Ascending) :: Nil
+ val orderAB = SortOrder(a, Ascending) :: SortOrder(b, Ascending) :: Nil
+
+ // SortExec([a, b]) <- Filter(a IS NOT NULL) <- SortExec([a]). A filter is a cardinality
+ // reducer, so it is only crossed when the reducer config is on.
+ val lower = SortExec(orderA, global = false, scan)
+ val filter = FilterExec(IsNotNull(a), lower)
+ val upper = SortExec(orderAB, global = false, filter)
+
+ withSQLConf(SQLConf.PUSH_DOWN_LOCAL_SORT_ENABLED.key -> "true") {
+ // Reducer traversal off (default): the filter is not crossed, both sorts survive.
+ withSQLConf(
+ SQLConf.PUSH_DOWN_LOCAL_SORT_THROUGH_CARDINALITY_REDUCER_ENABLED.key -> "false") {
+ assert(PushDownLocalSort(upper).fastEquals(upper), "must not cross the filter by default")
+ }
+ // Reducer traversal on: the sort is pushed below the filter and the two sorts become one.
+ withSQLConf(
+ SQLConf.PUSH_DOWN_LOCAL_SORT_THROUGH_CARDINALITY_REDUCER_ENABLED.key -> "true") {
+ val rewritten = PushDownLocalSort(upper)
+ assert(rewritten.collect { case s: SortExec => s }.length == 1)
+ }
+ }
+ }
+
+ test("Plan-level: push a wider sort down through order-preserving operators") {
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val scan = LocalTableScanExec(Seq(a, b), Nil, None)
+ val orderA = SortOrder(a, Ascending) :: Nil
+ val orderAB = SortOrder(a, Ascending) :: SortOrder(b, Ascending) :: Nil
+
+ // SortExec([a, b]) <- Filter <- Project <- SortExec([a]) : the wider sort above is pushed down
+ // to widen the lower [a] sort into [a, b], and the upper sort is dropped, leaving one sort.
+ // Crossing the `Filter` requires the cardinality-reducer config.
+ val lower = SortExec(orderA, global = false, scan)
+ val project = ProjectExec(Seq(a, b), lower)
+ val filter = FilterExec(IsNotNull(a), project)
+ val upper = SortExec(orderAB, global = false, filter)
+
+ withSQLConf(
+ SQLConf.PUSH_DOWN_LOCAL_SORT_ENABLED.key -> "true",
+ SQLConf.PUSH_DOWN_LOCAL_SORT_THROUGH_CARDINALITY_REDUCER_ENABLED.key -> "true") {
+ val rewritten = PushDownLocalSort(upper)
+ val sorts = rewritten.collect { case s: SortExec => s }
+ assert(sorts.length == 1, "the two sorts become one")
+ // The single remaining sort is the wide [a, b] sort, sitting at the bottom of the chain.
+ assert(sorts.head.sortOrder == orderAB)
+ assert(sorts.head.child.isInstanceOf[LocalTableScanExec])
+ // The rewritten plan still re-exposes the [a, b] ordering that the dropped upper sort gave.
+ assert(SortOrder.orderingSatisfies(rewritten.outputOrdering, orderAB))
+ // Every operator's required ordering is still satisfied.
+ assert(ValidateRequirements.validate(rewritten, UnspecifiedDistribution))
+ }
+ }
+
+ test("Plan-level: rewrite the ordering through a renaming project when pushing down") {
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val scan = LocalTableScanExec(Seq(a, b), Nil, None)
+ val orderA = SortOrder(a, Ascending) :: Nil
+ val bb = Alias(b, "bb")()
+ val orderABB = SortOrder(a, Ascending) :: SortOrder(bb.toAttribute, Ascending) :: Nil
+
+ // SortExec([a, bb]) <- Project([a, b AS bb]) <- SortExec([a]) : the upper sort is over the
+ // renamed `bb`. Pushing it below the project rewrites `bb` back to `b`, widening the lower sort
+ // to [a, b]; the upper sort is dropped and the project re-exposes [a, bb] above.
+ val lower = SortExec(orderA, global = false, scan)
+ val project = ProjectExec(Seq(a, bb), lower)
+ val upper = SortExec(orderABB, global = false, project)
+
+ withSQLConf(SQLConf.PUSH_DOWN_LOCAL_SORT_ENABLED.key -> "true") {
+ val rewritten = PushDownLocalSort(upper)
+ val sorts = rewritten.collect { case s: SortExec => s }
+ assert(sorts.length == 1, "the two sorts become one")
+ // The single remaining sort is the rewritten [a, b] sort in the project's input space.
+ assert(sorts.head.sortOrder == SortOrder(a, Ascending) :: SortOrder(b, Ascending) :: Nil)
+ assert(sorts.head.child.isInstanceOf[LocalTableScanExec])
+ // Above, the project re-exposes the [a, bb] ordering the dropped upper sort provided.
+ assert(SortOrder.orderingSatisfies(rewritten.outputOrdering, orderABB))
+ assert(ValidateRequirements.validate(rewritten, UnspecifiedDistribution))
+ }
+ }
+
+ test("Plan-level: push a wider sort down past three stacked sorts in a single pass") {
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ val c = AttributeReference("c", IntegerType)()
+ val scan = LocalTableScanExec(Seq(a, b, c), Nil, None)
+ val orderA = SortOrder(a, Ascending) :: Nil
+ val orderAB = SortOrder(a, Ascending) :: SortOrder(b, Ascending) :: Nil
+ val orderABC =
+ SortOrder(a, Ascending) :: SortOrder(b, Ascending) :: SortOrder(c, Ascending) :: Nil
+
+ // Sort([a,b,c]) <- Filter <- Sort([a,b]) <- Project <- Sort([a]) : the widest ordering is
+ // pushed all the way to the bottom sort, and both intermediate sorts are dropped, leaving one.
+ // Crossing the `Filter` requires the cardinality-reducer config.
+ val bottom = SortExec(orderA, global = false, scan)
+ val project = ProjectExec(Seq(a, b, c), bottom)
+ val middle = SortExec(orderAB, global = false, project)
+ val filter = FilterExec(IsNotNull(a), middle)
+ val top = SortExec(orderABC, global = false, filter)
+
+ withSQLConf(
+ SQLConf.PUSH_DOWN_LOCAL_SORT_ENABLED.key -> "true",
+ SQLConf.PUSH_DOWN_LOCAL_SORT_THROUGH_CARDINALITY_REDUCER_ENABLED.key -> "true") {
+ val rewritten = PushDownLocalSort(top)
+ val sorts = rewritten.collect { case s: SortExec => s }
+ assert(sorts.length == 1, "the three sorts become one")
+ assert(sorts.head.sortOrder == orderABC)
+ assert(sorts.head.child.isInstanceOf[LocalTableScanExec])
+ assert(SortOrder.orderingSatisfies(rewritten.outputOrdering, orderABC))
+ assert(ValidateRequirements.validate(rewritten, UnspecifiedDistribution))
+ }
+ }
+}
+
+class PushDownLocalSortSuite extends PushDownLocalSortSuiteBase
+ with DisableAdaptiveExecutionSuite
+
+class PushDownLocalSortSuiteAE extends PushDownLocalSortSuiteBase
+ with EnableAdaptiveExecutionSuite
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org