peter-toth commented on code in PR #57396: URL: https://github.com/apache/spark/pull/57396#discussion_r3622038437
########## 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 Review Comment: Widening the sort that feeds a window changes the row order the window sees *within ties of its own `ORDER BY`*. For order-sensitive window functions that's observable: - `collect_list`/`collect_set`/`first_value`/`last_value`/`nth_value` collect/pick in input order; - `lead`/`lag` and aggregates with a `ROWS` frame depend on physical adjacency; - a `row_number()`-based `WindowGroupLimitExec` keeps a *different set* of rows when ties reorder. Example: `collect_list(c) OVER (PARTITION BY a ORDER BY b)` stacked under a wider window that needs `[a, b, c]`. Today the collect_list window sorts by `[a, b]`; after this rule it sorts by `[a, b, c]`, so within equal `b` the list follows `c`. These are all non-deterministic over a non-unique `ORDER BY`, so it isn't *wrong* per Spark's contract -- but it's exactly the order-sensitivity you use to keep `CollectMetricsExec` out of `isOrderPreserving` (and the reason cited in the PR description). So the description's "query results are unchanged" holds for deterministic queries but not for these, and the line between "exclude CollectMetrics" and "push through a window carrying `first`/`last`/`collect_list`" is hard to defend. Two asks: (1) soften that claim in the description; (2) decide whether a window carrying an order-sensitive function should be treated like `CollectMetricsExec`. If you deliberately keep them in (they're non-deterministic anyway), a one-line comment here explaining why the window case is acceptable while CollectMetrics isn't would settle it for the next reader. ########## 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") Review Comment: Good case, but the data neutralizes the risk it looks like it's testing. Within any `(a, b)` group every row has the same `c`: ids sharing `(id%10, id%7)` differ by multiples of `lcm(10,7)=70`, and `70 % 5 == 0`, so `id % 5` is constant per `(a, b)`. Widening the window's sort from `[a, b]` to `[a, b, c]` therefore can't reorder within-`(a,b)` ties, so `row_number` comes out identical with the rule on and off and `checkAnswer` passes by construction. If `c` actually varied within `(a, b)`, the `[a,b,c]` sort would reorder ties and `row_number` would differ between the two configs -- which would make this comparison flaky rather than catch a bug. So this test can't observe the order change from the window comment. Worth either a note here, or a companion test that asserts the sort count drops while using a tie-safe function (`rank`/`dense_rank`) or a unique key, so the result stays stable even when ties genuinely reorder. ########## 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 Review Comment: I couldn't construct a plan where traversing a `SortAggregateExec` here actually enables a widening. To push `upperOrder` through it, line 118 requires it to satisfy the grouping-key ordering and line 119 requires every key to be a child column -- together that forces `upperOrder` to be a prefix-cover of the grouping keys and nothing wider, because any extra output column is an aggregate result that fails the `subsetOf(child.outputSet)` check. But the aggregate already requires its input sorted by the full grouping keys, so `EnsureRequirements` places that sort directly beneath it; there's no strictly-narrower lower sort left for `upperOrder` to widen. If that's right, the `SortAggregateExec` branch (here, plus the `resultExpressions` rewrite at line 104) never contributes to a real push-down, and dropping it would shrink the surface the rule must be trusted over. If you do have a firing case in mind, could you add a test? The `resultExpressions` rename path and the `WindowGroupLimitExec` traversal are both currently uncovered. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
