This is an automated email from the ASF dual-hosted git repository.
uros-b pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 6567619c0205 [SPARK-58233][SQL] Push down local sort to reduce
redundant sorts within a stage
6567619c0205 is described below
commit 6567619c0205f16d94633840eca5454fac70c9a0
Author: Xiduo You <[email protected]>
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 <[email protected]>
Signed-off-by: Uros Bojanic <[email protected]>
(cherry picked from commit 589adcb9b0b4ee247d37acadd851baac776a5a91)
Signed-off-by: Uros Bojanic <[email protected]>
---
.../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: [email protected]
For additional commands, e-mail: [email protected]
