fix: Various fixes for AQE join selection

[andygrove]

Which issue does this PR close?

Closes #2121.

Performance summary

At TPC-H SF1000 on the reference cluster (AQE on), the 21-query total excluding Q18 drops from 4661.0 s to 3936.8 s, about 16% faster. That puts Ballista at roughly the same speed as vanilla Spark on the same cluster and data:

Engine	Total (excl. Q18)
Spark 3.5.3 (vanilla)	3914.7 s
Comet	3705.0 s
Ballista (AQE on)	3936.8 s

Rationale for this change

Two related corrections to AQE join selection, plus the documentation needed to measure the result.

1. The byte-size checks did not reflect the side that is actually built

DynamicJoinSelectionExec::to_actual_join made two byte-size decisions against inputs that may not be the build side the executor actually runs. Both were raised by @phillipleblanc as non-blocking suggestions on #2084.

Broadcast eligibility was OR-ed across both sides. under_threshold was supports_collect_by_thresholds(left) || supports_collect_by_thresholds(right), so the join was promoted to CollectLeft when either side passed the byte threshold. Which side actually gets broadcast is decided separately, by SelectJoinRule::supports_swap_join_order, which falls back to row count when byte statistics are absent. So if only the right side passed the byte estimate while the wider-by-bytes left side had fewer rows, AQE would still broadcast the left, replicating more bytes to every probe task than the threshold intended.

The hash-join build-fit check measured the pre-swap input. max_per_partition_build_bytes(&self.left) measured self.left, but SelectJoinRule may swap the inputs, making self.right the real build side. The check could therefore validate a partition the build task never consumes, letting an oversized build side through.

Both defects share one root cause: the function reasons about self.left / self.right while the executor runs a possibly-swapped build side. It already computed the swap (for build_side_join_type) and threw the result away.

2. AQE was re-reading prefer_hash_join, a planner flag it should not consult

The (Repartitioned, Partitioned) arm was gated on config.optimizer.prefer_hash_join. That is a planner flag: DataFusion applies it when it builds the physical plan, and DelayJoinSelectionRule then folds both the resulting HashJoinExec and SortMergeJoinExec onto a single DynamicJoinSelectionExec. By the time AQE runs, the flag has already had its effect, so reading it again applies the same preference twice. It was also applied inconsistently: the two CollectLeft arms emit a hash join unconditionally, so the gate only ever covered one of the three paths that can produce one.

Setting prefer_hash_join=false therefore did not do what a user would expect under AQE. It did not force sort-merge (broadcast joins ignored it), and where it did bite, it sent the plan down a conversion path that discards the join's projection, so it measured a slower plan rather than a different join strategy. That mattered in practice: it is what the benchmarking guide told people to set, and the SF1000 numbers on that page were collected with it.

Removing the gate leaves hash_join_max_build_partition_bytes as the only thing that can lower a Partitioned join to the spillable SortMergeJoin, so its default of 0 (check disabled) would have made sort-merge unreachable on that arm. The default moves to 64 MiB, the value #2084's own benchmark section had already tuned and validated across the full suite.

What changes are included in this PR?

Join selection (dynamic_join.rs)

• Hoist the swap decision into a swap_inputs binding and derive build_side from it (self.right when the swap holds, self.left otherwise). build_side_join_type reuses it rather than calling supports_swap_join_order a second time.
• Broadcast eligibility now tests build_side alone instead of OR-ing both sides. This is strictly a tightening: it can only reject broadcasts the OR would have allowed, never introduce one.
• The build-fit check measures build_side instead of self.left.
• Drop the prefer_hash_join gate from the (Repartitioned, Partitioned) arm, with a comment recording why the flag is deliberately not consulted.
• The AQE join decision debug line reports which side was selected as build.

Configuration (config.rs, extension.rs)

• ballista.optimizer.hash_join_max_build_partition_bytes now defaults to 64 MiB instead of 0. 0 still disables the check, and the doc comment and config description now spell out that disabling it leaves the Partitioned arm unconditionally on hash join.
• Add SessionConfigExt::with_ballista_hash_join_max_build_partition_bytes, the fluent setter this key was missing.

Documentation (docs/source/contributors-guide/benchmarking.md)

• Stop instructing readers to pass prefer_hash_join=false, and explain what actually selects the join strategy under AQE.
• Note that the Comet replaceSortMergeJoin=false setting no longer pins the two engines to the same join strategy, so a Ballista-vs-Comet gap may partly reflect a strategy difference.
• Replace the SF1000 result column with numbers from a single 21-query suite run using --skip 18, and record that per-query run-to-run spread reached ~20% so small differences should be read as noise. Total (excl. Q18) moves from 4661.0 s to 3936.8 s.
• Drop the separate "hash join with a per-partition build-size fallback" section: it described prefer_hash_join=true as a distinct configuration, which is no longer a thing AQE distinguishes, and its numbers came from a different cluster and commit.
• Recommend --skip over a loop of single-query runs when recording a result.

Notes on what is deliberately not changed

Byte eligibility does not override the swap (i.e. "build the eligible side even though row count says otherwise"). SelectJoinRule re-derives the swap independently via supports_swap_join_order, so overriding it here would need the decision threaded through both rules, and the two would silently diverge if either drifts.

#2121 also asks for the budget to be applied post-coalesce. It is not, and the reason is ordering: SelectJoinRule runs inside replan_stages(), while CoalescePartitionsRule fires later in actionable_stages() on the stage the join was just placed into. At decision time there is no CoalescePlan to read, so a coalesce-aware check would have to re-derive the bin-pack over a leaf set that differs from the rule's own for chained joins. That is recorded as a known conservative gap in the doc comment on max_per_partition_build_bytes; ballista.planner.coalesce.enabled is false by default, so it is inert unless coalescing is explicitly turned on. Happy to split it out into a follow-up if reviewers would rather see it addressed.

Q18 is recorded as TBD rather than OOM in the results table. It was excluded from the suite run with --skip 18 and has not since been run to completion at this commit. The earlier OOM observation on this branch was made under prefer_hash_join=false, which no longer describes any supported configuration, so that result does not transfer.

Are there any user-facing changes?

Yes, though no breaking API change.

Behavior. Three changes, all narrowing toward what the existing thresholds already intended:

• Fewer joins are promoted to CollectLeft — only when the side that actually gets broadcast fits the byte threshold.
• The hash_join_max_build_partition_bytes fallback to sort-merge now triggers on the swapped-in build side.
• datafusion.optimizer.prefer_hash_join no longer affects join selection under AQE. Users who set it to false expecting sort-merge throughout will now get AQE's runtime choice instead; the supported lever is ballista.optimizer.hash_join_max_build_partition_bytes.

Defaults. ballista.optimizer.hash_join_max_build_partition_bytes changes from 0 to 64 MiB, so a Partitioned hash join whose largest build partition exceeds 64 MiB now falls back to SortMergeJoin where it previously did not. Setting the key back to 0 restores the old behavior, at the cost of making sort-merge unreachable on that path.

API. SessionConfigExt gains with_ballista_hash_join_max_build_partition_bytes. Additive for callers; external implementors of the trait would need to add the method.

Test notes

One pre-existing test, build_over_threshold_falls_back_to_sort_merge_join, had to change its expectation. Its fixture paired a 500 MB left ExchangeExec with a 16 MB StatisticsExec on the right, so under the fix the swap correctly moves the build onto that 16 MB side, which genuinely fits the 200 MB budget — HashJoinExec is now the right answer, and the old SortMergeJoinExec assertion was pinning the exact bug being fixed. The other two build-fit tests had the same fixture problem and were passing only because the swap landed on a non-ExchangeExec whose size is unknown, so nothing was measured at all. All three now use a two-ExchangeExec helper, which is the shape the (Repartitioned, Partitioned) arm always sees in practice.

The build-fit tests' shared config now sets prefer_hash_join = false deliberately, the hostile value: AQE must not consult it, so any test below asserting a HashJoinExec would fail if that ever regressed.

join_selection.rs's make_ctx_without_collect_left previously distinguished its hash-join and sort-merge cases by toggling prefer_hash_join. It now takes a build budget instead, above or below the 10 bytes per shuffle partition the mocks report, which is what actually selects the strategy.

New tests: broadcast eligibility following the build side; the fit check measuring the post-swap build; a no-swap mirror case pinning that selection follows the swap decision rather than blanket-switching to the right; and prefer_hash_join not affecting the decision under either setting. The existing default-value test is retargeted at 64 MiB, with a comment on why the default must stay non-zero.

[andygrove]

@phillipleblanc fyi

[phillipleblanc]

Thanks @andygrove!

[avantgardnerio]

I lack the big picture here, but each individual part seems like an obvious win. I hope to learn more about joins in Ballista soon: I think I can improve SMJ using OrderedRepartitionExec, and broadcast hash join using MPTs.