Skip to content

Commit bb578b2

Browse files
committed
[SPARK-59279][SQL] Don't re-read the sorted-merge config after GroupPartitionsExec is planned
### What changes were proposed in this pull request? `GroupPartitionsExec.canUseSortedMerge` is replaced by two members, and the config read moves into the planner's own method. - `kWayMergeIsFeasible` holds the two live terms, the child having an ordering and the child subtree being `SafeForKWayMerge`. - `usesSortedMerge` is what `doExecute`, `supportsColumnar` and `outputOrdering` ask, and it carries no config term. - `tryEnableSortedMerge` reads the config itself, so it appears once in the file, at the only place that decides anything with it. It reads it through `conf` rather than `SQLConf.get`. The `enableSortedMerge` scaladoc now states the contract, that the flag is the decision rather than a hint. EXPLAIN shows the flag too, since it is now the only thing that says whether a node k-way merges. ### Why are the changes needed? A sort-merge join silently drops rows when `spark.sql.sources.v2.bucketing.preserveOrderingOnCoalesce.enabled` is turned off between planning and execution. Both tables identity-partitioned on the join key, both reporting a two-column ordering, two splits per key so `GroupPartitionsExec` coalesces: val df = sql("SELECT /*+ MERGE(i, p) */ i.id, i.name FROM items i JOIN purchases p " + "ON p.item_id = i.id AND p.time = i.arrive_time") df.queryExecution.executedPlan // planned with the config on, no SortExec below the join spark.conf.set("spark.sql.sources.v2.bucketing.preserveOrderingOnCoalesce.enabled", "false") df.collect() // 3 rows instead of 5 At planning, `tryEnableSortedMerge()` finds the config on and returns `copy(enableSortedMerge = true)`. That copy reports the child's full ordering, so `EnsureRequirements` adds no `SortExec` under the join. The copy is a new instance, so its own `canUseSortedMerge` is still unevaluated. At execution `doExecute` forces it, now under the new config value, and builds a plain `CoalescedRDD` instead. The concatenated partitions are no longer sorted, and the sort-merge join walks them as if they were. `outputOrdering` read the same member, so the node and the plan above it ended up disagreeing about what the node delivers. `supportsColumnar` read it too, which is a second route to the same lost rows. Under AQE, `ApplyColumnarRulesAndInsertTransitions` runs when the result stage is created, so it sees the flipped config. A columnar child would then make `supportsColumnar` true and route to `doExecuteColumnar`, which only ever builds a plain `CoalescedRDD`, while the join above had already been planned against the merged ordering. That route is closed by the same change and is not covered by a test, because the suite has no columnar V2 source. Measured on `branch-4.3` and `branch-4.2` as well, same 3 rows of 5 in both, with the same two rows lost. `branch-4.1` has no `GroupPartitionsExec`, so the path does not exist there. `enableSortedMerge` is already the record of the planner's decision, and only `tryEnableSortedMerge` sets it, after checking the config. So the execution side does not need the config at all. `childIsSafeForKWayMerge` does have to stay live. `ApplyColumnarRulesAndInsertTransitions` and `CollapseCodegenStages` insert nodes under the child after `EnsureRequirements` ran. Those wrappers are all whitelisted, which is what keeps the answer stable, and the live check is the fallback if something outside the whitelist ever appears. Splitting the terms apart is what lets the config term go while that guard stays. Two alternatives were rejected. Snapshotting the config into a `private val` at construction, the way `SortExec` does with `enableRadixSort`, fixes the same case, but `SparkPlan.conf` is the live conf and a `val` freezes at construction rather than at the decision, so the window reopens for any node copied after `EnsureRequirements` ran. Dropping the `&& canUseSortedMerge` re-check outright would take `childIsSafeForKWayMerge` with it. Reading `conf` instead of `SQLConf.get` is not what fixes this, and that was measured rather than assumed. `SparkPlan.conf` is `session.sessionState.conf`, the session's live mutable `SQLConf`, which is the same object `spark.conf.set` and `withSQLConf` mutate. The read does switch to `conf` here, on separate grounds. It now happens only on the driver during planning, so the node's own session conf is the right one to ask, and `SparkPlan.conf` falls back to `SQLConf.get` when the node has no session. It also makes the file consistent, since `outputOrdering` already read `conf`. `outputOrdering`'s other config, `preserveKeyOrderingOnCoalesce`, is deliberately still a live read, and the file now says why. It gates whether to *report* an ordering that holds either way, so a late read only makes the node claim less than it delivers. The sorted-merge config gates whether the merge *happens*. ### Does this PR introduce _any_ user-facing change? Yes, it fixes wrong results. A query whose plan was built with the config on keeps the k-way merge and returns the right rows when the config is turned off before it runs. One consequence worth stating. The config is documented as a cost knob, and turning it off no longer stops a merge in a plan that is already built. That includes an `InMemoryRelation`'s cached plan, which can outlive many config changes in a session. Re-planning is what picks the new value up. ### How was this patch tested? A new test in `KeyGroupedPartitioningSuite`, "SPARK-59279: a planned k-way merge survives a later config change". It forces the plan with the config on, asserts there is no `SortExec` below the sort-merge join, then turns the config off and executes the same `DataFrame`. It covers both AQE modes, because they freeze the plan at different points. AQE builds it in `AdaptiveSparkPlanExec.initialPlan`, and without AQE `prepareForExecution` does. Without the fix both modes return 3 of the 5 rows, dropping `[1,aa]` and `[2,cc]`. One existing test changed, plus a new unit test beside it. `GroupPartitionsExecSuite`'s "SPARK-55715: coalescing with enableSortedMerge = true returns full child ordering" wrapped its flag assertion in `withSQLConf(preserveOrderingOnCoalesce -> true)`. That wrapper is inert now, because `outputOrdering` no longer reads that config, so it is dropped and the assertion runs at the config's default of `false`. A new "SPARK-59279: enableSortedMerge decides the k-way merge, not the config" states the invariant directly, as a grid over the config with and without the flag. The sibling test's name lost the words "sorted merge config enabled", for the same reason. All three fail without the fix. The new integration test's fixture was byte-identical to the two tests above it, so it is extracted into `createOrderedIdTables` plus `orderedIdJoinRows`, next to the suite's other fixture helpers. "SPARK-56549: k-way merge enabled only when parent requires ordering" now uses it too. The third copy, under "SPARK-55715", is left alone to keep this diff small. "SPARK-55992: GroupPartitions string in simple and extended explain" now expects the `SortedMerge` field. `KeyGroupedPartitioningSuite`, `GroupPartitionsExecSuite`, `EnsureRequirementsSuite`, `SortedMergeCoalescedRDDSuite`, `PlannerSuite` and `ProjectedOrderingAndPartitioningSuite`, 313 tests. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) Closes #58543 from peter-toth/SPARK-59279-sorted-merge-conf-freeze. Authored-by: Peter Toth <peter.toth@gmail.com> Signed-off-by: Peter Toth <peter.toth@gmail.com> (cherry picked from commit 6a23cb6) Signed-off-by: Peter Toth <peter.toth@gmail.com>
1 parent 4116a7e commit bb578b2

3 files changed

Lines changed: 166 additions & 67 deletions

File tree

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala

Lines changed: 33 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ import org.apache.spark.sql.catalyst.plans.QueryPlan
2929
import org.apache.spark.sql.catalyst.plans.physical.{IdentityReducer, KeyedPartitioning, KeyReducer, Partitioning, PartitioningCollection, UnknownPartitioning}
3030
import org.apache.spark.sql.catalyst.util.{truncatedString, InternalRowComparableWrapper}
3131
import org.apache.spark.sql.execution.{SafeForKWayMerge, SparkPlan, UnaryExecNode}
32-
import org.apache.spark.sql.internal.SQLConf
3332
import org.apache.spark.sql.types.DataType
3433
import org.apache.spark.sql.vectorized.ColumnarBatch
3534

@@ -57,6 +56,10 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
5756
* of the coalesced partitions, preserving the child's output ordering
5857
* end-to-end. Set by [[EnsureRequirements]] when a parent operator
5958
* requires the ordering that this node can satisfy via sorted merge.
59+
* This flag is the decision, not a hint. Nothing below re-reads the
60+
* config that produced it. A node planned with the merge would otherwise
61+
* concatenate instead, and a sort-merge join above it would lose rows
62+
* (SPARK-59279).
6063
*/
6164
case class GroupPartitionsExec(
6265
child: SparkPlan,
@@ -317,28 +320,33 @@ case class GroupPartitionsExec(
317320
//
318321
// `outputOrdering` is first evaluated during EnsureRequirements (which decides whether to
319322
// add SortExec), but the child plan tree changes afterwards when
320-
// ApplyColumnarRulesAndInsertTransitions and CollapseCodegenStages insert wrapper nodes. The
321-
// correctness of this code relies on all such insertable nodes (WholeStageCodegenExec,
322-
// InputAdapter, ColumnarToRowExec) being in the SafeForKWayMerge whitelist so the evaluation
323-
// stays consistent.
323+
// ApplyColumnarRulesAndInsertTransitions and CollapseCodegenStages insert wrapper nodes. Each
324+
// such change mints a fresh node through `withNewChildInternal`, so this is recomputed against
325+
// the new child. The correctness of this code relies on all such insertable nodes
326+
// (WholeStageCodegenExec, InputAdapter, ColumnarToRowExec) being in the SafeForKWayMerge
327+
// whitelist so the evaluation stays consistent.
324328
@transient private lazy val childIsSafeForKWayMerge: Boolean =
325329
!child.exists {
326330
case _: SafeForKWayMerge => false
327331
case _ => true
328332
}
329333

330-
@transient private lazy val canUseSortedMerge: Boolean =
331-
SQLConf.get.v2BucketingPreserveOrderingOnCoalesceEnabled &&
332-
child.outputOrdering.nonEmpty &&
333-
childIsSafeForKWayMerge
334+
/** Whether a k-way merge would work at all, leaving aside whether it is switched on. */
335+
@transient private lazy val kWayMergeIsFeasible: Boolean =
336+
child.outputOrdering.nonEmpty && childIsSafeForKWayMerge
337+
338+
/** Whether this node performs the k-way merge. No config term, see `enableSortedMerge`. */
339+
@transient private lazy val usesSortedMerge: Boolean =
340+
enableSortedMerge && hasCoalescing && kWayMergeIsFeasible
334341

335342
/**
336-
* Returns a copy of this node with k-way merge enabled if it is feasible: the config is on,
337-
* the child has an ordering, the child subtree is `SafeForKWayMerge`, and this node actually
338-
* coalesces partitions.
343+
* Returns a copy of this node with k-way merge enabled, when the config is on, this node
344+
* coalesces partitions and the merge is feasible. The only read of
345+
* `preserveOrderingOnCoalesce`.
339346
*/
340347
def tryEnableSortedMerge(): Option[GroupPartitionsExec] = {
341-
Option.when(hasCoalescing && canUseSortedMerge) {
348+
Option.when(conf.v2BucketingPreserveOrderingOnCoalesceEnabled && hasCoalescing &&
349+
kWayMergeIsFeasible) {
342350
val newGroupPartitions = copy(enableSortedMerge = true)
343351
newGroupPartitions.copyTagsFrom(this)
344352
newGroupPartitions
@@ -357,7 +365,7 @@ case class GroupPartitionsExec(
357365
override protected def doExecute(): RDD[InternalRow] = {
358366
if (groupedPartitions.isEmpty) {
359367
sparkContext.emptyRDD
360-
} else if (hasCoalescing && enableSortedMerge && canUseSortedMerge) {
368+
} else if (usesSortedMerge) {
361369
val partitionCoalescer = new GroupedPartitionCoalescer(groupedPartitions.map(_._2))
362370
val rowOrdering = new LazyCodeGenOrdering(kWayMergeOrdering, child.output)
363371
new SortedMergeCoalescedRDD[InternalRow](
@@ -371,8 +379,7 @@ case class GroupPartitionsExec(
371379
}
372380
}
373381

374-
override def supportsColumnar: Boolean =
375-
child.supportsColumnar && !(hasCoalescing && enableSortedMerge && canUseSortedMerge)
382+
override def supportsColumnar: Boolean = child.supportsColumnar && !usesSortedMerge
376383

377384
override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
378385
if (groupedPartitions.isEmpty) {
@@ -394,7 +401,7 @@ case class GroupPartitionsExec(
394401
// within-partition ordering is fully preserved (including any key-derived ordering that
395402
// `DataSourceV2ScanExecBase` already prepended).
396403
child.outputOrdering
397-
} else if (enableSortedMerge && canUseSortedMerge) {
404+
} else if (usesSortedMerge) {
398405
// Coalescing with sorted merge: SortedMergeCoalescedRDD performs a k-way merge using the
399406
// child's ordering, so the full within-partition ordering is preserved end-to-end.
400407
child.outputOrdering
@@ -405,6 +412,10 @@ case class GroupPartitionsExec(
405412
// sorted ascending by the data column), concatenating them yields (A,1),(A,3),(A,2),(A,5)
406413
// which is no longer sorted by the data column. Only sort orders over partition key
407414
// expressions remain valid -- they evaluate to the same value (A) in every merged partition.
415+
//
416+
// The config below stays a per-call read. It gates only whether to report an ordering that
417+
// holds either way, so a late read can never claim more than this node delivers. The merge
418+
// config is different, because what it gates changes what this node produces.
408419
outputPartitioning match {
409420
case p: Partitioning with Expression
410421
if reducers.isEmpty && conf.v2BucketingPreserveKeyOrderingOnCoalesceEnabled =>
@@ -440,8 +451,11 @@ case class GroupPartitionsExec(
440451
s"Reducers: ${truncatedString(names, "[", ", ", "]", joinKeyMaxFields)}"
441452
}
442453
val distributeStr = Iterator(s"DistributePartitions: $distributePartitions")
443-
joinKeyStr ++ expectedStr ++ reducersStr ++ distributeStr
444-
454+
// Rendered from the constructor field, as `DistributePartitions` above is. Not from
455+
// `usesSortedMerge`, because that forces `grouping`, which can throw, and this method feeds
456+
// `simpleString`, which `treeString` calls on error paths.
457+
val sortedMergeStr = Iterator(s"SortedMerge: $enableSortedMerge")
458+
joinKeyStr ++ expectedStr ++ reducersStr ++ distributeStr ++ sortedMergeStr
445459
}
446460
}
447461

sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala

Lines changed: 96 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -333,6 +333,40 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
333333
"JOIN testcat.ns.bucket8 b8 ON b12.id = b8.id " +
334334
s"JOIN testcat.ns.bucket$third b ON b12.id = b.id")
335335

336+
/**
337+
* Creates `items` and `purchases` identity-partitioned on the join key, each reporting a
338+
* two-column ordering. Keys 1 and 2 sit on two splits per side, so a join over them makes
339+
* `GroupPartitionsExec` coalesce. Rows are inserted so that concatenating a key's splits violates
340+
* the reported ordering, which is what a k-way merge has to repair.
341+
*/
342+
private def createOrderedIdTables(): Unit = {
343+
val itemOrdering = Array(
344+
sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST),
345+
sort(FieldReference("arrive_time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST))
346+
createTable(items, itemsColumns, Array(identity("id")), itemOrdering)
347+
sql(s"INSERT INTO testcat.ns.$items VALUES " +
348+
"(2, 'cc', 30.0, cast('2023-06-15' as timestamp)), " +
349+
"(1, 'bb', 20.0, cast('2022-03-10' as timestamp)), " +
350+
"(3, 'dd', 40.0, cast('2024-01-01' as timestamp)), " +
351+
"(1, 'aa', 10.0, cast('2021-05-20' as timestamp)), " +
352+
"(2, 'ee', 50.0, cast('2025-09-01' as timestamp))")
353+
354+
val purchaseOrdering = Array(
355+
sort(FieldReference("item_id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST),
356+
sort(FieldReference("time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST))
357+
createTable(purchases, purchasesColumns, Array(identity("item_id")), purchaseOrdering)
358+
sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
359+
"(2, 50.0, cast('2025-09-01' as timestamp)), " +
360+
"(1, 10.0, cast('2021-05-20' as timestamp)), " +
361+
"(3, 40.0, cast('2024-01-01' as timestamp)), " +
362+
"(2, 30.0, cast('2023-06-15' as timestamp)), " +
363+
"(1, 20.0, cast('2022-03-10' as timestamp))")
364+
}
365+
366+
/** What a join of the `createOrderedIdTables` tables on both ordering columns returns. */
367+
private val orderedIdJoinRows =
368+
Seq(Row(1, "aa"), Row(1, "bb"), Row(2, "cc"), Row(2, "ee"), Row(3, "dd"))
369+
336370
/** The `(id, ts)` rows the `withReducedTsJoinLegs` tables are filled from, one per year. */
337371
private val row2020 = "(0, cast('2020-01-01' as timestamp))"
338372
private val row2021 = "(1, cast('2021-01-03' as timestamp))"
@@ -4448,10 +4482,10 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
44484482
|""".stripMargin)
44494483
val simpleAndExtendedKeyword =
44504484
"GroupPartitions JoinKeyPositions: [0] ExpectedPartitionKeys: 2 " +
4451-
"Reducers: [BucketReducer(2)] DistributePartitions: false"
4485+
"Reducers: [BucketReducer(2)] DistributePartitions: false SortedMerge: false"
44524486
val formattedKeyword =
44534487
"Arguments: JoinKeyPositions: [0], ExpectedPartitionKeys: 2, " +
4454-
"Reducers: [BucketReducer(2)], DistributePartitions: false"
4488+
"Reducers: [BucketReducer(2)], DistributePartitions: false, SortedMerge: false"
44554489
checkKeywordsExistsInExplain(df, SimpleMode, simpleAndExtendedKeyword)
44564490
checkKeywordsExistsInExplain(df, ExtendedMode, simpleAndExtendedKeyword)
44574491
checkKeywordsExistsInExplain(df, FormattedMode, formattedKeyword)
@@ -4925,32 +4959,9 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
49254959
}
49264960

49274961
test("SPARK-56549: k-way merge enabled only when parent requires ordering") {
4928-
// Both tables are partitioned by id/item_id and report a two-column ordering.
4929-
// Key 1 appears on two splits on each side, so GroupPartitionsExec must coalesce.
4930-
//
49314962
// Dynamic gate: with the config enabled, k-way merge must be activated only when the parent
49324963
// actually requires ordering (SMJ), and must stay off when the parent does not (hash join).
4933-
val itemOrdering = Array(
4934-
sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST),
4935-
sort(FieldReference("arrive_time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST))
4936-
createTable(items, itemsColumns, Array(identity("id")), itemOrdering)
4937-
sql(s"INSERT INTO testcat.ns.$items VALUES " +
4938-
"(2, 'cc', 30.0, cast('2023-06-15' as timestamp)), " +
4939-
"(1, 'bb', 20.0, cast('2022-03-10' as timestamp)), " +
4940-
"(3, 'dd', 40.0, cast('2024-01-01' as timestamp)), " +
4941-
"(1, 'aa', 10.0, cast('2021-05-20' as timestamp)), " +
4942-
"(2, 'ee', 50.0, cast('2025-09-01' as timestamp))")
4943-
4944-
val purchaseOrdering = Array(
4945-
sort(FieldReference("item_id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST),
4946-
sort(FieldReference("time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST))
4947-
createTable(purchases, purchasesColumns, Array(identity("item_id")), purchaseOrdering)
4948-
sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
4949-
"(2, 50.0, cast('2025-09-01' as timestamp)), " +
4950-
"(1, 10.0, cast('2021-05-20' as timestamp)), " +
4951-
"(3, 40.0, cast('2024-01-01' as timestamp)), " +
4952-
"(2, 30.0, cast('2023-06-15' as timestamp)), " +
4953-
"(1, 20.0, cast('2022-03-10' as timestamp))")
4964+
createOrderedIdTables()
49544965

49554966
withSQLConf(
49564967
SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
@@ -4962,7 +4973,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
49624973
|FROM testcat.ns.$items i
49634974
|JOIN testcat.ns.$purchases p ON p.item_id = i.id AND p.time = i.arrive_time
49644975
|""".stripMargin)
4965-
checkAnswer(hashDf, Seq(Row(1, "aa"), Row(1, "bb"), Row(2, "cc"), Row(2, "ee"), Row(3, "dd")))
4976+
checkAnswer(hashDf, orderedIdJoinRows)
49664977
val hashPlan = hashDf.queryExecution.executedPlan
49674978
assert(collect(hashPlan) { case j: ShuffledHashJoinExec => j }.nonEmpty,
49684979
"expected ShuffledHashJoinExec")
@@ -4984,7 +4995,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
49844995
|FROM testcat.ns.$items i
49854996
|JOIN testcat.ns.$purchases p ON p.item_id = i.id AND p.time = i.arrive_time
49864997
|""".stripMargin)
4987-
checkAnswer(smjDf, Seq(Row(1, "aa"), Row(1, "bb"), Row(2, "cc"), Row(2, "ee"), Row(3, "dd")))
4998+
checkAnswer(smjDf, orderedIdJoinRows)
49884999
val smjPlan = smjDf.queryExecution.executedPlan
49895000
assert(collectAllShuffles(smjPlan).isEmpty, "should not shuffle for compatible partitioning")
49905001
val smjCoalescing =
@@ -4999,6 +5010,63 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
49995010
}
50005011
}
50015012

5013+
test("SPARK-59279: a planned k-way merge survives a later config change") {
5014+
// The plan is built with the config on, so the k-way merge delivers the child's full ordering
5015+
// and EnsureRequirements adds no SortExec below the sort-merge join. Turning the config off
5016+
// afterwards must not change what the already planned node does. If it did, the join would read
5017+
// concatenated partitions as if they were still sorted and silently drop rows.
5018+
//
5019+
// Both AQE modes are covered, because a fresh node instance is minted at a different point in
5020+
// each. Without AQE, `CollapseCodegenStages` puts a `WholeStageCodegenExec` under this node
5021+
// during `prepareForExecution`. With AQE, the wrappers go in when the result stage is created,
5022+
// which is after the flip below.
5023+
createOrderedIdTables()
5024+
5025+
Seq(true, false).foreach { aqeEnabled =>
5026+
withSQLConf(
5027+
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString,
5028+
SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
5029+
SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true") {
5030+
val df = sql(
5031+
s"""
5032+
|${selectWithMergeJoinHint("i", "p")}
5033+
|i.id, i.name
5034+
|FROM testcat.ns.$items i
5035+
|JOIN testcat.ns.$purchases p ON p.item_id = i.id AND p.time = i.arrive_time
5036+
|""".stripMargin)
5037+
// Force the plan without executing it, so the k-way merge is decided under this config.
5038+
val plan = df.queryExecution.executedPlan
5039+
assert(collectAllShuffles(plan).isEmpty, "should not shuffle for compatible partitioning")
5040+
val coalescing =
5041+
collectAllGroupPartitions(plan).filter(_.groupedPartitions.exists(_._2.size > 1))
5042+
assert(coalescing.nonEmpty, "expected coalescing GroupPartitionsExec")
5043+
coalescing.foreach { gp =>
5044+
assert(gp.enableSortedMerge,
5045+
"sort-merge join requires ordering: enableSortedMerge must be true")
5046+
}
5047+
val smjs = collect(plan) { case j: SortMergeJoinExec => j }
5048+
assert(smjs.nonEmpty, "expected SortMergeJoinExec")
5049+
assert(smjs.flatMap(_.children).forall(c => collect(c) { case s: SortExec => s }.isEmpty),
5050+
"the k-way merge satisfies the ordering, so no SortExec should be added")
5051+
5052+
// Flip only the ordering config. The plan above is already committed.
5053+
withSQLConf(SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "false") {
5054+
checkAnswer(df, orderedIdJoinRows)
5055+
// Re-collected, because under AQE `executedPlan` only reaches the final plan's nodes
5056+
// once the query has run.
5057+
val executed =
5058+
collectAllGroupPartitions(df.queryExecution.executedPlan)
5059+
.filter(_.groupedPartitions.exists(_._2.size > 1))
5060+
assert(executed.nonEmpty, "expected coalescing GroupPartitionsExec")
5061+
executed.foreach { gp =>
5062+
assert(gp.execute().isInstanceOf[SortedMergeCoalescedRDD[_]],
5063+
"the planned k-way merge must not be dropped by a config change")
5064+
}
5065+
}
5066+
}
5067+
}
5068+
}
5069+
50025070
test("SPARK-46367: partition key alias in subquery projects KeyedPartitioning") {
50035071
// A subquery that renames a partition key (id -> pk) creates a ProjectExec between the scan and
50045072
// the join. This test verifies that KeyedPartitioning expressions are correctly projected

0 commit comments

Comments
 (0)