peter-toth commented on code in PR #57986:
URL: https://github.com/apache/spark/pull/57986#discussion_r3805517394
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -4761,6 +4761,17 @@ object SQLConf {
"The threshold of window group limit must be -1, 0 or positive
integer.")
.createWithDefault(1000)
+ val COLLAPSE_WINDOW_WITH_EMPTY_ORDER_SPEC_IN_CHILD =
+ buildConf("spark.sql.optimizer.collapseWindowWithEmptyOrderSpecInChild")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .doc("When true, the optimizer collapses two adjacent windows with the
same partition " +
+ "spec into one when the window with the empty order spec is the child
(inner) window. " +
+ "This saves a WindowExec pass but can disable the WindowGroupLimit
optimization for " +
Review Comment:
**Finding 8.** The doc names only `WindowGroupLimit`, but the flag also
costs `LimitPushDownThroughWindow`. Measured on this head with `SELECT c1,
count(1) OVER () cnt, row_number() OVER (ORDER BY c2) rk FROM t3 LIMIT 5`:
flag off (default):
Project [c1, cnt, rk]
+- Window [row_number() ... AS rk], [c2 ASC]
+- GlobalLimit 5
+- LocalLimit 5
+- Sort [c2 ASC], true
+- Window [count(1) ... AS cnt]
flag on:
GlobalLimit 5
+- LocalLimit 5
+- Project [c1, cnt, rk]
+- Window [count(1) ..., row_number() ...], [c2 ASC]
So the top-5 `Limit + Sort` below the window is gone and the merged window
sorts and buffers the whole input in a single partition.
`LimitPushDownThroughWindow` sits *earlier* than `CollapseWindow` in the
operator-optimization rule list, which is why I ruled it out last round, but it
loses the race: the analyzed plan is
`LocalLimit(Project(Project(Window(Window(...)))))`, and neither of its
patterns (`LocalLimit(_, Window)`, `LocalLimit(_, Project(_, Window))`) matches
through two projects. `CollapseProject` and `CollapseWindow` both run later in
that same iteration, so by iteration 2 the pattern finally matches but
`supportsPushdownThroughWindow` fails on the merged-in `count`.
Nothing to fix in the rule -- both losses are the same "merged window is no
longer all-expanding" effect and the flag is off by default. Just worth naming
both, here and in the rule comment at `Optimizer.scala:1799`, so someone
flipping this for the saved `WindowExec` pass knows what else goes:
```suggestion
"This saves a WindowExec pass but can disable the WindowGroupLimit
and the LocalLimit " +
"push-down optimizations for top-k queries. " +
```
##########
sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala:
##########
@@ -79,6 +79,27 @@ class DataFrameWindowFunctionsSuite extends
SharedSparkSession
parameters = Map("wf_name" -> "row_number", "wf_expr" -> "row_number()"))
}
+ test("SPARK-58757: collapse window with an empty order spec into an ordered
sibling") {
+ val df = Seq(
+ (0, 0), (0, 2), (0, 4),
+ (1, 1), (1, 3), (1, 5)).toDF("k", "v")
+ val ordered = Window.partitionBy("k").orderBy("v")
+ val unordered = Window.partitionBy("k")
+ checkAnswer(
Review Comment:
**Finding 9.** `row_number` and `count` over a whole-partition frame are
both order-insensitive, and there is no plan assertion, so this test cannot
tell a merged plan from an unmerged one -- with
`spark.sql.optimizer.excludedRules=org.apache.spark.sql.catalyst.optimizer.CollapseWindow`
it still passes, on a plan that has two `Window` operators. As a regression
test for this PR it can only catch a crash.
The sibling collapse test in this file already has the idiom
(`DataFrameWindowFunctionsSuite.scala:1716`):
```scala
val windows = df.queryExecution.optimizedPlan.collect { case w:
LogicalWindow => w }
assert(windows.size === 1)
```
Adding that plus an order-dependent function makes the test pin both the
merge and the row order the merged-in expressions now see -- which is the
interesting property, since the parent-empty direction is supposed to preserve
it exactly. Verified this passes on this head:
```scala
val res = df.select(
$"k",
$"v",
row_number().over(ordered).as("rn"),
collect_list($"v").over(unordered).as("vs"),
first($"v").over(unordered).as("f"))
assert(res.queryExecution.optimizedPlan.collect { case w: LogicalWindow
=> w }.size === 1)
checkAnswer(res, Seq(
Row(0, 0, 1, Array(0, 2, 4), 0),
Row(0, 2, 2, Array(0, 2, 4), 0),
Row(0, 4, 3, Array(0, 2, 4), 0),
Row(1, 1, 1, Array(1, 3, 5), 1),
Row(1, 3, 2, Array(1, 3, 5), 1),
Row(1, 5, 3, Array(1, 3, 5), 1)))
```
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseWindowSuite.scala:
##########
@@ -168,4 +174,220 @@ 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 has a RANGE
whole-partition frame") {
+ // `RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING` covers the
whole partition just
+ // like `ROWS`, so it also collapses.
+ 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(RangeFrame, 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 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. This direction can disable
InferWindowGroupLimit, so it is
+ // gated by `spark.sql.optimizer.collapseWindowWithEmptyOrderSpecInChild`.
+ 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 = withSQLConf(
+ SQLConf.COLLAPSE_WINDOW_WITH_EMPTY_ORDER_SPEC_IN_CHILD.key -> "true") {
+ Optimize.execute(analyzed)
+ }
+ assert(analyzed.output === optimized.output)
+
+ val correctAnswer = testRelation
+ .window(Seq(cnt, rk), partitionSpec1, orderSpec1)
+
+ comparePlans(optimized, correctAnswer)
+ }
+
+ test("don't collapse the inner empty-order window by default") {
+ // Merging an empty-order child into an ordered parent can disable
InferWindowGroupLimit for
+ // top-k queries, so it is off by default.
+ 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 optimized = Optimize.execute(query.analyze)
+ val correctAnswer = query.analyze
+
+ comparePlans(optimized, correctAnswer)
+ }
+
+ test("collapse windows with a Project between them when one has an empty
order spec") {
Review Comment:
**Finding 10.** This covers the `Project`-between pattern only in the
config-gated direction, so the direction that is on by default has no test for
that second `apply` case. Worth mirroring it -- verified this merges on this
head with no config set:
```scala
val query = testRelation
.window(Seq(rk), partitionSpec1, orderSpec1)
.select($"a", $"b", $"c", $"rk")
.window(Seq(cnt), partitionSpec1, Nil)
.select($"a", $"b", $"c", $"rk", $"cnt")
```
Project [a#0, b#1, c#2, rk#4, cnt#6L]
+- Window [row_number() ... AS rk#4, count(c#2) ... AS cnt#6L], [c#2],
[c#2 ASC NULLS FIRST]
+- LocalRelation <empty>, [a#0, b#1, c#2]
--
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]
