Looks promising:

      "queries": [
        {
          "query_id": "TPCH-8",
          "has_errors": false,
          "has_warnings": false,
          "mean_s": 37.76327075066666,
          "min_s": 36.831046083,
          "max_s": 39.35304816,
          "median_s": 37.105718009,
          "stdev_s": 1.3836203688704116,
          "iterations": 3
        }
      ]

Reviewed this against the PR head (1340cc97). Baseline is green for me: 416/416 in cargo test -p ballista-scheduler --lib.

First, the part you asked people to push on — the one-sided shuffle safety argument holds. I traced it rather than taking it on faith: find_runnable_exchanges only pushes ExchangeExec nodes as stages and stops descending at a resolved one, so in the staged state the probe subtree creates no stage, and the join sits above the exchange rather than inside its subtree. Repartitioned is reachable only via to_partitioned() or inputs_already_partitioned(), and staging sets neither. to_broadcast sharing stage_id/shuffle_partitions checks out too.

Two things I'd want changed before merge.

1. Staging fires for join types that can never broadcast

The StageBuildSide guard in to_actual_join checks the config, !null_aware, the exchange guard and should_stage_build_side — but not collect_left_broadcast_safe(build_side_join_type), which is already computed ~80 lines above.

partition_mode is CollectLeft only when under_threshold && collect_left_broadcast_safe(..). That second term is a pure function of the join type, so no measurement can change it. For a non-broadcastable type the second pass can't reach CollectLeft, hits the (now failing) exchange guard, and falls to Repartition — one serialised stage boundary that never had a chance of changing the decision.

I confirmed it with a throwaway test on the q8 shape:

join_type=Full     collect_left_broadcast_safe=false  action=StageBuildSide
join_type=Left     collect_left_broadcast_safe=false  action=StageBuildSide
join_type=LeftSemi collect_left_broadcast_safe=false  action=StageBuildSide
join_type=LeftAnti collect_left_broadcast_safe=false  action=StageBuildSide

Full is affected unconditionally (Full.swap() == Full); Left/LeftSemi/LeftAnti/LeftMark whenever the build side is already the left input.

One line fixes it:

 (JoinInputState::Unknown, PartitionMode::Partitioned)
     if bc.stage_build_side_enabled()
         && !self.null_aware
+        // Measuring the build side can only change the outcome if
+        // CollectLeft is reachable at all, which is a property of the
+        // (post-swap) join type alone.
+        && collect_left_broadcast_safe(build_side_join_type)
         && !self.left.is::<ExchangeExec>()

Applied locally: 415/415 pass, no changes needed to any test in this PR — they all use Inner.

On the +3.1s band: this is a candidate partial explanation, not a diagnosis. q16 (NOT IN, which not_in_subquery_rewrite turns into a plain anti join) and q22 (NOT EXISTS) are both anti joins and both sit in the band. It doesn't explain q5/q7/q12, which are plain inner joins — those look more like the serialisation cost you already acknowledge. Re-running the six with stage_build_side=false per query would separate the two.

2. Nothing tests the multi-pass path

All the new coverage is single-pass: should_stage_build_side in isolation, one to_actual_join call, exchange_on. The mechanism this PR is actually about — stage → resolve with measured stats → re-decide → CollectLeft — isn't exercised, and neither is the EnsureRequirements/CoalescePartitionsExec behaviour you say you verified by hand.

The harness for exactly this already exists in aqe/test/join_selection.rs (AdaptivePlanner + mock_partitions_with_statistics + assert_plan!). Two assert_plan! tests would cover it: the q8 shape resolving to a broadcast CollectLeft, and the fallthrough where the measured size is still too large and exchange_on reuse yields one exchange per side rather than a nested pair.

Worth noting test_left_join_not_collected_left (the #1055 guard) passes today only because its MemTable stats don't match the staging shape — nothing currently pins "a Left join must not stage."

Smaller

• I'd split the broadcast_join_threshold_rows semantics change out. It's well argued and clearly flagged, but by your own account it isn't what fixes q8, and it replaces a test that asserted the old behaviour deliberately. Separate PRs means either half can be reverted alone if the SF1000 numbers move.
• The exchange guard only inspects immediate children, so a build side shaped like FilterExec(ExchangeExec(..)) gets a second boundary stacked on the first. Edge case.
• CoalescePartitionsRule's docs assume both join legs read upstream shuffle output, and attach the same CoalescePlan to every leaf to keep buckets aligned. In the staged state only one leg is an exchange. It's off by default so the risk is low, but does the leaf grouping handle a one-exchange join?

On defaulting to true

Fix #1 first, then yes. As it stands, default-on ships a path that's provably unprofitable for Full and left-sided semi/anti joins, which is a bad trade for queries that get nothing back. After the guard I'd have no objection.

Benchmark methodology is good — real cluster, row counts cross-checked, regressions disclosed rather than smoothed. Two caveats worth stating inline: one iteration per query, and the current totals include wasted stages on any anti/outer join matching the shape, so the post-fix number should come in a little under 696.09s.

I did not check what happens to a resolved exchange's shuffle_partitions when the executor holding that shuffle is lost. Looks pre-existing for all exchanges rather than something this PR introduces, but staging makes one more stage's output load-bearing.

Thanks for tracing the safety argument rather than taking it on faith — that was the part I most wanted someone to actually walk, and your reading of find_runnable_exchanges and the two Repartitioned entry points matches mine.

Both blocking items are addressed, plus the three inline ones. Pushed.

1. Join-type guard — fixed

You're right, and it's a real waste rather than a cosmetic one. Applied exactly as you wrote it.

One thing worth adding, because it was the only way I could see the guard being too strict: the swap decision can't rescue those cases either. supports_swap_join_order is left_bytes > right_bytes, so swap_inputs == false means the build side is already the smaller input — and measuring it only shrinks it further, since the whole premise is that the estimate overshoots. So the swap can never flip false → true and turn a Left into a broadcastable Right after staging. The guard loses nothing.

Both directions are pinned now, across all ten join types: does_not_stage_a_join_type_that_can_never_broadcast and stages_the_join_types_that_can_broadcast.

2. Multi-pass coverage — added

New file ballista/scheduler/src/state/aqe/test/stage_build_side.rs, using the AdaptivePlanner + StatsTable + assert_plan! harness you pointed at. It needed one new fixture helper (estimated_statistics, for an Inexact total_byte_size — sized_statistics is Exact, which never stages) and a parameterised mock_partitions_with_size, with mock_partitions_with_statistics delegating to it.

Three tests:

• stages_the_build_side_then_broadcasts_the_measured_result — the mechanism end to end. Pass one asserts an ExchangeExec on the build side and a bare scan on the probe; actionable_stages() returns exactly one stage and that stage is the build side alone; then after the shuffle reports its measured size, pass two resolves to HashJoinExec: mode=CollectLeft over broadcast=true, with the fact table still unshuffled.
• reuses_the_staged_exchange_when_the_build_side_measures_too_large — the fallthrough. Asserts one exchange per side, the build's being the resolved staged one (stage_id=0, stage_resolved=true) and the probe's fresh, rather than a nested pair.
• does_not_stage_a_left_join — pins the thing you noted nothing currently pins. You were right that test_left_join_not_collected_left passes only incidentally; its MemTable stats don't match the staging shape. These ones do.

418 passed, 0 failed. Clippy and rustfmt clean.

3. Splitting the row-threshold change — done

Now #2444. Agreed on the reasoning: it isn't what fixes q8, and it deliberately replaces a test that asserted the old semantics, so the two should be revertable independently.

Smaller

Exchange guard depth. Acknowledged, and left as is deliberately — I'd rather not widen it without a shape that actually motivates it. The fire-at-most-once bound is unaffected, because staging always produces a direct ExchangeExec child, which the guard does catch. And in the FilterExec(ExchangeExec(..)) case the staged measurement is of the filter's output, not the inner shuffle's, so it isn't redundant work — just an extra boundary. A subtree-wide search would also decline legitimately useful cases, like a build side that is itself a join over an already-resolved broadcast exchange. Noted as a follow-up in the description; happy to tighten if you think the edge case is likelier than I'm reading it.

CoalescePartitionsRule with a one-exchange join. I traced it. Short answer: the leaf grouping handles it fine, and the risk you're pointing at isn't reachable — but it's safe by scheduling order rather than by construction, which I think is worth recording.

The rule is entirely join-agnostic: it collects leaf ExchangeExecs with no notion of which leg they feed, and the only cardinality guard is leaves.is_empty(). A one-leaf group is already routine (any single-input stage reading one shuffle), and start_indices_to_partition_groups covers [0, M) exactly, so coalescing one leg changes task count and never row count.

The window you're describing — build leaf carrying a plan from a one-leaf pass while the probe leaf carries none — can't occur, because the rule only runs after a full replan. finalise_stage_internal calls replan_stages(), which runs SelectJoinRule and re-decides the join; CoalescePartitionsRule runs later, inside actionable_stages(). So the instant the staged exchange resolves, the join is already re-decided, and by the time the rule looks the plan is either broadcast (it bails on any leaf .broadcast) or two leaves with the probe unresolved (it bails on the unresolved leaf). m is also read from properties().partitioning, which set_coalesce never rewrites, so re-running on an already-coalesced leaf re-derives the same M and the heterogeneous-M guard doesn't misfire.

The thing I'd raise separately, which this PR doesn't introduce but does sit close to: nothing ever clears a CoalescePlan. set_coalesce is the only writer, and both bail paths return without clearing — including the degenerate-K path that says so explicitly in its own debug line. If a leaf ever acquires a plan and the group's byte profile later changes such that the rule bails, the stale K persists on one leg while its partner stays at M. Only the replan-before-rule ordering keeps that from being reachable here. Clearing the slot on the two bail paths, or asserting the group agrees on coalesce() before returning, would make it structural. I'd rather do that in its own PR than widen this one — let me know if you'd prefer it bundled.

Executor loss and a resolved exchange's shuffle_partitions. Agreed it looks pre-existing for all exchanges rather than something staging introduces, and agreed that staging makes one more stage's output load-bearing. I haven't checked it either; not planning to in this PR.

Defaults and benchmark numbers

Keeping stage_build_side at true, per your "fix #1 first, then yes".

Both caveats are now stated inline in the description rather than in a footnote: one iteration per query, and the totals were measured before the join-type guard, so they include stages that were staged and thrown away on anti and outer joins. The post-fix total should land slightly under 696.09s. I haven't re-measured on a cluster yet.

Your partial explanation for the +3.1s band holds up and I've written it into the description: q16 (NOT IN → anti join via not_in_subquery_rewrite) and q22 (NOT EXISTS) are both anti joins, and the guard stops both from staging at all now. It doesn't cover q5/q7/q12, which are plain inner joins — those look like the serialisation cost. Re-running the six with stage_build_side=false per query to separate the two is the next measurement I plan to take.

some s3 tests are failing

its the same minio issue observed in ci: datafusion-cli MinIO docker login failed · Issue #25215 · apache/datafusion
apache/datafusion#25215

UPD: #2446

Thanks both. @milenkovicm for the approval, and @comphead for the second pass — I took essentially all of it. Net −232 lines (386 insertions, 618 deletions) across the branch, and 410 passed / 0 failed, clippy and rustfmt clean.

The load-bearing check: the only snapshots insta accepted were the four in stage_build_side.rs, and their plan shapes and plan_ids are byte-identical to before the refactor. Nothing in test/join_selection.rs or the plan-stability files moved, which is the evidence that this round is behaviour-preserving.

I verified your three structural claims rather than taking them on faith, and all three hold: create_partitioning(self, ..) does consume (so the clone in join_key_partitioning was genuinely dead), partition_mode is Partitioned only when !null_aware, and the local selection_state can only be Unknown when self.selection_state already is — so with_selection_state(Unknown) really was copying eleven fields to change nothing.

Taken

Production. Dropped the join payload (the arm rebuilds from the visited node); build_side: JoinSide → build_is_right: bool, which deletes the unreachable JoinSide::None arm, its error string and the JoinSide import from both files; collapsed the two (Unknown, Partitioned) arms into one with the join-type condition moved inside requires_build_staging (which retires the "caller is responsible for the fourth condition" paragraph); removed with_selection_state and restored to_partitioned's inline body; BuildStagingLimits → plain params; probe_stats hoisted beside build_stats; join_key_partitioning returns a tuple by value over the owned Vec; exchange_on → SelectJoinRule::ensure_exchange.

On the naming: agreed exchange_on read as "on the join keys". The receiver keeps the fetch_add laziness without the generic, as you suggested. I also took the doc line about inputs_already_partitioned — the structural comparison is deliberate, since the only exchange worth reusing is one this rule just placed, and that is now written down so nobody unifies them later.

Tests. Six single-assert staging tests folded into two behind a stages(build, probe, threshold) predicate; the two hand-partitioned join-type lists merged into one loop over all ten with collect_left_broadcast_safe as the oracle (your point that a new DataFusion join type fails neither current list is exactly right, and the merged form does catch it); q8_shaped_action* un-forked from run_to_actual_join, which gained a stage_build_side param; the fourth StatisticsExec builder parameterised with the other three as wrappers; mock_partitions_with_statistics_no_data delegating to mock_partitions_with_size; estimated_statistics using Statistics::to_inexact.

You were right that the fixtures in stage_build_side.rs were a straight copy — ballista_ctx, narrow_schema, register_stats_table and MB now live in test/mod.rs and both files import them, so "the shipped session config" means one thing.

Dropped does_not_stage_a_left_join and exchange_on_wraps_a_non_exchange_child. On the first, your read was better than mine: Left is already a row of the merged loop with the same q8-shaped statistics, so the only unique content was a plan snapshot that needs re-blessing on every unrelated plan_id shift.

Also fixed the shadowed use datafusion::common::{ColumnStatistics, Statistics, stats::Precision}; — good catch, since as you say -D warnings never would have.

Not taken

The exchange-guard depth, deliberately. The fire-at-most-once bound is unaffected, because staging always produces a direct ExchangeExec child, which the guard does catch. In the FilterExec(ExchangeExec(..)) shape the staged measurement is of the filter's output rather than the inner shuffle's, so it is an extra boundary but not redundant work, and a subtree-wide search would decline legitimately useful cases — a build side that is itself a join over an already-resolved broadcast exchange, for one. Noted as a follow-up; happy to tighten if you think the shape is likelier than I am reading it.

Threading one StatisticsContext through to_actual_join. Agreed it is the natural moment, and agreed this PR adds no new statistics computation. But it is a change to a path five callers share, with no test that would catch a memo-sharing mistake, so I would rather it not ride along behind a behaviour-preserving cleanup. Separate PR.

@milenkovicm

Law of Demeter: taken, and it lines up with comphead's item — requires_build_staging now takes min_probe_ratio and max_estimate_multiple as plain params rather than reaching through &BallistaConfig, and broadcast_threshold_bytes is passed the same way, so the mix-and-match is gone.

Config wording: also taken. All three stage_build_side* entries now say explicitly that staging applies only where the build side's size is an estimate rather than a measurement — an exact size over the budget is a fact no further measurement can change. Docs regenerated.

And yes on #1916 and cardinality estimation as the follow-up. Better cardinalities would shrink the window this heuristic exists to cover, which is the right direction.

Benchmarks and docs

I reverted the benchmarking.md change entirely, per your point. It pinned a row to an unmerged PR commit behind an unowned TODO, and it was the only {note} admonition in docs/source. The deeper problem is the one you found in Q1 vs Q6: both are join-free, one improved 4% and the other regressed 7%, so a ±7% noise floor was being published as measured signal next to a claim that this PR accounts for almost all of the improvement.

Those numbers also predate the join-type guard, so they are stale regardless. I will re-measure on the cluster and land the docs separately after merge, with the iteration count and the regression band stated rather than rounded away.

CI

The S3 failures look like the MinIO issue you linked (#2446) rather than anything here.

Merged. Thanks @milenkovicm and @comphead