cloud-fan commented on code in PR #58144: URL: https://github.com/apache/spark/pull/58144#discussion_r4057251287
########## sql/core/src/test/scala/org/apache/spark/sql/execution/TransposeWindowQuerySuite.scala: ########## @@ -0,0 +1,143 @@ +/* + * 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, 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 the exchange-minimizing 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 2 exchanges are optimal. 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 actual = withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "true") { + val df = sql(query) + assert(numWindows(df) == 6) + assert(numExchanges(df) == 2) + df.collect().toSeq + } + // The reordered plan must produce the same result as the default (reorder off) plan. + val expected = withSQLConf(SQLConf.WINDOW_REORDER_ENABLED.key -> "false") { + sql(query).collect().toSeq + } + assert(actual == expected) Review Comment: **Non-blocking (P2):** These queries have no outer `ORDER BY`, while the optimization changes shuffle and sort placement. Direct `Seq` equality can fail for semantically identical rows returned in a different physical order, so these result-equivalence assertions should not make output order part of the contract. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala: ########## @@ -1841,12 +1841,339 @@ 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 so that the fewest possible exchanges are 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 minimum exchange count for a stack therefore equals the + * number of its minimal partition specs (under semantic subset), and it is achieved by + * grouping the windows by the minimal partition spec their own partition spec contains, + * smaller specs first within a group. 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 minimum equals the number of distinct exact + * partition specs, achieved by grouping identical specs adjacently. + * + * 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 an exchange keyed on `spec2` can satisfy the `ClusteredDistribution` of a window + * with partition spec `spec1`. By default (subset semantics) a partitioning satisfies the + * distribution if every partitioning key appears in the required clustering keys, so + * `spec1` being a subset of `spec2` 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 subsetOf(spec1: Seq[Expression], spec2: Seq[Expression]): Boolean = + if (conf.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION)) { + spec1.length == spec2.length && + spec1.zip(spec2).forall { case (e1, e2) => e1.semanticEquals(e2) } + } else { + spec1.forall(e1 => spec2.exists(e1.semanticEquals)) + } + + /** Semantic equivalence of two partition specs, i.e. mutual subset. */ + private def equivalent(spec1: Seq[Expression], spec2: Seq[Expression]): Boolean = + subsetOf(spec1, spec2) && subsetOf(spec2, spec1) + + /** + * 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), 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)) && Review Comment: **Non-blocking (P2):** This admits an empty-order `ROWS ... CURRENT ROW` frame even though that frame depends on the inherited row order. Regrouping the chain can place it above a differently ordered Window, changing the prefix attached to each row and producing different results when the opt-in rule is enabled. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala: ########## @@ -1841,12 +1841,339 @@ 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 so that the fewest possible exchanges are 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 minimum exchange count for a stack therefore equals the + * number of its minimal partition specs (under semantic subset), and it is achieved by + * grouping the windows by the minimal partition spec their own partition spec contains, + * smaller specs first within a group. 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 minimum equals the number of distinct exact + * partition specs, achieved by grouping identical specs adjacently. + * + * 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 an exchange keyed on `spec2` can satisfy the `ClusteredDistribution` of a window + * with partition spec `spec1`. By default (subset semantics) a partitioning satisfies the + * distribution if every partitioning key appears in the required clustering keys, so + * `spec1` being a subset of `spec2` 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 subsetOf(spec1: Seq[Expression], spec2: Seq[Expression]): Boolean = + if (conf.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION)) { + spec1.length == spec2.length && + spec1.zip(spec2).forall { case (e1, e2) => e1.semanticEquals(e2) } + } else { + spec1.forall(e1 => spec2.exists(e1.semanticEquals)) + } + + /** Semantic equivalence of two partition specs, i.e. mutual subset. */ + private def equivalent(spec1: Seq[Expression], spec2: Seq[Expression]): Boolean = + subsetOf(spec1, spec2) && subsetOf(spec2, spec1) + + /** + * 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), 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) && + 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 + } + } + } + + /** + * The exchange-minimal window order, as indices into `windows` (bottom-to-top): group the + * windows by the minimal partition spec their own partition spec contains, groups ordered + * by first appearance, smaller specs first within a group; the sort is stable, so windows + * with equal keys keep their original relative order. Each group leader pays for its + * exchange and every other member rides it, because a riding window never changes the + * partitioning seen by the windows above it, so the result needs one exchange per minimal + * partition spec. + */ + private def optimalOrder(windows: Seq[Window]): Seq[Int] = { + val specs = windows.map(_.partitionSpec) + val minimal = specs.indices.map { i => + !specs.indices.exists { j => + j != i && subsetOf(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) + } + } + val classOf = specs.indices.map { i => + minimalSpecs.indices.filter(j => subsetOf(minimalSpecs(j), specs(i))).min + } + specs.indices.sortBy(i => (classOf(i), specs(i).length)) + } + + /** + * 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 = { + if (conf.windowGroupLimitThreshold == -1) return false + 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) => Review Comment: **Non-blocking (P2):** The current Project passes `rn` through unchanged, so `aliasMap` is empty, and the assertion stops at the TransposeWindow output. A regression in remapping a shape such as `rn AS rank` could silently prevent InferWindowGroupLimit from firing while this test remains green. See **Shared repair plan 1** in the review body. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala: ########## @@ -1841,12 +1841,339 @@ 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 so that the fewest possible exchanges are 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 minimum exchange count for a stack therefore equals the + * number of its minimal partition specs (under semantic subset), and it is achieved by + * grouping the windows by the minimal partition spec their own partition spec contains, + * smaller specs first within a group. 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 minimum equals the number of distinct exact + * partition specs, achieved by grouping identical specs adjacently. + * + * 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 an exchange keyed on `spec2` can satisfy the `ClusteredDistribution` of a window + * with partition spec `spec1`. By default (subset semantics) a partitioning satisfies the + * distribution if every partitioning key appears in the required clustering keys, so + * `spec1` being a subset of `spec2` 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 subsetOf(spec1: Seq[Expression], spec2: Seq[Expression]): Boolean = + if (conf.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION)) { + spec1.length == spec2.length && + spec1.zip(spec2).forall { case (e1, e2) => e1.semanticEquals(e2) } + } else { + spec1.forall(e1 => spec2.exists(e1.semanticEquals)) + } + + /** Semantic equivalence of two partition specs, i.e. mutual subset. */ + private def equivalent(spec1: Seq[Expression], spec2: Seq[Expression]): Boolean = + subsetOf(spec1, spec2) && subsetOf(spec2, spec1) + + /** + * 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), 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) && + 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 + } + } + } + + /** + * The exchange-minimal window order, as indices into `windows` (bottom-to-top): group the + * windows by the minimal partition spec their own partition spec contains, groups ordered + * by first appearance, smaller specs first within a group; the sort is stable, so windows + * with equal keys keep their original relative order. Each group leader pays for its + * exchange and every other member rides it, because a riding window never changes the + * partitioning seen by the windows above it, so the result needs one exchange per minimal + * partition spec. + */ + private def optimalOrder(windows: Seq[Window]): Seq[Int] = { + val specs = windows.map(_.partitionSpec) + val minimal = specs.indices.map { i => + !specs.indices.exists { j => + j != i && subsetOf(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) + } + } + val classOf = specs.indices.map { i => + minimalSpecs.indices.filter(j => subsetOf(minimalSpecs(j), specs(i))).min + } + specs.indices.sortBy(i => (classOf(i), specs(i).length)) + } + + /** + * 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 = { + if (conf.windowGroupLimitThreshold == -1) return false + 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) } + } + } + + topWindow.windowExpressions.exists { + case alias @ Alias(WindowExpression(rankLikeFunction, _), _) + if InferWindowGroupLimit.support(rankLikeFunction) => + InferWindowGroupLimit.extractLimits(condition, alias.toAttribute) + .exists(_ <= conf.windowGroupLimitThreshold) Review Comment: **Non-blocking (P2):** This pins the top Window whenever any supported rank predicate is below the threshold, but InferWindowGroupLimit also checks `child.maxRows` and selects among competing rank predicates. When those gates reject the rewrite, reordering is still constrained and can retain an avoidable exchange for an optimization that never occurs. See **Shared repair plan 1** in the review body. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala: ########## @@ -1841,12 +1841,339 @@ 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 so that the fewest possible exchanges are 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 minimum exchange count for a stack therefore equals the + * number of its minimal partition specs (under semantic subset), and it is achieved by + * grouping the windows by the minimal partition spec their own partition spec contains, + * smaller specs first within a group. 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 minimum equals the number of distinct exact + * partition specs, achieved by grouping identical specs adjacently. + * + * 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 an exchange keyed on `spec2` can satisfy the `ClusteredDistribution` of a window + * with partition spec `spec1`. By default (subset semantics) a partitioning satisfies the + * distribution if every partitioning key appears in the required clustering keys, so + * `spec1` being a subset of `spec2` suffices; under Review Comment: **Nit (P3):** The parameter roles here are reversed: callers pass the exchange/minimal spec as `spec1` and the required Window spec as `spec2`, and satisfaction requires the exchange keys to be a subset of the required keys. The current Scaladoc teaches the opposite physical contract. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala: ########## @@ -4882,6 +4882,20 @@ 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 is reordered, grouping windows that share a partition spec so that " + + "the number of inserted exchanges (shuffles) and sorts is minimized. Windows are " + Review Comment: **Non-blocking (P2):** The implementation orders partition groups but never examines `orderSpec`. A non-prefix X, Y, X sequence therefore remains three sorts, while grouping X, X, Y lets the equal-order Windows share or collapse and needs only two. This public config description promises a guarantee the enabled rule does not provide. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala: ########## @@ -1841,12 +1841,339 @@ 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 so that the fewest possible exchanges are 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 minimum exchange count for a stack therefore equals the + * number of its minimal partition specs (under semantic subset), and it is achieved by + * grouping the windows by the minimal partition spec their own partition spec contains, + * smaller specs first within a group. 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 minimum equals the number of distinct exact + * partition specs, achieved by grouping identical specs adjacently. + * + * 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 an exchange keyed on `spec2` can satisfy the `ClusteredDistribution` of a window + * with partition spec `spec1`. By default (subset semantics) a partitioning satisfies the + * distribution if every partitioning key appears in the required clustering keys, so + * `spec1` being a subset of `spec2` 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 subsetOf(spec1: Seq[Expression], spec2: Seq[Expression]): Boolean = + if (conf.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION)) { + spec1.length == spec2.length && + spec1.zip(spec2).forall { case (e1, e2) => e1.semanticEquals(e2) } + } else { + spec1.forall(e1 => spec2.exists(e1.semanticEquals)) + } + + /** Semantic equivalence of two partition specs, i.e. mutual subset. */ + private def equivalent(spec1: Seq[Expression], spec2: Seq[Expression]): Boolean = + subsetOf(spec1, spec2) && subsetOf(spec2, spec1) + + /** + * 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), 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) && + 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 + } + } + } + + /** + * The exchange-minimal window order, as indices into `windows` (bottom-to-top): group the + * windows by the minimal partition spec their own partition spec contains, groups ordered + * by first appearance, smaller specs first within a group; the sort is stable, so windows + * with equal keys keep their original relative order. Each group leader pays for its + * exchange and every other member rides it, because a riding window never changes the + * partitioning seen by the windows above it, so the result needs one exchange per minimal + * partition spec. + */ + private def optimalOrder(windows: Seq[Window]): Seq[Int] = { + val specs = windows.map(_.partitionSpec) + val minimal = specs.indices.map { i => + !specs.indices.exists { j => + j != i && subsetOf(specs(j), specs(i)) && !equivalent(specs(j), specs(i)) + } + } + // The distinct minimal specs, ranked by first appearance, bottom-to-top. Review Comment: **Non-blocking (P2):** The first incomparable group is chosen only by appearance. If the child already satisfies B but the chain is ordered A then B, this inserts exchanges for both A and B; choosing B then A reuses the child partitioning and needs only A. The current ordering therefore does not establish the claimed exchange minimum. **Recommended change:** Carry the safe Window-chain alternatives from TransposeWindow to physical planning and finalize the order of incomparable minimal groups using the bottom child's actual outputPartitioning before exchanges are inserted. Add ordinary and adaptive physical-plan coverage for a pre-partitioned child. **Why this works:** Keep logical dependency, containment, and pinned-top constraints in TransposeWindow, but defer the otherwise stable first-appearance choice among incomparable minimal groups. At the physical distribution boundary, select a satisfying bottom group first when the child partitioning can be reused, then preserve stable order for remaining cost ties before EnsureRequirements materializes exchanges. **Scope:** Make the exchange-minimum decision at the common planner boundary that knows incoming partitioning. **Compatibility:** TransposeWindow continues to define which Window permutations are semantically safe; physical planning chooses only among those safe alternatives. **Risks:** Physical reordering can break Window expression dependencies or Project output identity if it is not restricted to the alternatives proven safe by TransposeWindow. A decision made only in the initial physical plan can become stale after adaptive rewrites change child partitioning. **Constraints:** Preserve the disabled adjacent-pair behavior and all TransposeWindow dependency guards. Preserve pinned rank topology and stable relative order when exchange cost is tied. Apply the same input-aware decision on ordinary and adaptive requirement-enforcement routes. **Success:** For incomparable minimal specs A and B above a child already partitioned by B, B is placed first and only the A exchange is inserted. The same reuse decision holds on both ordinary and adaptive planning paths. Plans without useful incoming partitioning retain a stable valid order and unchanged results and output identity. -- 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]
