feat: Improve TPC-H performance with AQE on - #2084
Merged
Merged
Conversation
Contributor
|
One note, while trying AQE on TPCDS Q72 for some reason bite size is not available in all joins (not sure why) hence AQE makes wrong decisions. |
Member
Author
Yes, that's exactly what I am seeing. I am getting up to speed with the statistics code now. |
The ballista.optimizer.broadcast_join_threshold_bytes config is only consumed by the static distributed planner (maybe_promote_to_broadcast). Under adaptive query planning, broadcast (CollectLeft) selection in DynamicJoinSelectionExec::to_actual_join instead used DataFusion's hash_join_single_partition_threshold (1 MiB default), so the Ballista key had no effect and the effective cutoff was silently a different value. Use broadcast_join_threshold_bytes as the byte threshold in the AQE join selection path, keeping DataFusion's row threshold as the absent-stats fallback. A value of 0 disables broadcast promotion, matching the static planner. This gives a single config key consistent behavior under both planners. Closes apache#2085
Follow-up within the same change: the AQE join-selection path also used DataFusion's hash_join_single_partition_threshold_rows as the row-count fallback. There was no Ballista equivalent, so the row threshold still escaped the single-config goal. Add ballista.optimizer.broadcast_join_threshold_rows (default 128K, mirroring DataFusion's previous default) plus SessionConfigExt accessors, and use it in DynamicJoinSelectionExec::to_actual_join instead of the DataFusion key. AQE broadcast selection now depends only on Ballista config. Document both broadcast thresholds in the AQE tuning guide.
…havior SessionConfig::new_with_ballista() already installs a 1,000,000 row threshold for DataFusion's hash_join_single_partition_threshold_rows. Default the new Ballista row key to the same value so consolidating AQE onto the Ballista keys does not silently lower the effective row-count broadcast cutoff.
…ig defaults Instead of hard-coding 10 MiB / 1M for the DataFusion hash_join_single_partition_threshold[_rows] session settings, read them from BallistaConfig::default().broadcast_join_threshold_bytes()/_rows(). The Ballista broadcast-threshold defaults are now the single source of truth for both DataFusion's built-in JoinSelection and Ballista's AQE join selection. Values are unchanged.
In the static planner, maybe_promote_to_broadcast trusted any broadcast-safe HashJoinExec(CollectLeft) that DataFusion's JoinSelection produced. DataFusion decides CollectLeft from its own session threshold, which can exceed a runtime override of ballista.optimizer.broadcast_join_threshold_bytes. Demote such a join back to Partitioned when its build side is not under the current Ballista threshold (or when broadcasts are disabled with threshold 0), so the Ballista key is authoritative in the static path too. Null-aware anti joins are never demoted since they require CollectLeft.
The join-selection repartition tests forced the repartition path by setting DataFusion's hash_join_single_partition_threshold[_rows] to 0. AQE join selection now reads the broadcast cutoff from the Ballista config (broadcast_join_threshold_bytes), so those DataFusion keys no longer gate CollectLeft promotion and the small test tables were promoted to broadcast. Set the Ballista broadcast byte threshold to 0 in the helper, which disables CollectLeft promotion and restores the repartitioned plans the snapshots assert.
`supports_collect_by_thresholds` compared a row count against `hash_join_single_partition_threshold_rows` whenever `total_byte_size` was unknown, so a build side of up to a million arbitrarily wide rows could be broadcast to every probe task without the byte threshold ever applying. Unknown `total_byte_size` is the common case, not an edge case: DataFusion discards it on every join, and rebuilding it in `Statistics::calculate_total_byte_size` only works when every column has a fixed width, so a single `Utf8` column loses it permanently. In TPC-H that covers most dimension-side join results. Estimate the size instead and hold it to the same byte threshold. Each column contributes its own `byte_size` statistic when present -- a total for that column's output, already scaled for filters and limits -- otherwise its fixed width times the row count, otherwise a default width mirroring Spark's `StringType`/`BinaryType` defaults. An overflowing estimate declines the broadcast rather than wrapping to a small number. The row threshold is retained as a ceiling, so this can only reject a broadcast the row rule would have allowed, never introduce a new one. Closes apache#2081.
The broadcast-vs-partitioned decision is a function of statistics, but the tests around it could only describe tables they were willing to materialise, so the sizes it actually turns on had no coverage: a build side of hundreds of thousands of rows, or one whose `total_byte_size` is unknown. The existing tests instead toggle the decision by zeroing the threshold, which shows the rule is self-consistent but not that the shipped thresholds behave. Add `StatsTable`, a table that declares its statistics and holds no rows, so a fixture can say "800,000 rows of unknown size" in one line. Its scan reports the declared figures and cannot be executed, which is enough for the planner tests, and it deliberately does not recompute `total_byte_size` on projection, since an unknown size is the case these fixtures exist to express. Add tests covering the decision at both edges -- wide rows of unknown size are not broadcast, while small dimensions, narrow rows, and known sizes under the threshold still are -- run under `SessionConfig::new_with_ballista` so they exercise the 10 MB / 1,000,000 row thresholds a deployment ships with rather than DataFusion's defaults, plus a test pinning those defaults directly. `wide_rows_of_unknown_size_are_not_broadcast` fails on the rule that preceded the previous commit and passes with it. The rest pass either way: they guard against the estimate rejecting broadcasts it should allow. Part of apache#2081.
andygrove
force-pushed
the
fix/2081-size-aware-broadcast
branch
from
July 18, 2026 14:23
3bec11b to
f33bef6
Compare
github-actions
Bot
added
the
documentation
github-actions
Bot
removed
the
documentation
andygrove
changed the title
Surface two previously debug-only diagnostics at levels visible under the default INFO filter, so join-strategy and memory-pressure behavior can be observed without enabling debug logging. - AQE dynamic join selection now logs each decision at INFO, naming the resolved action (CollectLeft/Hash/SortMerge/Repartition), the partition mode, and the size-aware inputs (per-side row/byte estimates and the byte/row broadcast thresholds). - Sort-shuffle write completion logs at INFO with row and spill counts; when a partition spills under memory pressure it logs at WARN with the spilled bytes, batches, and event count. Per-batch spill events stay at debug.
…ough serde The sort-shuffle writer's per-task buffered-bytes cap defaulted to 256 MB, far below the per-task memory-pool budget, so it spilled long before the pool was under pressure. Default the cap to 0, which disables it: spilling is then driven solely by memory-pool pressure. A non-zero value still adds a second spill trigger and is retained for tests and explicit tuning. The cap was also dropped during physical-plan serialization (the executor rebuilt the config with the default), so a configured value never reached the executor. Carry memory_limit_per_task_bytes through the SortShuffleWriterExecNode protobuf and apply it on decode, so an override via ballista.shuffle.sort_based.memory_limit_per_task_bytes takes effect on executors.
Reads the actual materialized per-partition byte sizes off the resolved ExchangeExec feeding a Partitioned hash join's build side (the same source CoalescePartitionsRule reads) and returns the MAX rather than the average, since a single oversized partition is enough to OOM even when the average partition is small (the Q18 failure shape).
Record a full 22-query SF1000 run with prefer_hash_join=true and the AQE hash-join build-size safety fallback (hash_join_max_build_partition_bytes=64 MiB) on a 2 executor x 16 core cluster at target_partitions=64: all 22 queries complete with no OOM, where a pure hash-join run fails on Q18. Remove the AQE-off column pending a re-run at a matched core count.
github-actions
Bot
added
the
documentation
Resolve conflicts from the merged sort-shuffle spill-cap work (apache#2091): adopt the optional proto field and Option<u64> serde handling from main, keep the per-task budget default at 0 (uncapped), and preserve this branch's INFO-level shuffle-write logging while folding in main's repart/spill/write timing breakdown. Take main's refreshed benchmarking.md results as the base for the pending SF1000 re-run.
andygrove
force-pushed
the
fix/2081-size-aware-broadcast
branch
from
July 20, 2026 23:29
f799293 to
becb376
Compare
Re-ran the TPC-H SF1000 suite (AQE on, target_partitions=64, prefer_hash_join=false, 1 iteration) on the 2x16-core reference cluster against the PR build (becb376). Q1-Q17 from a full-suite run, Q19-Q22 as individual jobs; Q18 still OOMs (Partitioned build side, unchanged). Ballista total (excl. Q18) improves 4817.8 -> 4661.0s, led by the join-heavy queries (Q7 -88s, Q8 -191s, Q9 -169s).
Drop this branch's change of the sort-shuffle per-task spill-cap default to 0 (uncapped); restore the 256 MiB default from main. The serde plumbing for the value stays (it landed on main via apache#2091). This PR no longer alters the shipped spill-cap default.
Make explicit that the reference AQE-on numbers were produced with the sort-shuffle per-task spill cap overridden to 0 (uncapped); the shipped default is 256 MiB. Point the Ballista row at the current branch commit.
andygrove
marked this pull request as ready for review
July 21, 2026 03:14
andygrove
changed the title
phillipleblanc
approved these changes
Jul 21, 2026
phillipleblanc
left a comment
phillipleblanc
left a comment
Contributor
There was a problem hiding this comment.
Some minor feedback below.
| ); | ||
| } else { | ||
| debug!( | ||
| info!( |
Contributor
There was a problem hiding this comment.
Looks like this got flipped back to info
Member
Author
There was a problem hiding this comment.
Fixed in b563844 — demoted to debug!. It fires once per shuffle partition; the spill case stays at warn!.
| return false; | ||
| } | ||
|
|
||
| estimate_output_byte_size(&plan.schema(), num_rows, &stats.column_statistics) |
Contributor
There was a problem hiding this comment.
This new byte estimate is not used when choosing the build side: to_actual_join ORs both sides, then build-side selection can fall back to row count. If only the right side passes this estimate but the wider left has fewer rows, AQE still broadcasts the left. It might be better to preserve the per-side eligibility when selecting the build input?
Member
Author
There was a problem hiding this comment.
Confirmed — the per-side eligibility is lost once under_threshold collapses to an OR of both sides. Tracking as a follow-up in #2121; it would need a re-benchmark so keeping it out of this PR.
| let threshold_collect_left_join_bytes = bc.broadcast_join_threshold_bytes(); | ||
| let threshold_collect_left_join_rows = bc.broadcast_join_threshold_rows(); | ||
| let max_build_bytes = bc.hash_join_max_build_partition_bytes(); | ||
| let build_max_partition_bytes = max_per_partition_build_bytes(&self.left); |
Contributor
There was a problem hiding this comment.
self.left is measured before SelectJoinRule may swap the hash inputs and before AQE coalescing groups partitions. The eventual build task can therefore consume a different or larger partition than this check validated. It might be better to apply the budget to the post-swap, post-coalesce input.
Member
Author
There was a problem hiding this comment.
Confirmed — the check measures the pre-swap, pre-coalesce self.left, which can differ from the partition the build task consumes. Tracking in #2121 to apply the budget to the post-swap, post-coalesce build input.
Co-authored-by: Phillip LeBlanc <phillip@spice.ai>
Demote the per-partition shuffle-write completion log to debug (it fires once per shuffle partition; spills stay at WARN). Update the dynamic_join log import to match the debug! call.
andygrove
changed the title
Member
Author
|
Thanks for the review @phillipleblanc. I'm going to work on the points you raised next (filed as #2121) |
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
3 participants
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.
Which issue does this PR close?
Closes #2081.
Related: #2025 (Q18 hash-join build-side OOM at SF1000).
Rationale for this change
Two AQE join behaviours concentrate work on a few tasks or fail outright at
SF1000:
could broadcast a build side that is actually large, or repartition one that
is actually small. Sizing the decision by bytes keeps the broadcast fast path
on genuinely small build sides.
Partitionedhash join whose largest build partition does not fit one taskslot's memory pool builds in memory and OOM-kills the executor —
DataFusion's hash-join build side does not spill — notably Q18 (Hash-join build side OOMs on TPC-H Q18 at SF1000 and does not shrink with target_partitions #2025). There
was no way to keep such a join running.
This PR makes the broadcast decision size-aware and adds an opt-in per-partition
build-size check that lowers an oversized hash join to
SortMergeJoin(whichspills) instead of aborting. It also adds INFO/WARN observability for these
runtime decisions so they are visible in executor logs.
What changes are included in this PR?
Size-aware broadcast join (AQE):
ballista.optimizer.broadcast_join_threshold_bytesandballista.optimizer.broadcast_join_threshold_rows; the AQE planner sizes thebroadcast decision by bytes, falling back to rows when byte statistics are
unavailable.
CollectLeftjoins whose build side exceeds the Ballistathreshold.
Hash-join build-size fallback (opt-in):
ballista.optimizer.hash_join_max_build_partition_bytes(default0=off). When a
Partitionedhash join's largest build partition exceeds thebudget, the join is lowered to
SortMergeJoin; the decision is logged at INFO.Observability:
write/spill activity; byte-size statistics for completed query stages.
Docs:
demonstration in
docs/source/contributors-guide/benchmarking.md.Benchmark results (TPC-H SF1000, 2×16-core reference cluster, AQE on, 1 iteration)
Default config — SMJ-planned (
prefer_hash_join=false, fallback off).Ballista total excl. Q18: 4817.8 → 4661.0 s. Largest movers are the
join-heavy queries:
mainQ18 still OOMs in this configuration — the build-size fallback is opt-in and off
here.
Hash-join mode + fallback (
prefer_hash_join=true,hash_join_max_build_partition_bytes=64 MiB). All 22 queries complete with noOOM; Q18 finishes at 744.6 s, where the same run with the fallback disabled
fails the job (#2025). Full tables are in
benchmarking.md.Caveats: single iteration. The "Prior" column is the previous
mainbaseline;this build also changes the sort-shuffle spill-cap default, so per-query deltas
reflect the combined change rather than broadcast sizing alone. The two result
sections use different
prefer_hash_joinsettings and are not directlycomparable to each other.
Are there any user-facing changes?
New session config keys (defaults preserve prior behaviour):
ballista.optimizer.broadcast_join_threshold_bytesballista.optimizer.broadcast_join_threshold_rowsballista.optimizer.hash_join_max_build_partition_bytes(default0, off)New INFO/WARN executor log lines for AQE join decisions and shuffle spill
activity. No breaking API changes.