peter-toth commented on code in PR #59306:
URL: https://github.com/apache/spark/pull/59306#discussion_r4258962852
##########
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:
Thanks. I'd keep the explicit failure. `SafeForKWayMerge` is what a node
under a merging node has to provide. A plugin that replaces such a node, but
keeps the `GroupPartitionsExec` above it, has to keep that guarantee too. If
its node does, it can mix in the trait. A sort in this node would hide the
plugin's mistake behind a slower plan. It would also need a sorter here, with
its own memory, spill and metrics. The error tells the plugin author what to
fix.
##########
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:
Done in
[`4c9a5377689`](https://github.com/apache/spark/pull/59306/commits/4c9a537768922dc56addf65078f718e561d41bef).
The scaladoc now says that `enableSortedMerge` is driver-only, like the other
accessors of `@transient` fields such as `distributePartitions`.
##########
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:
Done in
[`4c9a5377689`](https://github.com/apache/spark/pull/59306/commits/4c9a537768922dc56addf65078f718e561d41bef).
The test now also rebuilds the `DebugExec` wrapping the way `debug()` does,
and checks that the merging node's `execute()` returns a
`SortedMergeCoalescedRDD`. It fails if the merge is skipped silently.
##########
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:
Done in
[`4c9a5377689`](https://github.com/apache/spark/pull/59306/commits/4c9a537768922dc56addf65078f718e561d41bef).
##########
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:
Done in
[`4c9a5377689`](https://github.com/apache/spark/pull/59306/commits/4c9a537768922dc56addf65078f718e561d41bef).
Both tests now share `createKWayMergeWindowQuery`, `kWayMergeWindowConfs` and
`kWayMergeWindowRows`.
--
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]
