cloud-fan commented on code in PR #58144:
URL: https://github.com/apache/spark/pull/58144#discussion_r4253700800
##########
sql/core/benchmarks/WindowReorderBenchmark-results.txt:
##########
@@ -0,0 +1,12 @@
+================================================================================================
+WindowReorder: stack reordering of the window chain
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.18+8-LTS on Mac OS X 26.4.1
Review Comment:
**Nit (P3):** This file was generated locally (`OpenJDK 64-Bit Server VM
17.0.18+8-LTS on Mac OS X 26.4.1`, Apple M4 Pro). Only 2 of the 246 results
files under `sql/core/benchmarks` come from Mac OS X. The rest come from the
GitHub Actions benchmark workflow, and benchmarks such as `WindowBenchmark`
also commit `-jdk21-results.txt` and `-jdk25-results.txt`. Could you regenerate
the results with the benchmark workflow in your fork and commit all JDK
variants? Otherwise the 1.3X numbers cannot be compared with the other
committed results.
**Verification:**
- **Inspection:** The committed result headers show the workflow's reference
environment, and JDK 21 and JDK 25 variants exist.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -1845,12 +1845,417 @@ object CollapseWindow extends Rule[LogicalPlan] {
}
/**
- * Transpose Adjacent Window Expressions.
- * - If the partition spec of the parent Window expression is compatible with
the partition spec
- * of the child window expression, transpose them.
+ * Reorder a stack of `Window` operators to reduce the exchanges inserted for
it. Window
+ * operators only append their output columns and each of them
(re-)partitions/(re-)sorts
+ * its input independently of its position in the stack, so two windows can be
swapped
+ * freely as long as neither references the other's output. Physically, a
window rides the
+ * exchange created for the nearest window below it iff that exchange's key
set is a subset
+ * of the window's partition spec (`HashPartitioning` satisfying
`ClusteredDistribution`),
+ * and `EnsureRequirements` only creates exchanges keyed on exactly some
window's partition
+ * spec. The exchange count is at least the number of minimal partition specs
in the stack
+ * (under semantic subset); grouping windows by the minimal partition spec
their own
+ * partition spec contains reaches that lower bound when the bottom child does
not already
+ * satisfy any window's clustering. Incomparable minimal groups are ordered
without
+ * consulting the bottom child's output partitioning. Within a group, smaller
specs come
+ * first. Windows with the same complete required child ordering (partition
keys as ascending
+ * sort keys, then order spec) are made adjacent so they may share sorts; sort
minimization
+ * is not claimed. This subsumes the previous adjacent-pair transposition.
+ * Under `requireAllClusterKeysForDistribution`, a partitioning only satisfies
a window's
+ * `ClusteredDistribution` with the exact same keys in the same order, so the
subset
+ * relation degenerates to exact spec equality: the lower bound equals the
number of
+ * distinct exact partition specs, reached by grouping identical specs
adjacently when the
+ * bottom child does not already satisfy any window's clustering.
+ *
+ * The matched stack is a maximal run of `Window` operators connected by
transparent links,
+ * where a transparent link is a deterministic `Project` made only of plain
attributes and
+ * aliases. Besides pass-through, rename or drop, an alias may also compute a
derived,
+ * deterministic expression (e.g. `Alias(a + b, "s")`). Hoisting such a link
above the
+ * windows only changes where its expression is computed, not the values it
computes on
+ * (windows never rewrite their input), so the one unsafe case is a window
above the link
+ * referencing its aliased output, which `reorderable` excludes. The links are
stripped and
+ * re-applied above the reordered windows, and a chain-top `Project` restores
the original
+ * output attribute order when needed.
+ *
+ * If the stack is directly topped by a rank filter (e.g. `rn <= 1`, the
pattern that
+ * `InferWindowGroupLimit` rewrites later, in a `Once` batch), the top window
is pinned in
+ * place and only the windows below it are reordered, with the order-restoring
`Project`
+ * placed below the pinned window, so that the strict `Filter`-over-`Window`
match survives.
+ *
+ * The stack reordering is gated by
`spark.sql.optimizer.windowReorder.enabled` (default
+ * false); when disabled, the rule falls back to transposing only adjacent
window pairs whose
+ * upper partition spec is a proper subset of the lower one (see
`transposeAdjacentPairs`).
*/
object TransposeWindow extends Rule[LogicalPlan] {
- private def compatiblePartitions(ps1 : Seq[Expression], ps2:
Seq[Expression]): Boolean = {
+
+ /**
+ * A `Project` is a transparent link of a window chain if it only passes
through, renames
+ * or drops attributes, or adds deterministic aliases (which may compute
derived
+ * expressions), and is itself deterministic.
+ */
+ private def isTransparentLink(plan: LogicalPlan): Boolean = plan match {
+ case p: Project =>
+ p.projectList.forall(e => e.isInstanceOf[Attribute] ||
e.isInstanceOf[Alias]) &&
+ p.expressions.forall(_.deterministic)
+ case _ => false
+ }
+
+ /**
+ * A maximal run of adjacent `Window` operators connected by transparent
links.
+ *
+ * @param windows the windows of the chain, bottom-to-top
+ * @param betweenLinks `betweenLinks(i)` holds the transparent links between
`windows(i)`
+ * and `windows(i + 1)`, bottom-to-top; it has one
element less than
+ * `windows`
+ * @param topLinks the transparent links above the top window, bottom-to-top
+ * @param bottomChild the child of the bottom window
+ * @param topWindowChildOutput the output of the child of the top window, in
the original
+ * plan
+ * @param originalOutput the output of the whole chain, in the original plan
+ */
+ private case class WindowChain(
+ windows: Seq[Window],
+ betweenLinks: Seq[Seq[Project]],
+ topLinks: Seq[Project],
+ bottomChild: LogicalPlan,
+ topWindowChildOutput: Seq[Attribute],
+ originalOutput: Seq[Attribute])
+
+ /** Collect the maximal window chain topped by `plan`, if `plan` tops at
least 2 windows. */
+ private def collectChain(plan: LogicalPlan): Option[WindowChain] = {
+ // Collect the chain top-down: `windows` is filled from the top window
down,
+ // `betweenLinks(i)` holds the transparent links accumulated right above
`windows(i)`,
+ // and `pending` the links found since the last window. The buffers are
reversed into
+ // the bottom-to-top order of `WindowChain` at the end.
+ val windows = mutable.ArrayBuffer.empty[Window]
+ val betweenLinks = mutable.ArrayBuffer.empty[Seq[Project]]
+ var pending = Vector.empty[Project]
+ var cur = plan
+ var collecting = true
+ while (collecting) {
+ cur match {
+ case p: Project if isTransparentLink(p) =>
+ pending = pending :+ p
+ cur = p.child
+ case w: Window =>
+ betweenLinks += pending
+ windows += w
+ pending = Vector.empty
+ cur = w.child
+ case _ =>
+ collecting = false
+ }
+ }
+
+ if (windows.length < 2) {
+ None
+ } else {
+ val n = windows.length
+ Some(WindowChain(
+ windows = windows.reverse.toSeq,
+ betweenLinks = (0 until n - 1).map(i => betweenLinks(n - 1 -
i).reverse),
+ topLinks = betweenLinks.head.reverse,
+ bottomChild = windows.last.child,
+ topWindowChildOutput = windows.head.child.output,
+ originalOutput = plan.output))
+ }
+ }
+
+ /**
+ * Whether a window with partition spec `requiredKeys` can ride an exchange
keyed on
+ * `exchangeKeys`: by default (subset semantics) a partitioning satisfies a
+ * `ClusteredDistribution` if its keys are a subset of the required
clustering keys, so
+ * `exchangeKeys` being a subset of `requiredKeys` suffices; under
+ * `requireAllClusterKeysForDistribution` the partitioning must match the
required keys
+ * exactly and in order, so only identical specs ride each other's exchange.
+ */
+ private def canRideExchange(
+ exchangeKeys: Seq[Expression],
+ requiredKeys: Seq[Expression]): Boolean =
+ if (conf.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION)) {
+ exchangeKeys.length == requiredKeys.length &&
+ exchangeKeys.zip(requiredKeys).forall { case (e1, e2) =>
e1.semanticEquals(e2) }
+ } else {
+ exchangeKeys.forall(e1 => requiredKeys.exists(e1.semanticEquals))
+ }
+
+ /** Semantic equivalence of two partition specs, i.e. mutual subset. */
+ private def equivalent(spec1: Seq[Expression], spec2: Seq[Expression]):
Boolean =
+ canRideExchange(spec1, spec2) && canRideExchange(spec2, spec1)
+
+ /** Semantic equality of two sort order sequences, i.e. element-wise
`semanticEquals`. */
+ private def sameSortOrders(orders1: Seq[SortOrder], orders2:
Seq[SortOrder]): Boolean =
+ orders1.length == orders2.length &&
+ orders1.zip(orders2).forall { case (o1, o2) => o1.semanticEquals(o2) }
+
+ /** Required child ordering for a window, matching `WindowExecBase`. */
+ private def requiredChildOrdering(window: Window): Seq[SortOrder] =
+ window.partitionSpec.map(SortOrder(_, Ascending)) ++ window.orderSpec
+
+ private def sameRequiredChildOrdering(w1: Window, w2: Window): Boolean =
+ sameSortOrders(requiredChildOrdering(w1), requiredChildOrdering(w2))
+
+ /**
+ * The chain is reorderable if all windows are deterministic and have a
non-empty partition
+ * spec (an empty one requires `AllTuples` and must not move), every window
with an empty
+ * order spec is order-safe (see `orderSafeWithEmptyOrder`), no window
references another
+ * window's output, and no window references the aliased output of a link
below it (such a
+ * link could not be re-applied above that window). The aliased output may
be any
+ * deterministic expression; hoisting it above the windows only changes
where it is
+ * computed, never the values it sees, so this check is the real safety
property regardless
+ * of whether a link's alias is a bare rename or a derived expression.
+ */
+ private def reorderable(chain: WindowChain): Boolean = {
+ val windows = chain.windows
+ windows.forall(_.expressions.forall(_.deterministic)) &&
+ windows.forall(_.partitionSpec.nonEmpty) &&
+ // An empty-order window with a relative frame (e.g. `ROWS ... CURRENT
ROW`) is
+ // order-sensitive: moving it, or any window below it, changes results --
skip the chain.
+ windows.forall { w =>
+ w.orderSpec.nonEmpty ||
w.windowExpressions.forall(orderSafeWithEmptyOrder)
+ } &&
+ windows.indices.forall { i =>
+ windows.indices.forall { j =>
+ i == j ||
windows(i).references.intersect(windows(j).windowOutputSet).isEmpty
+ }
+ } &&
+ chain.betweenLinks.zipWithIndex.forall { case (links, i) =>
+ val aliasedOutputs = AttributeSet(links.flatMap(_.projectList).collect {
+ case a: Alias => a.toAttribute
+ })
+ // The windows above these links are `windows(i + 1)` and up.
+ aliasedOutputs.isEmpty ||
+ (i + 1 until windows.length).forall { k =>
+ windows(k).references.intersect(aliasedOutputs).isEmpty
+ }
+ }
+ }
+
+ /**
+ * Returns a window order as indices into `windows` (bottom-to-top) that
groups by minimal
+ * partition spec and makes equal complete required child orderings
adjacent. Does not use
+ * the bottom child's output partitioning when ordering incomparable minimal
groups.
+ */
+ private def optimalOrder(windows: Seq[Window]): Seq[Int] = {
+ val specs = windows.map(_.partitionSpec)
+ // A spec is minimal if no other spec in the stack is a strict subset of
it, i.e. no
+ // other window already keys an exchange that this window could ride.
+ val minimal = specs.indices.map { i =>
+ !specs.indices.exists { j =>
+ j != i && canRideExchange(specs(j), specs(i)) && !equivalent(specs(j),
specs(i))
+ }
+ }
+ // The distinct minimal specs, ranked by first appearance, bottom-to-top.
+ val minimalSpecs = mutable.ArrayBuffer.empty[Seq[Expression]]
+ specs.indices.filter(minimal).foreach { i =>
+ if (!minimalSpecs.exists(equivalent(_, specs(i)))) {
+ minimalSpecs += specs(i)
+ }
+ }
+ // Every window rides the exchange of the first minimal spec its partition
spec
+ // contains, so all windows of the same class share that exchange.
+ val classOf = specs.indices.map { i =>
+ minimalSpecs.indices.filter(j => canRideExchange(minimalSpecs(j),
specs(i))).min
+ }
+ // Each group's leader is the first member whose partition spec is the
group's minimal
+ // spec; it must come first within the group, so that its exchange is
keyed on the
+ // minimal spec.
+ val leaders = minimalSpecs.indices.map { g =>
+ specs.indices.find(i => classOf(i) == g && equivalent(specs(i),
minimalSpecs(g))).get
+ }
+ // The required-ordering classes of the windows that are not ranked -1, in
first-appearance
+ // order among those windows. Windows whose complete required child
ordering equals their
+ // group leader's get the special rank -1: they belong right after the
leader, where they
+ // ride its sort as well as its exchange (or merge with it). Ranking them
by first
+ // appearance like the other classes would not be stable: the leader is
forced to the
+ // front, which changes its class's first appearance, so re-applying the
rule would
+ // compute a different order. With the special rank the first application
is a fixed point.
+ val orderClass = Array.fill(windows.length)(-1)
+ val rankedOrders = mutable.ArrayBuffer.empty[Seq[SortOrder]]
+ windows.indices.foreach { i =>
+ val leader = leaders(classOf(i))
+ if (i != leader && !sameRequiredChildOrdering(windows(leader),
windows(i))) {
+ val spec = requiredChildOrdering(windows(i))
+ val rank = rankedOrders.indexWhere(sameSortOrders(_, spec))
+ if (rank < 0) {
+ rankedOrders += spec
+ orderClass(i) = rankedOrders.length - 1
+ } else {
+ orderClass(i) = rank
+ }
+ }
+ }
+ // The other members ride the leader's exchange in any order, so they are
ordered by
+ // orderClass (then by partition spec length, for stability): windows with
equal complete
+ // required child orderings become adjacent and share the sort inserted
for the first of
+ // them, without affecting the exchange count.
+ minimalSpecs.indices.flatMap { g =>
+ val members = specs.indices.filter(classOf(_) == g)
+ val leader = leaders(g)
+ leader +: members.filterNot(_ == leader).sortBy(i => (orderClass(i),
specs(i).length))
+ }
+ }
+
+ /**
+ * Whether the expression is safe in a window with an empty order spec, i.e.
contains no
+ * window expression whose frame has bounds relative to the current row:
only the
+ * whole-partition frame covers the same rows under any row order (the same
distinction
+ * `CollapseWindow` draws). A bare aggregate without an explicit frame also
evaluates
+ * over the whole partition, so it is safe.
+ */
+ private def orderSafeWithEmptyOrder(windowExpression: NamedExpression):
Boolean =
+ !windowExpression.exists {
+ case WindowExpression(_, WindowSpecDefinition(_, _, frame:
SpecifiedWindowFrame)) =>
+ !(frame.lower == UnboundedPreceding && frame.upper ==
UnboundedFollowing)
+ case _ => false
+ }
+
+ /**
+ * Whether `filter` above the chain top is a rank filter on the chain's top
window that
+ * `InferWindowGroupLimit` would rewrite later (e.g. `rn <= 1` under
+ * `windowGroupLimitThreshold`). The condition is looked through the chain's
top links, as
+ * the predicate pushdown rules may not have moved the filter below them yet.
+ */
+ private def isRankFilterOnTopWindow(filter: Filter, chain: WindowChain):
Boolean = {
+ val topWindow = chain.windows.last
+ if (topWindow.orderSpec.isEmpty ||
+
!topWindow.windowExpressions.forall(InferWindowGroupLimit.isExpandingWindow)) {
+ return false
+ }
+
+ // There might be Project(s) between Filter and Window now. Here we peel
off the
+ // aliases added by these Projects, so we can check whether Filter
contains conditions
+ // that is checking rank function result produced by the windows.
+ val condition = chain.topLinks.reverse.foldLeft(filter.condition) { (cond,
link) =>
+ val aliasMap = link.projectList.collect { case a: Alias => a.toAttribute
-> a.child }.toMap
+ if (aliasMap.isEmpty) {
+ cond
+ } else {
+ cond.transform { case a: Attribute if aliasMap.contains(a) =>
aliasMap(a) }
+ }
+ }
+
+ // Pin only when the group-limit rule would actually select a rank
candidate for the
+ // remapped condition, so that reordering is not constrained by an
optimization that
+ // would never fire (over the threshold, outcompeted by another rank
predicate, or
+ // unnecessary given the known child cardinality). `chain.bottomChild` has
the same
+ // `maxRows` as the pinned Window's child would have: the reorder is
+ // cardinality-preserving.
+ InferWindowGroupLimit
+ .selectRankCandidate(condition, topWindow.windowExpressions,
chain.bottomChild)
+ .isDefined
+ }
+
+ /**
+ * Reorder `chain` using `optimalOrder` and rebuild it, keeping the top
window in place if
+ * `pinTopWindow` (see `isRankFilterOnTopWindow`): only the windows below it
are reordered
+ * then and the order-restoring `Project` is placed below the pinned window,
so that the
+ * strict `Filter`-over-`Window` match of `InferWindowGroupLimit` survives.
The stripped
+ * links are re-applied above the reordered windows (below the pinned window
if pinned),
+ * each additionally passing through the attributes that anything above it
still references.
+ * Returns `None` if the chain is already in that target order.
+ */
+ private def reorderChain(chain: WindowChain, pinTopWindow: Boolean):
Option[LogicalPlan] = {
+ val windows = chain.windows
+ val toPermute = if (pinTopWindow) windows.dropRight(1) else windows
+ val order = optimalOrder(toPermute)
+ if (order == toPermute.indices) {
+ return None
+ }
+
+ val baseAttrs = chain.bottomChild.output ++
+ windows.flatMap(_.windowExpressions.map(_.toAttribute))
+ val hoistedLinks =
+ if (pinTopWindow) chain.betweenLinks.flatten
+ else chain.betweenLinks.flatten ++ chain.topLinks
+ val boundaryOutput =
+ if (pinTopWindow) chain.topWindowChildOutput else chain.originalOutput
+ val boundaryIds = boundaryOutput.map(_.exprId).toSet
+
+ // refsAbove(i): the attribute ids referenced by the links above
hoistedLinks(i).
+ val refsAbove = new Array[Set[ExprId]](hoistedLinks.length)
+ var acc = Set.empty[ExprId]
+ hoistedLinks.indices.reverse.foreach { i =>
+ refsAbove(i) = acc
+ acc = acc ++ hoistedLinks(i).references.map(_.exprId)
+ }
+
+ var cur: LogicalPlan = chain.bottomChild
+ order.foreach(i => cur = toPermute(i).withNewChildren(Seq(cur)))
+
+ var available = baseAttrs.map(_.exprId).toSet
+ hoistedLinks.indices.foreach { i =>
+ val link = hoistedLinks(i)
+ val required = boundaryIds ++ refsAbove(i)
+ val missing = baseAttrs.filter(a =>
+ required.contains(a.exprId) && available.contains(a.exprId) &&
+ !link.output.exists(_.exprId == a.exprId))
+ cur = if (missing.isEmpty) {
+ link.withNewChildren(Seq(cur))
+ } else {
+ link.copy(projectList = link.projectList ++ missing, child = cur)
+ }
+ available = cur.output.map(_.exprId).toSet
+ }
+
+ if (pinTopWindow) {
+ if (cur.output.map(_.exprId) !=
chain.topWindowChildOutput.map(_.exprId)) {
+ cur = Project(chain.topWindowChildOutput, cur)
+ }
+ cur = windows.last.withNewChildren(Seq(cur))
+ chain.topLinks.foreach(link => cur = link.withNewChildren(Seq(cur)))
+ } else if (cur.output.map(_.exprId) != chain.originalOutput.map(_.exprId))
{
+ cur = Project(chain.originalOutput, cur)
+ }
+ Some(cur)
+ }
+
+ /** Reorder the window chain topped by `plan`, if any, given the plan's
parent. */
+ private def reorderIfChainTop(
+ plan: LogicalPlan, parent: Option[LogicalPlan]): Option[LogicalPlan] = {
+ val chain = collectChain(plan).getOrElse(return None)
+ if (!reorderable(chain)) {
Review Comment:
**Non-blocking (P2):** With the flag on, `apply` sends every plan through
`reorderWindowStacks`, and this early return leaves a rejected chain untouched.
`reorderable` is all-or-nothing, so one blocking window anywhere in the chain
also cancels the pairwise transposition that the flag-off path still performs
on the other windows.
Example, bottom-to-top: `r = rank() OVER (PARTITION BY a, b ORDER BY x)`, `s
= sum(x) OVER (PARTITION BY a)`, `t = sum(r) OVER (PARTITION BY c)`. With the
flag off, `transposeAdjacentPairs` swaps `s` below `r`, so the plan shuffles on
`[a]` (which `r` rides) and on `[c]`: 2 exchanges. With the flag on, `t` reads
`r`'s output, `reorderable` is false, and the plan keeps exchanges on `[a, b]`,
`[a]` and `[c]`: 3. The same happens when a chain with a transposable pair also
contains a nondeterministic window, an empty partition spec, or an empty-order
window with a relative frame. This contradicts the overview's "This subsumes
the previous adjacent-pair transposition". Results are unchanged; only the
opt-in plan is worse than flag-off.
**Recommended change:** When the flag is on and reorderable rejects a chain,
apply the legacy adjacent-pair transposition to that chain's windows instead of
returning it unchanged. Limit the fallback to the rejected chain (not below
chain.bottomChild), so it never reorders windows already handled by the stack
reordering. Update the Scaladoc's 'subsumes' sentence to describe this
fallback. Add optimizer tests for flag-on chains that contain a blocking window
next to a transposable pair.
**Why this works:** Every chain the stack reordering declines receives
exactly the transformation the flag-off path would apply, and chains it accepts
are untouched by the legacy rule. The enabled mode is therefore never worse
than flag-off on rejected chains, and the documented subsumption holds.
**Scope:** Per-chain legacy fallback inside TransposeWindow's enabled path,
plus Scaladoc and window-optimizer test updates.
**Compatibility:** The flag-off plans and the reordered plans of reorderable
chains.
**Risks:** Applying the legacy rule beyond the rejected chain could undo or
oscillate with stack-reordered chains.
**Constraints:** The flag-off path must stay byte-for-byte the existing
transposeAdjacentPairs. Reorderable chains must keep their current optimalOrder
result.
**Success:** With the flag on, a chain such as r = rank() OVER (a, b), s =
sum OVER (a), t = sum(r) OVER (c) gets the same 2-exchange plan as with the
flag off. With the flag on, no rejected chain plans more exchanges than with
the flag off. The optimizer reaches a fixed point on such chains without
alternating between rewrites.
**Verification:**
- **Regression:** With the flag on, a chain containing a window that
references another window's output, next to a strict-subset pair, is transposed
exactly as with the flag off.
- **Regression:** With the flag on, chains blocked by a nondeterministic
window, an empty-order relative frame, or an empty partition spec still receive
the legacy pairwise transposition.
- **Compatibility:** With the flag off, all existing transposition plans are
unchanged, and reorderable chains keep their reordered plans.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/WindowReorderBenchmark.scala:
##########
@@ -0,0 +1,137 @@
+/*
+ * 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.benchmark
+
+import org.apache.spark.SparkConf
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Benchmark to measure the `TransposeWindow` stack reordering behind
+ * `spark.sql.optimizer.windowReorder.enabled` (default false). The query
+ * mirrors the window-function shape of a typical LLM-serving ETL: one wide
+ * SELECT that annotates every inference request with a batch of ROW_NUMBER /
+ * LAG ranking columns, keyed on a single user dimension. The window functions
+ * share one timestamp ORDER BY but are partitioned over two overlapping key
Review Comment:
**Nit (P3):** `windowedSelect` uses both `ORDER BY create_ts` and `ORDER BY
create_ts DESC` (`rn_user_desc`, `rn_user_tier_desc`, `rn_user_model_desc`).
Every complete required ordering in the query is distinct, so here the reorder
saves shuffles but cannot share sorts. The doc misdescribes the measured query
and where the 1.3X comes from.
**Verification:**
- **Inspection:** The described ORDER BY directions and claimed savings
match windowedSelect.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/TransposeWindowSuite.scala:
##########
@@ -161,4 +185,683 @@ class TransposeWindowSuite extends PlanTest {
comparePlans(optimized, correctAnswer.analyze)
}
+ // Data-backed so that maxRows (Some(3)) does not suppress the
WindowGroupLimit that
+ // InferWindowGroupLimit produces in the rank-filter tests below.
+ private val wideRelation = LocalRelation.fromExternalRows(
+ Seq($"k1".string, $"k2".string, $"k3".string, $"k4".string, $"v".int,
$"u".int),
+ 1.to(3).map(i => Row(s"a$i", s"b$i", s"c$i", s"d$i", i, i)))
+ private val k1 = wideRelation.output(0)
+ private val k2 = wideRelation.output(1)
+ private val k3 = wideRelation.output(2)
+ private val k4 = wideRelation.output(3)
+ private val v = wideRelation.output(4)
+ private val u = wideRelation.output(5)
+
+ // specF is a strict subset of specP; specS is incomparable to both.
+ private val specF = Seq(k1, k2)
+ private val specKF = Seq(k2, k1)
+ private val specP = Seq(k1, k2, k3)
+ private val specS = Seq(k1, k4)
+ // Partition spec with an extra key between two shared keys
(subset-compatible with specF).
+ private val specFExtra = Seq(k1, k4, k2)
+ private val order = Seq(k1.asc)
+
+ test("don't reorder window stacks with non-deterministic expressions") {
+ // Interleaved partition specs that are regrouped when stack reordering is
enabled (see the
+ // test below). `reorderable` must reject the whole chain when any window
is
+ // non-deterministic, not only the legacy adjacent-pair path.
+ val query = wideRelation
+ .window(Seq(sum(v).as("f_s")), specF, order)
+ .window(Seq(sum(v).as("p_s")), specP, order)
+ .window(Seq(sum(v).as("s_s")), specS, order)
+ .window(Seq(Rand(0).as("r"), sum(v).as("f_a")), specF, order)
+ .window(Seq(sum(v).as("p_a")), specP, order)
+ .window(Seq(sum(v).as("s_a")), specS, order)
+
+ val analyzed = query.analyze
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ comparePlans(Optimize.execute(analyzed), analyzed)
+ }
+ }
+
+ test("reorder a window stack by grouping minimal partition specs") {
+ val query = wideRelation
+ .window(Seq(sum(v).as("f_s")), specF, order)
+ .window(Seq(sum(v).as("p_s")), specP, order)
+ .window(Seq(sum(v).as("s_s")), specS, order)
+ .window(Seq(sum(v).as("f_a")), specF, order)
+ .window(Seq(sum(v).as("p_a")), specP, order)
+ .window(Seq(sum(v).as("s_a")), specS, order)
+
+ val analyzed = query.analyze
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("f_s")), specF, order)
+ .window(Seq(sum(v).as("f_a")), specF, order)
+ .window(Seq(sum(v).as("p_s")), specP, order)
+ .window(Seq(sum(v).as("p_a")), specP, order)
+ .window(Seq(sum(v).as("s_s")), specS, order)
+ .window(Seq(sum(v).as("s_a")), specS, order)
+ .select($"k1", $"k2", $"k3", $"k4", $"v", $"u",
+ $"f_s", $"p_s", $"s_s", $"f_a", $"p_a", $"s_a")
+
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ val optimized = Optimize.execute(analyzed)
+ comparePlans(optimized, correctAnswer.analyze)
+ }
+ }
+
+ test("already optimally ordered window stack is unchanged") {
+ // Distinct order specs so that CollapseWindow cannot merge the same-spec
windows, and
+ // ordered so that no window can move next to an equal-order window to
share a sort: the
+ // stack is optimal both in exchanges and in sorts, so the rule must leave
it unchanged.
+ val query = wideRelation
+ .window(Seq(sum(v).as("s1")), Seq(k1), order)
+ .window(Seq(sum(v).as("s2a")), specF, order)
+ .window(Seq(sum(v).as("s2b")), specF, Seq(k2.asc))
+ .window(Seq(sum(v).as("s3")), specP, Seq(k3.asc))
+
+ val analyzed = query.analyze
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ comparePlans(Optimize.execute(analyzed), analyzed)
+ }
+ }
+
+ test("windows with equal required child orderings become adjacent within a
group") {
+ // All three windows share ORDER BY k3, but WindowExec prepends partition
keys to the
+ // required sort. The middle window's extra partition key (k4) must not
sit between the
+ // two specF windows, which share the complete ordering k1, k2, k3.
+ val orderK3 = Seq(k3.asc)
+ val query = wideRelation
+ .window(Seq(sum(v).as("s1")), specF, orderK3)
+ .window(Seq(sum(v).as("s2")), specFExtra, orderK3)
+ .window(Seq(sum(v).as("s3")), specF, orderK3)
+
+ val analyzed = query.analyze
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("s1")), specF, orderK3)
+ .window(Seq(sum(v).as("s3")), specF, orderK3)
+ .window(Seq(sum(v).as("s2")), specFExtra, orderK3)
+ .select($"k1", $"k2", $"k3", $"k4", $"v", $"u", $"s1", $"s2", $"s3")
+
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ val optimized = Optimize.execute(analyzed)
+ comparePlans(optimized, correctAnswer.analyze)
+ }
+ }
+
+ test("reordering a window stack is idempotent") {
+ // The leader (the window whose partition spec is the minimal spec (k1))
carries
+ // ORDER BY k1, which is not the first-appearing order spec in the input:
pulling the
+ // leader to the front changes its class's first appearance, so unless the
leader's
+ // class is ranked specially (see `optimalOrder`) a second application
would reorder
+ // the stack again instead of being a no-op.
+ val query = wideRelation
+ .window(Seq(sum(v).as("b")), Seq(k1, k2), Seq(k2.asc))
+ .window(Seq(sum(v).as("a")), Seq(k1), Seq(k1.asc))
+ .window(Seq(sum(v).as("c")), Seq(k1, k2), Seq(k1.asc))
+
+ val analyzed = query.analyze
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ val reordered = TransposeWindow.apply(analyzed)
+ comparePlans(reordered, TransposeWindow.apply(reordered))
+ }
+ }
+
+ test("windows with identical partition specs become adjacent") {
+ val query = wideRelation
+ .window(Seq(sum(v).as("sum_f")), specF, order)
+ .window(Seq(sum(v).as("sum_p")), specP, order)
+ .window(Seq(count(v).as("count_f")), specF, order)
+
+ val analyzed = query.analyze
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("sum_f")), specF, order)
+ .window(Seq(count(v).as("count_f")), specF, order)
+ .window(Seq(sum(v).as("sum_p")), specP, order)
+ .select($"k1", $"k2", $"k3", $"k4", $"v", $"u", $"sum_f", $"sum_p",
$"count_f")
+
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ val optimized = Optimize.execute(analyzed)
+ comparePlans(optimized, correctAnswer.analyze)
+ }
+ }
+
+ test("requireAllClusterKeysForDistribution groups identical specs, not
subsets") {
+ // specF = (k1, k2) and specS = (k1, k4) both contain k1 but are
incomparable; the two
+ // (k1, k2) windows bookend the (k1, k4) window. Under the subset model
all windows group
+ // under the minimal spec (k1) and the stack is already optimally ordered.
Under
+ // `requireAllClusterKeysForDistribution` a window can only ride an
exchange with the
+ // exact same keys in the same order, so the two (k1, k2) windows must
become adjacent
+ // to share one exchange; the distinct order specs of the two (k1, k2)
windows keep
+ // `CollapseWindow` from merging them.
+ val query = wideRelation
+ .window(Seq(sum(v).as("s1")), Seq(k1), order)
+ .window(Seq(sum(v).as("sum_f")), specF, Seq(k1.asc))
+ .window(Seq(sum(v).as("sum_s")), specS, order)
+ .window(Seq(sum(v).as("sum_f2")), specF, Seq(k2.asc))
+ val analyzed = query.analyze
+
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ withSQLConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key ->
"true") {
+ val optimized = Optimize.execute(analyzed)
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("s1")), Seq(k1), order)
+ .window(Seq(sum(v).as("sum_f")), specF, Seq(k1.asc))
+ .window(Seq(sum(v).as("sum_f2")), specF, Seq(k2.asc))
+ .window(Seq(sum(v).as("sum_s")), specS, order)
+ .select($"k1", $"k2", $"k3", $"k4", $"v", $"u",
+ $"s1", $"sum_f", $"sum_s", $"sum_f2")
+
+ comparePlans(optimized, correctAnswer.analyze)
+ }
+
+ // With the default subset semantics the stack is already optimally
ordered, so the rule
+ // must leave it unchanged.
+ comparePlans(Optimize.execute(analyzed), analyzed)
+ }
+ }
+
+ test("requireAllClusterKeysForDistribution distinguishes partition key
order") {
+ // Under strict clustering, (k1, k2) and (k2, k1) are distinct specs. The
two (k1, k2)
+ // windows must be regrouped together without treating (k2, k1) as the
same class.
+ val query = wideRelation
+ .window(Seq(sum(v).as("sum_fk1")), specF, Seq(k1.asc))
+ .window(Seq(sum(v).as("sum_kf")), specKF, Seq(k2.asc))
+ .window(Seq(sum(v).as("sum_fk2")), specF, Seq(k2.asc))
+
+ val analyzed = query.analyze
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("sum_fk1")), specF, Seq(k1.asc))
+ .window(Seq(sum(v).as("sum_fk2")), specF, Seq(k2.asc))
+ .window(Seq(sum(v).as("sum_kf")), specKF, Seq(k2.asc))
+ .select($"k1", $"k2", $"k3", $"k4", $"v", $"u", $"sum_fk1", $"sum_kf",
$"sum_fk2")
+
+ withSQLConf(
+ SQLConf.WINDOW_REORDER_ENABLED.key -> "true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key -> "true") {
+ comparePlans(Optimize.execute(analyzed), correctAnswer.analyze)
+ }
+ }
+
+ test("windowReorder disabled still transposes compatible adjacent windows") {
+ // With the stack reordering off, the original adjacent-pair transposition
applies: the
+ // upper window's partition spec (k1, k2) is a proper subset of the lower
one's (k1, k2,
+ // k3), so the pair is swapped and the output order restored with a
Project.
+ val query = wideRelation
+ .window(Seq(sum(v).as("sum_p")), specP, order) // bottom (k1, k2, k3)
+ .window(Seq(sum(v).as("sum_f")), specF, order) // top (k1, k2)
+ val analyzed = query.analyze
+
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "false") {
+ val optimized = Optimize.execute(analyzed)
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("sum_f")), specF, order)
+ .window(Seq(sum(v).as("sum_p")), specP, order)
+ .select(k1, k2, k3, k4, v, u, $"sum_p", $"sum_f")
+
+ comparePlans(optimized, correctAnswer.analyze)
+ }
+ }
+
+ test("windowReorder disabled disables stack reordering") {
+ // The stack (k1), (k1, k2), (k1, k4), (k1, k2) is regrouped by the stack
reordering
+ // under requireAllClusterKeysForDistribution (the two (k1, k2) windows
become
+ // adjacent, see the test above). With the config disabled no adjacent
pair is
+ // compatible, so the rule must leave the plan untouched.
+ val query = wideRelation
+ .window(Seq(sum(v).as("s1")), Seq(k1), order)
+ .window(Seq(sum(v).as("sum_f")), specF, Seq(k1.asc))
+ .window(Seq(sum(v).as("sum_s")), specS, order)
+ .window(Seq(sum(v).as("sum_f2")), specF, Seq(k2.asc))
+ val analyzed = query.analyze
+
+ withSQLConf(
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key -> "true",
+ SQLConf.WINDOW_REORDER_ENABLED.key -> "false") {
+ comparePlans(Optimize.execute(analyzed), analyzed)
+ }
+ }
+
+ test("attribute-only projects between windows are transparent to
reordering") {
+ // The project between the windows drops the unused attribute `u`.
+ val query = wideRelation
+ .window(Seq(sum(v).as("sum_p")), specP, order)
+ .select(k1, k2, k3, k4, v, $"sum_p")
+ .window(Seq(sum(v).as("sum_f")), specF, order)
+
+ val analyzed = query.analyze
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("sum_f")), specF, order)
+ .window(Seq(sum(v).as("sum_p")), specP, order)
+ .select(k1, k2, k3, k4, v, $"sum_p", $"sum_f")
+
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ val optimized = Optimize.execute(analyzed)
+ comparePlans(optimized, correctAnswer.analyze)
+ }
+ }
+
+ test("derived aliases in transparent links do not block reordering") {
+ // The project between the windows adds the derived alias `s = v + u`.
Since neither
+ // window above it references `s`, the link is safe to hoist above the
reordered
+ // windows: windows only append columns, so `s` computes the same values
wherever it
+ // is re-applied.
+ val query = wideRelation
+ .window(Seq(sum(v).as("sum_a")), specF, order)
+ .select(k1, k2, k3, k4, v, u, $"sum_a",
+ (v + u).as("s"))
+ .window(Seq(sum(v).as("sum_b")), Seq(k1), order)
+
+ val analyzed = query.analyze
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("sum_b")), Seq(k1), order)
+ .window(Seq(sum(v).as("sum_a")), specF, order)
+ .select(k1, k2, k3, k4, v, u, $"sum_a",
+ (v + u).as("s"), $"sum_b")
+
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ val optimized = Optimize.execute(analyzed)
+ comparePlans(optimized, correctAnswer.analyze)
+ }
+ }
+
+ test("nondeterministic project between windows blocks stack reordering") {
+ // Same reorderable window pair as the derived-alias link test, but the
between-window
+ // project adds a nondeterministic alias. `isTransparentLink` must reject
it so the
+ // chain is not reordered when stack reordering is enabled.
+ val query = wideRelation
+ .window(Seq(sum(v).as("sum_a")), specF, order)
+ .select(k1, k2, k3, k4, v, u, $"sum_a", Rand(0).as("r"))
+ .window(Seq(sum(v).as("sum_b")), Seq(k1), order)
+
+ val analyzed = query.analyze
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ comparePlans(Optimize.execute(analyzed), analyzed)
+ }
+ }
+
+ test("stack reordering preserves attributes across two transparent links") {
+ // Two between-window projects: the lower link defines `s`, the upper link
references
+ // `s` but no window does. Rebuilding hoisted links must keep `s`
available for the
+ // upper link (`refsAbove` / missing-attribute propagation).
+ val query = wideRelation
+ .window(Seq(sum(v).as("sum_f")), specF, order)
+ .select(k1, k2, k3, k4, v, u, $"sum_f", (v + u).as("s"))
+ .window(Seq(sum(v).as("sum_k1")), Seq(k1), order)
+ .select(k1, k2, k3, k4, v, u, $"sum_f", $"sum_k1", $"s", ($"s" +
0).as("t"))
+ .window(Seq(sum(v).as("sum_p")), specP, order)
+
+ val analyzed = query.analyze
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("sum_k1")), Seq(k1), order)
+ .window(Seq(sum(v).as("sum_f")), specF, order)
+ .window(Seq(sum(v).as("sum_p")), specP, order)
+ .select(k1, k2, k3, k4, (v + u).as("s"), $"sum_f", $"sum_k1", $"sum_p",
u, v)
+ .select(k1, k2, k3, k4, v, u, $"sum_f", $"sum_k1", $"s", ($"s" +
0).as("t"), $"sum_p")
+
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ val optimized = Optimize.execute(analyzed)
+ comparePlans(optimized, correctAnswer.analyze)
+ }
+ }
+
+ test("transparent link alias referenced by an upper window blocks
reordering") {
+ // The derived alias `s = v + u` feeds the upper window, so hoisting the
link above it
+ // would hide `s` from that window; the chain must be left untouched.
+ val query = wideRelation
+ .window(Seq(sum(v).as("sum_a")), specF, order)
+ .select(k1, k2, k3, k4, v, u, $"sum_a",
+ (v + u).as("s"))
+ .window(Seq(sum($"s").as("sum_s")), Seq(k1), order)
+
+ val analyzed = query.analyze
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ comparePlans(Optimize.execute(analyzed), analyzed)
+ }
+ }
+
+ test("window with an empty partition spec blocks reordering") {
+ val query = wideRelation
+ .window(Seq(sum(v).as("sum_f")), specF, order)
+ .window(Seq(sum(v).as("sum_all")), Seq.empty, order)
+
+ val analyzed = query.analyze
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ comparePlans(Optimize.execute(analyzed), analyzed)
+ }
+ }
+
+ // A window with no ORDER BY and an explicit relative frame. Post-analysis
it keeps the
+ // `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW` frame it was given.
+ private def runningSumAlias = WindowExpression(
+ sum(v),
+ WindowSpecDefinition(Seq(k1), Seq.empty,
+ SpecifiedWindowFrame(RowFrame, UnboundedPreceding,
CurrentRow))).as("run_sum")
+
+ test("window with an empty order spec and a relative frame blocks
reordering") {
+ // The running-sum window has no ORDER BY, so its relative frame covers
whatever row
+ // prefix the operator receives, in the delivery order established by the
nearest window
+ // below it. Permuting the chain moves that boundary across different row
orders and
+ // changes results, so the whole chain must be skipped.
+ val query = wideRelation
+ .window(Seq(sum(v).as("sum_f")), specF, order)
+ .window(Seq(runningSumAlias), Seq(k1), Seq.empty)
+ .window(Seq(sum(v).as("sum_s")), specS, order)
+
+ val analyzed = query.analyze
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ comparePlans(Optimize.execute(analyzed), analyzed)
+ }
+ }
+
+ test("window with an empty order spec and the whole-partition frame is still
reorderable") {
+ // Without an explicit frame, analysis assigns the whole-partition frame
+ // (`UNBOUNDED PRECEDING` to `UNBOUNDED FOLLOWING`) to an empty-order
window, which
+ // covers the same rows under any row order, so it does not pin the
chain's order.
+ val query = wideRelation
+ .window(Seq(sum(v).as("sum_f")), specF, order)
+ .window(Seq(sum(v).as("sum_k1")), Seq(k1), Seq.empty)
+
+ val analyzed = query.analyze
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("sum_k1")), Seq(k1), Seq.empty)
+ .window(Seq(sum(v).as("sum_f")), specF, order)
+ .select($"k1", $"k2", $"k3", $"k4", $"v", $"u", $"sum_f", $"sum_k1")
+
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ val optimized = Optimize.execute(analyzed)
+ comparePlans(optimized, correctAnswer.analyze)
+ }
+ }
+
+ test("a single window is not a chain and is left unchanged") {
+ // `collectChain` requires at least two adjacent windows, so a lone window
is a no-op.
+ val query = wideRelation.window(Seq(sum(v).as("s1")), Seq(k1), order)
+
+ val analyzed = query.analyze
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ comparePlans(Optimize.execute(analyzed), analyzed)
+ }
+ }
+
+ private def rankAlias = WindowExpression(
+ RowNumber(),
+ WindowSpecDefinition(specP, order,
+ SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rn")
+
+ private def denseRankAlias = WindowExpression(
+ DenseRank(Seq.empty),
+ WindowSpecDefinition(specP, order,
+ SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("dr")
+
+ test("rank filter on the top window pins it and reorders the windows below")
{
Review Comment:
**Nit (P3):** In this test and the other pinned variants (attribute-only,
renaming, two renaming links, and both `TransposeWindowQuerySuite` rank-filter
tests), the unpinned `optimalOrder` would also keep the rank window on top. So
the catalyst tests tell pinned from unpinned only by where the order-restoring
`Project` lands, and only because their batches omit `PushDownPredicates`. The
SQL assertions would pass with no pin at all. A stack such as `sum [k1, k2]`,
`sum [k1, k4]`, `rn = row_number() [k1, k2]` with `rn <= 1` makes the pin
observable. Its unpinned order is `[0, 2, 1]`, so without the pin `rn` ends up
under the `[k1, k4]` window and `WindowGroupLimit` is lost.
See **Shared repair plan 1** in the review body.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -4883,6 +4883,24 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val WINDOW_REORDER_ENABLED =
+ buildConf("spark.sql.optimizer.windowReorder.enabled")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .doc("When true, a stack of adjacent Window operators connected by
transparent projections " +
Review Comment:
**Nit (P3):** This conf is public, so its doc shows up in `SET -v` and on
the generated configuration page, but it reads like the `TransposeWindow`
Scaladoc ("distinct minimal partition spec", "incomparable minimal groups",
"complete required child orderings", "rides the exchange"). It also points
users at `spark.sql.requireAllClusterKeysForDistribution`, which is
`.internal()` and so does not appear in the same docs. The sibling 4.4.0 conf
`collapseWindowWithEmptyOrderSpecInChild` states only the user-visible effect
and its trade-off. Could this doc say in plain words what changes when the flag
is enabled, keep the qualification from the last round (no guaranteed shuffle
minimum, the input's existing partitioning is not considered, sorts are not
minimized), and leave the algorithm to the Scaladoc?
**Verification:**
- **Inspection:** The public conf doc states the effect, the default and the
caveats, including that the shuffle count is not guaranteed minimal and sorts
are not minimized, without planner-internal terms or internal conf references.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/TransposeWindowQuerySuite.scala:
##########
@@ -0,0 +1,260 @@
+/*
+ * 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.QueryTest
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
+import org.apache.spark.sql.execution.window.{WindowExec, WindowGroupLimitExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+
+/**
+ * SQL end-to-end tests for window stack reordering done by the
`TransposeWindow` optimizer
+ * rule, run through a full SparkSession. The logical-rule transformations
themselves are
+ * covered in [[org.apache.spark.sql.catalyst.optimizer.TransposeWindowSuite]].
+ */
+class TransposeWindowQuerySuite extends QueryTest with SharedSparkSession {
+
+ private def withInput(f: => Unit): Unit = {
+ withTempView("t") {
+ spark.range(1000).selectExpr(
+ "cast(id % 10 as string) k1", "cast(id % 7 as string) k2",
+ "cast(id % 5 as string) k3", "cast(id % 3 as string) k4", "id v")
+ .createOrReplaceTempView("t")
+ f
+ }
+ }
+
+ private def numExchanges(df: DataFrame): Int =
+ df.queryExecution.executedPlan.collect { case _: ShuffleExchangeExec => ()
}.size
+
+ private def numWindows(df: DataFrame): Int =
+ df.queryExecution.executedPlan.collect { case _: WindowExec => () }.size
+
+ private def numSorts(df: DataFrame): Int =
+ df.queryExecution.executedPlan.collect { case _: SortExec => () }.size
+
+ test("stacked windows are regrouped to minimize exchanges") {
+ // Partition specs (k1, k2), (k1, k2, k3) and (k1, k4) interleaved in
select-list order;
+ // (k1, k2) and (k1, k4) are the minimal specs, so reordering targets 2
exchanges on a
+ // scan child. Distinct order specs keep CollapseWindow from merging the
same-spec windows.
+ val query =
+ """
+ |SELECT k1, k2, k3, k4, v,
+ | sum(v) OVER (PARTITION BY k1, k2 ORDER BY k1) AS f1,
+ | sum(v) OVER (PARTITION BY k1, k2, k3 ORDER BY k1) AS p1,
+ | sum(v) OVER (PARTITION BY k1, k4 ORDER BY k1) AS s1,
+ | sum(v) OVER (PARTITION BY k1, k2 ORDER BY k2) AS f2,
+ | sum(v) OVER (PARTITION BY k1, k2, k3 ORDER BY k2) AS p2,
+ | sum(v) OVER (PARTITION BY k1, k4 ORDER BY k2) AS s2
+ |FROM t
+ """.stripMargin
+
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ withInput {
+ val actualDf = withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key ->
"true") {
+ val df = sql(query)
+ assert(numWindows(df) == 6)
+ assert(numExchanges(df) == 2)
+ df
+ }
+ // The reordered plan must produce the same result as the default
(reorder off)
+ // plan. The queries have no outer ORDER BY, so `checkAnswer` compares
the answers
+ // without making physical row order part of the contract.
+ val expectedDf = withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key ->
"false") {
Review Comment:
**Nit (P3):** `sql(query)` only analyzes eagerly. The plan is optimized when
`checkAnswer(actualDf, expectedDf)` collects `expectedDf`, after this
`withSQLConf` scope has closed, so the baseline is optimized under whatever the
session value is at that point, not with the flag off. Today that is the
default `false`. If the default flips, or the suite runs with the flag enabled
at session level, this compares reorder-on with reorder-on, and a
result-changing reorder would pass. The regrouping, shared-partition-spec and
both rank-filter tests use the same pattern. Forcing
`expectedDf.queryExecution.executedPlan`, or collecting, inside the scope would
fix it, as the equal-order and relative-frame tests already do.
**Verification:**
- **Regression:** With the flag enabled at session level, the comparison
still uses a flag-off baseline plan.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/TransposeWindowSuite.scala:
##########
@@ -161,4 +185,683 @@ class TransposeWindowSuite extends PlanTest {
comparePlans(optimized, correctAnswer.analyze)
}
+ // Data-backed so that maxRows (Some(3)) does not suppress the
WindowGroupLimit that
+ // InferWindowGroupLimit produces in the rank-filter tests below.
+ private val wideRelation = LocalRelation.fromExternalRows(
+ Seq($"k1".string, $"k2".string, $"k3".string, $"k4".string, $"v".int,
$"u".int),
+ 1.to(3).map(i => Row(s"a$i", s"b$i", s"c$i", s"d$i", i, i)))
+ private val k1 = wideRelation.output(0)
+ private val k2 = wideRelation.output(1)
+ private val k3 = wideRelation.output(2)
+ private val k4 = wideRelation.output(3)
+ private val v = wideRelation.output(4)
+ private val u = wideRelation.output(5)
+
+ // specF is a strict subset of specP; specS is incomparable to both.
+ private val specF = Seq(k1, k2)
+ private val specKF = Seq(k2, k1)
+ private val specP = Seq(k1, k2, k3)
+ private val specS = Seq(k1, k4)
+ // Partition spec with an extra key between two shared keys
(subset-compatible with specF).
+ private val specFExtra = Seq(k1, k4, k2)
+ private val order = Seq(k1.asc)
+
+ test("don't reorder window stacks with non-deterministic expressions") {
+ // Interleaved partition specs that are regrouped when stack reordering is
enabled (see the
+ // test below). `reorderable` must reject the whole chain when any window
is
+ // non-deterministic, not only the legacy adjacent-pair path.
+ val query = wideRelation
+ .window(Seq(sum(v).as("f_s")), specF, order)
+ .window(Seq(sum(v).as("p_s")), specP, order)
+ .window(Seq(sum(v).as("s_s")), specS, order)
+ .window(Seq(Rand(0).as("r"), sum(v).as("f_a")), specF, order)
+ .window(Seq(sum(v).as("p_a")), specP, order)
+ .window(Seq(sum(v).as("s_a")), specS, order)
+
+ val analyzed = query.analyze
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ comparePlans(Optimize.execute(analyzed), analyzed)
+ }
+ }
+
+ test("reorder a window stack by grouping minimal partition specs") {
+ val query = wideRelation
+ .window(Seq(sum(v).as("f_s")), specF, order)
+ .window(Seq(sum(v).as("p_s")), specP, order)
+ .window(Seq(sum(v).as("s_s")), specS, order)
+ .window(Seq(sum(v).as("f_a")), specF, order)
+ .window(Seq(sum(v).as("p_a")), specP, order)
+ .window(Seq(sum(v).as("s_a")), specS, order)
+
+ val analyzed = query.analyze
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("f_s")), specF, order)
+ .window(Seq(sum(v).as("f_a")), specF, order)
+ .window(Seq(sum(v).as("p_s")), specP, order)
+ .window(Seq(sum(v).as("p_a")), specP, order)
+ .window(Seq(sum(v).as("s_s")), specS, order)
+ .window(Seq(sum(v).as("s_a")), specS, order)
+ .select($"k1", $"k2", $"k3", $"k4", $"v", $"u",
+ $"f_s", $"p_s", $"s_s", $"f_a", $"p_a", $"s_a")
+
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ val optimized = Optimize.execute(analyzed)
+ comparePlans(optimized, correctAnswer.analyze)
+ }
+ }
+
+ test("already optimally ordered window stack is unchanged") {
+ // Distinct order specs so that CollapseWindow cannot merge the same-spec
windows, and
+ // ordered so that no window can move next to an equal-order window to
share a sort: the
+ // stack is optimal both in exchanges and in sorts, so the rule must leave
it unchanged.
+ val query = wideRelation
+ .window(Seq(sum(v).as("s1")), Seq(k1), order)
+ .window(Seq(sum(v).as("s2a")), specF, order)
+ .window(Seq(sum(v).as("s2b")), specF, Seq(k2.asc))
+ .window(Seq(sum(v).as("s3")), specP, Seq(k3.asc))
+
+ val analyzed = query.analyze
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ comparePlans(Optimize.execute(analyzed), analyzed)
+ }
+ }
+
+ test("windows with equal required child orderings become adjacent within a
group") {
+ // All three windows share ORDER BY k3, but WindowExec prepends partition
keys to the
+ // required sort. The middle window's extra partition key (k4) must not
sit between the
+ // two specF windows, which share the complete ordering k1, k2, k3.
+ val orderK3 = Seq(k3.asc)
+ val query = wideRelation
+ .window(Seq(sum(v).as("s1")), specF, orderK3)
+ .window(Seq(sum(v).as("s2")), specFExtra, orderK3)
+ .window(Seq(sum(v).as("s3")), specF, orderK3)
+
+ val analyzed = query.analyze
+
+ val correctAnswer = wideRelation
+ .window(Seq(sum(v).as("s1")), specF, orderK3)
+ .window(Seq(sum(v).as("s3")), specF, orderK3)
+ .window(Seq(sum(v).as("s2")), specFExtra, orderK3)
+ .select($"k1", $"k2", $"k3", $"k4", $"v", $"u", $"s1", $"s2", $"s3")
+
+ withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") {
+ val optimized = Optimize.execute(analyzed)
+ comparePlans(optimized, correctAnswer.analyze)
+ }
+ }
+
+ test("reordering a window stack is idempotent") {
+ // The leader (the window whose partition spec is the minimal spec (k1))
carries
Review Comment:
**Nit (P3):** With classification by complete required child ordering, the
three windows here have the orderings `[k1, k2, k2]`, `[k1, k1]` and `[k1, k2,
k1]`, which are pairwise distinct. The special leader rank
(`!sameRequiredChildOrdering(windows(leader), windows(i))`) therefore never
applies, and the test passes unchanged without it. The comment describes the
older order-spec-only classification, in which `c`'s `ORDER BY k1` matched the
leader's. The test also covers only one group, so it misses the cross-group
case in the `optimalOrder` finding.
See **Shared repair plan 2** in the review body.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -1845,12 +1845,417 @@ object CollapseWindow extends Rule[LogicalPlan] {
}
/**
- * Transpose Adjacent Window Expressions.
- * - If the partition spec of the parent Window expression is compatible with
the partition spec
- * of the child window expression, transpose them.
+ * Reorder a stack of `Window` operators to reduce the exchanges inserted for
it. Window
+ * operators only append their output columns and each of them
(re-)partitions/(re-)sorts
+ * its input independently of its position in the stack, so two windows can be
swapped
+ * freely as long as neither references the other's output. Physically, a
window rides the
+ * exchange created for the nearest window below it iff that exchange's key
set is a subset
+ * of the window's partition spec (`HashPartitioning` satisfying
`ClusteredDistribution`),
+ * and `EnsureRequirements` only creates exchanges keyed on exactly some
window's partition
+ * spec. The exchange count is at least the number of minimal partition specs
in the stack
+ * (under semantic subset); grouping windows by the minimal partition spec
their own
+ * partition spec contains reaches that lower bound when the bottom child does
not already
+ * satisfy any window's clustering. Incomparable minimal groups are ordered
without
+ * consulting the bottom child's output partitioning. Within a group, smaller
specs come
+ * first. Windows with the same complete required child ordering (partition
keys as ascending
+ * sort keys, then order spec) are made adjacent so they may share sorts; sort
minimization
+ * is not claimed. This subsumes the previous adjacent-pair transposition.
+ * Under `requireAllClusterKeysForDistribution`, a partitioning only satisfies
a window's
+ * `ClusteredDistribution` with the exact same keys in the same order, so the
subset
+ * relation degenerates to exact spec equality: the lower bound equals the
number of
+ * distinct exact partition specs, reached by grouping identical specs
adjacently when the
+ * bottom child does not already satisfy any window's clustering.
+ *
+ * The matched stack is a maximal run of `Window` operators connected by
transparent links,
+ * where a transparent link is a deterministic `Project` made only of plain
attributes and
+ * aliases. Besides pass-through, rename or drop, an alias may also compute a
derived,
+ * deterministic expression (e.g. `Alias(a + b, "s")`). Hoisting such a link
above the
+ * windows only changes where its expression is computed, not the values it
computes on
+ * (windows never rewrite their input), so the one unsafe case is a window
above the link
+ * referencing its aliased output, which `reorderable` excludes. The links are
stripped and
+ * re-applied above the reordered windows, and a chain-top `Project` restores
the original
+ * output attribute order when needed.
+ *
+ * If the stack is directly topped by a rank filter (e.g. `rn <= 1`, the
pattern that
+ * `InferWindowGroupLimit` rewrites later, in a `Once` batch), the top window
is pinned in
+ * place and only the windows below it are reordered, with the order-restoring
`Project`
+ * placed below the pinned window, so that the strict `Filter`-over-`Window`
match survives.
+ *
+ * The stack reordering is gated by
`spark.sql.optimizer.windowReorder.enabled` (default
+ * false); when disabled, the rule falls back to transposing only adjacent
window pairs whose
+ * upper partition spec is a proper subset of the lower one (see
`transposeAdjacentPairs`).
*/
object TransposeWindow extends Rule[LogicalPlan] {
- private def compatiblePartitions(ps1 : Seq[Expression], ps2:
Seq[Expression]): Boolean = {
+
+ /**
+ * A `Project` is a transparent link of a window chain if it only passes
through, renames
+ * or drops attributes, or adds deterministic aliases (which may compute
derived
+ * expressions), and is itself deterministic.
+ */
+ private def isTransparentLink(plan: LogicalPlan): Boolean = plan match {
+ case p: Project =>
+ p.projectList.forall(e => e.isInstanceOf[Attribute] ||
e.isInstanceOf[Alias]) &&
+ p.expressions.forall(_.deterministic)
+ case _ => false
+ }
+
+ /**
+ * A maximal run of adjacent `Window` operators connected by transparent
links.
+ *
+ * @param windows the windows of the chain, bottom-to-top
+ * @param betweenLinks `betweenLinks(i)` holds the transparent links between
`windows(i)`
+ * and `windows(i + 1)`, bottom-to-top; it has one
element less than
+ * `windows`
+ * @param topLinks the transparent links above the top window, bottom-to-top
+ * @param bottomChild the child of the bottom window
+ * @param topWindowChildOutput the output of the child of the top window, in
the original
+ * plan
+ * @param originalOutput the output of the whole chain, in the original plan
+ */
+ private case class WindowChain(
+ windows: Seq[Window],
+ betweenLinks: Seq[Seq[Project]],
+ topLinks: Seq[Project],
+ bottomChild: LogicalPlan,
+ topWindowChildOutput: Seq[Attribute],
+ originalOutput: Seq[Attribute])
+
+ /** Collect the maximal window chain topped by `plan`, if `plan` tops at
least 2 windows. */
+ private def collectChain(plan: LogicalPlan): Option[WindowChain] = {
+ // Collect the chain top-down: `windows` is filled from the top window
down,
+ // `betweenLinks(i)` holds the transparent links accumulated right above
`windows(i)`,
+ // and `pending` the links found since the last window. The buffers are
reversed into
+ // the bottom-to-top order of `WindowChain` at the end.
+ val windows = mutable.ArrayBuffer.empty[Window]
+ val betweenLinks = mutable.ArrayBuffer.empty[Seq[Project]]
+ var pending = Vector.empty[Project]
+ var cur = plan
+ var collecting = true
+ while (collecting) {
+ cur match {
+ case p: Project if isTransparentLink(p) =>
+ pending = pending :+ p
+ cur = p.child
+ case w: Window =>
+ betweenLinks += pending
+ windows += w
+ pending = Vector.empty
+ cur = w.child
+ case _ =>
+ collecting = false
+ }
+ }
+
+ if (windows.length < 2) {
+ None
+ } else {
+ val n = windows.length
+ Some(WindowChain(
+ windows = windows.reverse.toSeq,
+ betweenLinks = (0 until n - 1).map(i => betweenLinks(n - 1 -
i).reverse),
+ topLinks = betweenLinks.head.reverse,
+ bottomChild = windows.last.child,
+ topWindowChildOutput = windows.head.child.output,
+ originalOutput = plan.output))
+ }
+ }
+
+ /**
+ * Whether a window with partition spec `requiredKeys` can ride an exchange
keyed on
+ * `exchangeKeys`: by default (subset semantics) a partitioning satisfies a
+ * `ClusteredDistribution` if its keys are a subset of the required
clustering keys, so
+ * `exchangeKeys` being a subset of `requiredKeys` suffices; under
+ * `requireAllClusterKeysForDistribution` the partitioning must match the
required keys
+ * exactly and in order, so only identical specs ride each other's exchange.
+ */
+ private def canRideExchange(
+ exchangeKeys: Seq[Expression],
+ requiredKeys: Seq[Expression]): Boolean =
+ if (conf.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION)) {
+ exchangeKeys.length == requiredKeys.length &&
+ exchangeKeys.zip(requiredKeys).forall { case (e1, e2) =>
e1.semanticEquals(e2) }
+ } else {
+ exchangeKeys.forall(e1 => requiredKeys.exists(e1.semanticEquals))
+ }
+
+ /** Semantic equivalence of two partition specs, i.e. mutual subset. */
+ private def equivalent(spec1: Seq[Expression], spec2: Seq[Expression]):
Boolean =
+ canRideExchange(spec1, spec2) && canRideExchange(spec2, spec1)
+
+ /** Semantic equality of two sort order sequences, i.e. element-wise
`semanticEquals`. */
+ private def sameSortOrders(orders1: Seq[SortOrder], orders2:
Seq[SortOrder]): Boolean =
+ orders1.length == orders2.length &&
+ orders1.zip(orders2).forall { case (o1, o2) => o1.semanticEquals(o2) }
+
+ /** Required child ordering for a window, matching `WindowExecBase`. */
+ private def requiredChildOrdering(window: Window): Seq[SortOrder] =
+ window.partitionSpec.map(SortOrder(_, Ascending)) ++ window.orderSpec
+
+ private def sameRequiredChildOrdering(w1: Window, w2: Window): Boolean =
+ sameSortOrders(requiredChildOrdering(w1), requiredChildOrdering(w2))
+
+ /**
+ * The chain is reorderable if all windows are deterministic and have a
non-empty partition
+ * spec (an empty one requires `AllTuples` and must not move), every window
with an empty
+ * order spec is order-safe (see `orderSafeWithEmptyOrder`), no window
references another
+ * window's output, and no window references the aliased output of a link
below it (such a
+ * link could not be re-applied above that window). The aliased output may
be any
+ * deterministic expression; hoisting it above the windows only changes
where it is
+ * computed, never the values it sees, so this check is the real safety
property regardless
+ * of whether a link's alias is a bare rename or a derived expression.
+ */
+ private def reorderable(chain: WindowChain): Boolean = {
+ val windows = chain.windows
+ windows.forall(_.expressions.forall(_.deterministic)) &&
+ windows.forall(_.partitionSpec.nonEmpty) &&
+ // An empty-order window with a relative frame (e.g. `ROWS ... CURRENT
ROW`) is
+ // order-sensitive: moving it, or any window below it, changes results --
skip the chain.
+ windows.forall { w =>
+ w.orderSpec.nonEmpty ||
w.windowExpressions.forall(orderSafeWithEmptyOrder)
+ } &&
+ windows.indices.forall { i =>
+ windows.indices.forall { j =>
+ i == j ||
windows(i).references.intersect(windows(j).windowOutputSet).isEmpty
+ }
+ } &&
+ chain.betweenLinks.zipWithIndex.forall { case (links, i) =>
+ val aliasedOutputs = AttributeSet(links.flatMap(_.projectList).collect {
+ case a: Alias => a.toAttribute
+ })
+ // The windows above these links are `windows(i + 1)` and up.
+ aliasedOutputs.isEmpty ||
+ (i + 1 until windows.length).forall { k =>
+ windows(k).references.intersect(aliasedOutputs).isEmpty
+ }
+ }
+ }
+
+ /**
+ * Returns a window order as indices into `windows` (bottom-to-top) that
groups by minimal
+ * partition spec and makes equal complete required child orderings
adjacent. Does not use
+ * the bottom child's output partitioning when ordering incomparable minimal
groups.
+ */
+ private def optimalOrder(windows: Seq[Window]): Seq[Int] = {
+ val specs = windows.map(_.partitionSpec)
+ // A spec is minimal if no other spec in the stack is a strict subset of
it, i.e. no
+ // other window already keys an exchange that this window could ride.
+ val minimal = specs.indices.map { i =>
+ !specs.indices.exists { j =>
+ j != i && canRideExchange(specs(j), specs(i)) && !equivalent(specs(j),
specs(i))
+ }
+ }
+ // The distinct minimal specs, ranked by first appearance, bottom-to-top.
+ val minimalSpecs = mutable.ArrayBuffer.empty[Seq[Expression]]
+ specs.indices.filter(minimal).foreach { i =>
+ if (!minimalSpecs.exists(equivalent(_, specs(i)))) {
+ minimalSpecs += specs(i)
+ }
+ }
+ // Every window rides the exchange of the first minimal spec its partition
spec
+ // contains, so all windows of the same class share that exchange.
+ val classOf = specs.indices.map { i =>
+ minimalSpecs.indices.filter(j => canRideExchange(minimalSpecs(j),
specs(i))).min
+ }
+ // Each group's leader is the first member whose partition spec is the
group's minimal
+ // spec; it must come first within the group, so that its exchange is
keyed on the
+ // minimal spec.
+ val leaders = minimalSpecs.indices.map { g =>
+ specs.indices.find(i => classOf(i) == g && equivalent(specs(i),
minimalSpecs(g))).get
+ }
+ // The required-ordering classes of the windows that are not ranked -1, in
first-appearance
+ // order among those windows. Windows whose complete required child
ordering equals their
+ // group leader's get the special rank -1: they belong right after the
leader, where they
+ // ride its sort as well as its exchange (or merge with it). Ranking them
by first
+ // appearance like the other classes would not be stable: the leader is
forced to the
+ // front, which changes its class's first appearance, so re-applying the
rule would
+ // compute a different order. With the special rank the first application
is a fixed point.
+ val orderClass = Array.fill(windows.length)(-1)
+ val rankedOrders = mutable.ArrayBuffer.empty[Seq[SortOrder]]
Review Comment:
**Nit (P3):** `rankedOrders` is shared by all minimal-spec groups, but
members are only sorted within their own group, so a class's rank can change
between passes. Example, bottom-to-top: `L0 [k3] ORDER BY k9`, `A [k1] ORDER BY
k8`, `B [k1] ORDER BY k4`, `C [k1] ORDER BY k3, k5`, `X1 [k1, k3] ORDER BY k5`.
`C` and `X1` share the complete ordering `[k1, k3, k5]`, but `X1` joins group
`(k3)`. The first application gives `L0, X1, A, B, C`. In the second, `X1`
appears first, so its class (shared with `C`) now outranks `B`'s, which gives
`L0, X1, A, C, B`. Only the third application is the identity.
So the comment's "With the special rank the first application is a fixed
point" does not hold. In the FixedPoint batch this costs extra rewrites;
exchange and sort counts are unaffected. A `Once` batch containing
`TransposeWindow` would fail the test-mode idempotence check on such a stack.
See **Shared repair plan 2** in the review body.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -1845,12 +1845,417 @@ object CollapseWindow extends Rule[LogicalPlan] {
}
/**
- * Transpose Adjacent Window Expressions.
- * - If the partition spec of the parent Window expression is compatible with
the partition spec
- * of the child window expression, transpose them.
+ * Reorder a stack of `Window` operators to reduce the exchanges inserted for
it. Window
+ * operators only append their output columns and each of them
(re-)partitions/(re-)sorts
+ * its input independently of its position in the stack, so two windows can be
swapped
+ * freely as long as neither references the other's output. Physically, a
window rides the
+ * exchange created for the nearest window below it iff that exchange's key
set is a subset
+ * of the window's partition spec (`HashPartitioning` satisfying
`ClusteredDistribution`),
+ * and `EnsureRequirements` only creates exchanges keyed on exactly some
window's partition
+ * spec. The exchange count is at least the number of minimal partition specs
in the stack
+ * (under semantic subset); grouping windows by the minimal partition spec
their own
+ * partition spec contains reaches that lower bound when the bottom child does
not already
+ * satisfy any window's clustering. Incomparable minimal groups are ordered
without
+ * consulting the bottom child's output partitioning. Within a group, smaller
specs come
Review Comment:
**Nit (P3):** A few of the new comments no longer match the code:
1. Here, "Within a group, smaller specs come first": non-leader members are
sorted by `(orderClass(i), specs(i).length)`, so a longer spec in an earlier
ordering class comes before a shorter one. The PR body repeats this claim.
2. `equivalent` is documented as "i.e. mutual subset". It is actually mutual
`canRideExchange`, which under `requireAllClusterKeysForDistribution` means
exact ordered equality, so `(k1, k2)` and `(k2, k1)` are not equivalent.
3. The Scaladoc "The original adjacent-pair transposition: ..." sits on
`compatiblePartitions` rather than on `transposeAdjacentPairs`.
4. In `isRankFilterOnTopWindow`, "conditions that is checking rank function
result" should read "conditions that check the rank function result".
**Verification:**
- **Inspection:** The within-group order, the equivalent semantics, the
legacy-path doc placement and the rank-filter comment each match the code they
describe.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InferWindowGroupLimit.scala:
##########
@@ -73,19 +73,57 @@ object InferWindowGroupLimit extends Rule[LogicalPlan] with
PredicateHelper {
* All window expressions should use the same expanding window and do not
contains
* `SizeBasedWindowFunction`, so that we can safely do the early stop.
*/
- private def isExpandingWindow(
+ private[optimizer] def isExpandingWindow(
windowExpression: NamedExpression): Boolean = windowExpression match {
case Alias(WindowExpression(windowFunction, WindowSpecDefinition(_, _,
SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))), _)
if !windowFunction.isInstanceOf[SizeBasedWindowFunction] => true
case _ => false
}
- private def support(windowFunction: Expression): Boolean = windowFunction
match {
+ private[optimizer] def support(windowFunction: Expression): Boolean =
windowFunction match {
Review Comment:
**Nit (P3):** `support` is only called from `selectRankCandidate` in this
object, and `TransposeWindow` uses only `isExpandingWindow` and
`selectRankCandidate`. It can stay `private`.
**Verification:**
- **Compile:** The catalyst optimizer compiles with support private, and
TransposeWindow still compiles against the remaining shared helpers.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -1845,12 +1845,417 @@ object CollapseWindow extends Rule[LogicalPlan] {
}
/**
- * Transpose Adjacent Window Expressions.
- * - If the partition spec of the parent Window expression is compatible with
the partition spec
- * of the child window expression, transpose them.
+ * Reorder a stack of `Window` operators to reduce the exchanges inserted for
it. Window
+ * operators only append their output columns and each of them
(re-)partitions/(re-)sorts
+ * its input independently of its position in the stack, so two windows can be
swapped
+ * freely as long as neither references the other's output. Physically, a
window rides the
+ * exchange created for the nearest window below it iff that exchange's key
set is a subset
+ * of the window's partition spec (`HashPartitioning` satisfying
`ClusteredDistribution`),
+ * and `EnsureRequirements` only creates exchanges keyed on exactly some
window's partition
+ * spec. The exchange count is at least the number of minimal partition specs
in the stack
+ * (under semantic subset); grouping windows by the minimal partition spec
their own
+ * partition spec contains reaches that lower bound when the bottom child does
not already
+ * satisfy any window's clustering. Incomparable minimal groups are ordered
without
+ * consulting the bottom child's output partitioning. Within a group, smaller
specs come
+ * first. Windows with the same complete required child ordering (partition
keys as ascending
+ * sort keys, then order spec) are made adjacent so they may share sorts; sort
minimization
+ * is not claimed. This subsumes the previous adjacent-pair transposition.
+ * Under `requireAllClusterKeysForDistribution`, a partitioning only satisfies
a window's
+ * `ClusteredDistribution` with the exact same keys in the same order, so the
subset
+ * relation degenerates to exact spec equality: the lower bound equals the
number of
+ * distinct exact partition specs, reached by grouping identical specs
adjacently when the
+ * bottom child does not already satisfy any window's clustering.
+ *
+ * The matched stack is a maximal run of `Window` operators connected by
transparent links,
+ * where a transparent link is a deterministic `Project` made only of plain
attributes and
+ * aliases. Besides pass-through, rename or drop, an alias may also compute a
derived,
+ * deterministic expression (e.g. `Alias(a + b, "s")`). Hoisting such a link
above the
+ * windows only changes where its expression is computed, not the values it
computes on
+ * (windows never rewrite their input), so the one unsafe case is a window
above the link
+ * referencing its aliased output, which `reorderable` excludes. The links are
stripped and
+ * re-applied above the reordered windows, and a chain-top `Project` restores
the original
+ * output attribute order when needed.
+ *
+ * If the stack is directly topped by a rank filter (e.g. `rn <= 1`, the
pattern that
+ * `InferWindowGroupLimit` rewrites later, in a `Once` batch), the top window
is pinned in
+ * place and only the windows below it are reordered, with the order-restoring
`Project`
+ * placed below the pinned window, so that the strict `Filter`-over-`Window`
match survives.
+ *
+ * The stack reordering is gated by
`spark.sql.optimizer.windowReorder.enabled` (default
+ * false); when disabled, the rule falls back to transposing only adjacent
window pairs whose
+ * upper partition spec is a proper subset of the lower one (see
`transposeAdjacentPairs`).
*/
object TransposeWindow extends Rule[LogicalPlan] {
- private def compatiblePartitions(ps1 : Seq[Expression], ps2:
Seq[Expression]): Boolean = {
+
+ /**
+ * A `Project` is a transparent link of a window chain if it only passes
through, renames
+ * or drops attributes, or adds deterministic aliases (which may compute
derived
+ * expressions), and is itself deterministic.
+ */
+ private def isTransparentLink(plan: LogicalPlan): Boolean = plan match {
+ case p: Project =>
+ p.projectList.forall(e => e.isInstanceOf[Attribute] ||
e.isInstanceOf[Alias]) &&
+ p.expressions.forall(_.deterministic)
+ case _ => false
+ }
+
+ /**
+ * A maximal run of adjacent `Window` operators connected by transparent
links.
+ *
+ * @param windows the windows of the chain, bottom-to-top
+ * @param betweenLinks `betweenLinks(i)` holds the transparent links between
`windows(i)`
+ * and `windows(i + 1)`, bottom-to-top; it has one
element less than
+ * `windows`
+ * @param topLinks the transparent links above the top window, bottom-to-top
+ * @param bottomChild the child of the bottom window
+ * @param topWindowChildOutput the output of the child of the top window, in
the original
+ * plan
+ * @param originalOutput the output of the whole chain, in the original plan
+ */
+ private case class WindowChain(
+ windows: Seq[Window],
+ betweenLinks: Seq[Seq[Project]],
+ topLinks: Seq[Project],
+ bottomChild: LogicalPlan,
+ topWindowChildOutput: Seq[Attribute],
+ originalOutput: Seq[Attribute])
+
+ /** Collect the maximal window chain topped by `plan`, if `plan` tops at
least 2 windows. */
+ private def collectChain(plan: LogicalPlan): Option[WindowChain] = {
+ // Collect the chain top-down: `windows` is filled from the top window
down,
+ // `betweenLinks(i)` holds the transparent links accumulated right above
`windows(i)`,
+ // and `pending` the links found since the last window. The buffers are
reversed into
+ // the bottom-to-top order of `WindowChain` at the end.
+ val windows = mutable.ArrayBuffer.empty[Window]
+ val betweenLinks = mutable.ArrayBuffer.empty[Seq[Project]]
+ var pending = Vector.empty[Project]
+ var cur = plan
+ var collecting = true
+ while (collecting) {
+ cur match {
+ case p: Project if isTransparentLink(p) =>
+ pending = pending :+ p
+ cur = p.child
+ case w: Window =>
+ betweenLinks += pending
+ windows += w
+ pending = Vector.empty
+ cur = w.child
+ case _ =>
+ collecting = false
+ }
+ }
+
+ if (windows.length < 2) {
+ None
+ } else {
+ val n = windows.length
+ Some(WindowChain(
+ windows = windows.reverse.toSeq,
+ betweenLinks = (0 until n - 1).map(i => betweenLinks(n - 1 -
i).reverse),
+ topLinks = betweenLinks.head.reverse,
+ bottomChild = windows.last.child,
+ topWindowChildOutput = windows.head.child.output,
+ originalOutput = plan.output))
+ }
+ }
+
+ /**
+ * Whether a window with partition spec `requiredKeys` can ride an exchange
keyed on
+ * `exchangeKeys`: by default (subset semantics) a partitioning satisfies a
+ * `ClusteredDistribution` if its keys are a subset of the required
clustering keys, so
+ * `exchangeKeys` being a subset of `requiredKeys` suffices; under
+ * `requireAllClusterKeysForDistribution` the partitioning must match the
required keys
+ * exactly and in order, so only identical specs ride each other's exchange.
+ */
+ private def canRideExchange(
+ exchangeKeys: Seq[Expression],
+ requiredKeys: Seq[Expression]): Boolean =
+ if (conf.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION)) {
+ exchangeKeys.length == requiredKeys.length &&
+ exchangeKeys.zip(requiredKeys).forall { case (e1, e2) =>
e1.semanticEquals(e2) }
+ } else {
+ exchangeKeys.forall(e1 => requiredKeys.exists(e1.semanticEquals))
+ }
+
+ /** Semantic equivalence of two partition specs, i.e. mutual subset. */
+ private def equivalent(spec1: Seq[Expression], spec2: Seq[Expression]):
Boolean =
+ canRideExchange(spec1, spec2) && canRideExchange(spec2, spec1)
+
+ /** Semantic equality of two sort order sequences, i.e. element-wise
`semanticEquals`. */
+ private def sameSortOrders(orders1: Seq[SortOrder], orders2:
Seq[SortOrder]): Boolean =
+ orders1.length == orders2.length &&
+ orders1.zip(orders2).forall { case (o1, o2) => o1.semanticEquals(o2) }
+
+ /** Required child ordering for a window, matching `WindowExecBase`. */
+ private def requiredChildOrdering(window: Window): Seq[SortOrder] =
+ window.partitionSpec.map(SortOrder(_, Ascending)) ++ window.orderSpec
+
+ private def sameRequiredChildOrdering(w1: Window, w2: Window): Boolean =
+ sameSortOrders(requiredChildOrdering(w1), requiredChildOrdering(w2))
+
+ /**
+ * The chain is reorderable if all windows are deterministic and have a
non-empty partition
+ * spec (an empty one requires `AllTuples` and must not move), every window
with an empty
+ * order spec is order-safe (see `orderSafeWithEmptyOrder`), no window
references another
+ * window's output, and no window references the aliased output of a link
below it (such a
+ * link could not be re-applied above that window). The aliased output may
be any
+ * deterministic expression; hoisting it above the windows only changes
where it is
+ * computed, never the values it sees, so this check is the real safety
property regardless
+ * of whether a link's alias is a bare rename or a derived expression.
+ */
+ private def reorderable(chain: WindowChain): Boolean = {
+ val windows = chain.windows
+ windows.forall(_.expressions.forall(_.deterministic)) &&
+ windows.forall(_.partitionSpec.nonEmpty) &&
+ // An empty-order window with a relative frame (e.g. `ROWS ... CURRENT
ROW`) is
+ // order-sensitive: moving it, or any window below it, changes results --
skip the chain.
+ windows.forall { w =>
+ w.orderSpec.nonEmpty ||
w.windowExpressions.forall(orderSafeWithEmptyOrder)
+ } &&
+ windows.indices.forall { i =>
+ windows.indices.forall { j =>
+ i == j ||
windows(i).references.intersect(windows(j).windowOutputSet).isEmpty
+ }
+ } &&
+ chain.betweenLinks.zipWithIndex.forall { case (links, i) =>
+ val aliasedOutputs = AttributeSet(links.flatMap(_.projectList).collect {
+ case a: Alias => a.toAttribute
+ })
+ // The windows above these links are `windows(i + 1)` and up.
+ aliasedOutputs.isEmpty ||
+ (i + 1 until windows.length).forall { k =>
+ windows(k).references.intersect(aliasedOutputs).isEmpty
+ }
+ }
+ }
+
+ /**
+ * Returns a window order as indices into `windows` (bottom-to-top) that
groups by minimal
+ * partition spec and makes equal complete required child orderings
adjacent. Does not use
+ * the bottom child's output partitioning when ordering incomparable minimal
groups.
+ */
+ private def optimalOrder(windows: Seq[Window]): Seq[Int] = {
+ val specs = windows.map(_.partitionSpec)
+ // A spec is minimal if no other spec in the stack is a strict subset of
it, i.e. no
+ // other window already keys an exchange that this window could ride.
+ val minimal = specs.indices.map { i =>
+ !specs.indices.exists { j =>
+ j != i && canRideExchange(specs(j), specs(i)) && !equivalent(specs(j),
specs(i))
+ }
+ }
+ // The distinct minimal specs, ranked by first appearance, bottom-to-top.
+ val minimalSpecs = mutable.ArrayBuffer.empty[Seq[Expression]]
+ specs.indices.filter(minimal).foreach { i =>
+ if (!minimalSpecs.exists(equivalent(_, specs(i)))) {
+ minimalSpecs += specs(i)
+ }
+ }
+ // Every window rides the exchange of the first minimal spec its partition
spec
+ // contains, so all windows of the same class share that exchange.
+ val classOf = specs.indices.map { i =>
+ minimalSpecs.indices.filter(j => canRideExchange(minimalSpecs(j),
specs(i))).min
+ }
+ // Each group's leader is the first member whose partition spec is the
group's minimal
+ // spec; it must come first within the group, so that its exchange is
keyed on the
+ // minimal spec.
+ val leaders = minimalSpecs.indices.map { g =>
+ specs.indices.find(i => classOf(i) == g && equivalent(specs(i),
minimalSpecs(g))).get
+ }
+ // The required-ordering classes of the windows that are not ranked -1, in
first-appearance
+ // order among those windows. Windows whose complete required child
ordering equals their
+ // group leader's get the special rank -1: they belong right after the
leader, where they
+ // ride its sort as well as its exchange (or merge with it). Ranking them
by first
+ // appearance like the other classes would not be stable: the leader is
forced to the
+ // front, which changes its class's first appearance, so re-applying the
rule would
+ // compute a different order. With the special rank the first application
is a fixed point.
+ val orderClass = Array.fill(windows.length)(-1)
+ val rankedOrders = mutable.ArrayBuffer.empty[Seq[SortOrder]]
+ windows.indices.foreach { i =>
+ val leader = leaders(classOf(i))
+ if (i != leader && !sameRequiredChildOrdering(windows(leader),
windows(i))) {
+ val spec = requiredChildOrdering(windows(i))
+ val rank = rankedOrders.indexWhere(sameSortOrders(_, spec))
+ if (rank < 0) {
+ rankedOrders += spec
+ orderClass(i) = rankedOrders.length - 1
+ } else {
+ orderClass(i) = rank
+ }
+ }
+ }
+ // The other members ride the leader's exchange in any order, so they are
ordered by
+ // orderClass (then by partition spec length, for stability): windows with
equal complete
+ // required child orderings become adjacent and share the sort inserted
for the first of
+ // them, without affecting the exchange count.
+ minimalSpecs.indices.flatMap { g =>
+ val members = specs.indices.filter(classOf(_) == g)
+ val leader = leaders(g)
+ leader +: members.filterNot(_ == leader).sortBy(i => (orderClass(i),
specs(i).length))
+ }
+ }
+
+ /**
+ * Whether the expression is safe in a window with an empty order spec, i.e.
contains no
+ * window expression whose frame has bounds relative to the current row:
only the
+ * whole-partition frame covers the same rows under any row order (the same
distinction
+ * `CollapseWindow` draws). A bare aggregate without an explicit frame also
evaluates
+ * over the whole partition, so it is safe.
+ */
+ private def orderSafeWithEmptyOrder(windowExpression: NamedExpression):
Boolean =
+ !windowExpression.exists {
+ case WindowExpression(_, WindowSpecDefinition(_, _, frame:
SpecifiedWindowFrame)) =>
+ !(frame.lower == UnboundedPreceding && frame.upper ==
UnboundedFollowing)
+ case _ => false
+ }
+
+ /**
+ * Whether `filter` above the chain top is a rank filter on the chain's top
window that
+ * `InferWindowGroupLimit` would rewrite later (e.g. `rn <= 1` under
+ * `windowGroupLimitThreshold`). The condition is looked through the chain's
top links, as
+ * the predicate pushdown rules may not have moved the filter below them yet.
+ */
+ private def isRankFilterOnTopWindow(filter: Filter, chain: WindowChain):
Boolean = {
+ val topWindow = chain.windows.last
+ if (topWindow.orderSpec.isEmpty ||
+
!topWindow.windowExpressions.forall(InferWindowGroupLimit.isExpandingWindow)) {
+ return false
+ }
+
+ // There might be Project(s) between Filter and Window now. Here we peel
off the
+ // aliases added by these Projects, so we can check whether Filter
contains conditions
+ // that is checking rank function result produced by the windows.
+ val condition = chain.topLinks.reverse.foldLeft(filter.condition) { (cond,
link) =>
+ val aliasMap = link.projectList.collect { case a: Alias => a.toAttribute
-> a.child }.toMap
+ if (aliasMap.isEmpty) {
+ cond
+ } else {
+ cond.transform { case a: Attribute if aliasMap.contains(a) =>
aliasMap(a) }
+ }
+ }
+
+ // Pin only when the group-limit rule would actually select a rank
candidate for the
+ // remapped condition, so that reordering is not constrained by an
optimization that
+ // would never fire (over the threshold, outcompeted by another rank
predicate, or
+ // unnecessary given the known child cardinality). `chain.bottomChild` has
the same
+ // `maxRows` as the pinned Window's child would have: the reorder is
+ // cardinality-preserving.
+ InferWindowGroupLimit
+ .selectRankCandidate(condition, topWindow.windowExpressions,
chain.bottomChild)
Review Comment:
**Non-blocking (P2):** The pin is decided once, on the predicate and window
shape as they look when `TransposeWindow` runs. `InferWindowGroupLimit`,
however, matches the shape left after the rest of the operator-optimization
batch. With the flag on, this goes wrong in two ways:
1. `extractLimits` only recognizes `IntegerLiteral` bounds, and
`TransposeWindow` runs before `ConstantFolding` and
`UnwrapCastInBinaryComparison`. Take windows `a [k1] ORDER BY v`, `s [k2] ORDER
BY v` and `rn = row_number() [k1] ORDER BY v DESC` with the DataFrame filter
`col("rn") <= lit(1L)`. The analyzed condition is `CAST(rn AS BIGINT) <= 1L`,
so nothing is pinned, and `optimalOrder` gives `a, rn, s`, which puts `s`
(RANGE frame) on top. After the predicate is normalized to `rn <= 1`, the
Filter cannot be pushed through `s`, the pin check fails on `s`, and
`WindowGroupLimit` is never inserted.
2. The pinned order can put a window with the rank window's partition and
order spec directly under it. For `b [k3] ORDER BY x`, `a [k1] ORDER BY k2`, `c
[k3] ORDER BY y` and `rn [k1] ORDER BY k2` with `rn <= 1`, the pinned order is
`b, c, a`. `CollapseWindow`'s Window-Project-Window case then merges `a` into
`rn`, and the merged `sum` with its default RANGE frame fails
`isExpandingWindow`.
With the flag off, both queries keep `WindowGroupLimit`. That breaks the
`selectRankCandidate` contract that the top Window is pinned if and only if the
group-limit rule selects a candidate. Results are unchanged, but
top-k-per-group queries lose the pre-shuffle row reduction.
See **Shared repair plan 1** in the review body.
--
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]
