dongjoon-hyun commented on code in PR #59306:
URL: https://github.com/apache/spark/pull/59306#discussion_r4247346868
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -337,8 +349,23 @@ case class GroupPartitionsExec(
}
}
+ /**
+ * Fails a planned merge that would not run. The parent's required ordering
was checked against
+ * the merge, so no sort was added above it. A rule after
`EnsureRequirements` can still put a
+ * node outside the `SafeForKWayMerge` whitelist below this one. This node
would then concatenate
+ * and hand the parent unsorted rows. No built-in rule does that, see
`childIsSafeForKWayMerge`.
+ * Like `checkChildStillMatches`, this makes the failure explicit if one
starts.
+ */
+ private def checkPlannedMergeRuns(): Unit = {
Review Comment:
Thank you for adding this check. It closes the silent case I raised earlier.
One consequence is worth a look before merging, though.
`ApplyColumnarRulesAndInsertTransitions` runs after `EnsureRequirements`,
and the `preColumnarTransitions` rules of a columnar plugin (for example RAPIDS
or Gluten) run inside it. If such a rule replaces a whitelisted `ProjectExec`,
`FilterExec` or `BatchScanExec` under a merging node with its own operator,
`childIsSafeForKWayMerge` becomes false. Every such query then fails here with
an internal error once `preserveOrderingOnCoalesce` is on. On the current
master the same queries run but return rows in the wrong order. So this is
strictly better, but it turns the plugin case into a hard failure.
The node now holds the frozen ordering, so it could still deliver that
ordering without the merge. It could concatenate as usual and then sort each
output partition by `sortedMergeOrdering`, as a `SortExec` above it would. That
keeps such a query correct instead of failing it, and the k-way merge stays the
fast path. Would that fallback be worth it here, or do you prefer to keep the
explicit failure and handle plugins separately?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -246,30 +250,38 @@ case class GroupPartitionsExec(
// such change mints a fresh node through `withNewChildInternal`, so this is
recomputed against
// the new child. The correctness of this code relies on all such insertable
nodes
// (WholeStageCodegenExec, InputAdapter, ColumnarToRowExec) being in the
SafeForKWayMerge
- // whitelist so the evaluation stays consistent.
+ // whitelist so the evaluation stays consistent. So is the `DebugExec` that
the `debug()` helper
+ // (`org.apache.spark.sql.execution.debug`) wraps every node in without AQE.
A rule that breaks
+ // this under a planned merge fails the query at execution, see
`checkPlannedMergeRuns`.
@transient private lazy val childIsSafeForKWayMerge: Boolean =
!child.exists {
case _: SafeForKWayMerge => false
case _ => true
}
- /** Whether a k-way merge would work at all, leaving aside whether it is
switched on. */
- @transient private lazy val kWayMergeIsFeasible: Boolean =
- child.outputOrdering.nonEmpty && childIsSafeForKWayMerge
+ // A merge by an empty ordering would run but sort nothing.
+ require(sortedMergeOrdering.forall(_.nonEmpty), "A k-way merge needs an
ordering to merge by")
- /** Whether this node performs the k-way merge. No config term, see
`enableSortedMerge`. */
+ /** Whether the k-way merge was planned for this node. See
`sortedMergeOrdering`. */
+ def enableSortedMerge: Boolean = sortedMergeOrdering.isDefined
Review Comment:
nit. `sortedMergeOrdering` is `@transient`, so it is `null` on a
deserialized copy of this node. Then `enableSortedMerge`, and through it
`usesSortedMerge`, `supportsColumnar` and the `SortedMerge:` rendering, would
throw an NPE. The previous `@transient enableSortedMerge: Boolean` read `false`
there instead.
I did not find a path that serializes this node today, and
`distributePartitions` already has the same shape with `ungroupingOrigin`. So
this is only about keeping the safe default the Boolean had. For example,
`Option(sortedMergeOrdering).exists(_.nonEmpty)`, or a short comment that these
accessors are driver-only.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -337,8 +349,23 @@ case class GroupPartitionsExec(
}
}
+ /**
+ * Fails a planned merge that would not run. The parent's required ordering
was checked against
+ * the merge, so no sort was added above it. A rule after
`EnsureRequirements` can still put a
+ * node outside the `SafeForKWayMerge` whitelist below this one. This node
would then concatenate
+ * and hand the parent unsorted rows. No built-in rule does that, see
`childIsSafeForKWayMerge`.
+ * Like `checkChildStillMatches`, this makes the failure explicit if one
starts.
+ */
+ private def checkPlannedMergeRuns(): Unit = {
+ if (enableSortedMerge && hasCoalescing && !usesSortedMerge) {
Review Comment:
nit. `usesSortedMerge` is `enableSortedMerge && hasCoalescing &&
childIsSafeForKWayMerge`, so this condition is the same as `enableSortedMerge
&& hasCoalescing && !childIsSafeForKWayMerge`. Writing it that way names the
one cause this check is about, the whitelist, which also matches the error
message and the scaladoc above.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -6793,6 +6794,84 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-60044: a planned k-way merge keeps its ordering after a later
config change") {
+ // The table reports no ordering. The scan derives [id, name] from its
partition keys and drops
+ // the transform between them. The splits are sorted by the full key, so
within id 1 the 'bb'
+ // split comes before the 'aa' one. The window groups by id and orders by
name. The k-way merge
+ // delivers that ordering, so EnsureRequirements adds no SortExec. Turning
partitionKeyOrdering
+ // off afterwards empties the scan's ordering. The planned merge must
still merge by [id, name].
+ // Otherwise the window numbers the rows of id 1 in split order.
+ //
+ // Both AQE modes are covered, see the SPARK-59279 test above.
+ createTable(items, itemsColumns, Array(identity("id"),
days("arrive_time"), identity("name")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ "(1, 'bb', 10.0, cast('2021-01-01' as timestamp)), " +
+ "(1, 'aa', 20.0, cast('2022-01-01' as timestamp)), " +
+ "(2, 'cc', 30.0, cast('2021-01-01' as timestamp))")
+ // The query selects `arrive_time`, so the transform's column stays in the
scan output.
+ val query =
+ s"""SELECT id, name, arrive_time, ROW_NUMBER() OVER (PARTITION BY id
ORDER BY name) AS rn
+ |FROM testcat.ns.$items
+ |""".stripMargin
+ val expected = Seq(
+ Row(1, "aa", Timestamp.valueOf("2022-01-01 00:00:00"), 1),
+ Row(1, "bb", Timestamp.valueOf("2021-01-01 00:00:00"), 2),
+ Row(2, "cc", Timestamp.valueOf("2021-01-01 00:00:00"), 1))
+
+ Seq(true, false).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString,
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key ->
"true") {
+ val df = sql(query)
+ // Force the plan without executing it, so the k-way merge is decided
under this config.
+ val plan = df.queryExecution.executedPlan
+ val merging =
collectAllGroupPartitions(plan).filter(_.enableSortedMerge)
+ assert(merging.length == 1, "expected one k-way merge")
+ assert(merging.head.outputOrdering.map(_.child.sql) == Seq("id",
"name"))
+ assert(collect(plan) { case s: SortExec => s }.isEmpty,
+ "the k-way merge satisfies the window's ordering, so no SortExec
should be added")
+
+ // Flip only the key ordering config. The plan above is already
committed.
+ withSQLConf(SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key ->
"false") {
+ assert(merging.head.child.outputOrdering.isEmpty,
+ "test setup: the scan no longer derives an ordering from its keys")
+ checkAnswer(df, expected)
+ // Re-collected, because under AQE `executedPlan` only reaches the
final plan's nodes
+ // once the query has run.
+ val executed =
collectAllGroupPartitions(df.queryExecution.executedPlan)
+ assert(executed.map(_.outputOrdering.map(_.child.sql)) ==
Seq(Seq("id", "name")),
+ "the planned k-way merge must keep the ordering it was planned
with")
+ }
+ }
+ }
+ }
+
+ test("SPARK-60044: the debug() helper runs a planned k-way merge") {
Review Comment:
nit. This test copies the table, the three-row `INSERT`, the confs and the
`ROW_NUMBER` query from the test above it. If the data or query in one test
changes later, for example to cover a different split order, the other test can
stop producing a merge and fail only on its setup assertion. Could the two
tests share a small helper for the table, rows and query?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -6793,6 +6794,84 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-60044: a planned k-way merge keeps its ordering after a later
config change") {
+ // The table reports no ordering. The scan derives [id, name] from its
partition keys and drops
+ // the transform between them. The splits are sorted by the full key, so
within id 1 the 'bb'
+ // split comes before the 'aa' one. The window groups by id and orders by
name. The k-way merge
+ // delivers that ordering, so EnsureRequirements adds no SortExec. Turning
partitionKeyOrdering
+ // off afterwards empties the scan's ordering. The planned merge must
still merge by [id, name].
+ // Otherwise the window numbers the rows of id 1 in split order.
+ //
+ // Both AQE modes are covered, see the SPARK-59279 test above.
+ createTable(items, itemsColumns, Array(identity("id"),
days("arrive_time"), identity("name")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ "(1, 'bb', 10.0, cast('2021-01-01' as timestamp)), " +
+ "(1, 'aa', 20.0, cast('2022-01-01' as timestamp)), " +
+ "(2, 'cc', 30.0, cast('2021-01-01' as timestamp))")
+ // The query selects `arrive_time`, so the transform's column stays in the
scan output.
+ val query =
+ s"""SELECT id, name, arrive_time, ROW_NUMBER() OVER (PARTITION BY id
ORDER BY name) AS rn
+ |FROM testcat.ns.$items
+ |""".stripMargin
+ val expected = Seq(
+ Row(1, "aa", Timestamp.valueOf("2022-01-01 00:00:00"), 1),
+ Row(1, "bb", Timestamp.valueOf("2021-01-01 00:00:00"), 2),
+ Row(2, "cc", Timestamp.valueOf("2021-01-01 00:00:00"), 1))
+
+ Seq(true, false).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString,
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key ->
"true") {
+ val df = sql(query)
+ // Force the plan without executing it, so the k-way merge is decided
under this config.
+ val plan = df.queryExecution.executedPlan
+ val merging =
collectAllGroupPartitions(plan).filter(_.enableSortedMerge)
+ assert(merging.length == 1, "expected one k-way merge")
+ assert(merging.head.outputOrdering.map(_.child.sql) == Seq("id",
"name"))
+ assert(collect(plan) { case s: SortExec => s }.isEmpty,
+ "the k-way merge satisfies the window's ordering, so no SortExec
should be added")
+
+ // Flip only the key ordering config. The plan above is already
committed.
+ withSQLConf(SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key ->
"false") {
+ assert(merging.head.child.outputOrdering.isEmpty,
+ "test setup: the scan no longer derives an ordering from its keys")
+ checkAnswer(df, expected)
+ // Re-collected, because under AQE `executedPlan` only reaches the
final plan's nodes
+ // once the query has run.
+ val executed =
collectAllGroupPartitions(df.queryExecution.executedPlan)
+ assert(executed.map(_.outputOrdering.map(_.child.sql)) ==
Seq(Seq("id", "name")),
+ "the planned k-way merge must keep the ordering it was planned
with")
+ }
+ }
+ }
+ }
+
+ test("SPARK-60044: the debug() helper runs a planned k-way merge") {
+ // Without AQE, `debug()` wraps every node of the executed plan in a
`DebugExec`, including the
+ // nodes under a merging `GroupPartitionsExec`. That wrapper must not stop
the planned merge.
+ import org.apache.spark.sql.execution.debug._
+ createTable(items, itemsColumns, Array(identity("id"),
days("arrive_time"), identity("name")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ "(1, 'bb', 10.0, cast('2021-01-01' as timestamp)), " +
+ "(1, 'aa', 20.0, cast('2022-01-01' as timestamp)), " +
+ "(2, 'cc', 30.0, cast('2021-01-01' as timestamp))")
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key ->
"true") {
+ val df = sql(
+ s"""SELECT id, name, arrive_time, ROW_NUMBER() OVER (PARTITION BY id
ORDER BY name) AS rn
+ |FROM testcat.ns.$items
+ |""".stripMargin)
+
assert(collectAllGroupPartitions(df.queryExecution.executedPlan).exists(_.enableSortedMerge),
+ "test setup: the plan has a k-way merge")
+ df.debug()
Review Comment:
This test checks only that `debug()` does not throw. `debug()` returns
`Unit` and prints the row count and the column stats to stdout, so the test
cannot tell whether the debug plan actually ran the merge or produced the
expected rows.
For example, the merge would be skipped without reaching
`checkPlannedMergeRuns` if a later change made the wrapped node stop
coalescing. The test would still pass.
Could we also assert on the outcome? For example, build the same
`DebugExec`-wrapped plan and check that the merging node's `execute()` returns
a `SortedMergeCoalescedRDD`. Or at least run `checkAnswer(df, expected)` in the
same config after `debug()`.
--
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]
