You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ul...@apache.org on 2026/07/31 05:21:24 UTC
(spark) branch branch-4.x updated: [SPARK-58405][SQL] Pull up project alias through window to eliminate redundant shuffle and sort
This is an automated email from the ASF dual-hosted git repository.
ulysses-you pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 4525284e193c [SPARK-58405][SQL] Pull up project alias through window to eliminate redundant shuffle and sort
4525284e193c is described below
commit 4525284e193caa15143f00cb3d01c870e4612628
Author: Xiduo You <ul...@gmail.com>
AuthorDate: Fri Jul 31 13:20:55 2026 +0800
[SPARK-58405][SQL] Pull up project alias through window to eliminate redundant shuffle and sort
### What changes were proposed in this pull request?
This PR adds a new logical optimizer rule `PullUpProjectAliasThroughWindow`.
When a window function's partition/order key is an alias of an underlying
column, the rename lives in a `Project` *below* the window, while the
analyzer-inserted `Project` *above* the window references the renamed column
as a bare `Attribute` (empty alias map). Since `Window.outputPartitioning`
/`outputOrdering` are pure pass-throughs of the child, the upper project
cannot translate the child's `HashPartitioning(k)`/`SortOrder(k)` into the
renamed attribute a downstream operator requires.
The new rule pulls such aliases up from the bottom `Project` into the top
`Project`, across a chain of one or more adjacent `Window` operators
(`Project - Window+ - Project`), reusing the same exprId. After the rewrite,
the top project's alias map maps `k -> a`, so the existing
`PartitioningPreservingUnaryExecNode`/`OrderPreservingUnaryExecNode`
machinery that `ProjectExec` mixes in projects the window's partitioning and
ordering up through the alias. The `Window` operators themselves are left
untouched.
Example:
```sql
SELECT userid, count(*), sum(size(vset))
FROM (
SELECT key AS userid, collect_set(value) OVER (PARTITION BY key) AS vset
FROM testData
) u
GROUP BY 1
```
Before: the window shuffles on `key`, then a second, redundant shuffle on
`userid` is inserted for the aggregate. After: a single shuffle.
An entry of the bottom project is pulled up only when: it is an `Alias`; no
window in the chain references it; all of its input attributes remain
produced by the pruned bottom project; and the top project consumes its
output solely as a bare pass-through attribute.
Besides removing the redundant shuffle/sort, this also narrows the data
crossing the window's shuffle: the alias no longer flows through the window
as a separate column and is recomputed cheaply above the exchange.
### Why are the changes needed?
To avoid a redundant shuffle (and possibly a redundant sort) whenever a
downstream consumer (aggregate, repartition, chained window, ...) requires
distribution/ordering on an alias of a window partition/order key. The window
has already shuffled/sorted on the underlying key; the second exchange is
pure overhead.
### Does this PR introduce _any_ user-facing change?
No. It is a plan optimization; results are unchanged.
### How was this patch tested?
New unit tests:
- `PullUpProjectAliasThroughWindowSuite` (catalyst): 14 tests covering pull-up
of partition/order key renames, multi-key windows, computed aliases whose
inputs are all window keys, window chains, and the negative/idempotency/
schema-preservation cases.
- `ProjectedOrderingAndPartitioningSuite` (sql/core): 4 end-to-end tests
asserting a single shuffle (and single sort where applicable) plus
`checkAnswer` correctness, for aggregate/repartition/stacked-window/
window-chain consumers.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 4.8
Closes #57603 from ulysses-you/SPARK-58405.
Authored-by: Xiduo You <ul...@gmail.com>
Signed-off-by: Xiduo You <ul...@apache.org>
(cherry picked from commit 0cb45c2cc03c3b245df4d8acf934805234b1fdd4)
Signed-off-by: Xiduo You <ul...@apache.org>
---
.../spark/sql/catalyst/optimizer/Optimizer.scala | 1 +
.../PullUpProjectAliasThroughWindow.scala | 140 +++++++++
.../sql/catalyst/rules/RuleIdCollection.scala | 1 +
.../PullUpProjectAliasThroughWindowSuite.scala | 315 +++++++++++++++++++++
.../ProjectedOrderingAndPartitioningSuite.scala | 136 +++++++++
5 files changed, 593 insertions(+)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index c058b54d06a7..68bb2aa14529 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -125,6 +125,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
OptimizeRepartition,
EliminateWindowPartitions,
TransposeWindow,
+ PullUpProjectAliasThroughWindow,
NullPropagation,
// NullPropagation may introduce Exists subqueries, so RewriteNonCorrelatedExists must run
// after.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindow.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindow.scala
new file mode 100644
index 000000000000..d58404cc2a1a
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindow.scala
@@ -0,0 +1,140 @@
+/*
+ * 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.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, AttributeSet, NamedExpression}
+import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project, Window}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern.WINDOW
+
+/**
+ * A [[Window]] operator does not project its output partitioning or ordering through aliases: its
+ * physical `outputPartitioning`/`outputOrdering` are pure pass-throughs of the child's. So when a
+ * window partitioned/ordered by `k` is followed by a consumer (aggregate, repartition, a chained
+ * window, ...) that requires distribution/ordering on a rename `k AS a`, a redundant shuffle and/or
+ * sort is inserted: the window already shuffled/sorted on `k`, but `HashPartitioning(k)` does not
+ * satisfy `ClusteredDistribution(a)` because the rename lives in a [[Project]] *below* the window
+ * while the analyzer-inserted [[Project]] *above* references `a` as a bare [[Attribute]] (empty
+ * alias map), and `k` and `a`, though value-identical, have distinct expr ids.
+ *
+ * This rule pulls such aliases up from the bottom [[Project]] into the top [[Project]], across the
+ * chain of one or more [[Window]] operators in between (`Project - Window+ - Project`; adjacent
+ * windows that cannot be collapsed, e.g. with different order specs, leave no [[Project]] between
+ * them). The top project's alias map then maps `k -> a`, and the
+ * `PartitioningPreservingUnaryExecNode` / `OrderPreservingUnaryExecNode` machinery that
+ * `ProjectExec` mixes in projects the windows' `HashPartitioning(k)`/`SortOrder(k)` up through the
+ * alias, satisfying the consumer. The [[Window]] operators themselves are left untouched.
+ *
+ * Besides removing the redundant shuffle/sort, this also narrows the data crossing the window's
+ * shuffle: the alias no longer flows through the window as a separate column (only its input `k`,
+ * already needed for partitioning, does) and is recomputed cheaply in the top project above the
+ * exchange.
+ *
+ * An entry of the bottom project is pulled up only when:
+ * - it is an [[Alias]] (bare pass-through attributes stay below so the windows keep producing
+ * them for the top project to reference);
+ * - it is deterministic: a nondeterministic alias (e.g. `rand()`, `spark_partition_id()`)
+ * evaluated above the window's exchange/sort instead of below it could produce different
+ * values, so it must stay below;
+ * - no window in the chain references it, so window semantics are unaffected;
+ * - all of its input attributes remain produced by the pruned bottom project (i.e. they are
+ * referenced by some window and thus retained), so a computed alias whose inputs would be
+ * dropped is not lifted above windows that no longer output them;
+ * - the top project consumes its output solely as a bare pass-through attribute (never inside a
+ * larger expression), so replacing that attribute with the alias fully preserves it.
+ *
+ * The rewrite is a no-op on its own output (a moved alias is no longer a bare attribute above nor
+ * present below), so it converges immediately at the fixed point.
+ */
+object PullUpProjectAliasThroughWindow extends Rule[LogicalPlan] {
+
+ /**
+ * Matches a non-empty chain of [[Window]] operators bottomed out by a [[Project]], returning the
+ * windows top-to-bottom and that bottom project.
+ */
+ private object WindowChain {
+ def unapply(plan: LogicalPlan): Option[(Seq[Window], Project)] = plan match {
+ case w @ Window(_, _, _, lower: Project, _) => Some((Seq(w), lower))
+ case w @ Window(_, _, _, WindowChain(windows, lower), _) => Some((w +: windows, lower))
+ case _ => None
+ }
+ }
+
+ override def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning(
+ _.containsPattern(WINDOW), ruleId) {
+ // Match the `Project - Window+ - Project` shape: the top project is the windows' parent
+ // scaffolding, and the bottom project is where the rename (`key AS userid`) is defined.
+ case p @ Project(projectList, WindowChain(windows, lower)) =>
+ // Entries any window references must stay below: they define what the bottom project must
+ // keep producing, and in particular carry the partition/order key attributes the pulled-up
+ // aliases reuse.
+ val windowRefs = AttributeSet(windows.flatMap(_.references))
+ val (retained, candidates) =
+ lower.projectList.partition(e => windowRefs.contains(e.toAttribute))
+ val retainedAttrs = AttributeSet(retained.map(_.toAttribute))
+ // Attributes the top project consumes as bare pass-through entries, and those it consumes
+ // inside a larger expression (an alias child, a window-output reference, ...).
+ val bareAttrs = AttributeSet(projectList.collect { case a: Attribute => a })
+ val referencedByExpr = AttributeSet(projectList.flatMap {
+ case _: Attribute => Nil
+ case other => other.references
+ })
+ // Pull up an alias only when: it is deterministic (a nondeterministic alias must not move
+ // across the window's exchange/sort, which would change its per-partition/per-row values);
+ // its inputs survive in the pruned bottom project; the top project passes its output through
+ // as a bare attribute; and the top project does not also consume that output inside an
+ // expression (which would require the windows to keep it).
+ val pullUp = candidates.collect {
+ case a: Alias
+ if a.deterministic &&
+ a.references.subsetOf(retainedAttrs) &&
+ bareAttrs.contains(a.toAttribute) &&
+ !referencedByExpr.contains(a.toAttribute) => a
+ }
+ if (pullUp.isEmpty) {
+ p
+ } else {
+ val pullUpSet: Set[NamedExpression] = pullUp.toSet
+ // Key the lookup by expr id (not by `Attribute.equals`, which also compares qualifier):
+ // the top project may reference these attributes with a different qualifier than the alias
+ // carries below (e.g. across a subquery alias).
+ val pullUpMap = AttributeMap(pullUp.map(a => a.toAttribute -> a))
+ val newProjectList = projectList.map {
+ // Rebuild the alias so it keeps the lower alias's child and expr id but adopts the top
+ // attribute's full identity (name, qualifier, metadata). The lookup is keyed by expr id
+ // only, so a resolved top attribute may carry a different name/qualifier/metadata than
+ // the lower alias (e.g. a different qualifier across a subquery alias); taking them from
+ // the top attribute keeps the output schema byte-for-byte identical.
+ case attr: Attribute =>
+ pullUpMap.get(attr) match {
+ case Some(a) => Alias(a.child, attr.name)(
+ exprId = a.exprId,
+ qualifier = attr.qualifier,
+ explicitMetadata = Some(attr.metadata))
+ case None => attr
+ }
+ case other => other
+ }
+ val newLower = lower.copy(projectList = lower.projectList.filterNot(pullUpSet.contains))
+ val newChild = windows.foldRight(newLower: LogicalPlan) { (w, child) =>
+ w.withNewChildren(child :: Nil)
+ }
+ Project(newProjectList, newChild)
+ }
+ }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala
index 2953c09e183d..a71b4bc6380f 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala
@@ -162,6 +162,7 @@ object RuleIdCollection {
"org.apache.spark.sql.catalyst.optimizer.Optimizer$OptimizeSubqueries" ::
"org.apache.spark.sql.catalyst.optimizer.PropagateEmptyRelation" ::
"org.apache.spark.sql.catalyst.optimizer.PruneFilters" ::
+ "org.apache.spark.sql.catalyst.optimizer.PullUpProjectAliasThroughWindow" ::
"org.apache.spark.sql.catalyst.optimizer.PushDownJoinThroughUnion" ::
"org.apache.spark.sql.catalyst.optimizer.PushDownLeftSemiAntiJoin" ::
"org.apache.spark.sql.catalyst.optimizer.PushExtraPredicateThroughJoin" ::
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindowSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindowSuite.scala
new file mode 100644
index 000000000000..25d62cb016b2
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindowSuite.scala
@@ -0,0 +1,315 @@
+/*
+ * 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.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans._
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.rules._
+import org.apache.spark.sql.types.MetadataBuilder
+
+class PullUpProjectAliasThroughWindowSuite extends PlanTest {
+
+ private object Optimize extends RuleExecutor[LogicalPlan] {
+ val batches =
+ Batch("Pull up project alias through window", FixedPoint(20),
+ PullUpProjectAliasThroughWindow) :: Nil
+ }
+
+ private val testRelation = LocalRelation($"key".int, $"key2".int, $"value".string)
+ private val key = testRelation.output(0)
+ private val key2 = testRelation.output(1)
+ private val value = testRelation.output(2)
+ private val windowFrame = SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow)
+
+ // Builds `Window [row_number() ... AS <name>], <partitionSpec>, <orderSpec>` over `child`.
+ private def windowOver(
+ child: LogicalPlan,
+ name: String,
+ partitionSpec: Seq[Expression],
+ orderSpec: Seq[SortOrder] = Nil): Window = {
+ val spec = windowSpec(partitionSpec, orderSpec, windowFrame)
+ val winExpr = windowExpr(RowNumber(), spec).as(name)
+ Window(Seq(winExpr), partitionSpec, orderSpec, child)
+ }
+
+ test("pull up a rename of the window partition key") {
+ // Project [userid, w] <- both bare attributes
+ // +- Window [... AS w], [key]
+ // +- Project [key AS userid, value, key]
+ val userid = Alias(key, "userid")()
+ val bottom = Project(Seq(userid, value, key), testRelation)
+ val window = windowOver(bottom, "w", key :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(userid.toAttribute, w), window)
+
+ // `userid` is pulled up into the parent project; the lower project keeps the rest.
+ val prunedBottom = Project(Seq(value, key), testRelation)
+ val expected = Project(
+ Seq(Alias(key, "userid")(exprId = userid.exprId), w),
+ window.copy(child = prunedBottom))
+ comparePlans(Optimize.execute(originalQuery), expected)
+ }
+
+ test("pull up every applicable key in a multi-key window") {
+ val u1 = Alias(key, "u1")()
+ val u2 = Alias(key2, "u2")()
+ val bottom = Project(Seq(u1, u2, value, key, key2), testRelation)
+ val window = windowOver(bottom, "w", key :: key2 :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(u1.toAttribute, u2.toAttribute, w), window)
+
+ val prunedBottom = Project(Seq(value, key, key2), testRelation)
+ val expected = Project(
+ Seq(
+ Alias(key, "u1")(exprId = u1.exprId),
+ Alias(key2, "u2")(exprId = u2.exprId),
+ w),
+ window.copy(child = prunedBottom))
+ comparePlans(Optimize.execute(originalQuery), expected)
+ }
+
+ test("pull up a rename of the window order key") {
+ // Window ordered by `value`; the top Project passes `tstamp` (a rename of `value`) through as
+ // a bare attribute. It must be pulled up so the window's output ordering projects through it.
+ val tstamp = Alias(value, "tstamp")()
+ val bottom = Project(Seq(key, tstamp, value), testRelation)
+ val window = windowOver(bottom, "w", key :: Nil, SortOrder(value, Ascending) :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(key, tstamp.toAttribute, w), window)
+
+ val prunedBottom = Project(Seq(key, value), testRelation)
+ val expected = Project(
+ Seq(key, Alias(value, "tstamp")(exprId = tstamp.exprId), w),
+ window.copy(child = prunedBottom))
+ comparePlans(Optimize.execute(originalQuery), expected)
+ }
+
+ test("pull up renames of both partition and order keys") {
+ val userid = Alias(key, "userid")()
+ val tstamp = Alias(value, "tstamp")()
+ val bottom = Project(Seq(userid, tstamp, key, value), testRelation)
+ val window = windowOver(bottom, "w", key :: Nil, SortOrder(value, Ascending) :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(userid.toAttribute, tstamp.toAttribute, w), window)
+
+ val prunedBottom = Project(Seq(key, value), testRelation)
+ val expected = Project(
+ Seq(
+ Alias(key, "userid")(exprId = userid.exprId),
+ Alias(value, "tstamp")(exprId = tstamp.exprId),
+ w),
+ window.copy(child = prunedBottom))
+ comparePlans(Optimize.execute(originalQuery), expected)
+ }
+
+ test("keep a bare pass-through column the window does not reference below the window") {
+ // `key2` is a bare pass-through that the window neither partitions nor orders by. It must stay
+ // below the window so the window keeps producing it for the parent project to reference; only
+ // the rename `key AS userid` is pulled up.
+ val userid = Alias(key, "userid")()
+ val bottom = Project(Seq(userid, key2, value, key), testRelation)
+ val window = windowOver(bottom, "w", key :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(userid.toAttribute, key2, w), window)
+
+ val prunedBottom = Project(Seq(key2, value, key), testRelation)
+ val expected = Project(
+ Seq(Alias(key, "userid")(exprId = userid.exprId), key2, w),
+ window.copy(child = prunedBottom))
+ comparePlans(Optimize.execute(originalQuery), expected)
+ }
+
+ test("no rewrite when the renamed column is not a window key") {
+ // `userid` renames `key`, but the window is partitioned by `value`.
+ val userid = Alias(key, "userid")()
+ val bottom = Project(Seq(userid, value, key), testRelation)
+ val window = windowOver(bottom, "w", value :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(userid.toAttribute, w), window)
+
+ comparePlans(Optimize.execute(originalQuery), originalQuery)
+ }
+
+ test("no rewrite for a computed alias whose input is not a window key") {
+ // `(key2 + 1) AS z` depends on `key2`, which the window does not produce once pruned, so it
+ // cannot be pulled above the window and must stay below.
+ val z = Alias(key2 + 1, "z")()
+ val bottom = Project(Seq(z, value, key), testRelation)
+ val window = windowOver(bottom, "w", key :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(z.toAttribute, w), window)
+
+ comparePlans(Optimize.execute(originalQuery), originalQuery)
+ }
+
+ test("pull up a computed alias whose inputs are all window keys") {
+ // `(key + 1) AS z` is computed rather than a pure rename, but its only input `key` is the
+ // window partition key and is retained below. It is pulled up so the expression is evaluated
+ // in the top project, above the window's shuffle rather than being carried across it. This
+ // narrows the shuffled data even though it yields no partitioning benefit (`HashPartitioning`
+ // on `key` does not project through `key + 1`).
+ val z = Alias(key + 1, "z")()
+ val bottom = Project(Seq(z, value, key), testRelation)
+ val window = windowOver(bottom, "w", key :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(z.toAttribute, w), window)
+
+ val prunedBottom = Project(Seq(value, key), testRelation)
+ val expected = Project(
+ Seq(Alias(key + 1, "z")(exprId = z.exprId), w),
+ window.copy(child = prunedBottom))
+ comparePlans(Optimize.execute(originalQuery), expected)
+ }
+
+ test("pull up a computed alias combining a partition key and an order key") {
+ // `(key + key2) AS z` combines the partition key `key` and the order key `key2`; both are
+ // retained below, so the whole expression lifts above the window.
+ val z = Alias(key + key2, "z")()
+ val bottom = Project(Seq(z, key, key2), testRelation)
+ val window = windowOver(bottom, "w", key :: Nil, SortOrder(key2, Ascending) :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(z.toAttribute, w), window)
+
+ val prunedBottom = Project(Seq(key, key2), testRelation)
+ val expected = Project(
+ Seq(Alias(key + key2, "z")(exprId = z.exprId), w),
+ window.copy(child = prunedBottom))
+ comparePlans(Optimize.execute(originalQuery), expected)
+ }
+
+ test("no rewrite when the window child is not a Project") {
+ val window = windowOver(testRelation, "w", key :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(key, w), window)
+
+ comparePlans(Optimize.execute(originalQuery), originalQuery)
+ }
+
+ test("pull up an alias across a chain of windows") {
+ // Project [userid, w1, w2]
+ // +- Window [... AS w2], [key], [value DESC]
+ // +- Window [... AS w1], [key], [value] <- adjacent, no Project between
+ // +- Project [key AS userid, value, key]
+ // `userid` renames `key` and is referenced by neither window, so it is pulled all the way up
+ // to the top project, through both windows.
+ val userid = Alias(key, "userid")()
+ val bottom = Project(Seq(userid, value, key), testRelation)
+ val w1 = windowOver(bottom, "w1", key :: Nil, SortOrder(value, Ascending) :: Nil)
+ val w1a = w1.windowExpressions.head.toAttribute
+ val w2 = windowOver(w1, "w2", key :: Nil, SortOrder(value, Descending) :: Nil)
+ val w2a = w2.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(userid.toAttribute, w1a, w2a), w2)
+
+ val prunedBottom = Project(Seq(value, key), testRelation)
+ val newW1 = w1.copy(child = prunedBottom)
+ val newW2 = w2.copy(child = newW1)
+ val expected = Project(
+ Seq(Alias(key, "userid")(exprId = userid.exprId), w1a, w2a), newW2)
+ comparePlans(Optimize.execute(originalQuery), expected)
+ }
+
+ test("keep an alias referenced by an inner window below the chain") {
+ // The chain's inner window orders by `tstamp` (a rename of `value`), so `tstamp` is referenced
+ // by a window and must stay below; only `userid` (referenced by no window) is pulled up.
+ val userid = Alias(key, "userid")()
+ val tstamp = Alias(value, "tstamp")()
+ val bottom = Project(Seq(userid, tstamp, key, value), testRelation)
+ val w1 = windowOver(bottom, "w1", key :: Nil, SortOrder(tstamp.toAttribute, Ascending) :: Nil)
+ val w1a = w1.windowExpressions.head.toAttribute
+ val w2 = windowOver(w1, "w2", key :: Nil)
+ val w2a = w2.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(userid.toAttribute, tstamp.toAttribute, w1a, w2a), w2)
+
+ // `tstamp` stays below (referenced by w1's order spec); `userid` is pulled up.
+ val prunedBottom = Project(Seq(tstamp, key, value), testRelation)
+ val newW1 = w1.copy(child = prunedBottom)
+ val newW2 = w2.copy(child = newW1)
+ val expected = Project(
+ Seq(Alias(key, "userid")(exprId = userid.exprId), tstamp.toAttribute, w1a, w2a), newW2)
+ comparePlans(Optimize.execute(originalQuery), expected)
+ }
+
+ test("no rewrite for a nondeterministic alias") {
+ // `spark_partition_id() AS pid` is a leaf with no references, so it would pass the input-
+ // survival check vacuously, but moving it above the window's exchange/sort would change its
+ // per-partition value. It must stay below.
+ val pid = Alias(SparkPartitionID(), "pid")()
+ val bottom = Project(Seq(pid, value, key), testRelation)
+ val window = windowOver(bottom, "w", key :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(pid.toAttribute, w), window)
+
+ comparePlans(Optimize.execute(originalQuery), originalQuery)
+ }
+
+ test("rewrite preserves the top attribute's name, qualifier, and metadata") {
+ // The lookup is keyed by expr id only, so a resolved top attribute can carry a different name,
+ // qualifier, and metadata than the lower alias (same expr id). The rebuilt alias must adopt the
+ // top attribute's full identity so the output schema is byte-for-byte unchanged.
+ val userid = Alias(key, "userid")()
+ val bottom = Project(Seq(userid, value, key), testRelation)
+ val window = windowOver(bottom, "w", key :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ // The top project references `userid` under a different name, qualifier, and metadata.
+ val metadata = new MetadataBuilder().putString("comment", "external").build()
+ val topAttr = AttributeReference("external", key.dataType, key.nullable, metadata)(
+ exprId = userid.exprId, qualifier = Seq("sub"))
+ val originalQuery = Project(Seq(topAttr, w), window)
+
+ val prunedBottom = Project(Seq(value, key), testRelation)
+ val expected = Project(
+ Seq(
+ Alias(key, "external")(
+ exprId = userid.exprId, qualifier = Seq("sub"), explicitMetadata = Some(metadata)),
+ w),
+ window.copy(child = prunedBottom))
+ val optimized = Optimize.execute(originalQuery)
+ comparePlans(optimized, expected)
+ // The output attribute must match the top attribute exactly, not the lower alias's identity.
+ val out = optimized.output.head
+ assert(out.name == "external")
+ assert(out.qualifier == Seq("sub"))
+ assert(out.metadata == metadata)
+ }
+
+ test("rewrite is idempotent") {
+ val userid = Alias(key, "userid")()
+ val bottom = Project(Seq(userid, value, key), testRelation)
+ val window = windowOver(bottom, "w", key :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(userid.toAttribute, w), window)
+
+ val once = Optimize.execute(originalQuery)
+ val twice = Optimize.execute(once)
+ comparePlans(once, twice)
+ }
+
+ test("rewrite preserves the output schema (exprId, name, type, nullability)") {
+ val userid = Alias(key, "userid")()
+ val bottom = Project(Seq(userid, value, key), testRelation)
+ val window = windowOver(bottom, "w", key :: Nil)
+ val w = window.windowExpressions.head.toAttribute
+ val originalQuery = Project(Seq(userid.toAttribute, w), window)
+
+ val optimized = Optimize.execute(originalQuery)
+ // The output must be byte-for-byte identical so that downstream references still resolve.
+ assert(optimized.output === originalQuery.output)
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala
index 42f5e217ff4a..f1f2bb7e0833 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala
@@ -23,6 +23,7 @@ import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeRef
import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, HashPartitioning, KeyedPartitioning, Partitioning, PartitioningCollection, UnknownPartitioning}
import org.apache.spark.sql.connector.catalog.functions.{BucketFunction, YearsFunction}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, TimestampType}
@@ -701,6 +702,141 @@ class ProjectedOrderingAndPartitioningSuite
case other => fail(s"Expected UnknownPartitioning, got $other")
}
}
+
+ test("SPARK-58405: eliminate redundant shuffle when aggregate groups by an alias of " +
+ "the window partition key") {
+ // Register the `testData` view (scoped to just these tests that need it).
+ testData
+ // The window output `vset` is consumed downstream so the window is not eliminated.
+ // The window is hash-partitioned by `key`; the outer aggregate groups by `userid`, which is
+ // an alias of `key`. Since `key` is already a partition key of the window, the outer
+ // aggregate's shuffle on `userid` is redundant.
+ val df = sql(
+ """
+ |SELECT userid, count(*), sum(size(vset))
+ |FROM (
+ | SELECT key AS userid,
+ | collect_set(value) OVER (PARTITION BY key) AS vset
+ | FROM testData
+ |) u
+ |GROUP BY 1
+ """.stripMargin)
+
+ // Correctness: results must match grouping the same window output by `key` directly.
+ val expected = sql(
+ """
+ |SELECT key, count(*), sum(size(vset))
+ |FROM (
+ | SELECT key,
+ | collect_set(value) OVER (PARTITION BY key) AS vset
+ | FROM testData
+ |)
+ |GROUP BY 1
+ """.stripMargin)
+ checkAnswer(df, expected.collect())
+
+ // `collect` from `AdaptiveSparkPlanHelper` descends into the finalized AQE plan, so AQE can
+ // stay enabled. `checkAnswer` above has already materialized the query.
+ val plan = df.queryExecution.executedPlan
+ val shuffles = collect(plan) { case e: ShuffleExchangeExec => e }
+ assert(shuffles.size == 1,
+ s"Expected 1 shuffle but found ${shuffles.size}:\n$plan")
+ }
+
+ test("SPARK-58405: eliminate redundant shuffle for a repartition consumer of a window " +
+ "partitioned by an aliased key") {
+ // Register the `testData` view (scoped to just these tests that need it).
+ testData
+ // Consumer-agnostic: a repartition on `userid` (an alias of the window key `key`) reuses the
+ // window's shuffle rather than adding its own.
+ val df = sql(
+ """
+ |SELECT /*+ REPARTITION(userid) */ userid, vset
+ |FROM (
+ | SELECT key AS userid,
+ | collect_set(value) OVER (PARTITION BY key) AS vset
+ | FROM testData
+ |) u
+ """.stripMargin)
+
+ df.collect()
+ val plan = df.queryExecution.executedPlan
+ val shuffles = collect(plan) { case e: ShuffleExchangeExec => e }
+ assert(shuffles.size == 1,
+ s"Expected 1 shuffle but found ${shuffles.size}:\n$plan")
+ }
+
+ test("SPARK-58405: eliminate redundant shuffle and sort for a window consumer partitioned " +
+ "and ordered by aliased keys") {
+ // Register the `testData` view (scoped to just these tests that need it).
+ testData
+ // Stacked windows: the outer window is PARTITION BY userid ORDER BY tstamp, where `userid`
+ // and `tstamp` are aliases of the inner window's partition key `key` and order key `value`.
+ // The inner window's child is already partitioned by `key` and ordered by `[key, value]`;
+ // pulling both `key AS userid` and `value AS tstamp` above the inner window projects the
+ // partitioning and the full ordering up through the aliases, so the outer window needs
+ // neither a redundant shuffle nor a redundant sort. Without the rule the outer window adds
+ // one of each (2 shuffles, 2 sorts). `size(vset)` keeps the inner window output live so it
+ // is not pruned away.
+ val df = sql(
+ """
+ |SELECT userid, tstamp,
+ | sum(size(vset)) OVER (PARTITION BY userid ORDER BY tstamp) AS s
+ |FROM (
+ | SELECT key AS userid, value AS tstamp,
+ | collect_set(value) OVER (PARTITION BY key ORDER BY value) AS vset
+ | FROM testData
+ |) u
+ """.stripMargin)
+
+ df.collect()
+ val plan = df.queryExecution.executedPlan
+ val shuffles = collect(plan) { case e: ShuffleExchangeExec => e }
+ val sorts = collect(plan) { case s: SortExec => s }
+ assert(shuffles.size == 1,
+ s"Expected 1 shuffle but found ${shuffles.size}:\n$plan")
+ assert(sorts.size == 1,
+ s"Expected 1 sort but found ${sorts.size}:\n$plan")
+ }
+
+ test("SPARK-58405: eliminate redundant shuffle across a chain of windows over an aliased key") {
+ // Register the `testData` view (scoped to just these tests that need it).
+ testData
+ // Two adjacent windows (different order specs, so they are not collapsed and leave no Project
+ // between them) both partition by `key`; the aggregate downstream groups by `userid`, an alias
+ // of `key`. Pulling `key AS userid` up across the whole window chain lets the parent project's
+ // `HashPartitioning(userid)` satisfy the aggregate, so no redundant shuffle is inserted.
+ // `sum(r1)`/`sum(r2)` keep both window outputs live so the chain is not pruned away.
+ val df = sql(
+ """
+ |SELECT userid, count(*), sum(r1), sum(r2)
+ |FROM (
+ | SELECT key AS userid,
+ | row_number() OVER (PARTITION BY key ORDER BY value) AS r1,
+ | rank() OVER (PARTITION BY key ORDER BY value DESC) AS r2
+ | FROM testData
+ |) u
+ |GROUP BY 1
+ """.stripMargin)
+
+ val expected = sql(
+ """
+ |SELECT key, count(*), sum(r1), sum(r2)
+ |FROM (
+ | SELECT key,
+ | row_number() OVER (PARTITION BY key ORDER BY value) AS r1,
+ | rank() OVER (PARTITION BY key ORDER BY value DESC) AS r2
+ | FROM testData
+ |)
+ |GROUP BY 1
+ """.stripMargin)
+ checkAnswer(df, expected.collect())
+
+ val plan = df.queryExecution.executedPlan
+ val shuffles = collect(plan) { case e: ShuffleExchangeExec => e }
+ assert(shuffles.size == 1,
+ s"Expected 1 shuffle but found ${shuffles.size}:\n$plan")
+ }
}
private case class DummyLeafExecWithPartitioning(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org