From 334258cd70b76010b4b7d7352adc0687640438c2 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Thu, 13 Aug 2026 16:57:17 +0800 Subject: [PATCH 1/2] [SPARK-58757][SQL] Allow CollapseWindow to merge windows with an empty order spec Two adjacent Window operators can currently be collapsed only when their order specs are identical. This relaxes CollapseWindow to also merge a window whose order spec is empty into a sibling window with a non-empty order spec, provided every window expression of the empty-order window is order-insensitive. A window expression is order-insensitive when its frame is the whole partition (UNBOUNDED PRECEDING to UNBOUNDED FOLLOWING): the frame always covers every row of the partition regardless of ordering, so aggregates like count or sum give the same value under any ordering, and functions whose result depends on the row order (e.g. collect_list, first) are non-deterministic when the order spec is empty, so evaluating them under any ordering yields a valid result. A bounded frame is order-sensitive and is therefore never merged. The merged window keeps the non-empty order spec of the other window, so a query like SELECT c2, c1, row_number() OVER (PARTITION BY c1 ORDER BY c2) AS rk, count(1) OVER (PARTITION BY c1) FROM t3 now runs with a single window operator instead of two, saving one WindowExec pass (the sort on [c1, c2] is shared in both plans). Co-Authored-By: Claude --- .../sql/catalyst/optimizer/Optimizer.scala | 41 ++++- .../optimizer/CollapseWindowSuite.scala | 165 ++++++++++++++++++ 2 files changed, 203 insertions(+), 3 deletions(-) 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 ea78b8be49bee..921ad3073bf2b 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 @@ -1761,9 +1761,38 @@ object CollapseWindow extends Rule[LogicalPlan] { s1.zip(s2).forall(e => e._1.semanticEquals(e._2)) } + /** + * Returns true if the given window expression can be evaluated under any ordering of the rows + * within a partition without changing the result, so that it can be merged into another window + * with a different (non-empty) order spec. + * + * The frame determines whether the ordering matters. When the frame is the whole partition + * (`UNBOUNDED PRECEDING` to `UNBOUNDED FOLLOWING`), it always covers all the rows of the + * partition regardless of the ordering, so the ordering does not affect the result: aggregates + * such as `count` or `sum` give the same value under any ordering, and functions whose result + * does depend on the row order, such as `collect_list` or `first`, are non-deterministic when + * the order spec is empty, so evaluating them under any ordering yields a valid result. On the + * other hand, a bounded frame (e.g. `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW`) is + * order-sensitive: which rows are in the frame depends on the ordering, so even `count` or + * `sum` would change value, and such a window must not be merged. + */ + private def orderInsensitive(windowExpression: NamedExpression): Boolean = + windowExpression match { + case Alias(WindowExpression(_, WindowSpecDefinition(_, _, + SpecifiedWindowFrame(_, UnboundedPreceding, UnboundedFollowing))), _) => true + case _ => false + } + private def windowsCompatible(w1: Window, w2: Window): Boolean = { specCompatible(w1.partitionSpec, w2.partitionSpec) && - specCompatible(w1.orderSpec, w2.orderSpec) && + // The order specs can differ when one of them is empty, as long as the window expressions + // of the window with the empty order spec are insensitive to the row order. In that case, + // they can be evaluated under the non-empty order spec of the other window. + (specCompatible(w1.orderSpec, w2.orderSpec) || + (w1.orderSpec.isEmpty && w2.orderSpec.nonEmpty && + w1.windowExpressions.forall(orderInsensitive)) || + (w2.orderSpec.isEmpty && w1.orderSpec.nonEmpty && + w2.windowExpressions.forall(orderInsensitive))) && w1.references.intersect(w2.windowOutputSet).isEmpty && w1.windowExpressions.nonEmpty && w2.windowExpressions.nonEmpty && // This assumes Window contains the same type of window expressions. This is ensured @@ -1776,13 +1805,19 @@ object CollapseWindow extends Rule[LogicalPlan] { _.containsPattern(WINDOW), ruleId) { case w1 @ Window(we1, _, _, w2 @ Window(we2, _, _, grandChild, _), _) if windowsCompatible(w1, w2) => - w1.copy(windowExpressions = we2 ++ we1, child = grandChild) + w1.copy( + orderSpec = if (w1.orderSpec.nonEmpty) w1.orderSpec else w2.orderSpec, + windowExpressions = we2 ++ we1, + child = grandChild) case w1 @ Window(we1, _, _, Project(pl, w2 @ Window(we2, _, _, grandChild, _)), _) if windowsCompatible(w1, w2) && w1.references.subsetOf(grandChild.outputSet) => Project( pl ++ w1.windowOutputSet, - w1.copy(windowExpressions = we2 ++ we1, child = grandChild)) + w1.copy( + orderSpec = if (w1.orderSpec.nonEmpty) w1.orderSpec else w2.orderSpec, + windowExpressions = we2 ++ we1, + child = grandChild)) } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseWindowSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseWindowSuite.scala index 515203da7caf6..4fe7596e124ce 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseWindowSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseWindowSuite.scala @@ -19,6 +19,11 @@ package org.apache.spark.sql.catalyst.optimizer import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ +import org.apache.spark.sql.catalyst.expressions.{ + CurrentRow, RowFrame, RowNumber, SpecifiedWindowFrame, + UnboundedFollowing, UnboundedPreceding} +import org.apache.spark.sql.catalyst.expressions.aggregate.{ + AggregateExpression, Complete, Count, First, Sum} import org.apache.spark.sql.catalyst.plans.PlanTest import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan} import org.apache.spark.sql.catalyst.rules.RuleExecutor @@ -168,4 +173,164 @@ class CollapseWindowSuite extends PlanTest { comparePlans(optimized, correctAnswer) } + + test("collapse windows when one has an empty order spec " + + "(row_number + count over the whole partition)") { + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + + val query = testRelation + .window(Seq(rk), partitionSpec1, orderSpec1) + .window(Seq(cnt), partitionSpec1, Nil) + + val analyzed = query.analyze + val optimized = Optimize.execute(analyzed) + assert(analyzed.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(rk, cnt), partitionSpec1, orderSpec1) + + comparePlans(optimized, correctAnswer) + } + + test("collapse windows when the empty-order window has multiple window expressions") { + // Every window expression of the empty-order window must be order-insensitive for the merge. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + val sm = windowExpr( + AggregateExpression(Sum(b), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("sm") + + val query = testRelation + .window(Seq(rk), partitionSpec1, orderSpec1) + .window(Seq(cnt, sm), partitionSpec1, Nil) + + val analyzed = query.analyze + val optimized = Optimize.execute(analyzed) + assert(analyzed.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(rk, cnt, sm), partitionSpec1, orderSpec1) + + comparePlans(optimized, correctAnswer) + } + + test("collapse windows when the empty-order window has first() over the whole partition") { + // `first` is non-deterministic when the order is not determined by the query, so evaluating it + // under the other window's order spec yields a valid result. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val fr = windowExpr( + First(a, ignoreNulls = true).toAggregateExpression(), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("fr") + + val query = testRelation + .window(Seq(rk), partitionSpec1, orderSpec1) + .window(Seq(fr), partitionSpec1, Nil) + + val analyzed = query.analyze + val optimized = Optimize.execute(analyzed) + assert(analyzed.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(rk, fr), partitionSpec1, orderSpec1) + + comparePlans(optimized, correctAnswer) + } + + test("don't collapse windows when the empty-order window has a bounded frame") { + // The frame `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW` is order-sensitive: which rows + // fall in the frame depends on the ordering, so the window cannot be evaluated under the other + // window's order spec. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("cnt") + + val query = testRelation + .window(Seq(rk), partitionSpec1, orderSpec1) + .window(Seq(cnt), partitionSpec1, Nil) + + val optimized = Optimize.execute(query.analyze) + val correctAnswer = query.analyze + + comparePlans(optimized, correctAnswer) + } + + test("collapse windows when the empty-order window is the inner window") { + // The empty-order window can also be the child of the ordered window. In that case its + // expressions are evaluated under the ordered window's order spec, which is valid because all + // of them are order-insensitive. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + + val query = testRelation + .window(Seq(cnt), partitionSpec1, Nil) + .window(Seq(rk), partitionSpec1, orderSpec1) + + val analyzed = query.analyze + val optimized = Optimize.execute(analyzed) + assert(analyzed.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(cnt, rk), partitionSpec1, orderSpec1) + + comparePlans(optimized, correctAnswer) + } + + test("collapse windows with a Project between them when one has an empty order spec") { + // The same merge applies when a Project sits between the two windows and only passes through + // columns that are available below the inner window (SPARK-34565 shape). + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + + val query = testRelation + .window(Seq(cnt), partitionSpec1, Nil) + .select($"a", $"b", $"c", $"cnt") + .window(Seq(rk), partitionSpec1, orderSpec1) + .select($"a", $"b", $"c", $"cnt", $"rk") + + val analyzed = query.analyze + val optimized = Optimize.execute(analyzed) + assert(analyzed.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(cnt, rk), partitionSpec1, orderSpec1) + .select(a, b, c, $"cnt", $"rk") + .analyze + + comparePlans(optimized, correctAnswer) + } } From a7ee6d56bf9fa89171fd9cf51033b52edda81440 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Fri, 14 Aug 2026 11:20:30 +0800 Subject: [PATCH 2/2] [SPARK-58757][SQL] Allow CollapseWindow to merge windows with an empty order spec Address review comments: rename `orderInsensitive` to `canEvaluateUnderAnyOrder` and clarify the doc comment. The value of order-dependent expressions (e.g. `first`, `collect_list`, floating-point `sum`/`avg`) may differ, but since the order spec is empty the query does not fix the row order, so evaluating them under any ordering yields a valid result. Co-Authored-By: Claude --- .../sql/catalyst/optimizer/Optimizer.scala | 33 ++++++++++--------- 1 file changed, 17 insertions(+), 16 deletions(-) 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 921ad3073bf2b..76545a8f16a47 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 @@ -1762,21 +1762,22 @@ object CollapseWindow extends Rule[LogicalPlan] { } /** - * Returns true if the given window expression can be evaluated under any ordering of the rows - * within a partition without changing the result, so that it can be merged into another window - * with a different (non-empty) order spec. + * Returns true if the given window expression can still be evaluated correctly when the rows + * of the partition are reordered, so that it can be merged into another window with a different + * (non-empty) order spec. * - * The frame determines whether the ordering matters. When the frame is the whole partition + * The frame determines whether reordering is safe. When the frame is the whole partition * (`UNBOUNDED PRECEDING` to `UNBOUNDED FOLLOWING`), it always covers all the rows of the - * partition regardless of the ordering, so the ordering does not affect the result: aggregates - * such as `count` or `sum` give the same value under any ordering, and functions whose result - * does depend on the row order, such as `collect_list` or `first`, are non-deterministic when - * the order spec is empty, so evaluating them under any ordering yields a valid result. On the - * other hand, a bounded frame (e.g. `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW`) is - * order-sensitive: which rows are in the frame depends on the ordering, so even `count` or - * `sum` would change value, and such a window must not be merged. + * partition regardless of the ordering, so reordering changes only the order in which the rows + * are seen, never which rows are in the frame. Since the order spec of the window is empty, + * the query does not fix the row order, so evaluating its expressions under any ordering + * yields a valid result, even though the value may differ for order-dependent expressions + * such as `first`, `collect_list`, or floating-point `sum`/`avg`. On the other hand, a bounded + * frame (e.g. `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW`) is order-sensitive: which + * rows are in the frame depends on the ordering, so even `count` or `sum` would change value, + * and such a window must not be merged. */ - private def orderInsensitive(windowExpression: NamedExpression): Boolean = + private def canEvaluateUnderAnyOrder(windowExpression: NamedExpression): Boolean = windowExpression match { case Alias(WindowExpression(_, WindowSpecDefinition(_, _, SpecifiedWindowFrame(_, UnboundedPreceding, UnboundedFollowing))), _) => true @@ -1786,13 +1787,13 @@ object CollapseWindow extends Rule[LogicalPlan] { private def windowsCompatible(w1: Window, w2: Window): Boolean = { specCompatible(w1.partitionSpec, w2.partitionSpec) && // The order specs can differ when one of them is empty, as long as the window expressions - // of the window with the empty order spec are insensitive to the row order. In that case, - // they can be evaluated under the non-empty order spec of the other window. + // of the window with the empty order spec are safe to evaluate under any row order. In that + // case, they can be evaluated under the non-empty order spec of the other window. (specCompatible(w1.orderSpec, w2.orderSpec) || (w1.orderSpec.isEmpty && w2.orderSpec.nonEmpty && - w1.windowExpressions.forall(orderInsensitive)) || + w1.windowExpressions.forall(canEvaluateUnderAnyOrder)) || (w2.orderSpec.isEmpty && w1.orderSpec.nonEmpty && - w2.windowExpressions.forall(orderInsensitive))) && + w2.windowExpressions.forall(canEvaluateUnderAnyOrder))) && w1.references.intersect(w2.windowOutputSet).isEmpty && w1.windowExpressions.nonEmpty && w2.windowExpressions.nonEmpty && // This assumes Window contains the same type of window expressions. This is ensured