[SPARK-45352][SQL] Eliminate foldable window partitions #43144

zml1206 · 2023-09-27T08:37:15Z

What changes were proposed in this pull request?

This PR add a new optimizer rule EliminateWindowPartitions, it remove window partition if partition expressions are foldable.
sql1:
select row_number() over(order by a) b from t
sql2:
select row_number() over(partition by 1 order by a) b from t
After this PR, the optimizedPlan for sql1 and sql2 is the same.

Why are the changes needed?

Foldable partition is redundant, remove it not only can simplify plan, but some rules can also take effect when the partitions are all foldable, such as LimitPushDownThroughWindow.

Does this PR introduce any user-facing change?

No

How was this patch tested?

UT

Was this patch authored or co-authored using generative AI tooling?

No

beliefer · 2023-09-27T10:56:30Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

OptimizeWindowPartitions -> EliminateWindowPartitions

beliefer · 2023-09-27T11:02:08Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

Suggested change

plan.transformAllExpressionsWithPruning(_.containsAnyPattern(WINDOW_EXPRESSION), ruleId) {

plan.transformAllExpressionsWithPruning(_.containsPattern(WINDOW_EXPRESSION), ruleId) {

beliefer · 2023-09-27T11:05:03Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

Suggested change

case we @ WindowExpression(_, ws @ WindowSpecDefinition(ps, _, _)) =>

case we @ WindowExpression(_, ws @ WindowSpecDefinition(ps, _, _)) if ps.exists(_.foldable) =>

beliefer · 2023-09-27T11:08:48Z

sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFramesSuite.scala

Please add more test cases:

partitions only unfoldable.

Mix unfoldable and foldable partition specs.

@beliefer Thanks, updated all.

beliefer · 2023-10-04T09:45:09Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

It seems we can simplify the code here with transformWithPruning only, then we can remove the removeWindowExpressionPartitions.

If do as you say, this will cause the partitionSpec in WindowExpression to be inconsistent with the partitionSpec in Window, which may cause hidden dangers.

What I mean is merge the logic of removeWindowExpressionPartitions into transformWithPruning.

beliefer · 2023-10-07T02:51:17Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

case w @ Window(windowExpressions, partitionSpec, _, _) if partitionSpec.exists(_.foldable) => val newWindowExpressions = windowExpressions.map(_.transformWithPruning(_.containsPattern(WINDOW_EXPRESSION)) { case wsd @ WindowSpecDefinition(ps, _, _) if ps.exists(_.foldable) => wsd.copy(partitionSpec = ps.filter(!_.foldable)) }.asInstanceOf[NamedExpression])

Should this be done? WINDOW_EXPRESSION mismatch WindowSpecDefinition.

case w @ Window(we, ps, _, _) if ps.exists(_.foldable) => val newWe = we.map(_.transformWithPruning(_.containsPattern(WINDOW_EXPRESSION)) { case _we @ WindowExpression(_, wsd @ WindowSpecDefinition(_ps, _, _)) if _ps.exists(_.foldable) => val newWsd = wsd.copy(partitionSpec = _ps.filter(!_.foldable)) _we.copy(windowSpec = newWsd) }.asInstanceOf[NamedExpression]) w.copy(windowExpressions = newWe, partitionSpec = ps.filter(!_.foldable))

I haven't actually built it, just for reference. You should ensure that it is correct.

Could we make the variable name more readable? e.g. we => windowExprs or windowExpressions and others.

I adjusted the above code

case w @ Window(windowExprs, partitionSpec, _, _) if partitionSpec.exists(_.foldable) => val newWe = windowExprs.map(_.transformWithPruning(_.containsPattern(WINDOW_EXPRESSION)) { case windowExpr @ WindowExpression(_, wsd @ WindowSpecDefinition(ps, _, _)) if ps.exists(_.foldable) => val newWsd = wsd.copy(partitionSpec = ps.filter(!_.foldable)) windowExpr.copy(windowSpec = newWsd) }.asInstanceOf[NamedExpression]) w.copy(windowExpressions = newWe, partitionSpec = partitionSpec.filter(!_.foldable))

beliefer · 2023-10-07T05:56:55Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

Uh, I'm missing. Please rename newWe to newWindowExprs.

zml1206 · 2023-10-30T01:59:54Z

cc @wangyum

beliefer

LGTM except one comment.

beliefer · 2023-10-30T12:25:38Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

Please use withNewChildren instead.

Which object uses withNewChildren? What is the purpose of using withNewChildren?

reduce object copy.

Updated, thanks.

beliefer · 2023-10-30T12:28:36Z

cc @cloud-fan

cloud-fan · 2023-12-27T12:16:32Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

+/**
+ * Remove window partition if partition expressions are foldable
+ */
+object EliminateWindowPartitions extends Rule[LogicalPlan] {


nit: let's put it in a new file

cloud-fan · 2023-12-27T12:18:33Z

sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFramesSuite.scala

      Row(1) :: Row(1) :: Nil)
  }
+
+  test("SPARK-45352: Remove window partition if partition expression are foldable") {


for a new optimizer rule, we should add unit tests, like LimitPushdownSuite

cloud-fan · 2023-12-27T12:20:45Z

sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFramesSuite.scala

+    val window1 = new WindowSpec(Seq(), Seq(sortOrder), UnspecifiedFrame)
+    val window2 = new WindowSpec(Seq(lit(1).expr), Seq(sortOrder), UnspecifiedFrame)
+    val df1 = ds.select(row_number().over(window1).alias("num"))
+    val df2 = ds.select(row_number().over(window2).alias("num"))


for end-to-end tests, let's use end-to-end APIs, e.g. Window.partitionBy($"key").orderBy($"value")

Updated, thanks.

zml1206 · 2024-01-02T06:13:28Z

@cloud-fan Already fixed unit test, could you have time to take a look？Thanks.

cloud-fan · 2024-01-03T06:21:29Z

...alyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/EliminateWindowPartitions.scala

+        case windowExpr @ WindowExpression(_, wsd @ WindowSpecDefinition(ps, _, _))
+          if ps.exists(_.foldable) =>
+          val newWsd = wsd.copy(partitionSpec = ps.filter(!_.foldable))
+          windowExpr.withNewChildren(Seq(windowExpr.windowFunction, newWsd))


nit: windowExpr.copy(windowSpec = newWsd)

cloud-fan · 2024-01-03T06:23:02Z

.../src/test/scala/org/apache/spark/sql/catalyst/optimizer/EliminateWindowPartitionsSuite.scala

+      testRelation
+        .select(a, b,
+          windowExpr(RowNumber(),
+            windowSpec(Nil, b.desc :: Nil, windowFrame)).as("rn"))


@beliefer can you help confirm that no partition column means a single partition in the window operator?

Yes. All the data rows in a single partition.

cloud-fan · 2024-01-03T06:23:32Z

.../src/test/scala/org/apache/spark/sql/catalyst/optimizer/EliminateWindowPartitionsSuite.scala

+          windowExpr(RowNumber(),
+            windowSpec(a :: Nil, b.desc :: Nil, windowFrame)).as("rn"))
+
+    val correctAnswer =


nit: val correctAnswer = originalQuery

cloud-fan · 2024-01-03T06:24:52Z

sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFramesSuite.scala

      Row(1) :: Row(1) :: Nil)
  }
+
+  test("SPARK-45352: Eliminate foldable window partitions") {


shall we turn on/off the new optimizer rule and run the test twice to make sure the result is the same?

zml1206 · 2024-01-03T09:46:50Z

Updated all. @cloud-fan

cloud-fan · 2024-01-08T09:16:18Z

thanks, merging to master!

zml1206 · 2024-01-08T09:20:45Z

Thank you for review. @cloud-fan @beliefer

### What changes were proposed in this pull request? This PR add a new optimizer rule `EliminateWindowPartitions`, it remove window partition if partition expressions are foldable. sql1: `select row_number() over(order by a) b from t ` sql2: `select row_number() over(partition by 1 order by a) b from t ` After this PR, the `optimizedPlan` for sql1 and sql2 is the same. ### Why are the changes needed? Foldable partition is redundant, remove it not only can simplify plan, but some rules can also take effect when the partitions are all foldable, such as `LimitPushDownThroughWindow`. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? UT ### Was this patch authored or co-authored using generative AI tooling? No Closes apache#43144 from zml1206/SPARK-45352. Authored-by: zml1206 <zhuml1206@gmail.com> Signed-off-by: Wenchen Fan <wenchen@databricks.com>

github-actions bot added the SQL label Sep 27, 2023

beliefer reviewed Sep 27, 2023

View reviewed changes

beliefer reviewed Oct 4, 2023

View reviewed changes

beliefer reviewed Oct 7, 2023

View reviewed changes

zml1206 changed the title ~~[SPARK-45352][SQL] Remove window partition if partition expressions are foldable~~ [SPARK-45352][SQL] Eliminate foldable window partitions Oct 8, 2023

beliefer reviewed Oct 30, 2023

View reviewed changes

zml1206 added 9 commits December 26, 2023 23:14

Remove window partition if partition expression are foldable

d6cabaa

update

6c4ceef

update

d4d53ac

update

5495be9

update

84cd78a

update

71dd250

update

7cb785f

use withNewChildren instead

77bc92e

update rule order

7854e80

zml1206 force-pushed the SPARK-45352 branch from e16daea to 7854e80 Compare December 26, 2023 16:25

beliefer approved these changes Dec 27, 2023

View reviewed changes

cloud-fan reviewed Dec 27, 2023

View reviewed changes

update unit tests

334ceed

cloud-fan reviewed Jan 3, 2024

View reviewed changes

update

6829480

zml1206 requested a review from cloud-fan January 4, 2024 13:35

cloud-fan approved these changes Jan 8, 2024

View reviewed changes

cloud-fan closed this in 8facc99 Jan 8, 2024

	plan.transformAllExpressionsWithPruning(_.containsAnyPattern(WINDOW_EXPRESSION), ruleId) {
	plan.transformAllExpressionsWithPruning(_.containsPattern(WINDOW_EXPRESSION), ruleId) {

	case we @ WindowExpression(_, ws @ WindowSpecDefinition(ps, _, _)) =>
	case we @ WindowExpression(_, ws @ WindowSpecDefinition(ps, _, _)) if ps.exists(_.foldable) =>

[SPARK-45352][SQL] Eliminate foldable window partitions #43144

[SPARK-45352][SQL] Eliminate foldable window partitions #43144

Uh oh!

Conversation

zml1206 commented Sep 27, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

beliefer Oct 5, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zml1206 Oct 7, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zml1206 commented Oct 30, 2023

Uh oh!

beliefer left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

beliefer commented Oct 30, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zml1206 commented Jan 2, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

zml1206 commented Sep 27, 2023 •

edited

Loading

beliefer Oct 5, 2023 •

edited

Loading

zml1206 Oct 7, 2023 •

edited

Loading

zml1206 commented Jan 3, 2024 •

edited

Loading