uros-b commented on code in PR #58144:
URL: https://github.com/apache/spark/pull/58144#discussion_r3873004883
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -1834,12 +1834,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] = {
Review Comment:
optimalOrder sorts only by partition-spec length. For windows [a] → [a,b,c]
→ [a,b] with empty order specs, the original needs two sorts because [a,b,c]
satisfies [a,b]. Reordering to [a] → [a,b] → [a,b,c] needs three, without
reducing exchanges.
--
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]
