Skip to content

[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
apache:masterfrom
peter-toth:SPARK-59279-sorted-merge-conf-freeze
Open

[SPARK-59279][SQL] Don't re-read the sorted-merge config after GroupPartitionsExec is planned#58543
peter-toth wants to merge 1 commit into
apache:masterfrom
peter-toth:SPARK-59279-sorted-merge-conf-freeze

Conversation

@peter-toth

Copy link
Copy Markdown
Contributor

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)

…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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant