Skip to content

Commit a3a1889

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. ### Why this shape `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)
1 parent 2a7cfea commit a3a1889

3 files changed

Lines changed: 162 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}
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,
@@ -272,28 +275,33 @@ case class GroupPartitionsExec(
272275
//
273276
// `outputOrdering` is first evaluated during EnsureRequirements (which decides whether to
274277
// add SortExec), but the child plan tree changes afterwards when
275-
// ApplyColumnarRulesAndInsertTransitions and CollapseCodegenStages insert wrapper nodes. The
276-
// correctness of this code relies on all such insertable nodes (WholeStageCodegenExec,
277-
// InputAdapter, ColumnarToRowExec) being in the SafeForKWayMerge whitelist so the evaluation
278-
// stays consistent.
278+
// ApplyColumnarRulesAndInsertTransitions and CollapseCodegenStages insert wrapper nodes. Each
279+
// such change mints a fresh node through `withNewChildInternal`, so this is recomputed against
280+
// the new child. The correctness of this code relies on all such insertable nodes
281+
// (WholeStageCodegenExec, InputAdapter, ColumnarToRowExec) being in the SafeForKWayMerge
282+
// whitelist so the evaluation stays consistent.
279283
@transient private lazy val childIsSafeForKWayMerge: Boolean =
280284
!child.exists {
281285
case _: SafeForKWayMerge => false
282286
case _ => true
283287
}
284288

285-
@transient private lazy val canUseSortedMerge: Boolean =
286-
SQLConf.get.v2BucketingPreserveOrderingOnCoalesceEnabled &&
287-
child.outputOrdering.nonEmpty &&
288-
childIsSafeForKWayMerge
289+
/** Whether a k-way merge would work at all, leaving aside whether it is switched on. */
290+
@transient private lazy val kWayMergeIsFeasible: Boolean =
291+
child.outputOrdering.nonEmpty && childIsSafeForKWayMerge
292+
293+
/** Whether this node performs the k-way merge. No config term, see `enableSortedMerge`. */
294+
@transient private lazy val usesSortedMerge: Boolean =
295+
enableSortedMerge && hasCoalescing && kWayMergeIsFeasible
289296

290297
/**
291-
* Returns a copy of this node with k-way merge enabled if it is feasible: the config is on,
292-
* the child has an ordering, the child subtree is `SafeForKWayMerge`, and this node actually
293-
* coalesces partitions.
298+
* Returns a copy of this node with k-way merge enabled, when the config is on, this node
299+
* coalesces partitions and the merge is feasible. The only read of
300+
* `preserveOrderingOnCoalesce`.
294301
*/
295302
def tryEnableSortedMerge(): Option[GroupPartitionsExec] = {
296-
Option.when(hasCoalescing && canUseSortedMerge) {
303+
Option.when(conf.v2BucketingPreserveOrderingOnCoalesceEnabled && hasCoalescing &&
304+
kWayMergeIsFeasible) {
297305
val newGroupPartitions = copy(enableSortedMerge = true)
298306
newGroupPartitions.copyTagsFrom(this)
299307
newGroupPartitions
@@ -312,7 +320,7 @@ case class GroupPartitionsExec(
312320
override protected def doExecute(): RDD[InternalRow] = {
313321
if (groupedPartitions.isEmpty) {
314322
sparkContext.emptyRDD
315-
} else if (hasCoalescing && enableSortedMerge && canUseSortedMerge) {
323+
} else if (usesSortedMerge) {
316324
val partitionCoalescer = new GroupedPartitionCoalescer(groupedPartitions.map(_._2))
317325
val rowOrdering = new LazyCodeGenOrdering(kWayMergeOrdering, child.output)
318326
new SortedMergeCoalescedRDD[InternalRow](
@@ -326,8 +334,7 @@ case class GroupPartitionsExec(
326334
}
327335
}
328336

329-
override def supportsColumnar: Boolean =
330-
child.supportsColumnar && !(hasCoalescing && enableSortedMerge && canUseSortedMerge)
337+
override def supportsColumnar: Boolean = child.supportsColumnar && !usesSortedMerge
331338

332339
override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
333340
if (groupedPartitions.isEmpty) {
@@ -349,7 +356,7 @@ case class GroupPartitionsExec(
349356
// within-partition ordering is fully preserved (including any key-derived ordering that
350357
// `DataSourceV2ScanExecBase` already prepended).
351358
child.outputOrdering
352-
} else if (enableSortedMerge && canUseSortedMerge) {
359+
} else if (usesSortedMerge) {
353360
// Coalescing with sorted merge: SortedMergeCoalescedRDD performs a k-way merge using the
354361
// child's ordering, so the full within-partition ordering is preserved end-to-end.
355362
child.outputOrdering
@@ -360,6 +367,10 @@ case class GroupPartitionsExec(
360367
// sorted ascending by the data column), concatenating them yields (A,1),(A,3),(A,2),(A,5)
361368
// which is no longer sorted by the data column. Only sort orders over partition key
362369
// expressions remain valid -- they evaluate to the same value (A) in every merged partition.
370+
//
371+
// The config below stays a per-call read. It gates only whether to report an ordering that
372+
// holds either way, so a late read can never claim more than this node delivers. The merge
373+
// config is different, because what it gates changes what this node produces.
363374
outputPartitioning match {
364375
case p: Partitioning with Expression
365376
if reducers.isEmpty && conf.v2BucketingPreserveKeyOrderingOnCoalesceEnabled =>
@@ -395,8 +406,11 @@ case class GroupPartitionsExec(
395406
s"Reducers: ${truncatedString(names, "[", ", ", "]", joinKeyMaxFields)}"
396407
}
397408
val distributeStr = Iterator(s"DistributePartitions: $distributePartitions")
398-
joinKeyStr ++ expectedStr ++ reducersStr ++ distributeStr
399-
409+
// Rendered from the constructor field, as `DistributePartitions` above is. Not from
410+
// `usesSortedMerge`, because that forces `grouping`, which can throw, and this method feeds
411+
// `simpleString`, which `treeString` calls on error paths.
412+
val sortedMergeStr = Iterator(s"SortedMerge: $enableSortedMerge")
413+
joinKeyStr ++ expectedStr ++ reducersStr ++ distributeStr ++ sortedMergeStr
400414
}
401415
}
402416

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
@@ -591,6 +591,40 @@ class KeyGroupedPartitioningSuite
591591
"JOIN testcat.ns.bucket8 b8 ON b12.id = b8.id " +
592592
s"JOIN testcat.ns.bucket$third b ON b12.id = b.id")
593593

594+
/**
595+
* Creates `items` and `purchases` identity-partitioned on the join key, each reporting a
596+
* two-column ordering. Keys 1 and 2 sit on two splits per side, so a join over them makes
597+
* `GroupPartitionsExec` coalesce. Rows are inserted so that concatenating a key's splits violates
598+
* the reported ordering, which is what a k-way merge has to repair.
599+
*/
600+
private def createOrderedIdTables(): Unit = {
601+
val itemOrdering = Array(
602+
sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST),
603+
sort(FieldReference("arrive_time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST))
604+
createTable(items, itemsColumns, Array(identity("id")), itemOrdering)
605+
sql(s"INSERT INTO testcat.ns.$items VALUES " +
606+
"(2, 'cc', 30.0, cast('2023-06-15' as timestamp)), " +
607+
"(1, 'bb', 20.0, cast('2022-03-10' as timestamp)), " +
608+
"(3, 'dd', 40.0, cast('2024-01-01' as timestamp)), " +
609+
"(1, 'aa', 10.0, cast('2021-05-20' as timestamp)), " +
610+
"(2, 'ee', 50.0, cast('2025-09-01' as timestamp))")
611+
612+
val purchaseOrdering = Array(
613+
sort(FieldReference("item_id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST),
614+
sort(FieldReference("time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST))
615+
createTable(purchases, purchasesColumns, Array(identity("item_id")), purchaseOrdering)
616+
sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
617+
"(2, 50.0, cast('2025-09-01' as timestamp)), " +
618+
"(1, 10.0, cast('2021-05-20' as timestamp)), " +
619+
"(3, 40.0, cast('2024-01-01' as timestamp)), " +
620+
"(2, 30.0, cast('2023-06-15' as timestamp)), " +
621+
"(1, 20.0, cast('2022-03-10' as timestamp))")
622+
}
623+
624+
/** What a join of the `createOrderedIdTables` tables on both ordering columns returns. */
625+
private val orderedIdJoinRows =
626+
Seq(Row(1, "aa"), Row(1, "bb"), Row(2, "cc"), Row(2, "ee"), Row(3, "dd"))
627+
594628
/** The `(id, ts)` rows the `withReducedTsJoinLegs` tables are filled from, one per year. */
595629
private val row2020 = "(0, cast('2020-01-01' as timestamp))"
596630
private val row2021 = "(1, cast('2021-01-03' as timestamp))"
@@ -4124,10 +4158,10 @@ class KeyGroupedPartitioningSuite
41244158
|""".stripMargin)
41254159
val simpleAndExtendedKeyword =
41264160
"GroupPartitions JoinKeyPositions: [0] ExpectedPartitionKeys: 2 " +
4127-
"Reducers: [BucketReducer(2)] DistributePartitions: false"
4161+
"Reducers: [BucketReducer(2)] DistributePartitions: false SortedMerge: false"
41284162
val formattedKeyword =
41294163
"Arguments: JoinKeyPositions: [0], ExpectedPartitionKeys: 2, " +
4130-
"Reducers: [BucketReducer(2)], DistributePartitions: false"
4164+
"Reducers: [BucketReducer(2)], DistributePartitions: false, SortedMerge: false"
41314165
checkKeywordsExistsInExplain(df, SimpleMode, simpleAndExtendedKeyword)
41324166
checkKeywordsExistsInExplain(df, ExtendedMode, simpleAndExtendedKeyword)
41334167
checkKeywordsExistsInExplain(df, FormattedMode, formattedKeyword)
@@ -4601,32 +4635,9 @@ class KeyGroupedPartitioningSuite
46014635
}
46024636

46034637
test("SPARK-56549: k-way merge enabled only when parent requires ordering") {
4604-
// Both tables are partitioned by id/item_id and report a two-column ordering.
4605-
// Key 1 appears on two splits on each side, so GroupPartitionsExec must coalesce.
4606-
//
46074638
// Dynamic gate: with the config enabled, k-way merge must be activated only when the parent
46084639
// actually requires ordering (SMJ), and must stay off when the parent does not (hash join).
4609-
val itemOrdering = Array(
4610-
sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST),
4611-
sort(FieldReference("arrive_time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST))
4612-
createTable(items, itemsColumns, Array(identity("id")), itemOrdering)
4613-
sql(s"INSERT INTO testcat.ns.$items VALUES " +
4614-
"(2, 'cc', 30.0, cast('2023-06-15' as timestamp)), " +
4615-
"(1, 'bb', 20.0, cast('2022-03-10' as timestamp)), " +
4616-
"(3, 'dd', 40.0, cast('2024-01-01' as timestamp)), " +
4617-
"(1, 'aa', 10.0, cast('2021-05-20' as timestamp)), " +
4618-
"(2, 'ee', 50.0, cast('2025-09-01' as timestamp))")
4619-
4620-
val purchaseOrdering = Array(
4621-
sort(FieldReference("item_id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST),
4622-
sort(FieldReference("time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST))
4623-
createTable(purchases, purchasesColumns, Array(identity("item_id")), purchaseOrdering)
4624-
sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
4625-
"(2, 50.0, cast('2025-09-01' as timestamp)), " +
4626-
"(1, 10.0, cast('2021-05-20' as timestamp)), " +
4627-
"(3, 40.0, cast('2024-01-01' as timestamp)), " +
4628-
"(2, 30.0, cast('2023-06-15' as timestamp)), " +
4629-
"(1, 20.0, cast('2022-03-10' as timestamp))")
4640+
createOrderedIdTables()
46304641

46314642
withSQLConf(
46324643
SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
@@ -4638,7 +4649,7 @@ class KeyGroupedPartitioningSuite
46384649
|FROM testcat.ns.$items i
46394650
|JOIN testcat.ns.$purchases p ON p.item_id = i.id AND p.time = i.arrive_time
46404651
|""".stripMargin)
4641-
checkAnswer(hashDf, Seq(Row(1, "aa"), Row(1, "bb"), Row(2, "cc"), Row(2, "ee"), Row(3, "dd")))
4652+
checkAnswer(hashDf, orderedIdJoinRows)
46424653
val hashPlan = hashDf.queryExecution.executedPlan
46434654
assert(collect(hashPlan) { case j: ShuffledHashJoinExec => j }.nonEmpty,
46444655
"expected ShuffledHashJoinExec")
@@ -4660,7 +4671,7 @@ class KeyGroupedPartitioningSuite
46604671
|FROM testcat.ns.$items i
46614672
|JOIN testcat.ns.$purchases p ON p.item_id = i.id AND p.time = i.arrive_time
46624673
|""".stripMargin)
4663-
checkAnswer(smjDf, Seq(Row(1, "aa"), Row(1, "bb"), Row(2, "cc"), Row(2, "ee"), Row(3, "dd")))
4674+
checkAnswer(smjDf, orderedIdJoinRows)
46644675
val smjPlan = smjDf.queryExecution.executedPlan
46654676
assert(collectAllShuffles(smjPlan).isEmpty, "should not shuffle for compatible partitioning")
46664677
val smjCoalescing =
@@ -4675,6 +4686,63 @@ class KeyGroupedPartitioningSuite
46754686
}
46764687
}
46774688

4689+
test("SPARK-59279: a planned k-way merge survives a later config change") {
4690+
// The plan is built with the config on, so the k-way merge delivers the child's full ordering
4691+
// and EnsureRequirements adds no SortExec below the sort-merge join. Turning the config off
4692+
// afterwards must not change what the already planned node does. If it did, the join would read
4693+
// concatenated partitions as if they were still sorted and silently drop rows.
4694+
//
4695+
// Both AQE modes are covered, because a fresh node instance is minted at a different point in
4696+
// each. Without AQE, `CollapseCodegenStages` puts a `WholeStageCodegenExec` under this node
4697+
// during `prepareForExecution`. With AQE, the wrappers go in when the result stage is created,
4698+
// which is after the flip below.
4699+
createOrderedIdTables()
4700+
4701+
Seq(true, false).foreach { aqeEnabled =>
4702+
withSQLConf(
4703+
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString,
4704+
SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
4705+
SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true") {
4706+
val df = sql(
4707+
s"""
4708+
|${selectWithMergeJoinHint("i", "p")}
4709+
|i.id, i.name
4710+
|FROM testcat.ns.$items i
4711+
|JOIN testcat.ns.$purchases p ON p.item_id = i.id AND p.time = i.arrive_time
4712+
|""".stripMargin)
4713+
// Force the plan without executing it, so the k-way merge is decided under this config.
4714+
val plan = df.queryExecution.executedPlan
4715+
assert(collectAllShuffles(plan).isEmpty, "should not shuffle for compatible partitioning")
4716+
val coalescing =
4717+
collectAllGroupPartitions(plan).filter(_.groupedPartitions.exists(_._2.size > 1))
4718+
assert(coalescing.nonEmpty, "expected coalescing GroupPartitionsExec")
4719+
coalescing.foreach { gp =>
4720+
assert(gp.enableSortedMerge,
4721+
"sort-merge join requires ordering: enableSortedMerge must be true")
4722+
}
4723+
val smjs = collect(plan) { case j: SortMergeJoinExec => j }
4724+
assert(smjs.nonEmpty, "expected SortMergeJoinExec")
4725+
assert(smjs.flatMap(_.children).forall(c => collect(c) { case s: SortExec => s }.isEmpty),
4726+
"the k-way merge satisfies the ordering, so no SortExec should be added")
4727+
4728+
// Flip only the ordering config. The plan above is already committed.
4729+
withSQLConf(SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "false") {
4730+
checkAnswer(df, orderedIdJoinRows)
4731+
// Re-collected, because under AQE `executedPlan` only reaches the final plan's nodes
4732+
// once the query has run.
4733+
val executed =
4734+
collectAllGroupPartitions(df.queryExecution.executedPlan)
4735+
.filter(_.groupedPartitions.exists(_._2.size > 1))
4736+
assert(executed.nonEmpty, "expected coalescing GroupPartitionsExec")
4737+
executed.foreach { gp =>
4738+
assert(gp.execute().isInstanceOf[SortedMergeCoalescedRDD[_]],
4739+
"the planned k-way merge must not be dropped by a config change")
4740+
}
4741+
}
4742+
}
4743+
}
4744+
}
4745+
46784746
test("SPARK-46367: partition key alias in subquery projects KeyedPartitioning") {
46794747
// A subquery that renames a partition key (id -> pk) creates a ProjectExec between the scan and
46804748
// the join. This test verifies that KeyedPartitioning expressions are correctly projected

0 commit comments

Comments
 (0)