[SPARK-59279][SQL] Don't re-read the sorted-merge config after GroupPartitionsExec is planned - #58543
Open
peter-toth wants to merge 1 commit into
Open
Conversation
…artitionsExec 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)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
GroupPartitionsExec.canUseSortedMergeis replaced by two members, and the config read moves into the planner's own method.kWayMergeIsFeasibleholds the two live terms, the child having an ordering and the child subtree beingSafeForKWayMerge.usesSortedMergeis whatdoExecute,supportsColumnarandoutputOrderingask, and it carries no config term.tryEnableSortedMergereads the config itself, so it appears once in the file, at the only place that decides anything with it. It reads it throughconfrather thanSQLConf.get.The
enableSortedMergescaladoc 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.enabledis 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
GroupPartitionsExeccoalesces:At planning,
tryEnableSortedMerge()finds the config on and returnscopy(enableSortedMerge = true). That copy reports the child's full ordering, soEnsureRequirementsadds noSortExecunder the join. The copy is a new instance, so its owncanUseSortedMergeis still unevaluated. At executiondoExecuteforces it, now under the new config value, and builds a plainCoalescedRDDinstead. The concatenated partitions are no longer sorted, and the sort-merge join walks them as if they were.outputOrderingread the same member, so the node and the plan above it ended up disagreeing about what the node delivers.supportsColumnarread it too, which is a second route to the same lost rows. Under AQE,ApplyColumnarRulesAndInsertTransitionsruns when the result stage is created, so it sees the flipped config. A columnar child would then makesupportsColumnartrue and route todoExecuteColumnar, which only ever builds a plainCoalescedRDD, 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.3andbranch-4.2as well, same 3 rows of 5 in both, with the same two rows lost.branch-4.1has noGroupPartitionsExec, so the path does not exist there.Why this shape
enableSortedMergeis already the record of the planner's decision, and onlytryEnableSortedMergesets it, after checking the config. So the execution side does not need the config at all.childIsSafeForKWayMergedoes have to stay live.ApplyColumnarRulesAndInsertTransitionsandCollapseCodegenStagesinsert nodes under the child afterEnsureRequirementsran. 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 valat construction, the waySortExecdoes withenableRadixSort, fixes the same case, butSparkPlan.confis the live conf and avalfreezes at construction rather than at the decision, so the window reopens for any node copied afterEnsureRequirementsran. Dropping the&& canUseSortedMergere-check outright would takechildIsSafeForKWayMergewith it.Reading
confinstead ofSQLConf.getis not what fixes this, and that was measured rather than assumed.SparkPlan.confissession.sessionState.conf, the session's live mutableSQLConf, which is the same objectspark.conf.setandwithSQLConfmutate. The read does switch toconfhere, 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, andSparkPlan.conffalls back toSQLConf.getwhen the node has no session. It also makes the file consistent, sinceoutputOrderingalready readconf.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 noSortExecbelow the sort-merge join, then turns the config off and executes the sameDataFrame.It covers both AQE modes, because they freeze the plan at different points. AQE builds it in
AdaptiveSparkPlanExec.initialPlan, and without AQEprepareForExecutiondoes. 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 inwithSQLConf(preserveOrderingOnCoalesce -> true). That wrapper is inert now, becauseoutputOrderingno longer reads that config, so it is dropped and the assertion runs at the config's default offalse. 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
createOrderedIdTablesplusorderedIdJoinRows, 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
SortedMergefield.KeyGroupedPartitioningSuite,GroupPartitionsExecSuite,EnsureRequirementsSuite,SortedMergeCoalescedRDDSuite,PlannerSuiteandProjectedOrderingAndPartitioningSuite, 313 tests.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 5)