[SPARK-58084][SQL] Support converting sort merge join to shuffled hash join in AQE physical plan

[ulysses-you]

What changes were proposed in this pull request?

This PR lets adaptive query execution (AQE) convert a SortMergeJoinExec into a ShuffledHashJoinExec on the physical plan, so the conversion can see through (aggregate, project, filter, window, left-existence join) that sit between the join and its input shuffle.

Concretely:

1. New rule ReplaceSortMergeJoinToShuffledHashJoin (in queryStagePreparationRules, after ReplaceHashWithSortAgg). Once the join's input shuffles have materialized, it reaches each side's ShuffleQueryStageExec through the safelist of non-shuffle operators above it and, if a build side's per-partition sizes all fit spark.sql.adaptive.maxShuffledHashJoinLocalMapThreshold, rewrites the join to a shuffled hash join. The swap is shuffle-free (both are ShuffledJoin with the same distribution/partitioning); only the child sorts become unnecessary. Because a shuffled hash join loses the sort merge join's output ordering, EnsureRequirements is re-run so any ordering an ancestor still needs is re-established. Gated by spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled (default false).

2. SimpleCostEvaluator gains an optional local-sort cost (spark.sql.adaptive.costEvaluator.countLocalSort.enabled, default false). The cost is packed as [-numSkewJoins | numShuffles | numLocalSorts], so local sorts are a lowest-priority tiebreaker below shuffle count (and skew-join count). AQE's cost comparison then adopts the converted plan only when it does not add local sorts elsewhere.

3. preferShuffledHashJoin is lifted into JoinSelectionHelper and shared by DynamicJoinSelection and the new rule.

Class hierarchy note: SortMergeJoinExec and ShuffledHashJoinExec both extend ShuffledJoin, which is why the swap needs no new shuffle. ShuffledHashJoinExec (via HashJoin.outputOrdering) only preserves the streamed side's ordering, whereas SortMergeJoinExec (inner) keeps both sides'; the local-sort cost handles cases where that difference forces a new sort upstream.

Why are the changes needed?

DynamicJoinSelection already prefers a shuffled hash join over a sort merge join, but it works by adding a join hint on the logical plan and only fires when the join child is a shuffle stage directly. When non-inflating operators (e.g. an aggregate, or a filter from a HAVING) sit between the join and its input shuffle, the hint is never added and the join stays a sort merge join even though a shuffled hash join would be cheaper. Doing the selection on the physical plan — where each side's input shuffle is an explicit materialized ShuffleQueryStageExec — removes that restriction and reuses the correct runtime statistics regardless of the operators above the shuffle.

Does this PR introduce any user-facing change?

Yes, two new configurations, both documented in docs/sql-performance-tuning.md:

• spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled (default false)
• spark.sql.adaptive.costEvaluator.countLocalSort.enabled (default false)

Both default to false, so there is no behavior change unless a user opts in.

How was this patch tested?

New unit tests in AdaptiveQueryExecSuite:

• Converting a sort merge join to a shuffled hash join when an aggregate / HAVING filter sits above the shuffle.
• The converted plan stays valid when an ancestor needs the join's ordering (EnsureRequirements re-adds the sort).
• SimpleCostEvaluator orders plans by skew joins, then shuffles, then local sorts (including a real skew-join node).
• With countLocalSort enabled, a conversion that would add a local sort (both sides are sort aggregates, a parent window partitions by the build-side key) is rejected and the sort merge join is kept.

Full AdaptiveQueryExecSuite passes.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code

[ulysses-you]

cc @dongjoon-hyun @viirya @sunchao @cloud-fan if you have time to take o look, thank you!

[sunchao]

Summary

Prior state and problem

AQE can already prefer a shuffled hash join through logical re-optimization, but that path only sees a shuffle stage directly beneath the logical join. This PR adds a physical-plan conversion so AQE can reach materialized shuffle statistics through operators such as aggregates, projections, filters, and windows.

Design approach

The new preparation rule finds the underlying ShuffleQueryStageExec, checks each side against maxShuffledHashJoinLocalMapThreshold, and replaces an eligible SortMergeJoinExec with ShuffledHashJoinExec. It reruns EnsureRequirements after the replacement and optionally includes local sorts as a lower-priority AQE cost component.

Correctness / compatibility analysis

The general conversion path and ordering repair are coherent, and the added SPARK-58084 tests pass. I found two correctness gaps: the underlying shuffle bytes do not safely bound the hash-table input after row-widening operators, and the physical rewrite bypasses an explicit MERGE strategy hint.

Key design decisions

The behavior is opt-in, excludes skewed sort-merge joins, and limits traversal to a safelist of physical operators. Both issues arise at those boundaries: several safelisted operators can increase byte width, and the physical match does not retain the reason the sort-merge join was selected.

Implementation sketch

The patch adds the configuration and documentation, installs ReplaceSortMergeJoinToShuffledHashJoin in AQE's query-stage preparation rules, shares the shuffled-hash preference helper, and extends SimpleCostEvaluator with optional local-sort counting. The rule strips obsolete input sorts and lets the normal adaptive cost comparison choose whether to adopt the rewritten plan.

Behavioral changes worth calling out

When enabled, the rule can replace a spillable sort-merge join using a size estimate taken before a widening operator, which can make the hash build exceed the configured threshold and available memory. It can also replace a join that the query explicitly requested as sort-merge with MERGE.

Suggested improvements

Please restrict shuffle-stat traversal to operators proven not to increase the build input's byte size, or obtain statistics after the traversed subtree. Also preserve explicit join-strategy hints during the physical rewrite and add regression coverage for a widening aggregate/project and for MERGE with the new configuration enabled.

[sunchao]

[P1] Do not use pre-operator shuffle bytes as the hash-table size bound

This traversal assumes the stage's mapStats still upper-bounds the rows passed to the shuffled hash join, but ProjectExec, BaseAggregateExec, and WindowExecBase can increase each row's byte width. A regression probe on this head using repeat(max(c2), 10000) above the aggregate, with a 500-byte local-map threshold, still produced ShuffledHashJoin ... BuildRight: the shuffle below the aggregate was small, but each post-aggregate build row was about 10 KB. UnsafeHashedRelation builds from those widened rows and cannot spill, so this can turn a working spillable SMJ into an out-of-memory failure. Please only traverse operators proven byte-non-increasing, or obtain a size bound after the operator, and add a widening aggregate/project regression test.

[sunchao]

[P2] Preserve explicit MERGE hints

This physical rewrite converts every matching SMJ without checking whether it was selected by an explicit strategy hint. With the new configuration enabled, a regression probe using SELECT /*+ MERGE(t1) */ ... FROM t1 JOIN t2 ... finished as ShuffledHashJoin ... BuildRight, even though Spark's existing dynamic join-selection path deliberately avoids overriding an existing strategy hint. The current MERGE-hint test does not catch this because the new configuration defaults to false. Please retain/check the logical JoinHint before replacing the SMJ and add an enabled-configuration regression test.

[cloud-fan]

1 blocking, 1 non-blocking, 0 nits.
One net-new blocking correctness gap (missing binary-stability guard) plus a design direction on consolidating the two SHJ-over-SMJ paths. Two prior issues raised by @sunchao (P1, P2) remain open and are not duplicated here.

Correctness (1)

• ReplaceSortMergeJoinToShuffledHashJoin.scala:82: converts an SMJ to an SHJ without the hashJoinSupported (binary-stable key) guard both peer paths enforce — wrong results on collated / non-binary-stable string keys. See inline.

Design / architecture (1)

• SHJ-over-SMJ preference now lives in two places (DynamicJoinSelection, default-on logical; this rule, opt-in physical) sharing preferShuffledHashJoin. The SHJ-over-SMJ half (PREFER_SHUFFLE_HASH/SHUFFLE_HASH) can consolidate fully onto this physical rule, which is strictly more capable; only NO_BROADCAST_HASH broadcast demotion must remain logical. This is doable in this PR — the only thing gating it is that going default-on requires closing all the correctness gates first (the binary-stability finding + P1 + P2, and probably countLocalSort default-on). Worth deciding explicitly: consolidate now (default-on, single path) vs. ship transitional (two paths, dormant feature, cleanup later). I'd lean toward consolidating here. See PR-description suggestions.

Two prior issues raised by @sunchao remain open — P1 (hash-table size bound after row-widening operators, at line 121) and P2 (explicit MERGE hint bypassed, at line 78); this review defers to those threads rather than duplicating them.

Verification

I traced the SMJ→SHJ swap for result-equivalence. It is row-equivalent for ordinary keys (both are ShuffledJoin equi-joins; EnsureRequirements repairs the lost ordering, which the tests exercise). The one input dimension where it is not equivalent is join-key binary stability: collated string keys are orderable (so SMJ was planned) but not binary-stable, and both existing SHJ-construction paths gate on hashJoinSupported for exactly this reason. This rule omits that gate — hence the blocking finding.

PR description suggestions

• Document why this is a new physical rule rather than an extension of DynamicJoinSelection (the ordering-loss tradeoff needs a physical plan; true per-partition build-side sizes come only from the materialized ShuffleQueryStageExec, not logical subtree stats).
• Document how the new rule relates to DynamicJoinSelection, and state the scope choice: consolidate SHJ-over-SMJ onto the physical rule here (default-on, remove PREFER_SHUFFLE_HASH/SHUFFLE_HASH, leave only NO_BROADCAST_HASH) once the correctness gates are closed, or explain how the transitional two-path overlap will be resolved.

[cloud-fan]

This builds a ShuffledHashJoinExec as soon as a build side fits the size threshold, but never checks that the join keys are hash-join-compatible. Both peer SHJ-construction paths gate on hashJoinSupported first — JoinSelection.createShuffleHashJoin (SparkStrategies.scala:240,262) returns None for non-binary-stable keys and falls through to SMJ, and the DynamicJoinSelection hint path reaches SHJ only through that same guard.

A join on collated / non-binary-stable string keys (e.g. UTF8_LCASE) is orderable — so SortMergeJoinExec is chosen (OrderUtils.isOrderable treats every AtomicType as orderable) — but not binary-stable (UnsafeRowUtils.isBinaryStable is false), which is exactly why SHJ was rejected on those paths. Converting such an SMJ here matches keys by UnsafeRow binary equality instead of collation-aware equality, so it silently misses matches and returns wrong results once the config is enabled.

Gate each side (or the whole conversion) on hashJoinSupported / isBinaryStable(keys) before converting, and add a collated-key regression test.

[viirya]

Thank you @ulysses-you for working on this. The overall approach looks reasonable to me - doing the selection on the physical plan where the materialized shuffle stats are visible is the right direction, and the opt-in + cost-based adoption design is clean. I agree with the correctness issues already raised (the missing hashJoinSupported guard, the size bound after row-widening operators, and the bypassed MERGE hint); those need to be addressed before this can go in. A few additional comments inline.

[viirya]

The new rule runs after ReplaceHashWithSortAgg. That rule replaces a hash aggregate with a sort aggregate when the child is already sorted (e.g. on top of a sort merge join), and this rule then destroys exactly that ordering. When both are enabled: with countLocalSort off we adopt a strictly worse plan (sort aggregate plus the sort that EnsureRequirements re-inserts, instead of the original hash aggregate); with countLocalSort on, the extra sort makes the cost evaluator reject the whole candidate plan, so we also lose an otherwise beneficial SMJ-to-SHJ conversion. Placing this rule before ReplaceHashWithSortAgg avoids both cases, since the ordering is already gone when ReplaceHashWithSortAgg makes its decision.

[viirya]

(numShuffles.toLong << SHUFFLE_SHIFT) | sortCost is not masked, so the shuffle count only safely fits in the 16 bits between SHUFFLE_SHIFT and SKEW_SHIFT. The previous encoding was safe by construction (numShuffles is an Int in the low 32 bits), while the new one silently wraps if a plan ever has 65536+ shuffles - the overflow bits get absorbed into the sign-extended skew term. Not a practical concern for any real plan, but since we're now packing three components, I'd suggest replacing the bit tricks with a small Cost implementation holding (numSkewJoins, numShuffles, numLocalSorts) and comparing lexicographically - it also removes the need to keep the shift comment in sync (the comment currently says the shuffle band has 32 bits, but it has 16). Clamping the shuffle count is fine too if you prefer keeping the packed Long.

[viirya]

Enabling only convertSortMergeJoinToShuffledHashJoin.enabled without costEvaluator.countLocalSort.enabled means a conversion that pushes extra sorts elsewhere is still adopted (your own test shows the 5-sorts-vs-4 case). Is there a reason to keep them independent? If the cost-evaluator change is kept behind its own flag for safety, the docs should at least recommend enabling them together - or countLocalSort could arguably default to true, since it is a no-op for cost comparison when nothing else changes the sort count.

[viirya]

Nit: the rule name ReplaceSortMergeJoinToShuffledHashJoin mixes the two naming patterns nearby (ReplaceHashWithSortAgg vs the config's convert...To...) - ReplaceSortMergeJoinWithShuffledHashJoin or ConvertSortMergeJoinToShuffledHashJoin would be more consistent.

[viirya]

Nit: this entry only mentions the per-partition size condition, but preferShuffledHashJoin also requires spark.sql.adaptive.advisoryPartitionSizeInBytes <= maxShuffledHashJoinLocalMapThreshold, which is worth mentioning like the existing 3.2.0 entry above does.