cloud-fan commented on code in PR #58144: URL: https://github.com/apache/spark/pull/58144#discussion_r4346078525
########## sql/core/src/test/scala/org/apache/spark/sql/execution/TransposeWindowQuerySuite.scala: ########## @@ -0,0 +1,776 @@ +/* + * 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.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.window.{WindowExec, WindowGroupLimitExec} +import org.apache.spark.sql.expressions.Window +import org.apache.spark.sql.functions.{col, lit, row_number, sum} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{ArrayType, MapType} + +/** + * 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 + with AdaptiveSparkPlanHelper { + + 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") Review Comment: **Nit (P3):** `k3 = id % 5` is a function of `k1 = id % 10`, so `PARTITION BY k1, k3` and `k1, k2, k3` group rows exactly like `k1` and `k1, k2`. Most compared windows also order by one of their own partition keys, so each partition is all peers under the default RANGE frame. The windows whose placement the tests check therefore compute equal columns (`f1 = p1 = f2 = p2` and `s1 = s2` in the regrouping and AQE tests, `a = b`, `a1 = a2 = s1`, and `l = x` in the prefix-rider tests). `checkAnswer` against the flag-off plan would still pass if the reorder routed one moved window's output to another or dropped a partition key; only the counters would be left, and they do not see output wiring. Independent key columns (for example coprime moduli) and orderings over non-key columns would make the answer comparison discriminating. **Verification:** - **Regression:** A deliberately mis-wired reorder (two moved windows' outputs exchanged) fails the suite's answer comparison. - **Compatibility:** Every asserted exchange, sort and window count is unchanged with the new data and orderings. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala: ########## @@ -5074,6 +5074,22 @@ object SQLConf { .booleanConf .createWithDefault(false) + val WINDOW_REORDER_ENABLED = + buildConf("spark.sql.optimizer.windowReorder.enabled") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .doc("When true, the optimizer reorders a stack of adjacent window operators to reduce " + + "shuffles when their PARTITION BY columns nest or repeat, and places windows with " + + "compatible sort orders next to each other to share a sort. The shuffle count is " + + "not guaranteed to be minimal, and over input already partitioned for one of the " + + "windows this can even add an exchange. Among orders with the same number of " + Review Comment: **Nit (P3):** This caveat names one case where enabling the flag can add an exchange, input already partitioned for one of the windows. The rank-filter pin is another, on plain input: a `row_number`/`rank`/`dense_rank` filter keeps its window on top so `InferWindowGroupLimit` still applies, and when the disabled rule would swap that window below a window partitioned by a superset of its keys, the flag-on plan keeps an extra exchange and sort. Example: `SELECT * FROM (SELECT *, sum(v) OVER (PARTITION BY k1, k2, k3 ORDER BY ts) AS s, row_number() OVER (PARTITION BY k1, k2 ORDER BY ts) AS rn FROM t) WHERE rn = 1` plans two exchanges and three sorts with the flag on, and one exchange and two sorts with it off. The rejected-chain and derived-cut fallbacks make the same trade by design ('rank filter pins the top window of a rejected chain', 'rank filter pins a one-window chain over a derived-alias cut'). The `TransposeWindow` overview says the pin can cost an exchange, but a user reading only this doc cannot tell. Could the doc say, in plain words, that a filter on a `row_number`/`rank`/`dense_rank` result can keep a shuffle that leaving the flag off would remove? If you narrow the pin in answer to the question under Decision challenges, describe the narrowed rule instead. This extends an [existing discussion](https://github.com/apache/spark/pull/58144#discussion_r4312084105). **Verification:** - **Inspection:** The published windowReorder.enabled doc states, without planner terms, that a filter on a rank-like result can keep a shuffle the disabled rule removes, consistent with shouldPinTopWindow and the pinned fallback paths as shipped. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala: ########## @@ -2159,12 +2161,834 @@ 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, so independent windows can be regrouped by + * partition spec to share exchanges. + * + * - Scope: only maximal runs connected by transparent, deterministic pass-through/rename + * `Project`s; derived projects cut a chain so their expressions are not moved across a + * reordered exchange. Reordering preserves dependency order and does not move windows + * with order-sensitive frames. + * - A rank filter can pin the top window so `InferWindowGroupLimit` retains its + * `Filter`-over-`Window` pattern. The pin requires a group-limit candidate for the + * normalized filter condition and is skipped when the unpinned order would already keep + * the rank window on top riding a lower exchange (see `shouldPinTopWindow`); preserving + * the rewrite can therefore cost the rank window a shared exchange. Only the windows + * below the pin are reordered, in an order whose later collapse cannot merge into the + * pin and defeat the rewrite (see `groupLimitSafeOrder`). + * - The stack reordering is gated by `spark.sql.optimizer.windowReorder.enabled` (default + * false). When disabled, the rule only transposes adjacent window pairs whose upper + * partition spec is a proper subset of the lower one (see `transposeAdjacentPairs`). + * When enabled, reorderable chains are stack-reordered; rejected chains, one-window + * runs, and chain bottoms at a derived-alias cut still get that same pairwise + * transposition. A pinned top window is not swapped down on any of these paths, because + * that would lose the `InferWindowGroupLimit` rewrite (see `shouldPinTopWindow`). */ 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 + * (`Alias` of an attribute), or drops attributes, and is itself deterministic. Derived + * aliases are excluded so they cut the chain (see the `TransposeWindow` overview). + */ + private def isTransparentLink(plan: LogicalPlan): Boolean = plan match { + case p: Project => + p.projectList.forall { + case _: Attribute => true + case Alias(_: Attribute, _) => true + case _ => false + } && 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 one + * window. One-window chains are returned too: they have nothing to reorder, but their + * bottom window may sit on a derived `Project` over another window, where the legacy + * pairwise swap still applies (see `reorderIfChainTop`). + */ + 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.isEmpty) { + 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: mutual `canRideExchange`. Under default + * subset semantics that is unordered key-set equality; under + * `requireAllClusterKeysForDistribution` it is exact ordered equality, so `(k1, k2)` and + * `(k2, k1)` are not equivalent. + */ + 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 + + /** + * Whether `provider`'s required child ordering satisfies `rider`'s: `rider`'s equals + * `provider`'s or is a prefix of it. Mirrors the prefix semantics of + * `SortOrder.orderingSatisfies`, which `EnsureRequirements` uses to reuse a child's + * sort, so a window emitted directly above `provider` (no sort in between) adds no + * sort of its own. + */ + private def requiredOrderingSatisfies(provider: Window, rider: Window): Boolean = { + val (providerOrdering, riderOrdering) = + (requiredChildOrdering(provider), requiredChildOrdering(rider)) + riderOrdering.length <= providerOrdering.length && + sameSortOrders(providerOrdering.take(riderOrdering.length), riderOrdering) + } + + /** + * Whether every window of the chain can move. Each conjunct of the guard below states + * its own reason inline; the shared notions it leans on are order-safety for + * empty-order windows (see `orderSafeWithEmptyOrder`) and link transparency (see + * `isTransparentLink`). + */ + private def reorderable(chain: WindowChain): Boolean = { + val windows = chain.windows + // Reordering must not change the evaluation order of nondeterministic expressions. + windows.forall(_.expressions.forall(_.deterministic)) && + // Empty partition specs require AllTuples and have no exchange key to regroup. + 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) + } && + // Direct window-output dependencies make stack permutation unsafe. + windows.indices.forall { i => + windows.indices.forall { j => + i == j || windows(i).references.intersect(windows(j).windowOutputSet).isEmpty + } + } && + // Links are re-applied above the reordered stack, so no upper window may consume an + // alias introduced by a link below it. + 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 + } + } + } + + /** + * Whether a window appends a wide output: a collection-typed column such as + * collect_list's array, whose per-row payload is unbounded. Only collection types are + * detected (wide strings and many-column outputs are not): the proxy bounds the worst + * case of moving a window below another group's exchange rather than measuring width. + */ + private def hasWideOutput(w: Window): Boolean = + w.windowExpressions.exists { e => + e.dataType.existsRecursively(t => t.isInstanceOf[ArrayType] || t.isInstanceOf[MapType]) Review Comment: **Non-blocking (P2):** `hasWideOutput` treats only array and map results as wide, but `listagg`/`string_agg` (`ListAgg` extends `Collect` and returns `child.dataType`) builds the same frame-growing payload as `collect_list`, just as a string. With the flag on, the wide tier therefore never ranks a listagg group last. Example, bottom to top: `s1 = sum(v) OVER (PARTITION BY k1 ORDER BY v)`, `s4 = sum(v) OVER (PARTITION BY k4 ORDER BY v)`, `l = listagg(s, ',') OVER (PARTITION BY k1, k2 ORDER BY v)`. With the flag off, no adjacent pair swaps, `l` stays on top, and its strings cross no exchange. With the flag on, both groups are narrow and first appearance emits `s1, l, s4`, so `Exchange hash(k4)` and `s4`'s sort carry every row's running string, about p^2/2 characters per `(k1, k2)` partition of p rows. The Scaladoc's exclusion of wide strings fits bounded wide columns, not this aggregate. Checking for a `Collect` aggregate in the window expression, alongside the type check, would cover it. **Verification:** - **Regression:** With the flag on, a running listagg window that rides one of two incomparable minimal-spec groups is emitted in the last group, and the plan keeps the reorder's exchange saving. - **Behavior:** A whole-partition listagg window (no ORDER BY) is also treated as wide. - **Compatibility:** Existing collect_list and map wide-tier plans are unchanged, and one application remains a fixed point. -- 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]
