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..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 @@ -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 0000000000000..161e779a0c9ae --- /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 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..42ea16048d663 --- /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