From 0b2f14fe893d55f93a1201cfb1e4610529e452f2 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Tue, 21 Jul 2026 14:46:52 +0800 Subject: [PATCH 1/5] [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. Order-preserving operators traversed: `ProjectExec`, `FilterExec`, `SortAggregateExec`, `WindowExecBase`, `WindowGroupLimitExec`. When one renames an ordering column in its output (`b AS x`), the ordering is rewritten from the operator's output space back to its child's space as it is pushed through (plain renames only). The rule never crosses a shuffle or a non-order-preserving operator, so the pushed-down sort keeps the same meaning it had above. `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`). ### Why are the changes needed? Stacked windows / 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; query results are unchanged. 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, window to sort aggregate, renaming project), negative SQL cases (disjoint orderings, shuffle between sorts, direction mismatch, computed ordering column), and plan-level cases. Verified 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) Co-Authored-By: Claude Opus 4.8 (1M context) --- .../apache/spark/sql/internal/SQLConf.scala | 13 + .../sql/execution/PushDownLocalSort.scala | 136 ++++++++ .../spark/sql/execution/QueryExecution.scala | 4 + .../adaptive/AdaptiveSparkPlanExec.scala | 4 + .../execution/PushDownLocalSortSuite.scala | 320 ++++++++++++++++++ 5 files changed, 477 insertions(+) create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/PushDownLocalSort.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/PushDownLocalSortSuite.scala 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 4758d063ca03c..48d7a50b22266 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,19 @@ 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 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 0000000000000..7cac79abfe2e9 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/PushDownLocalSort.scala @@ -0,0 +1,136 @@ +/* + * 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.aggregate.SortAggregateExec +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 an operator on the path renames an ordering column in its output (a `ProjectExec` with + * `b AS x`, or a `SortAggregateExec` whose result renames a grouping key), the ordering is + * rewritten from the operator'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 + } + + plan.transform { + case upper @ SortExec(upperOrder, false, child, _) => + pushDown(child, upperOrder).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. + */ + private def pushDown( + plan: SparkPlan, + upperOrder: Seq[SortOrder]): 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) => + // Some order-preserving operators rename ordering columns in their output (a `ProjectExec` + // with `b AS x`, or a `SortAggregateExec` whose result renames a grouping key). Rewrite + // `upperOrder` from the operator's output space back to its child's space before pushing + // further down. Only plain renames are followed; an expression alias leaves the sort key + // referencing an output attribute the child does not produce, so the check below rejects it. + val outputExprs = plan match { + case p: ProjectExec => p.projectList + case a: SortAggregateExec => a.resultExpressions + 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).map(newChild => op.withNewChildren(Seq(newChild))) + } else { + None + } + + case _ => None + } + + private def isOrderPreserving(plan: UnaryExecNode): Boolean = plan match { + case _: ProjectExec => true + case _: FilterExec => true + case _: SortAggregateExec => true + case _: WindowExecBase => true + case _: WindowGroupLimitExec => 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 f70eafc70af35..0a1185aa4ea57 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 1905b31b49f97..bfe6a9a3f6332 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 0000000000000..db1fa3fc25364 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PushDownLocalSortSuite.scala @@ -0,0 +1,320 @@ +/* + * 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, 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 = { + 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") { + spark.range(200).selectExpr("id % 10 as a", "id % 7 as b", "id % 5 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(rn) AS cl + |FROM (SELECT a, b, c, ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) AS rn 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("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. + 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") { + 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. + 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") { + 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 From 2b05aa93ee0e709b198ca9ec81dffc2a370cc01a Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Wed, 22 Jul 2026 09:53:24 +0800 Subject: [PATCH 2/5] [SPARK-58233][SQL][FOLLOWUP] Address review: drop unreachable SortAggregate branch, harden window test Addresses review feedback on the rule: - Remove `SortAggregateExec` from the traversed order-preserving operators (and the `resultExpressions` rename path). Pushing a sort *through* a sort aggregate can never fire: the aggregate's child is already sorted by the full grouping keys, and any ordering the consumer above can carry down references only grouping-key columns, so it can never be strictly wider than that sort. The canonical window-to-sort-aggregate case is unaffected -- there the aggregate is the top consumer and the push-down traverses the `Window`, not the aggregate. - Document at the `WindowExecBase`/`WindowGroupLimitExec` entries why pushing a wider sort under a window is acceptable (it only refines the tie order for a non-unique window `ORDER BY`, already non-deterministic per Spark's contract) while `CollectMetricsExec` is excluded (its observed metric is an expected-stable side output). - Harden the window-to-sort-aggregate test: the previous data had `c` constant within every `(a, b)` group, so widening `[a,b]` to `[a,b,c]` could not reorder window ties. Use data where `c` varies within `(a, b)` and a tie-safe `RANK()` so the sort count still drops while the result stays stable even though ties genuinely reorder. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../sql/execution/PushDownLocalSort.scala | 30 +++++++++++-------- .../execution/PushDownLocalSortSuite.scala | 10 +++++-- 2 files changed, 24 insertions(+), 16 deletions(-) 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 index 7cac79abfe2e9..6712ff9f4ba6c 100644 --- 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 @@ -19,7 +19,6 @@ 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.aggregate.SortAggregateExec import org.apache.spark.sql.execution.window.{WindowExecBase, WindowGroupLimitExec} import org.apache.spark.sql.internal.SQLConf @@ -53,11 +52,10 @@ import org.apache.spark.sql.internal.SQLConf * Exchange(hashpartitioning([a])) * }}} * - * When an operator on the path renames an ordering column in its output (a `ProjectExec` with - * `b AS x`, or a `SortAggregateExec` whose result renames a grouping key), the ordering is - * rewritten from the operator'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. + * 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] { @@ -94,14 +92,13 @@ object PushDownLocalSort extends Rule[SparkPlan] { Some(SortExec(upperOrder, global = false, child = lower.child)) case op: UnaryExecNode if isOrderPreserving(op) => - // Some order-preserving operators rename ordering columns in their output (a `ProjectExec` - // with `b AS x`, or a `SortAggregateExec` whose result renames a grouping key). Rewrite - // `upperOrder` from the operator's output space back to its child's space before pushing - // further down. Only plain renames are followed; an expression alias leaves the sort key - // referencing an output attribute the child does not produce, so the check below rejects it. + // 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 a: SortAggregateExec => a.resultExpressions case _ => Nil } val rewrittenUpperOrder = if (outputExprs.isEmpty) { @@ -128,7 +125,14 @@ object PushDownLocalSort extends Rule[SparkPlan] { private def isOrderPreserving(plan: UnaryExecNode): Boolean = plan match { case _: ProjectExec => true case _: FilterExec => true - case _: SortAggregateExec => true + // 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) and a `row_number`-based + // `WindowGroupLimitExec`, 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 _: WindowGroupLimitExec => true case _ => false 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 index db1fa3fc25364..3049c847ff11a 100644 --- 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 @@ -110,7 +110,11 @@ abstract class PushDownLocalSortSuiteBase test("Push a wider sort down through a window to feed a sort aggregate above it") { withTempView("t") { - spark.range(200).selectExpr("id % 10 as a", "id % 7 as b", "id % 5 as c") + // `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 @@ -121,8 +125,8 @@ abstract class PushDownLocalSortSuiteBase withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "false") { val query = """ - |SELECT a, b, c, collect_list(rn) AS cl - |FROM (SELECT a, b, c, ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) AS rn FROM t) + |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) From dd42f7318f7ee2cea9e77df758850eae4e3a1bb8 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Wed, 22 Jul 2026 12:01:06 +0800 Subject: [PATCH 3/5] [SPARK-58233][SQL][FOLLOWUP] Guard determinism and gate cardinality-reducer traversal Addresses review feedback (sunchao): - [P1, correctness] Do not push a local sort through a non-deterministic `ProjectExec` or `FilterExec`. Moving the sort below them 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). `ProjectExec` is now gated on `projectList.forall(_.deterministic)` and `FilterExec` on `condition.deterministic`, mirroring `EliminateSorts.canEliminateSort`. - [P2, performance] `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. They are now crossed only when the new internal config `spark.sql.execution.pushDownLocalSort.throughCardinalityReducer` is enabled (default false). Adds plan-level tests: non-deterministic project/filter are not crossed, and a filter is crossed only when the reducer config is on. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../apache/spark/sql/internal/SQLConf.scala | 13 +++ .../sql/execution/PushDownLocalSort.scala | 29 ++++--- .../execution/PushDownLocalSortSuite.scala | 80 ++++++++++++++++++- 3 files changed, 110 insertions(+), 12 deletions(-) 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 48d7a50b22266..db832153facd7 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 @@ -3014,6 +3014,19 @@ object SQLConf { .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 index 6712ff9f4ba6c..8825052cbf8c0 100644 --- 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 @@ -123,18 +123,29 @@ object PushDownLocalSort extends Rule[SparkPlan] { } private def isOrderPreserving(plan: UnaryExecNode): Boolean = plan match { - case _: ProjectExec => true - case _: FilterExec => true + // 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) and a `row_number`-based - // `WindowGroupLimitExec`, 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. + // `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 _: WindowGroupLimitExec => true case _ => false } + + private def throughCardinalityReducer: Boolean = + conf.getConf(SQLConf.PUSH_DOWN_LOCAL_SORT_THROUGH_CARDINALITY_REDUCER_ENABLED) } 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 index 3049c847ff11a..9472878e517b2 100644 --- 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 @@ -18,7 +18,7 @@ package org.apache.spark.sql.execution import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, AttributeReference, IsNotNull, SortOrder} +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 @@ -231,6 +231,74 @@ abstract class PushDownLocalSortSuiteBase } } + 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)() @@ -240,12 +308,15 @@ abstract class PushDownLocalSortSuiteBase // 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") { + 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") @@ -299,13 +370,16 @@ abstract class PushDownLocalSortSuiteBase // 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") { + 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") From 4b5520cae9e98b445b412e0619a1cc4f6033211f Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Wed, 22 Jul 2026 12:12:39 +0800 Subject: [PATCH 4/5] [SPARK-58233][SQL][FOLLOWUP] Thread throughCardinalityReducer as a parameter Read the cardinality-reducer flag once in `apply` and pass it down through `pushDown`/`isOrderPreserving` as an explicit parameter instead of re-reading the SQL config from a helper on every operator check. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../sql/execution/PushDownLocalSort.scala | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) 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 index 8825052cbf8c0..161e779a0c9ae 100644 --- 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 @@ -63,10 +63,12 @@ object PushDownLocalSort extends Rule[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).getOrElse(upper) + pushDown(child, upperOrder, throughCardinalityReducer).getOrElse(upper) } } @@ -76,11 +78,14 @@ object PushDownLocalSort extends Rule[SparkPlan] { * 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. + * 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]): Option[SparkPlan] = plan match { + 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 @@ -91,7 +96,7 @@ object PushDownLocalSort extends Rule[SparkPlan] { AttributeSet(upperOrder.flatMap(_.references)).subsetOf(lower.child.outputSet) => Some(SortExec(upperOrder, global = false, child = lower.child)) - case op: UnaryExecNode if isOrderPreserving(op) => + 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 @@ -114,7 +119,8 @@ object PushDownLocalSort extends Rule[SparkPlan] { } if (SortOrder.orderingSatisfies(rewrittenUpperOrder, op.requiredChildOrdering.head) && AttributeSet(rewrittenUpperOrder.flatMap(_.references)).subsetOf(op.child.outputSet)) { - pushDown(op.child, rewrittenUpperOrder).map(newChild => op.withNewChildren(Seq(newChild))) + pushDown(op.child, rewrittenUpperOrder, throughCardinalityReducer) + .map(newChild => op.withNewChildren(Seq(newChild))) } else { None } @@ -122,7 +128,9 @@ object PushDownLocalSort extends Rule[SparkPlan] { case _ => None } - private def isOrderPreserving(plan: UnaryExecNode): Boolean = plan match { + 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 @@ -145,7 +153,4 @@ object PushDownLocalSort extends Rule[SparkPlan] { case _: WindowExecBase => true case _ => false } - - private def throughCardinalityReducer: Boolean = - conf.getConf(SQLConf.PUSH_DOWN_LOCAL_SORT_THROUGH_CARDINALITY_REDUCER_ENABLED) } From b6503d817a007385fbf1fd0f477edb5439fc1dbc Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Wed, 22 Jul 2026 12:50:29 +0800 Subject: [PATCH 5/5] [SPARK-58233][SQL][FOLLOWUP] Check the final adaptive plan in sort-count assertions `checkNumSorts` inspected `executedPlan` before executing the query. Under AQE that is the initial plan, before query-stage materialization and replanning where this rule also runs. Execute the query first so the assertion inspects the final adaptive plan as well. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../apache/spark/sql/execution/PushDownLocalSortSuite.scala | 3 +++ 1 file changed, 3 insertions(+) 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 index 9472878e517b2..42ea16048d663 100644 --- 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 @@ -32,6 +32,9 @@ abstract class PushDownLocalSortSuiteBase 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) }