feat: make broadcast_join_threshold_bytes/rows authoritative under AQE and the static planner

[andygrove]

Which issue does this PR close?

Closes #2085.

Rationale for this change

Broadcast (BHJ / CollectLeft) join selection in Ballista was governed by two disconnected sets of thresholds:

• The Ballista key ballista.optimizer.broadcast_join_threshold_bytes (default 10 MiB), consumed only by the static distributed planner's maybe_promote_to_broadcast.
• DataFusion's optimizer.hash_join_single_partition_threshold[_rows], which SessionConfig::new_with_ballista() installs (10 MiB / 1M) and which drove both DataFusion's built-in JoinSelection and Ballista's AQE join selection (DynamicJoinSelectionExec::to_actual_join).

So the Ballista broadcast key had no effect under AQE, and even in the static planner a runtime override of it could be silently ignored (a CollectLeft join DataFusion produced under its own session threshold was trusted unconditionally). This PR consolidates broadcast selection onto the Ballista keys so a single set of config is authoritative under both planners.

What changes are included in this PR?

• AQE join selection reads the Ballista keys. DynamicJoinSelectionExec::to_actual_join now uses BallistaConfig::broadcast_join_threshold_bytes() for the byte cutoff (instead of hash_join_single_partition_threshold), and a new ballista.optimizer.broadcast_join_threshold_rows for the row-count fallback (instead of hash_join_single_partition_threshold_rows). A byte threshold of 0 disables broadcast promotion entirely.
• New config key ballista.optimizer.broadcast_join_threshold_rows (default 1,000,000, matching the value new_with_ballista() already installed), with SessionConfigExt accessors (ballista_broadcast_join_threshold_rows / with_ballista_broadcast_join_threshold_rows).
• Single source of truth for defaults. SessionConfig::new_with_ballista() now derives the DataFusion hash_join_single_partition_threshold[_rows] session settings from BallistaConfig::default().broadcast_join_threshold_bytes()/_rows() instead of hard-coding 10 MiB / 1M. (Those DataFusion settings still exist to tune DataFusion's built-in JoinSelection in the static planner path.)
• Static planner honors a runtime override. maybe_promote_to_broadcast now demotes a broadcast-safe HashJoinExec(CollectLeft) that DataFusion produced back to Partitioned when its build side is not under the current Ballista threshold (or the threshold is 0). Null-aware anti joins are never demoted (they require CollectLeft). This makes the Ballista key authoritative even when overridden below the DataFusion session value.
• Docs: documented both broadcast thresholds in the AQE section of the tuning guide, and added a "Broadcast join selection" bullet to what AQE does today.
• Tests: AQE — side under the byte threshold collects; byte threshold below the side size repartitions; row-fallback (byte stats absent) driven by the new row key; 0 disables broadcast. Static planner — a DataFusion-promoted CollectLeft is kept under a high Ballista threshold and demoted to Partitioned under a low one.

The separate concern that shuffle total_byte_size is an on-disk Arrow IPC size (rather than an in-memory build-side size) is intentionally left as a follow-up.

Are there any user-facing changes?

Yes:

• New config key ballista.optimizer.broadcast_join_threshold_rows (default 1,000,000), with SessionConfigExt accessors.
• ballista.optimizer.broadcast_join_threshold_bytes now takes effect under AQE (previously ignored) and is authoritative in the static planner even when overridden at runtime below the DataFusion session value. Effective default thresholds are unchanged (10 MiB / 1M).

No breaking public API changes (the additions are new trait methods with provided implementations on SessionConfigExt).

[andygrove]

SF1000 benchmark (AQE on) — broadcast threshold raised to 128 MiB

Ran this branch on the k3s TPC-H SF1000 cluster with AQE on. Since the effective default thresholds are unchanged (10 MiB / 1M), a default run just reproduces the baseline — so to actually exercise the now-authoritative-under-AQE path I raised ballista.optimizer.broadcast_join_threshold_bytes to 128 MiB. Partial run (Q1–Q8, single iteration — lost the cluster before Q9), but the signal is already clear and positive.

Setup

Branch	andygrove/datafusion-ballista@6809fa5 (this PR)
Baseline	published benchmarking-guide SF1000 AQE-on numbers (main, same cluster)
Scale / data	SF1000, node-local Parquet
Cluster	2 executors, one per node, 8 cores / 56 GiB, --memory-pool-size=48GB
Config	target_partitions=32, prefer_hash_join=false, enable_dynamic_filter_pushdown=false, AQE on, broadcast_join_threshold_bytes=134217728

The PR does what it says: the Ballista key now drives AQE join selection

Under the old code these ballista.* broadcast keys were ignored under AQE. With this PR at 128 MiB, the scheduler plans broadcast joins over multi-table intermediates, not just tiny dim tables — e.g. in Q2:

• stage_id=5 broadcast = the part ⋈ partsupp result ([p_partkey, p_mfgr, ps_suppkey, ps_supplycost])
• stage_id=6 broadcast = the wider part⋈partsupp⋈supplier⋈nation⋈region intermediate

Those are ~10⁵–10⁶-row intermediates (well over the 10 MiB default), so at the default threshold AQE keeps them as Partitioned (shuffle) joins. Broadcasting them is only reachable because the raised Ballista key is now authoritative under AQE — which is exactly the behavior this PR adds. (Broadcasts of region/nation and the small part filter would happen at 10 MiB too.)

Results (Q1–Q8, seconds; ratio = baseline / PR, >1 = PR faster)

Query	PR (128 MiB)	main (AQE on)	ratio
1 (no joins — control)	56.7	36.2	0.64×	🔴
2	71.6	63.6	0.89×	🔴
3	240.5	255.0	1.06×	🟢
4	65.0	53.9	0.83×	🔴
5	643.3	700.5	1.09×	🟢
6 (no joins — control)	15.5	12.7	0.82×	🔴
7	505.8	531.4	1.05×	🟢
8	396.2	801.2	2.02×	🟢
Σ 1–8	1994.6	2454.5	1.23×

Q1–Q8 cumulative: ~23% faster than the AQE-on baseline, driven by the join-heavy queries — Q8 (an 8-way dimension-heavy join) halved (−405 s), with Q5 (−57 s) and Q7 (−26 s) also positive. No OOM through Q8.

Caveats (read the wins as understated, not overstated)

• The no-join control queries regressed (Q1 0.64×, Q6 0.82×). Broadcast selection can't touch them, so this is base-commit drift vs the documented main@26b29391 baseline plus single-iteration noise — a headwind the join queries are overcoming. A same-commit main baseline would tighten this.
• 128 MiB over-promotes on some shapes (Q2 0.89×, Q4 0.83×): it broadcasts ~100 MB intermediates whose replication to every task costs more than the shuffle it avoids. Suggests the sweet spot is lower than 128 MiB — a threshold sweep (e.g. 32 / 64 / 128 MiB) would find it.
• Single iteration, partial suite (no Q9–Q22 yet).

Next

• Finish the 22-query suite at 128 MiB.
• Sweep the threshold (32/64/128 MiB) to locate the point that keeps the dimension-join wins without over-broadcasting the mid-size intermediates.
• Re-baseline against main at this PR's base commit to remove the control-query drift.

Net: making broadcast_join_threshold_bytes authoritative under AQE is a clear, measurable win on the dimension-heavy joins (Q8 2× being the headline), and it exposes the tuning knob needed to avoid over-broadcasting — good change.

---

Benchmarking, plan/log analysis, and this write-up were done with LLM assistance (Claude).