peter-toth commented on code in PR #59306:
URL: https://github.com/apache/spark/pull/59306#discussion_r4282260479
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -360,6 +390,7 @@ case class GroupPartitionsExec(
override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
checkChildStillMatches()
+ checkPlannedMergeRuns()
Review Comment:
Done in
[`34f2672de31`](https://github.com/apache/spark/pull/59306/commits/34f2672de316aad274f3f555af3c04b9cc745a44).
`supportsColumnar` now follows the plan, `!(enableSortedMerge &&
hasCoalescing)`, and the columnar path fails for any planned merge. A new
assertion covers a columnar child outside the whitelist. It fails if
`supportsColumnar` reads `usesSortedMerge` again.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -2920,6 +2920,19 @@ class EnsureRequirementsSuite extends SharedSparkSession
{
}
}
+ test("SPARK-60044: pushing join key positions into a merging node keeps its
ordering") {
+ val keys = Seq(InternalRow(1, 1), InternalRow(1, 2))
+ val leaf = DummySparkPlan(
Review Comment:
Done in
[`34f2672de31`](https://github.com/apache/spark/pull/59306/commits/34f2672de316aad274f3f555af3c04b9cc745a44).
The child is now `DummyLeafSafeForKWayMerge`, and the test also checks that
the rebuilt node reports the frozen ordering.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -337,8 +352,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 && !childIsSafeForKWayMerge) {
+ throw SparkException.internalError(
+ "GroupPartitionsExec was planned with a k-way merge that its child no
longer allows")
Review Comment:
Done in
[`34f2672de31`](https://github.com/apache/spark/pull/59306/commits/34f2672de316aad274f3f555af3c04b9cc745a44).
The message names the node and `SafeForKWayMerge`, and the `SafeForKWayMerge`
scaladoc mentions this failure.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala:
##########
@@ -144,21 +144,66 @@ class GroupPartitionsExecSuite extends SharedSparkSession
{
}
test("SPARK-58324: k-way merge ordering drops sameOrderExpressions") {
- // The child ordering carries sameOrderExpressions (planner metadata). The
k-way merge
+ // The planned ordering carries sameOrderExpressions (planner metadata).
The k-way merge
// comparator only needs the sort key, so kWayMergeOrdering keeps
child/direction/nullOrdering
// but drops sameOrderExpressions, so LazyRowOrdering does not serialize
them with the RDD.
- val childOrdering = Seq(SortOrder(exprA, Ascending, Seq(exprB, exprC)))
+ val plannedOrdering = Seq(SortOrder(exprA, Ascending, Seq(exprB, exprC)))
val child = DummySparkPlan(
- outputPartitioning = KeyedPartitioning(Seq(exprA), Seq(row(1), row(2),
row(1))),
- outputOrdering = childOrdering)
- val gpe = GroupPartitionsExec(child)
+ outputPartitioning = KeyedPartitioning(Seq(exprA), Seq(row(1), row(2),
row(1))))
+ val gpe = GroupPartitionsExec(child, sortedMergeOrdering =
Some(plannedOrdering))
- assert(child.outputOrdering.head.sameOrderExpressions.nonEmpty, "test
setup")
val merged = gpe.kWayMergeOrdering
assert(merged.map(so => (so.child, so.direction)) === Seq((exprA,
Ascending)))
assert(merged.forall(_.sameOrderExpressions.isEmpty))
}
+ test("SPARK-60044: a merging node is not rebuilt over another child") {
+ // Its frozen ordering was checked against its own child, not against the
new one.
+ val leaf = DummySparkPlan(
+ outputPartitioning = KeyedPartitioning(Seq(exprA), Seq(row(1), row(1))))
+ assert(GroupPartitionsExec(leaf).withKeyPositionsFor(leaf).isDefined,
"test setup")
+ val merging =
+ GroupPartitionsExec(leaf, sortedMergeOrdering =
Some(Seq(SortOrder(exprA, Ascending))))
+ assert(merging.withKeyPositionsFor(leaf).isEmpty)
Review Comment:
Done in
[`34f2672de31`](https://github.com/apache/spark/pull/59306/commits/34f2672de316aad274f3f555af3c04b9cc745a44),
with a case over `InputAdapter(leaf)`.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -6793,6 +6797,98 @@ class KeyGroupedPartitioningSuite
}
}
+ /**
+ * The setup of the SPARK-60044 tests. Creates `items`, which reports no
ordering, so 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. Returns a
+ * window query that groups by id and orders by name. Under
`kWayMergeWindowConfs` a k-way merge
+ * delivers that ordering, so EnsureRequirements adds no SortExec. The query
selects
+ * `arrive_time`, so the transform's column stays in the scan output.
+ */
+ private def createKWayMergeWindowQuery(): String = {
+ 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))")
+ s"""SELECT id, name, arrive_time, ROW_NUMBER() OVER (PARTITION BY id ORDER
BY name) AS rn
+ |FROM testcat.ns.$items
+ |""".stripMargin
+ }
+
+ private val kWayMergeWindowConfs = Seq(
+ 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")
+
+ private val kWayMergeWindowRows = 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))
+
+ test("SPARK-60044: a planned k-way merge keeps its ordering after a later
config change") {
+ // Turning partitionKeyOrdering off after planning 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.
+ val query = createKWayMergeWindowQuery()
+ Seq(true, false).foreach { aqeEnabled =>
+ val confs = (SQLConf.ADAPTIVE_EXECUTION_ENABLED.key ->
aqeEnabled.toString) +:
+ kWayMergeWindowConfs
+ withSQLConf(confs: _*) {
+ 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, kWayMergeWindowRows)
+ // 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._
+ val query = createKWayMergeWindowQuery()
+ val confs = (SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") +:
kWayMergeWindowConfs
+ withSQLConf(confs: _*) {
+ val df = sql(query)
+ val plan = df.queryExecution.executedPlan
+ assert(collectAllGroupPartitions(plan).exists(_.enableSortedMerge),
+ "test setup: the plan has a k-way merge")
+ df.debug()
+
+ // `debug()` returns nothing, so its wrapping is rebuilt here, the same
way, to check that
+ // the merge runs.
+ val visited = mutable.HashSet.empty[TreeNodeRef]
+ val debugPlan = plan.transform {
Review Comment:
Done in
[`34f2672de31`](https://github.com/apache/spark/pull/59306/commits/34f2672de316aad274f3f555af3c04b9cc745a44).
The test now wraps only the merging node's child in a `DebugExec`, and the two
imports are gone.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala:
##########
@@ -448,17 +492,22 @@ class GroupPartitionsExecSuite extends SharedSparkSession
{
assert(!GroupPartitionsExec(child).groupedPartitions.forall(_._2.size <=
1),
"expected coalescing")
withSQLConf(SQLConf.V2_BUCKETING_PRESERVE_KEY_ORDERING_ON_COALESCE_ENABLED.key
-> "true") {
- // Even though enableSortedMerge = true, the child is not safe for k-way
merge,
- // so only key-expression orders survive (non-key exprC is dropped).
- val ordering = GroupPartitionsExec(child, enableSortedMerge =
true).outputOrdering
+ val merging = GroupPartitionsExec(child, sortedMergeOrdering =
Some(childOrdering))
+ val ordering = merging.outputOrdering
assert(ordering.length === 1)
assert(ordering.head.child === exprA)
+ // A planned merge that cannot run fails rather than concatenating
(SPARK-60044). Both
+ // execute paths have their own check.
+ Seq(() => merging.execute(), () => merging.executeColumnar()).foreach {
run =>
+ val error = intercept[SparkException](run())
+ assert(error.getMessage.contains("planned with a k-way merge"))
+ }
}
}
- test("SPARK-55715: coalescing with enableSortedMerge = true returns full
child ordering") {
+ test("SPARK-55715: coalescing with a planned merge returns full child
ordering") {
// Key 1 appears on partitions 0 and 2, causing coalescing. The child is a
LeafExecNode so
- // childIsSafeForKWayMerge = true. With enableSortedMerge = true the node
performs the k-way
+ // childIsSafeForKWayMerge = true. With a planned merge ordering the node
performs the k-way
Review Comment:
Done in
[`34f2672de31`](https://github.com/apache/spark/pull/59306/commits/34f2672de316aad274f3f555af3c04b9cc745a44).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -246,30 +250,41 @@ 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`. Driver-only, like
+ * the other accessors of `@transient` fields: a deserialized copy holds no
ordering.
+ */
+ def enableSortedMerge: Boolean = sortedMergeOrdering.isDefined
+
+ /** Whether this node performs the k-way merge. No config term, see
`sortedMergeOrdering`. */
@transient private lazy val usesSortedMerge: Boolean =
- enableSortedMerge && hasCoalescing && kWayMergeIsFeasible
+ enableSortedMerge && hasCoalescing && childIsSafeForKWayMerge
/**
* Returns a copy of this node with k-way merge enabled, when the config is
on, this node
- * coalesces partitions and the merge is feasible. The only read of
- * `preserveOrderingOnCoalesce`.
+ * coalesces partitions, the child reports an ordering and its subtree is
safe for the merge. The
+ * only read of `preserveOrderingOnCoalesce`. The copy keeps the ordering
checked here, see
+ * `sortedMergeOrdering`. A node that already merges is left as it is, so
the merge is decided
+ * once.
*/
def tryEnableSortedMerge(): Option[GroupPartitionsExec] = {
- Option.when(conf.v2BucketingPreserveOrderingOnCoalesceEnabled &&
hasCoalescing &&
- kWayMergeIsFeasible) {
- val newGroupPartitions = copy(enableSortedMerge = true)
+ // Read once, so that the ordering frozen is the one checked.
+ lazy val ordering = child.outputOrdering
Review Comment:
Done in
[`34f2672de31`](https://github.com/apache/spark/pull/59306/commits/34f2672de316aad274f3f555af3c04b9cc745a44),
with your shape.
--
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]
