peter-toth commented on code in PR #58775:
URL: https://github.com/apache/spark/pull/58775#discussion_r4007344283
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimits.scala:
##########
@@ -18,20 +18,56 @@
package org.apache.spark.sql.execution
import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.datasources.v2.GroupPartitionsExec
import org.apache.spark.sql.execution.window.{Final, Partial,
WindowGroupLimitExec}
/**
* Remove redundant partial WindowGroupLimitExec node from the spark plan. A
partial
- * WindowGroupLimitExec node is redundant when its child satisfies its
required child distribution.
+ * WindowGroupLimitExec node is redundant when the child left behind by
removing it satisfies the
+ * final node's required child distribution: the partial node only
pre-filters, keeping the rows
+ * whose rank within its own partition is within the limit, and a partition of
the final node is a
+ * union of partitions of the partial node, so it drops no row the final node
would keep. A
+ * [[GroupPartitionsExec]] between the two nodes, and a local sort, are looked
through.
*/
object RemoveRedundantWindowGroupLimits extends Rule[SparkPlan] {
def apply(plan: SparkPlan): SparkPlan = plan transform {
- case outer @ WindowGroupLimitExec(
- _, _, _, _, Final, WindowGroupLimitExec(_, _, _, _, Partial, child))
- if
child.outputPartitioning.satisfies(outer.requiredChildDistribution.head) =>
- val newOuter = outer.withNewChildren(Seq(child))
- newOuter
+ case outer @ WindowGroupLimitExec(_, _, _, _, Final, child) =>
+ val newChild = removePartialLimit(child, hasUpperSort = false)
+ if
(newChild.outputPartitioning.satisfies(outer.requiredChildDistribution.head)) {
+ outer.withNewChildren(Seq(newChild))
+ } else {
+ outer
+ }
+ }
+
+ /**
+ * Removes the partial [[WindowGroupLimitExec]] at the top of `plan`,
together with the local sort
+ * feeding it when `hasUpperSort` is set, or returns `plan` unchanged when
it does not hold one.
+ *
+ * A [[GroupPartitionsExec]] and a local sort are looked through:
`EnsureRequirements` adds the
+ * grouping between the two nodes when the child needs its partitions
coalesced or its partition
+ * keys projected to satisfy the final node's distribution, and adds the
sort between them to give
+ * the final node its ordering. That sort orders the rows the final node
ranks as a whole, which
+ * is what makes the sort feeding the partial node dead; on its own that one
is not redundant,
+ * having been added to give the partial node its required ordering.
+ *
+ * @param hasUpperSort whether a local sort between the final node and
`plan` orders the rows the
+ * final node ranks.
+ */
+ private def removePartialLimit(plan: SparkPlan, hasUpperSort: Boolean):
SparkPlan = plan match {
+ case WindowGroupLimitExec(_, _, _, _, Partial, child) =>
+ if (hasUpperSort) stripLocalSort(child) else child
Review Comment:
**Finding 1.** `stripLocalSort` returns `child` unchanged when the partial
node holds no local sort, so in that shape the partial node goes and nothing
goes with it. The sort above the grouping then has to sort everything the scan
produces instead of what the partial node let through. That is exactly the
third case of the new `KeyGroupedPartitioningSuite` test: the source reports
`(id, name, price)`, which covers the partial node's required ordering, so
there is no sort under it to remove, and the grouping still forces one above.
Measured on this branch against `b87b39c35699` - one partition key,
1,000,000 rows over 100 splits of 10,000, `rk = 1`, source reporting the
window's ordering:
| | rows into the surviving sort | `sortTime` | wall (steady) |
|---|---|---|---|
| `b87b39c35699` | 200, the partial node's output | 0 ms | 372-415 ms |
| `0066cc2b3a1` | 1,000,000 | 11-13 ms | 400-427 ms |
Wall clock only moves 4-7% here because the `InMemoryTable` scan dominates
and 1M narrow rows fit in memory - `spillSize` is 0 in both arms. On a real
source the same plan change is a 200-row sort against a spilling one.
`PushDownLocalSort`, in this package, declines this trade by default. Its
`isOrderPreserving` lists `WindowGroupLimitExec` as a cardinality reducer and
only crosses it behind
`spark.sql.execution.pushDownLocalSort.throughCardinalityReducer`, whose doc
says "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". That is
the same sentence as this shape.
The other two shapes are fine, so the removal is worth doing exactly when
something goes with it. Where a sort does go away with the partial node the
measurement came out slightly ahead (one 100k sort at 3 ms against ten 10k
sorts at 5 ms), and where no sort sits above the grouping the sort that stays
is unaffected.
```suggestion
if (hasUpperSort) {
// With no sort to remove, the partial node is the only cardinality
reducer between its
// child and the sort above, and dropping it leaves that sort with
the whole input -
// `PushDownLocalSort.isOrderPreserving` declines the same trade by
default.
child match {
case sort: SortExec if !sort.global => sort.child
case _ => plan
}
} else {
child
}
```
I ran that: `byKey` and `byPrice` still pass and the `reported` case keeps
both nodes, so the third block of the new test would want to assert that
instead. `stripLocalSort` then has no other caller, and the `hasUpperSort`
scaladoc needs a word about the sort being the condition rather than a bonus.
One more thing this shape changes, which argues the same way: while the
partial node sat under the grouping,
`GroupPartitionsExec.childIsSafeForKWayMerge` was false (`WindowGroupLimitExec`
is not a `SafeForKWayMerge`), so `EnsureRequirements.tryEnableSortedMerge`
could never enable the merge and always fell through to the `SortExec` at
`EnsureRequirements.scala:327`. After the removal the grouping's child is the
scan and the merge becomes feasible, but `enableSortedMerge` was already
decided and this rule runs afterwards, so nothing reconsiders it. I tried
reclaiming it (`newGroup.tryEnableSortedMerge().getOrElse(newGroup)` here) and
it does produce a sort-free plan - `GroupPartitions ... SortedMerge: true` with
`numSorts=0`, `RemoveRedundantSorts` dropping the sort - but it measured
606-608 ms against the 372-415 ms of master, so a 100-way merge over 1M rows is
no bargain either. Keeping the reducer is the cheap answer in this shape.
Finally, the description reads as a pure win ("drops an operator and a sort
pass from the plan"). Worth saying what the surviving sort now sees.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala:
##########
@@ -69,6 +69,36 @@ abstract class RemoveRedundantWindowGroupLimitsSuiteBase
checkWindowGroupLimits(query2, 2)
}
}
+
+ test("SPARK-59492: the local sort of a removed WindowGroupLimit stays where
it was") {
+ withTempView("t") {
+ spark.range(0, 100).withColumn("value",
lit(1)).createOrReplaceTempView("t")
+ // The child is already clustered on the window's PARTITION BY, so no
exchange lands between
+ // the two limit nodes and nothing above them orders the rows the final
node ranks: the local
+ // sort that fed the partial node is the one carrying that ordering, and
it stays where the
+ // partial node was.
+ val df = sql(
+ """
+ |SELECT *
+ |FROM (
+ | SELECT id, rank() OVER w AS rn
+ | FROM t
+ | GROUP BY id
+ | WINDOW w AS (PARTITION BY id ORDER BY max(value))
+ |)
+ |WHERE rn < 3
+ |""".stripMargin)
+ val plan = df.queryExecution.executedPlan
+ val limits = collectWithSubqueries(plan) { case w: WindowGroupLimitExec
=> w.mode }
+ assert(limits == Seq(Final), s"expected the final limit alone, got
$limits:\n$plan")
+ val sorts = collectWithSubqueries(plan) { case s: SortExec => s
}.filter(!_.global)
+ assert(sorts.length == 1, s"expected one local sort, got
${sorts.length}:\n$plan")
+ assert(!sorts.head.child.isInstanceOf[WindowGroupLimitExec],
Review Comment:
**Finding 3.** In this plan the partial node sits *above* the sort, never
below it, so `sorts.head.child` is the `HashAggregate` whether or not the rule
fires - on master too. The assertion therefore holds on every outcome and pins
nothing; `sorts.length == 1` on the line above is what carries the test. I
would drop this line, or, if the intent is that the sort ended up where the
partial node was, assert on the final node's child after unwrapping the codegen
wrapper (`SortExec` supports codegen here, so a bare `isInstanceOf` on it would
not hold).
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala:
##########
@@ -69,6 +69,36 @@ abstract class RemoveRedundantWindowGroupLimitsSuiteBase
checkWindowGroupLimits(query2, 2)
}
}
+
+ test("SPARK-59492: the local sort of a removed WindowGroupLimit stays where
it was") {
Review Comment:
**Finding 2.** This test passes on `b87b39c35699` in both AQE modes - I ran
it there. The query is `query1` from `remove redundant WindowGroupLimits`
verbatim, and on master the old rule already rewrote
`Final(Partial(Sort(...)))` into `Final(Sort(...))`, so every assertion here
holds without the patch.
It is not worthless, since it would catch `hasUpperSort` defaulting to
`true`. But as a separate test it re-runs an existing query and builds the same
plan twice. Folding the two sort assertions into `remove redundant
WindowGroupLimits`, right next to `checkWindowGroupLimits(query1, 1)`, pins the
same property without the duplicate.
--
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]
