dongjoon-hyun commented on code in PR #59306:
URL: https://github.com/apache/spark/pull/59306#discussion_r4264395551
##########
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:
This builds the merging node over `DummySparkPlan`, which is not
`SafeForKWayMerge`, and its two keys do not coalesce. So `tryEnableSortedMerge`
could not have produced it, and the node that `withJoinKeyPositions` returns is
a planned merge that cannot run. It coalesces after `Seq(0)`, but its
`outputOrdering` takes the key-only branch (`Nil` here), and `execute()` would
fail in `checkPlannedMergeRuns`. The test passes anyway, since it checks only
the field.
Could it use `DummyLeafSafeForKWayMerge` (L3105) with its ordering left
empty, and also assert `g.outputOrdering == frozen`? Then it would also catch a
rebuild that keeps the field but can no longer merge by it.
##########
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:
nit. The previous thread kept this failure because the error tells the
plugin author what to fix. But the message names neither the node that left the
whitelist nor `SafeForKWayMerge`, so the user only sees `[INTERNAL_ERROR]
GroupPartitionsExec was planned with a k-way merge that its child no longer
allows`. The scaladoc of `SafeForKWayMerge` does not mention this case either.
Could the message name both? For example:
```scala
val blocker = child.find(!_.isInstanceOf[SafeForKWayMerge]).map(_.nodeName)
throw SparkException.internalError(
"GroupPartitionsExec was planned with a k-way merge, but its child holds "
+
s"${blocker.getOrElse("a node")}, which does not mix in
SafeForKWayMerge")
```
##########
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:
nit. `leaf` is the merging node's own child here, so this does not set up
the case that the test name and the comment above describe. It passes because
L479 of `GroupPartitionsExec` refuses every merging node without looking at
`newChild`. This assertion is still useful, since it shows that a merging node
is refused even over its own child. Could we also add one over another child,
e.g. `merging.withKeyPositionsFor(InputAdapter(leaf))`? `InputAdapter` reports
the same partitioning, so the non-merging node above would still be rebuilt
over it.
##########
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:
nit. `ordering.nonEmpty` is evaluated before the cached
`childIsSafeForKWayMerge`, so for a subtree outside the whitelist each call
computes the child's ordering and drops it. The cost is small, but checking the
whitelist first also removes the need for the local `lazy val` and its comment:
```scala
Option.when(!enableSortedMerge &&
conf.v2BucketingPreserveOrderingOnCoalesceEnabled &&
hasCoalescing && childIsSafeForKWayMerge)(child.outputOrdering)
.filter(_.nonEmpty)
.map { ordering =>
val newGroupPartitions = copy(sortedMergeOrdering = Some(ordering))
newGroupPartitions.copyTagsFrom(this)
newGroupPartitions
}
```
The ordering is still read once, and the frozen value is still the one
checked.
##########
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:
nit. Since this comment is being updated anyway: the child counts as safe
because `DummyLeafSparkPlan` mixes in `SafeForKWayMerge`, not because it is a
`LeafExecNode`. `childIsSafeForKWayMerge` only looks for the trait, so a leaf
like `LocalTableScanExec` counts as unsafe. How about "The child,
`DummyLeafSparkPlan`, mixes in `SafeForKWayMerge`, so childIsSafeForKWayMerge =
true"?
##########
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:
Thank you for keeping this check on both execute paths. There is still one
gap on this path, though. Here and in `supportsColumnar` (L389), a planned
merge is judged by the live `childIsSafeForKWayMerge`, while
`doExecuteColumnar` never merges.
A planned merge reaches this method whenever its child was a columnar node
outside the whitelist when the transitions were inserted. `supportsColumnar` is
then `child.supportsColumnar`, so `ensureOutputsRowBased` puts the
`ColumnarToRowExec` above this node instead of below it. If the child is safe
again by execution, this check passes and `CoalescedRDD` concatenates.
For example, take a `ColumnarRule` whose `preColumnarTransitions` replaces
the scan with its own columnar node, and whose `postColumnarTransitions` swaps
in a node that mixes in `SafeForKWayMerge`, as suggested for plugins in the
previous thread. The window or sort-merge join above then reads unsorted rows
without an error. Since SPARK-52484, `ColumnarToRowExec` no longer asserts
`child.supportsColumnar`, so nothing else catches it.
Since this path can never merge, could the columnar side follow the frozen
plan only? For example:
```scala
override def supportsColumnar: Boolean =
child.supportsColumnar && !(enableSortedMerge && hasCoalescing)
```
and fail here for any planned merge. Then the transition always lands below
a planned merge, and `doExecute` either merges or fails in
`checkPlannedMergeRuns` as intended.
##########
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)
Review Comment:
nit. `checkAnswer` first runs `materializedRdd.count()`, which plans a new
`QueryExecution` under the flipped config. That plan has a `SortExec` and no
merge, so this part passes with or without the fix, and only the `collect()`
that follows runs the committed plan. `QueryTest.checkAnswer(df,
kWayMergeWindowRows, checkToRDD = false)` would run only the plan this test is
about. The SPARK-59279 test has the same pattern at L6784, if you want to keep
the two aligned.
##########
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:
nit. This copies the private wrapping of `debug()`, and the `mutable` and
`TreeNodeRef` imports are added only for it. So the `SortedMergeCoalescedRDD`
assertion checks this copy rather than what `debug()` runs, and a later change
to how `debug()` wraps would leave it green. It follows my earlier suggestion
to build the same wrapped plan, so this one is on me.
Wrapping only the merging node's child would check the same whitelist
decision without the copy, e.g.
`planned.withNewChildren(Seq(DebugExec(planned.child)))`, where `planned` is
the merging node of `plan`. Or a small `private[sql]` helper next to
`codegenSubtrees` could be shared by `debug()` and this test. What do you think?
--
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]
