comphead commented on code in PR #2434:
URL: 
https://github.com/apache/datafusion-ballista/pull/2434#discussion_r3995290061


##########
ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs:
##########
@@ -212,6 +212,16 @@ impl DisplayAs for DynamicJoinSelectionExec {
 }
 
 pub enum JoinSelectionAction {
+    /// Shuffle only the prospective build side, leaving the probe side
+    /// untouched, and re-decide the join once the build side's *measured* size
+    /// is known. See [`requires_build_staging`].
+    StageBuildSide {
+        join: Arc<DynamicJoinSelectionExec>,
+        /// Which child to give an exchange. Only ever [`JoinSide::Left`] or
+        /// [`JoinSide::Right`] — the resolver picks the side the join would
+        /// build from, which is never absent.
+        build_side: JoinSide,
+    },

Review Comment:
   Two things collapse here.
   
   `join` is a field-identical copy of the node `SelectJoinRule` is already 
visiting (see the `with_selection_state` call below), so the payload can go and 
the arm can rebuild from the visited `node`.
   
   `JoinSide` is a three-variant type for a two-state fact, which is what 
forces the unreachable `JoinSide::None` arm in `join_selection.rs`. The 
`BuildSide { Left, Right }` I originally suggested would have made that 
unrepresentable. A plain `build_is_right: bool` does too, and the producer 
below already computes exactly that as `swap_inputs`.
   
   ~13 lines.



##########
ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs:
##########
@@ -601,6 +657,107 @@ impl DynamicJoinSelectionExec {
     }
 }
 
+/// The configured limits that decide whether staging a join's build side on 
its
+/// own is worth an extra round trip. See [`requires_build_staging`].
+#[derive(Debug, Clone, Copy)]
+struct BuildStagingLimits {
+    /// `ballista.optimizer.broadcast_join_threshold_bytes`: the budget a build
+    /// side must come in under to be broadcast. `0` disables broadcast
+    /// promotion, and with it any reason to measure the build side again.
+    broadcast_threshold_bytes: usize,
+    /// `ballista.optimizer.stage_build_side_min_probe_ratio`: how many times
+    /// larger the probe side must be before staging pays off.
+    ///
+    /// Staging serialises the two shuffles that `Repartition` would otherwise
+    /// run concurrently, so the deferral costs at most the *smaller* side's
+    /// runtime. Requiring an order of magnitude keeps that cost well under the
+    /// probe-side shuffle it stands to avoid entirely.
+    min_probe_ratio: usize,
+    /// `ballista.optimizer.stage_build_side_max_estimate_multiple`: how far 
over
+    /// the broadcast budget an *estimated* build side may sit and still be 
worth
+    /// measuring.
+    ///
+    /// TPC-H q8's filtered `part` scan estimates roughly 10x over budget,
+    /// because the planner falls back to `default_filter_selectivity` for
+    /// `p_type = '...'`, and measures three orders of magnitude under it. A 
raw
+    /// fact-table scan sits far beyond this multiple and is shuffled without 
the
+    /// extra round trip.
+    max_estimate_multiple: usize,
+}
+
+impl BuildStagingLimits {
+    fn new(bc: &BallistaConfig, broadcast_threshold_bytes: usize) -> Self {
+        Self {
+            broadcast_threshold_bytes,
+            min_probe_ratio: bc.stage_build_side_min_probe_ratio(),
+            max_estimate_multiple: bc.stage_build_side_max_estimate_multiple(),
+        }
+    }
+}
+
+/// Whether to shuffle only the prospective build side now and re-decide the
+/// join once its measured size is known, rather than shuffling both sides.
+///
+/// [`JoinSelectionAction::Repartition`] shuffles both inputs at once, which 
for
+/// a fact-table probe side commits to the single most expensive stage in the
+/// query *before* any measurement exists. When the build side's size is only a
+/// guess, that guess is the sole reason the join is not a broadcast, and the
+/// probe side dwarfs it, one cheap stage buys an exact number and often flips
+/// the join to `CollectLeft`, leaving the probe side never shuffled at all.
+///
+/// Each of the three conditions is necessary:
+///
+/// * the build estimate must be **inexact**. An exact size over budget is a
+///   fact rather than a guess, so measuring it again cannot change the outcome
+///   and would only serialise two shuffles that could run concurrently.
+/// * the estimate must be **plausibly wrong enough to flip**. Past
+///   [`BuildStagingLimits::max_estimate_multiple`] the side is large on any
+///   reading, and no measurement brings it under budget.
+/// * the probe side must be **much larger**, per
+///   [`BuildStagingLimits::min_probe_ratio`], which is what bounds the cost of
+///   being wrong.

Review Comment:
   These three bullets map 1:1 onto the three early returns immediately below.
   
   Worth keeping: L698-706 (what `Repartition` actually costs, which is the 
real *why*) and L724-725 (why a missing probe size declines, not derivable from 
`else { return false }`).



##########
ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs:
##########
@@ -385,6 +395,41 @@ impl DynamicJoinSelectionExec {
                 JoinInputState::Repartitioned,
                 PartitionMode::Partitioned | PartitionMode::CollectLeft,
             ) => self.to_sort_merge_join().map(JoinSelectionAction::Sort),
+            (JoinInputState::Unknown, PartitionMode::Partitioned)
+                if bc.stage_build_side_enabled()
+                    && !self.null_aware
+                    // Measuring the build side can only change the decision if
+                    // `CollectLeft` is reachable at all, and that is a 
property
+                    // of the (post-swap) join type alone — no measurement 
makes
+                    // a `Full`, or a left-sided semi/anti/outer, join
+                    // broadcastable. `partition_mode` above is `CollectLeft`
+                    // only when this same predicate holds, so without it here
+                    // staging would buy a serialised stage boundary and no
+                    // decision.
+                    && collect_left_broadcast_safe(build_side_join_type)
+                    // Firing at most once per join is what keeps this from
+                    // ping-ponging: once it fires, one child is an
+                    // `ExchangeExec` and this guard never passes again.
+                    && !self.left.is::<ExchangeExec>()
+                    && !self.right.is::<ExchangeExec>()
+                    && requires_build_staging(
+                        build_stats,
+                        if swap_inputs { &stats_left } else { &stats_right },
+                        BuildStagingLimits::new(
+                            &bc,
+                            threshold_collect_left_join_bytes,
+                        ),
+                    ) =>
+            {
+                Ok(JoinSelectionAction::StageBuildSide {
+                    join: 
Arc::new(self.with_selection_state(JoinInputState::Unknown)),
+                    build_side: if swap_inputs {

Review Comment:
   `with_selection_state(Unknown)` is a no-op here. `JoinInputState` has two 
variants, and this arm is reachable only when the local `selection_state` 
(L366-372) is `Unknown`, which requires `self.selection_state == Unknown` 
already. So this copies all 11 fields and changes nothing.
   
   Combined with dropping the payload, `with_selection_state` can go entirely 
and `to_partitioned` can return to its inline body. Caller counts today: 
`to_partitioned` 1, `with_selection_state` 2, one of which is `to_partitioned` 
itself.



##########
ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs:
##########
@@ -385,6 +395,41 @@ impl DynamicJoinSelectionExec {
                 JoinInputState::Repartitioned,
                 PartitionMode::Partitioned | PartitionMode::CollectLeft,
             ) => self.to_sort_merge_join().map(JoinSelectionAction::Sort),
+            (JoinInputState::Unknown, PartitionMode::Partitioned)
+                if bc.stage_build_side_enabled()
+                    && !self.null_aware
+                    // Measuring the build side can only change the decision if
+                    // `CollectLeft` is reachable at all, and that is a 
property
+                    // of the (post-swap) join type alone — no measurement 
makes
+                    // a `Full`, or a left-sided semi/anti/outer, join
+                    // broadcastable. `partition_mode` above is `CollectLeft`
+                    // only when this same predicate holds, so without it here
+                    // staging would buy a serialised stage boundary and no
+                    // decision.
+                    && collect_left_broadcast_safe(build_side_join_type)
+                    // Firing at most once per join is what keeps this from
+                    // ping-ponging: once it fires, one child is an
+                    // `ExchangeExec` and this guard never passes again.
+                    && !self.left.is::<ExchangeExec>()
+                    && !self.right.is::<ExchangeExec>()
+                    && requires_build_staging(
+                        build_stats,
+                        if swap_inputs { &stats_left } else { &stats_right },

Review Comment:
   Three things in one guard.
   
   - `!self.null_aware` (L400) is already implied. `partition_mode` is `if 
self.null_aware || (..) { CollectLeft } else { Partitioned }`, so matching 
`Partitioned` guarantees it.
   - The 8-line comment at L401-408 sits inside a match *pattern*. The same 
rationale is already on `requires_build_staging`, on the join-type test, and in 
`stage_build_side.rs`. Keep one copy.
   - `if swap_inputs { &stats_left } else { &stats_right }` (L417) is the third 
`swap_inputs` ternary in this function. L292 and L328 already do it for 
`build_side` and `build_stats`.
   
   Structural: this arm and the fallthrough at L433 have the **identical** 
pattern, separated only by the guard, so a reader scans 25 lines past the 
pattern before learning what it matches. Hoisting `probe_stats` next to 
`build_stats` and folding the conditions into one `should_stage_build_side(..)` 
predicate collapses both into one arm, and retires the "the caller is 
responsible for the fourth condition" paragraph at L720.



##########
ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs:
##########
@@ -1166,6 +1323,306 @@ mod tests {
         );
     }
 
+    /// A source reporting `total_byte_size` with the given precision, shaped
+    /// like one side of the TPC-H q8 `part`/`lineitem` join.
+    fn sized_stats_exec(total_byte_size: Precision<usize>) -> Arc<dyn 
ExecutionPlan> {
+        Arc::new(StatisticsExec::new(
+            Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size,
+                column_statistics: vec![ColumnStatistics::new_unknown()],
+            },
+            Schema::new(vec![Field::new("k", DataType::Int32, false)]),
+        ))
+    }

Review Comment:
   Fourth `StatisticsExec` builder in this module, alongside 
`sizeless_stats_exec` (L919), `stats_exec` (L1055) and `stats_exec_rows_only` 
(L1068). All four wrap the same one-column `Int32` schema and differ only in 
how `total_byte_size` is filled.
   
   One parameterised pair, with the three existing ones as thin wrappers, would 
also let the six tests below drop their inline `Statistics { .. }` literals.



##########
ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs:
##########
@@ -294,31 +332,64 @@ impl PhysicalOptimizerRule for SelectJoinRule {
 
                                 Ok(Transformed::yes(exec.data))
                             }
-                            JoinSelectionAction::Repartition(dynamic_join) => {
-                                let partition_count = 
config.execution.target_partitions;
-                                let partitioning = dynamic_join
-                                    ._required_input_distribution()
-                                    .iter()
-                                    .map(|d| {
-                                        
d.clone().create_partitioning(partition_count)
-                                    })
-                                    .collect::<Vec<_>>();
-
-                                let left = dynamic_join.left.clone();
-                                let right = dynamic_join.right.clone();
-
-                                let left = Arc::new(ExchangeExec::new(
-                                    left,
-                                    Some(partitioning[0].clone()),
-                                    self.plan_id(),
-                                ));
+                            JoinSelectionAction::StageBuildSide { join, 
build_side } => {
+                                // Only the build side gets an exchange. The
+                                // probe side is left as it is, so no stage is
+                                // created for it and nothing of it is shuffled
+                                // until the join is decided for real.
+                                let partitioning = join_key_partitioning(
+                                    &join,
+                                    config.execution.target_partitions,
+                                );
+                                let build_idx = match build_side {
+                                    JoinSide::Left => 0,
+                                    JoinSide::Right => 1,
+                                    JoinSide::None => {
+                                        return Err(DataFusionError::Internal(
+                                            "StageBuildSide requires a build 
side"
+                                                .to_owned(),
+                                        ));
+                                    }
+                                };

Review Comment:
   Unreachable by construction. The producer computes this from a `bool` two 
lines earlier, and the field's own doc says it is "never absent". Switching the 
variant to `build_is_right: bool`, or to a two-variant `BuildSide`, deletes 
this arm, the error string and the three-line doc, and drops the `JoinSide` 
import from both files.
   
   If `JoinSide` stays: the repo precedent at 
`physical_optimizer/join_selection.rs:517` is `unreachable!()`, and the sibling 
arms in this same match use `DataFusionError::Execution`, not `Internal`.



##########
ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs:
##########
@@ -1166,6 +1323,306 @@ mod tests {
         );
     }
 
+    /// A source reporting `total_byte_size` with the given precision, shaped
+    /// like one side of the TPC-H q8 `part`/`lineitem` join.
+    fn sized_stats_exec(total_byte_size: Precision<usize>) -> Arc<dyn 
ExecutionPlan> {
+        Arc::new(StatisticsExec::new(
+            Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size,
+                column_statistics: vec![ColumnStatistics::new_unknown()],
+            },
+            Schema::new(vec![Field::new("k", DataType::Int32, false)]),
+        ))
+    }
+
+    /// The q8 shape: a filtered dimension scan whose size is a guess sitting
+    /// just over the broadcast budget, against a fact table that dwarfs it.
+    const STAGE_THRESHOLD: usize = 10 * MB;
+    const GUESSED_BUILD_BYTES: usize = 100 * MB;
+    const FACT_PROBE_BYTES: usize = 2000 * MB;
+
+    /// The staging limits for a given byte budget, carrying the shipped ratio
+    /// and multiple. Reading those from `BallistaConfig::default()` rather 
than
+    /// restating them means these tests pin the values a deployment runs with.
+    fn stage_limits(broadcast_threshold_bytes: usize) -> BuildStagingLimits {
+        BuildStagingLimits::new(&BallistaConfig::default(), 
broadcast_threshold_bytes)
+    }
+
+    #[test]
+    fn stages_the_build_side_when_its_size_is_only_a_guess() {
+        assert!(requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(FACT_PROBE_BYTES),
+                column_statistics: vec![],
+            },
+            stage_limits(STAGE_THRESHOLD),
+        ));
+    }
+
+    // An exact size over budget is a fact, not a guess. Measuring it again
+    // cannot flip the join, so both sides shuffle concurrently as before.
+    #[test]
+    fn does_not_stage_an_exactly_sized_build_side() {
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Exact(1_000),
+                total_byte_size: Precision::Exact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(FACT_PROBE_BYTES),
+                column_statistics: vec![],
+            },
+            stage_limits(STAGE_THRESHOLD),
+        ));
+    }
+
+    // Far enough over the budget that no measurement brings it back under, so
+    // the extra round trip would only add latency.
+    #[test]
+    fn does_not_stage_a_build_side_far_past_the_budget() {
+        let limits = stage_limits(STAGE_THRESHOLD);
+        let far_over = STAGE_THRESHOLD * (limits.max_estimate_multiple + 1);
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(far_over),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(far_over * 1_000),
+                column_statistics: vec![],
+            },
+            limits,
+        ));
+    }
+
+    // Without a probe side much larger than the build side, serialising the 
two
+    // shuffles costs about as much as it could save.
+    #[test]
+    fn does_not_stage_when_the_probe_side_is_comparable() {
+        let limits = stage_limits(STAGE_THRESHOLD);
+        let probe_bytes = GUESSED_BUILD_BYTES * (limits.min_probe_ratio - 1);
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(probe_bytes),
+                column_statistics: vec![],
+            },
+            limits,
+        ));
+    }
+
+    // No probe-side size is no evidence the round trip pays off.
+    #[test]
+    fn does_not_stage_without_a_probe_side_size() {
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Absent,
+                column_statistics: vec![],
+            },
+            stage_limits(STAGE_THRESHOLD),
+        ));
+    }
+
+    // With broadcast promotion off there is no decision to revisit.
+    #[test]
+    fn does_not_stage_when_broadcast_promotion_is_disabled() {
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(FACT_PROBE_BYTES),
+                column_statistics: vec![],
+            },
+            stage_limits(0),
+        ));
+    }
+
+    // The shipped defaults are what the benchmark numbers were measured under,
+    // so pin them here rather than only in the generated config docs.
+    #[test]
+    fn stage_build_side_ships_the_documented_defaults() {
+        let bc = BallistaConfig::default();
+
+        assert!(bc.stage_build_side_enabled());
+        assert_eq!(bc.stage_build_side_min_probe_ratio(), 10);
+        assert_eq!(bc.stage_build_side_max_estimate_multiple(), 32);
+    }
+
+    /// Runs the resolver over the q8 shape, with staging on or off.
+    fn q8_shaped_action(stage_build_side: bool) -> JoinSelectionAction {
+        q8_shaped_action_with_children(
+            sized_stats_exec(Precision::Inexact(GUESSED_BUILD_BYTES)),
+            sized_stats_exec(Precision::Exact(FACT_PROBE_BYTES)),
+            stage_build_side,
+        )
+    }
+
+    fn q8_shaped_action_with_children(
+        left: Arc<dyn ExecutionPlan>,
+        right: Arc<dyn ExecutionPlan>,
+        stage_build_side: bool,
+    ) -> JoinSelectionAction {
+        q8_shaped_action_with_join_type(left, right, JoinType::Inner, 
stage_build_side)
+    }
+
+    fn q8_shaped_action_with_join_type(
+        left: Arc<dyn ExecutionPlan>,
+        right: Arc<dyn ExecutionPlan>,
+        join_type: JoinType,
+        stage_build_side: bool,
+    ) -> JoinSelectionAction {
+        let on: JoinOn =
+            vec![(Arc::new(Column::new("k", 0)), Arc::new(Column::new("k", 
0)))];
+        let dj = DynamicJoinSelectionExec {
+            properties: Arc::clone(left.properties()),
+            left,
+            right,
+            on,
+            filter: None,
+            join_type,
+            projection: None,
+            null_equality: NullEquality::NullEqualsNothing,
+            selection_state: JoinInputState::Unknown,
+            null_aware: false,
+            plan_id: 0,
+        };
+        let mut config = ConfigOptions::new();
+        let mut bc = BallistaConfig::default();
+        bc.set(
+            "optimizer.broadcast_join_threshold_bytes",
+            &STAGE_THRESHOLD.to_string(),
+        )
+        .unwrap();
+        bc.set("optimizer.stage_build_side", &stage_build_side.to_string())
+            .unwrap();
+        config.extensions.insert(bc);
+        dj.to_actual_join(&config).unwrap()
+    }
+
+    // The whole point: a guessed-large build side against a fact-table probe
+    // side stages the build alone instead of committing to both shuffles.

Review Comment:
   This forks `run_to_actual_join` (L1081), documented in-file as *"Core 
driver: runs the join-strategy decision for the given children."* Same `on: 
JoinOn`, same 11-field struct literal, same `ConfigOptions` / `BallistaConfig` 
/ `insert` / `to_actual_join` plumbing. Only the keys being `set` differ, so 
adding a field to `DynamicJoinSelectionExec` now breaks four test constructors 
instead of fewer.
   
   Give `run_to_actual_join` a `stage_build_side` param and delete the fork, 
along with the middle link `q8_shaped_action_with_children` (L1486), which 
exists only to default `join_type` to `Inner` for two callers. ~42 lines.



##########
ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs:
##########
@@ -391,6 +462,69 @@ mod tests {
         physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner},
     };
 
+    fn key_partitioning(name: &str, schema: &Schema) -> Partitioning {
+        Partitioning::Hash(
+            vec![Arc::new(
+                
datafusion::physical_plan::expressions::Column::new_with_schema(
+                    name, schema,
+                )
+                .unwrap(),
+            )],
+            4,
+        )
+    }
+
+    // A build side that `StageBuildSide` already shuffled on the join key must
+    // not be shuffled again when the join later falls back to `Repartition`.
+    // Wrapping it a second time would nest one stage boundary inside another
+    // and move the same rows twice.
+    #[test]
+    fn exchange_on_reuses_a_matching_exchange() {
+        let schema = Schema::new(vec![Field::new("k", DataType::Int32, 
false)]);
+        let partitioning = key_partitioning("k", &schema);
+        let staged: Arc<dyn ExecutionPlan> = Arc::new(ExchangeExec::new(
+            Arc::new(EmptyExec::new(Arc::new(schema))),
+            Some(partitioning.clone()),
+            0,
+        ));
+
+        let reused = exchange_on(Arc::clone(&staged), &partitioning, || 1);
+
+        assert!(Arc::ptr_eq(&staged, &reused));
+    }
+
+    // An exchange on a different key is not reusable, so the join key's own
+    // exchange still has to be added on top.
+    #[test]
+    fn exchange_on_wraps_an_exchange_on_a_different_key() {
+        let schema = Schema::new(vec![
+            Field::new("k", DataType::Int32, false),
+            Field::new("other", DataType::Int32, false),
+        ]);
+        let staged: Arc<dyn ExecutionPlan> = Arc::new(ExchangeExec::new(
+            Arc::new(EmptyExec::new(Arc::new(schema.clone()))),
+            Some(key_partitioning("other", &schema)),
+            0,
+        ));
+
+        let wrapped = exchange_on(staged, &key_partitioning("k", &schema), || 
1);
+        let outer = wrapped.downcast_ref::<ExchangeExec>().unwrap();
+
+        assert!(outer.input().is::<ExchangeExec>());
+    }
+
+    // A plain child always gets an exchange.
+    #[test]
+    fn exchange_on_wraps_a_non_exchange_child() {
+        let schema = Schema::new(vec![Field::new("k", DataType::Int32, 
false)]);
+        let child: Arc<dyn ExecutionPlan> =
+            Arc::new(EmptyExec::new(Arc::new(schema.clone())));
+
+        let wrapped = exchange_on(child, &key_partitioning("k", &schema), || 
1);
+
+        assert!(wrapped.is::<ExchangeExec>());
+    }

Review Comment:
   This asserts that `Arc::new(ExchangeExec::new(..))` is an `ExchangeExec`. 
The unconditional-wrap behaviour predates the PR and is pinned end-to-end by 
four snapshots in `test/join_selection.rs` (L190, L217, L251, L278).
   
   Keep the other two: `exchange_on_reuses_a_matching_exchange` guards the 
load-bearing new branch, and `exchange_on_wraps_an_exchange_on_a_different_key` 
guards a case no integration test can produce.



##########
ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs:
##########
@@ -134,6 +135,43 @@ impl SelectJoinRule {
     }
 }
 
+/// The per-child shuffle partitioning a join's inputs need to be 
co-partitioned
+/// on its keys.
+fn join_key_partitioning(
+    join: &DynamicJoinSelectionExec,
+    partition_count: usize,
+) -> Vec<Partitioning> {
+    join._required_input_distribution()
+        .iter()
+        .map(|d| d.clone().create_partitioning(partition_count))
+        .collect()
+}

Review Comment:
   The clone here is dead. `_required_input_distribution()` returns an 
**owned** `Vec<Distribution>` (`dynamic_join.rs:236`), and 
`create_partitioning(self, ..)` **consumes** it 
(`datafusion-physical-expr-55.0.0/src/partitioning.rs:713`). `.into_iter()` 
removes two `Vec<Arc<dyn PhysicalExpr>>` allocations and 2N refcount 
round-trips.
   
   While here: the `StageBuildSide` arm builds this two-element `Vec` and uses 
exactly one element, then `exchange_on` clones the `&Partitioning` again at 
L170. Returning `(Partitioning, Partitioning)` and taking it **by value** in 
`exchange_on` removes all three copies.
   
   The clone predates the PR, but extracting the helper is the moment to drop 
it.



##########
ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs:
##########
@@ -294,31 +332,64 @@ impl PhysicalOptimizerRule for SelectJoinRule {
 
                                 Ok(Transformed::yes(exec.data))
                             }
-                            JoinSelectionAction::Repartition(dynamic_join) => {
-                                let partition_count = 
config.execution.target_partitions;
-                                let partitioning = dynamic_join
-                                    ._required_input_distribution()
-                                    .iter()
-                                    .map(|d| {
-                                        
d.clone().create_partitioning(partition_count)
-                                    })
-                                    .collect::<Vec<_>>();
-
-                                let left = dynamic_join.left.clone();
-                                let right = dynamic_join.right.clone();
-
-                                let left = Arc::new(ExchangeExec::new(
-                                    left,
-                                    Some(partitioning[0].clone()),
-                                    self.plan_id(),
-                                ));
+                            JoinSelectionAction::StageBuildSide { join, 
build_side } => {
+                                // Only the build side gets an exchange. The
+                                // probe side is left as it is, so no stage is
+                                // created for it and nothing of it is shuffled
+                                // until the join is decided for real.
+                                let partitioning = join_key_partitioning(
+                                    &join,
+                                    config.execution.target_partitions,
+                                );
+                                let build_idx = match build_side {
+                                    JoinSide::Left => 0,
+                                    JoinSide::Right => 1,
+                                    JoinSide::None => {
+                                        return Err(DataFusionError::Internal(
+                                            "StageBuildSide requires a build 
side"
+                                                .to_owned(),
+                                        ));
+                                    }
+                                };
 
-                                let right = Arc::new(ExchangeExec::new(
-                                    right,
-                                    Some(partitioning[1].clone()),
+                                let mut children =
+                                    vec![join.left.clone(), 
join.right.clone()];
+                                children[build_idx] = 
Arc::new(ExchangeExec::new(
+                                    children[build_idx].clone(),
+                                    Some(partitioning[build_idx].clone()),
                                     self.plan_id(),
                                 ));
 
+                                let join = join.replace_children(
+                                    children,
+                                    ReplaceChildrenOptions::new(
+                                        ChildrenPropertiesMode::Recompute,
+                                    ),
+                                )?;
+
+                                Ok(Transformed::yes(join))
+                            }
+                            JoinSelectionAction::Repartition(dynamic_join) => {
+                                let partitioning = join_key_partitioning(
+                                    &dynamic_join,
+                                    config.execution.target_partitions,
+                                );
+
+                                // A build side already staged by
+                                // `StageBuildSide` carries an exchange on the
+                                // join key, so reuse it rather than shuffling
+                                // that side a second time.

Review Comment:
   This says what `exchange_on`'s own doc at L153-157 already says, where the 
behaviour actually lives. The same rationale appears a third time at L475-478.



##########
docs/source/contributors-guide/benchmarking.md:
##########
@@ -135,29 +147,29 @@ mean of 3 iterations (cold iteration dropped by the 
harness).
 
 |     Query | Ballista (s) | Spark 3.4 (s) |
 | --------: | -----------: | ------------: |
-|         1 |        17.56 |         67.58 |
-|         2 |        27.13 |         29.80 |
-|         3 |        33.03 |         25.13 |
-|         4 |        18.34 |         21.19 |
-|         5 |        60.51 |         54.12 |
-|         6 |        14.25 |          1.23 |
-|         7 |        49.28 |         19.57 |
-|         8 |        96.02 |         48.60 |
-|         9 |       107.89 |         69.38 |
-|        10 |        55.78 |         35.92 |
-|        11 |        13.46 |         30.88 |
-|        12 |        16.49 |         10.78 |
-|        13 |        14.58 |         20.45 |
-|        14 |        17.99 |          7.00 |
-|        15 |        18.46 |         23.75 |
-|        16 |        14.29 |         23.41 |
-|        17 |        35.20 |         82.30 |
-|        18 |        53.64 |        129.40 |
-|        19 |        18.87 |         11.26 |
-|        20 |        29.29 |         19.22 |
-|        21 |        95.63 |        101.53 |
-|        22 |         9.37 |         12.71 |
-| **Total** |   **817.07** |    **845.21** |
+|         1 |        16.81 |         67.58 |
+|         2 |        30.24 |         29.80 |
+|         3 |        31.01 |         25.13 |
+|         4 |        20.50 |         21.19 |
+|         5 |        63.62 |         54.12 |
+|         6 |        15.28 |          1.23 |
+|         7 |        52.46 |         19.57 |
+|         8 |        40.68 |         48.60 |
+|         9 |        51.94 |         69.38 |
+|        10 |        56.01 |         35.92 |
+|        11 |        18.24 |         30.88 |
+|        12 |        19.60 |         10.78 |
+|        13 |        14.19 |         20.45 |
+|        14 |        15.70 |          7.00 |
+|        15 |        19.30 |         23.75 |
+|        16 |        17.40 |         23.41 |
+|        17 |        31.15 |         82.30 |
+|        18 |        51.41 |        129.40 |
+|        19 |        18.99 |         11.26 |
+|        20 |        33.49 |         19.22 |
+|        21 |        65.69 |        101.53 |
+|        22 |        12.40 |         12.71 |
+| **Total** |   **696.09** |    **845.21** |

Review Comment:
   Q6 is join-free (single-table scan, filter, aggregate on `lineitem`) and 
regressed 7% (14.25 to 15.28), while the equally join-free Q1 improved 4%. At 
one iteration per query that is a noise floor of roughly plus or minus 7% being 
published as measured signal, right next to a note saying this PR "accounts for 
almost all of the improvement".
   
   The three headline wins are an order of magnitude above that band and 
survive it fine. The 13 small regressions do not, and the doc says nothing 
about them while the PR description does.



##########
ballista/scheduler/src/state/aqe/test/stage_build_side.rs:
##########
@@ -0,0 +1,229 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Staging a join's build side, across the two passes it takes to pay off.
+//!
+//! The unit tests in `dynamic_join` pin the *decision* in isolation: given
+//! these statistics, does the resolver choose to stage? What they cannot show
+//! is the mechanism the feature exists for, which only appears across two
+//! passes — shuffle the build side alone, read back what that shuffle actually
+//! measured, then decide the join for real against a number instead of a 
guess.
+//! These tests drive `AdaptivePlanner` through both passes and assert on the
+//! plan that comes out of each.
+//!
+//! Sizes are declared ([`StatsTable`]) rather than materialised, because the
+//! shape staging turns on is a build side whose size is an *estimate* sitting
+//! over the broadcast budget. No table a test is willing to build produces 
that.
+
+use crate::state::aqe::planner::AdaptivePlanner;
+use crate::state::aqe::test::stats_table::{
+    StatsTable, estimated_statistics, sized_statistics,
+};
+use crate::state::aqe::test::{
+    mock_partitions_with_size, mock_partitions_with_statistics,
+};
+use ballista_core::assert_plan;
+use ballista_core::extension::SessionConfigExt;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::common::Statistics;
+use datafusion::execution::{
+    SessionStateBuilder, config::SessionConfig, context::SessionContext,
+};
+use std::sync::Arc;
+
+const MB: usize = 1024 * 1024;
+
+/// A build side whose size is only a guess, sitting over the shipped 128 MB
+/// broadcast budget but well inside the 32x window past which no measurement
+/// could bring it back under. This is the TPC-H q8 `part` shape: a filtered
+/// dimension scan the planner sizes with `default_filter_selectivity`.
+const GUESSED_BUILD_BYTES: usize = 200 * MB;
+
+/// A fact-table probe side, far past the 10x the ratio requires.
+const FACT_PROBE_BYTES: usize = 40 * 1024 * MB;
+
+/// A measured build size that still cannot be broadcast, for the fallthrough.
+const MEASURED_TOO_LARGE_BYTES: u64 = (200 * MB) as u64;
+
+fn join_schema() -> Arc<Schema> {
+    Arc::new(Schema::new(vec![
+        Field::new("id", DataType::Int32, false),
+        Field::new("val", DataType::Int32, false),
+    ]))
+}
+
+/// A context carrying Ballista's shipped configuration, so these tests run
+/// against the thresholds a deployment uses — including
+/// `ballista.optimizer.stage_build_side`, which defaults to on.
+fn ballista_ctx() -> SessionContext {
+    let config = SessionConfig::new_with_ballista()
+        .with_target_partitions(4)
+        .with_round_robin_repartition(false);
+    let state = SessionStateBuilder::new_with_default_features()
+        .with_config(config)
+        .build();
+    SessionContext::new_with_state(state)
+}
+
+fn register(ctx: &SessionContext, name: &str, stats: Statistics) {
+    ctx.register_table(name, Arc::new(StatsTable::new(join_schema(), stats, 
4)))
+        .unwrap();
+}
+
+/// The q8 shape: a guessed-large dimension side against a fact table that
+/// dwarfs it, joined on the key.
+fn q8_shaped_ctx() -> SessionContext {
+    let ctx = ballista_ctx();
+    let schema = join_schema();
+    register(
+        &ctx,
+        "dim",
+        estimated_statistics(&schema, 1_000_000, GUESSED_BUILD_BYTES),
+    );
+    register(
+        &ctx,
+        "fact",
+        sized_statistics(&schema, 1_000_000_000, FACT_PROBE_BYTES),
+    );
+    ctx
+}
+
+async fn planner_for(ctx: &SessionContext, sql: &str) -> AdaptivePlanner {
+    let lp = ctx.sql(sql).await.unwrap().into_optimized_plan().unwrap();
+    AdaptivePlanner::try_new(ctx, &lp, "test_job".into())
+        .await
+        .unwrap()
+}
+
+const Q8_SHAPED_JOIN: &str = "SELECT dim.val FROM dim JOIN fact ON dim.id = 
fact.id";
+
+/// The mechanism end to end: the first pass shuffles only the build side and
+/// leaves the fact table alone, and once that one cheap stage reports what it
+/// measured, the join resolves to a broadcast `CollectLeft` — so the fact 
table
+/// is never shuffled at all.
+#[tokio::test]
+async fn stages_the_build_side_then_broadcasts_the_measured_result() {
+    let ctx = q8_shaped_ctx();
+    let mut planner = planner_for(&ctx, Q8_SHAPED_JOIN).await;
+
+    // Pass one: an exchange on the build side only. The fact table is still a
+    // bare scan, so no stage is created for it.
+    assert_plan!(planner.current_plan(), @ "
+    AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending, 
stage_resolved=false
+      ProjectionExec: expr=[val@1 as val]
+        DynamicJoinSelectionExec: plan_id=0, join_type=Inner, on=[(id@0, 
id@0)] repartitioned=false
+          ExchangeExec: partitioning=Hash([id@0], 4), plan_id=1, 
stage_id=pending, stage_resolved=false
+            CooperativeExec
+              StatsExec: partitions=4, rows=Inexact(1000000), 
bytes=Inexact(209715200)
+          CooperativeExec
+            StatsExec: partitions=4, rows=Exact(1000000000), 
bytes=Exact(42949672960)
+    ");
+
+    let (stages, cancellable) = planner.actionable_stages().unwrap();
+    let stages = stages.unwrap();
+    assert_eq!(1, stages.len(), "only the build side should be staged");
+    assert_eq!(0, cancellable.len());
+
+    // And that one stage is the build side alone.
+    assert_plan!(stages.first().unwrap().plan.as_ref(), @ "
+    SortShuffleWriterExec: partitioning=Hash([id@0], 4)
+      CooperativeExec
+        StatsExec: partitions=4, rows=Inexact(1000000), 
bytes=Inexact(209715200)
+    ");
+
+    // The staged shuffle reports 10 bytes, three orders of magnitude under the
+    // estimate that kept the join partitioned.
+    planner
+        .finalise_stage_internal(0, mock_partitions_with_statistics())
+        .unwrap();
+
+    // Pass two: decided against the measurement, not the guess.
+    assert_plan!(planner.current_plan(), @ "
+    AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending, 
stage_resolved=false
+      HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)], 
projection=[val@1]
+        ExchangeExec: partitioning=None, plan_id=3, stage_id=0, 
stage_resolved=true, broadcast=true
+          CooperativeExec
+            StatsExec: partitions=4, rows=Inexact(1000000), 
bytes=Inexact(209715200)
+        CooperativeExec
+          StatsExec: partitions=4, rows=Exact(1000000000), 
bytes=Exact(42949672960)
+    ");
+}
+
+/// The other half of the decision: when the measurement confirms the build 
side
+/// really is too large, the join falls back to a partitioned one — and the
+/// staged exchange is *reused* rather than wrapped, so each side ends up 
behind
+/// exactly one exchange instead of a nested pair.
+#[tokio::test]
+async fn reuses_the_staged_exchange_when_the_build_side_measures_too_large() {
+    let ctx = q8_shaped_ctx();
+    let mut planner = planner_for(&ctx, Q8_SHAPED_JOIN).await;
+
+    let (stages, _) = planner.actionable_stages().unwrap();
+    assert_eq!(1, stages.unwrap().len());
+
+    // This time the shuffle confirms the build side is over the budget.
+    planner
+        .finalise_stage_internal(
+            0,
+            mock_partitions_with_size(1_000_000, MEASURED_TOO_LARGE_BYTES),
+        )
+        .unwrap();
+
+    // One exchange per side. The build side's is the staged one, already on 
the
+    // join key and already resolved; nothing is nested inside anything.
+    assert_plan!(planner.current_plan(), @ "
+    AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending, 
stage_resolved=false
+      ProjectionExec: expr=[val@1 as val]
+        DynamicJoinSelectionExec: plan_id=0, join_type=Inner, on=[(id@0, 
id@0)] repartitioned=true
+          ExchangeExec: partitioning=Hash([id@0], 4), plan_id=1, stage_id=0, 
stage_resolved=true
+            CooperativeExec
+              StatsExec: partitions=4, rows=Inexact(1000000), 
bytes=Inexact(209715200)
+          ExchangeExec: partitioning=Hash([id@0], 4), plan_id=3, 
stage_id=pending, stage_resolved=false
+            CooperativeExec
+              StatsExec: partitions=4, rows=Exact(1000000000), 
bytes=Exact(42949672960)
+    ");
+}
+
+/// A `Left` join can never be lowered to `CollectLeft` — the build side would
+/// emit its unmatched rows once per probe task (#1055) — so measuring it 
cannot
+/// change the decision, and staging would buy a serialised stage boundary for
+/// nothing. Both sides must be shuffled on the first pass.
+///
+/// Worth pinning separately: `test_left_join_not_collected_left` covers the
+/// broadcast half of this, but passes today only because its statistics do not
+/// match the staging shape. These ones do.
+#[tokio::test]
+async fn does_not_stage_a_left_join() {
+    let ctx = q8_shaped_ctx();
+    let planner = planner_for(
+        &ctx,
+        "SELECT dim.val FROM dim LEFT JOIN fact ON dim.id = fact.id",
+    )
+    .await;
+
+    assert_plan!(planner.current_plan(), @ "
+    AdaptiveDatafusionExec: is_final=false, plan_id=3, stage_id=pending, 
stage_resolved=false
+      ProjectionExec: expr=[val@1 as val]
+        DynamicJoinSelectionExec: plan_id=0, join_type=Left, on=[(id@0, id@0)] 
repartitioned=true
+          ExchangeExec: partitioning=Hash([id@0], 4), plan_id=1, 
stage_id=pending, stage_resolved=false
+            CooperativeExec
+              StatsExec: partitions=4, rows=Inexact(1000000), 
bytes=Inexact(209715200)
+          ExchangeExec: partitioning=Hash([id@0], 4), plan_id=2, 
stage_id=pending, stage_resolved=false
+            CooperativeExec
+              StatsExec: partitions=4, rows=Exact(1000000000), 
bytes=Exact(42949672960)
+    ");
+}

Review Comment:
   The justification in the doc does not quite hold up. 
`test_left_join_not_collected_left` (`test/join_selection.rs:190-198`) already 
snapshots the identical two-exchange `Left` shape, and `JoinType::Left` is 
already a case in `does_not_stage_a_join_type_that_can_never_broadcast`. The 
genuinely unique content is "with q8-shaped statistics too", which is one row 
of the merged loop.
   
   Judgment call rather than a clear cut, since `broadcast_thresholds.rs` does 
pair integration positives and negatives. But it is a ~30-line plan snapshot 
that will need re-blessing on every unrelated `plan_id` shift.



##########
ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs:
##########
@@ -601,6 +657,107 @@ impl DynamicJoinSelectionExec {
     }
 }
 
+/// The configured limits that decide whether staging a join's build side on 
its
+/// own is worth an extra round trip. See [`requires_build_staging`].
+#[derive(Debug, Clone, Copy)]
+struct BuildStagingLimits {
+    /// `ballista.optimizer.broadcast_join_threshold_bytes`: the budget a build
+    /// side must come in under to be broadcast. `0` disables broadcast
+    /// promotion, and with it any reason to measure the build side again.
+    broadcast_threshold_bytes: usize,
+    /// `ballista.optimizer.stage_build_side_min_probe_ratio`: how many times
+    /// larger the probe side must be before staging pays off.
+    ///
+    /// Staging serialises the two shuffles that `Repartition` would otherwise
+    /// run concurrently, so the deferral costs at most the *smaller* side's
+    /// runtime. Requiring an order of magnitude keeps that cost well under the
+    /// probe-side shuffle it stands to avoid entirely.
+    min_probe_ratio: usize,
+    /// `ballista.optimizer.stage_build_side_max_estimate_multiple`: how far 
over
+    /// the broadcast budget an *estimated* build side may sit and still be 
worth
+    /// measuring.
+    ///
+    /// TPC-H q8's filtered `part` scan estimates roughly 10x over budget,
+    /// because the planner falls back to `default_filter_selectivity` for
+    /// `p_type = '...'`, and measures three orders of magnitude under it. A 
raw
+    /// fact-table scan sits far beyond this multiple and is shuffled without 
the
+    /// extra round trip.
+    max_estimate_multiple: usize,
+}
+
+impl BuildStagingLimits {
+    fn new(bc: &BallistaConfig, broadcast_threshold_bytes: usize) -> Self {
+        Self {
+            broadcast_threshold_bytes,
+            min_probe_ratio: bc.stage_build_side_min_probe_ratio(),
+            max_estimate_multiple: bc.stage_build_side_max_estimate_multiple(),
+        }
+    }
+}

Review Comment:
   37 lines to bundle three `usize`s for one call site, and it diverges from 
the convention in this same file. `supports_collect_by_thresholds(plan, 
threshold_byte_size, threshold_num_rows)` (L509), `hash_build_fits(..)` (L766) 
and `broadcast_build_side_pays_off(..)` (`join_selection.rs:176`) all take 
plain params. I grepped `ballista/scheduler` and `ballista/core`: this is the 
only limits bundle in the tree.
   
   It is also internally inconsistent. Two limits come from the 
`&BallistaConfig` it receives, the third is passed by hand even though it is 
just `bc.broadcast_join_threshold_bytes()` from L274.
   
   Plain params drop ~31 lines and take the `stage_limits` test helper with 
them. Worth keeping: the q8 `part`-scan anecdote at L680-684, relocated as one 
sentence onto `requires_build_staging`. That one is measured evidence rather 
than prose.



##########
ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs:
##########
@@ -134,6 +135,43 @@ impl SelectJoinRule {
     }
 }
 
+/// The per-child shuffle partitioning a join's inputs need to be 
co-partitioned
+/// on its keys.
+fn join_key_partitioning(
+    join: &DynamicJoinSelectionExec,
+    partition_count: usize,
+) -> Vec<Partitioning> {
+    join._required_input_distribution()
+        .iter()
+        .map(|d| d.clone().create_partitioning(partition_count))
+        .collect()
+}
+
+/// Wraps `child` in an `ExchangeExec` shuffling on `partitioning`, unless it
+/// already is one that does.
+///
+/// `StageBuildSide` shuffles the build side on the join key ahead of the
+/// decision, so when the measured size then turns out to be too large to
+/// broadcast and the join falls back to `Repartition`, that side is already
+/// where it needs to be. Wrapping it again would nest one stage boundary
+/// inside another and move the same rows twice.
+fn exchange_on(
+    child: Arc<dyn ExecutionPlan>,
+    partitioning: &Partitioning,
+    plan_id: impl FnOnce() -> usize,
+) -> Arc<dyn ExecutionPlan> {
+    if let Some(exchange) = child.downcast_ref::<ExchangeExec>()
+        && exchange.partitioning.as_ref() == Some(partitioning)
+    {
+        return child;
+    }
+    Arc::new(ExchangeExec::new(
+        child,
+        Some(partitioning.clone()),
+        plan_id(),
+    ))
+}

Review Comment:
   Not a defect first: the `impl FnOnce() -> usize` laziness **is** 
load-bearing, since `plan_id()` is `fetch_add` on an `AtomicUsize` (L104). Do 
not swap it for a plain `usize`.
   
   But a method receiver gives the same laziness without the generic. `fn 
exchange_on(&self, child, partitioning)` calling `self.plan_id()` inline 
matches every other id consumer in this file (L220, L247, L308, L314, L360). 
`SelectJoinRule` derives `Default`, so the three tests cost one line each.
   
   On the name: in DataFusion vocabulary "on" means join keys (`JoinOn`, and 
the `on=[(id@0, id@0)]` in this PR's own snapshots). `exchange_on(child, 
partitioning)` reads as "exchange on the join keys" but means "ensure an 
exchange with this partitioning". `ensure_exchange` is unambiguous and matches 
the `add_sort_above` convention.
   
   One more: `DynamicJoinSelectionExec::inputs_already_partitioned` 
(`dynamic_join.rs:552`) answers the same "is this input already shuffled the 
way the join needs" question using the equivalence-aware 
`Partitioning::satisfaction`, where this uses struct equality. They agree 
today. If the plainer comparison is deliberate, one line in the doc saying so 
would stop someone unifying them later.



##########
ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs:
##########
@@ -1166,6 +1323,306 @@ mod tests {
         );
     }
 
+    /// A source reporting `total_byte_size` with the given precision, shaped
+    /// like one side of the TPC-H q8 `part`/`lineitem` join.
+    fn sized_stats_exec(total_byte_size: Precision<usize>) -> Arc<dyn 
ExecutionPlan> {
+        Arc::new(StatisticsExec::new(
+            Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size,
+                column_statistics: vec![ColumnStatistics::new_unknown()],
+            },
+            Schema::new(vec![Field::new("k", DataType::Int32, false)]),
+        ))
+    }
+
+    /// The q8 shape: a filtered dimension scan whose size is a guess sitting
+    /// just over the broadcast budget, against a fact table that dwarfs it.
+    const STAGE_THRESHOLD: usize = 10 * MB;
+    const GUESSED_BUILD_BYTES: usize = 100 * MB;
+    const FACT_PROBE_BYTES: usize = 2000 * MB;
+
+    /// The staging limits for a given byte budget, carrying the shipped ratio
+    /// and multiple. Reading those from `BallistaConfig::default()` rather 
than
+    /// restating them means these tests pin the values a deployment runs with.
+    fn stage_limits(broadcast_threshold_bytes: usize) -> BuildStagingLimits {
+        BuildStagingLimits::new(&BallistaConfig::default(), 
broadcast_threshold_bytes)
+    }
+
+    #[test]
+    fn stages_the_build_side_when_its_size_is_only_a_guess() {
+        assert!(requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(FACT_PROBE_BYTES),
+                column_statistics: vec![],
+            },
+            stage_limits(STAGE_THRESHOLD),
+        ));
+    }
+
+    // An exact size over budget is a fact, not a guess. Measuring it again
+    // cannot flip the join, so both sides shuffle concurrently as before.
+    #[test]
+    fn does_not_stage_an_exactly_sized_build_side() {
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Exact(1_000),
+                total_byte_size: Precision::Exact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(FACT_PROBE_BYTES),
+                column_statistics: vec![],
+            },
+            stage_limits(STAGE_THRESHOLD),
+        ));
+    }
+
+    // Far enough over the budget that no measurement brings it back under, so
+    // the extra round trip would only add latency.
+    #[test]
+    fn does_not_stage_a_build_side_far_past_the_budget() {
+        let limits = stage_limits(STAGE_THRESHOLD);
+        let far_over = STAGE_THRESHOLD * (limits.max_estimate_multiple + 1);
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(far_over),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(far_over * 1_000),
+                column_statistics: vec![],
+            },
+            limits,
+        ));
+    }
+
+    // Without a probe side much larger than the build side, serialising the 
two
+    // shuffles costs about as much as it could save.
+    #[test]
+    fn does_not_stage_when_the_probe_side_is_comparable() {
+        let limits = stage_limits(STAGE_THRESHOLD);
+        let probe_bytes = GUESSED_BUILD_BYTES * (limits.min_probe_ratio - 1);
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(probe_bytes),
+                column_statistics: vec![],
+            },
+            limits,
+        ));
+    }
+
+    // No probe-side size is no evidence the round trip pays off.
+    #[test]
+    fn does_not_stage_without_a_probe_side_size() {
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Absent,
+                column_statistics: vec![],
+            },
+            stage_limits(STAGE_THRESHOLD),
+        ));
+    }
+
+    // With broadcast promotion off there is no decision to revisit.
+    #[test]
+    fn does_not_stage_when_broadcast_promotion_is_disabled() {
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(FACT_PROBE_BYTES),
+                column_statistics: vec![],
+            },
+            stage_limits(0),
+        ));
+    }
+
+    // The shipped defaults are what the benchmark numbers were measured under,

Review Comment:
   Six single-assert tests, ~115 lines, each differing in **one** `Precision` 
value buried in identical boilerplate.
   
   This file already has the shape for it. `supports_collect(&plan)` (L934) and 
`is_collected(&action)` (L1162) are predicate helpers, and 
`known_total_byte_size_still_decides` (L996), 
`broadcast_threshold_bytes_drives_collect_decision` (L1223) and 
`to_actual_join_uses_post_swap_join_type` (L1207) each bundle positive and 
negative cases with explanatory messages.
   
   Add `fn stages(build: Precision<usize>, probe: Precision<usize>, threshold: 
usize) -> bool` and fold to two or three tests. ~80 lines.
   
   Not suggesting a named-case table harness. The file has none, and adding one 
would be the churn.



##########
ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs:
##########
@@ -294,31 +332,64 @@ impl PhysicalOptimizerRule for SelectJoinRule {
 
                                 Ok(Transformed::yes(exec.data))
                             }
-                            JoinSelectionAction::Repartition(dynamic_join) => {
-                                let partition_count = 
config.execution.target_partitions;
-                                let partitioning = dynamic_join
-                                    ._required_input_distribution()
-                                    .iter()
-                                    .map(|d| {
-                                        
d.clone().create_partitioning(partition_count)
-                                    })
-                                    .collect::<Vec<_>>();
-
-                                let left = dynamic_join.left.clone();
-                                let right = dynamic_join.right.clone();
-
-                                let left = Arc::new(ExchangeExec::new(
-                                    left,
-                                    Some(partitioning[0].clone()),
-                                    self.plan_id(),
-                                ));
+                            JoinSelectionAction::StageBuildSide { join, 
build_side } => {
+                                // Only the build side gets an exchange. The
+                                // probe side is left as it is, so no stage is
+                                // created for it and nothing of it is shuffled
+                                // until the join is decided for real.
+                                let partitioning = join_key_partitioning(
+                                    &join,
+                                    config.execution.target_partitions,
+                                );
+                                let build_idx = match build_side {
+                                    JoinSide::Left => 0,
+                                    JoinSide::Right => 1,
+                                    JoinSide::None => {
+                                        return Err(DataFusionError::Internal(
+                                            "StageBuildSide requires a build 
side"
+                                                .to_owned(),
+                                        ));
+                                    }
+                                };
 
-                                let right = Arc::new(ExchangeExec::new(
-                                    right,
-                                    Some(partitioning[1].clone()),
+                                let mut children =
+                                    vec![join.left.clone(), 
join.right.clone()];
+                                children[build_idx] = 
Arc::new(ExchangeExec::new(
+                                    children[build_idx].clone(),
+                                    Some(partitioning[build_idx].clone()),

Review Comment:
   Clones a `Vec` element it immediately overwrites. `std::mem::replace(&mut 
children[build_idx], staged)`, or building the exchange first and placing it, 
reads as intended. Cosmetic, but as written it looks like a slip.



##########
ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs:
##########
@@ -1166,6 +1323,306 @@ mod tests {
         );
     }
 
+    /// A source reporting `total_byte_size` with the given precision, shaped
+    /// like one side of the TPC-H q8 `part`/`lineitem` join.
+    fn sized_stats_exec(total_byte_size: Precision<usize>) -> Arc<dyn 
ExecutionPlan> {
+        Arc::new(StatisticsExec::new(
+            Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size,
+                column_statistics: vec![ColumnStatistics::new_unknown()],
+            },
+            Schema::new(vec![Field::new("k", DataType::Int32, false)]),
+        ))
+    }
+
+    /// The q8 shape: a filtered dimension scan whose size is a guess sitting
+    /// just over the broadcast budget, against a fact table that dwarfs it.
+    const STAGE_THRESHOLD: usize = 10 * MB;
+    const GUESSED_BUILD_BYTES: usize = 100 * MB;
+    const FACT_PROBE_BYTES: usize = 2000 * MB;
+
+    /// The staging limits for a given byte budget, carrying the shipped ratio
+    /// and multiple. Reading those from `BallistaConfig::default()` rather 
than
+    /// restating them means these tests pin the values a deployment runs with.
+    fn stage_limits(broadcast_threshold_bytes: usize) -> BuildStagingLimits {
+        BuildStagingLimits::new(&BallistaConfig::default(), 
broadcast_threshold_bytes)
+    }
+
+    #[test]
+    fn stages_the_build_side_when_its_size_is_only_a_guess() {
+        assert!(requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(FACT_PROBE_BYTES),
+                column_statistics: vec![],
+            },
+            stage_limits(STAGE_THRESHOLD),
+        ));
+    }
+
+    // An exact size over budget is a fact, not a guess. Measuring it again
+    // cannot flip the join, so both sides shuffle concurrently as before.
+    #[test]
+    fn does_not_stage_an_exactly_sized_build_side() {
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Exact(1_000),
+                total_byte_size: Precision::Exact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(FACT_PROBE_BYTES),
+                column_statistics: vec![],
+            },
+            stage_limits(STAGE_THRESHOLD),
+        ));
+    }
+
+    // Far enough over the budget that no measurement brings it back under, so
+    // the extra round trip would only add latency.
+    #[test]
+    fn does_not_stage_a_build_side_far_past_the_budget() {
+        let limits = stage_limits(STAGE_THRESHOLD);
+        let far_over = STAGE_THRESHOLD * (limits.max_estimate_multiple + 1);
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(far_over),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(far_over * 1_000),
+                column_statistics: vec![],
+            },
+            limits,
+        ));
+    }
+
+    // Without a probe side much larger than the build side, serialising the 
two
+    // shuffles costs about as much as it could save.
+    #[test]
+    fn does_not_stage_when_the_probe_side_is_comparable() {
+        let limits = stage_limits(STAGE_THRESHOLD);
+        let probe_bytes = GUESSED_BUILD_BYTES * (limits.min_probe_ratio - 1);
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(probe_bytes),
+                column_statistics: vec![],
+            },
+            limits,
+        ));
+    }
+
+    // No probe-side size is no evidence the round trip pays off.
+    #[test]
+    fn does_not_stage_without_a_probe_side_size() {
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Absent,
+                column_statistics: vec![],
+            },
+            stage_limits(STAGE_THRESHOLD),
+        ));
+    }
+
+    // With broadcast promotion off there is no decision to revisit.
+    #[test]
+    fn does_not_stage_when_broadcast_promotion_is_disabled() {
+        assert!(!requires_build_staging(
+            &Statistics {
+                num_rows: Precision::Inexact(1_000),
+                total_byte_size: Precision::Inexact(GUESSED_BUILD_BYTES),
+                column_statistics: vec![],
+            },
+            &Statistics {
+                num_rows: Precision::Exact(1_000_000),
+                total_byte_size: Precision::Exact(FACT_PROBE_BYTES),
+                column_statistics: vec![],
+            },
+            stage_limits(0),
+        ));
+    }
+
+    // The shipped defaults are what the benchmark numbers were measured under,
+    // so pin them here rather than only in the generated config docs.
+    #[test]
+    fn stage_build_side_ships_the_documented_defaults() {
+        let bc = BallistaConfig::default();
+
+        assert!(bc.stage_build_side_enabled());
+        assert_eq!(bc.stage_build_side_min_probe_ratio(), 10);
+        assert_eq!(bc.stage_build_side_max_estimate_multiple(), 32);
+    }
+
+    /// Runs the resolver over the q8 shape, with staging on or off.
+    fn q8_shaped_action(stage_build_side: bool) -> JoinSelectionAction {
+        q8_shaped_action_with_children(
+            sized_stats_exec(Precision::Inexact(GUESSED_BUILD_BYTES)),
+            sized_stats_exec(Precision::Exact(FACT_PROBE_BYTES)),
+            stage_build_side,
+        )
+    }
+
+    fn q8_shaped_action_with_children(
+        left: Arc<dyn ExecutionPlan>,
+        right: Arc<dyn ExecutionPlan>,
+        stage_build_side: bool,
+    ) -> JoinSelectionAction {
+        q8_shaped_action_with_join_type(left, right, JoinType::Inner, 
stage_build_side)
+    }
+
+    fn q8_shaped_action_with_join_type(
+        left: Arc<dyn ExecutionPlan>,
+        right: Arc<dyn ExecutionPlan>,
+        join_type: JoinType,
+        stage_build_side: bool,
+    ) -> JoinSelectionAction {
+        let on: JoinOn =
+            vec![(Arc::new(Column::new("k", 0)), Arc::new(Column::new("k", 
0)))];
+        let dj = DynamicJoinSelectionExec {
+            properties: Arc::clone(left.properties()),
+            left,
+            right,
+            on,
+            filter: None,
+            join_type,
+            projection: None,
+            null_equality: NullEquality::NullEqualsNothing,
+            selection_state: JoinInputState::Unknown,
+            null_aware: false,
+            plan_id: 0,
+        };
+        let mut config = ConfigOptions::new();
+        let mut bc = BallistaConfig::default();
+        bc.set(
+            "optimizer.broadcast_join_threshold_bytes",
+            &STAGE_THRESHOLD.to_string(),
+        )
+        .unwrap();
+        bc.set("optimizer.stage_build_side", &stage_build_side.to_string())
+            .unwrap();
+        config.extensions.insert(bc);
+        dj.to_actual_join(&config).unwrap()
+    }
+
+    // The whole point: a guessed-large build side against a fact-table probe
+    // side stages the build alone instead of committing to both shuffles.
+    #[test]
+    fn to_actual_join_stages_the_build_side_for_the_q8_shape() {
+        assert!(matches!(
+            q8_shaped_action(true),
+            JoinSelectionAction::StageBuildSide {
+                build_side: JoinSide::Left,
+                ..
+            }
+        ));
+    }
+
+    // Staging only pays off when `CollectLeft` is reachable, and that is a
+    // property of the (post-swap) join type alone. `Full` is never
+    // broadcastable, and `Left`/`LeftSemi`/`LeftAnti`/`LeftMark` are not when
+    // the build side is already the left input — as it is here, the left being
+    // the smaller side. Staging any of these would serialise a stage boundary
+    // that no measurement could ever cash in, so they must repartition 
instead.
+    #[test]
+    fn does_not_stage_a_join_type_that_can_never_broadcast() {
+        for join_type in [
+            JoinType::Full,
+            JoinType::Left,
+            JoinType::LeftSemi,
+            JoinType::LeftAnti,
+            JoinType::LeftMark,
+        ] {
+            let action = q8_shaped_action_with_join_type(
+                sized_stats_exec(Precision::Inexact(GUESSED_BUILD_BYTES)),
+                sized_stats_exec(Precision::Exact(FACT_PROBE_BYTES)),
+                join_type,
+                true,
+            );
+
+            assert!(
+                matches!(action, JoinSelectionAction::Repartition(_)),
+                "{join_type:?} can never reach CollectLeft, so it must 
repartition \
+                 rather than pay for a staging round trip"
+            );
+        }
+    }
+
+    // The counterweight: the broadcast-safe join types still stage, so the
+    // guard above rejects only what it is meant to.
+    #[test]
+    fn stages_the_join_types_that_can_broadcast() {
+        for join_type in [
+            JoinType::Inner,
+            JoinType::Right,
+            JoinType::RightSemi,
+            JoinType::RightAnti,
+            JoinType::RightMark,
+        ] {
+            let action = q8_shaped_action_with_join_type(
+                sized_stats_exec(Precision::Inexact(GUESSED_BUILD_BYTES)),
+                sized_stats_exec(Precision::Exact(FACT_PROBE_BYTES)),
+                join_type,
+                true,
+            );
+
+            assert!(
+                matches!(action, JoinSelectionAction::StageBuildSide { .. }),
+                "{join_type:?} can reach CollectLeft, so measuring its build 
side \
+                 is worth a staging round trip"
+            );
+        }
+    }
+

Review Comment:
   Two hand-partitioned five-element lists. 
`to_actual_join_collects_only_broadcast_safe_join_types` (L1176) already solves 
this exact problem 370 lines up: it loops **all ten** types and compares 
against `collect_left_broadcast_safe(join_type)` as the oracle.
   
   ```rust
   for join_type in [/* all ten */] {
       let action = q8_shaped_action_with_join_type(..., join_type, true);
       assert_eq!(
           matches!(action, JoinSelectionAction::StageBuildSide { .. }),
           collect_left_broadcast_safe(join_type),
           "join_type {join_type:?}",
       );
   }
   ```
   
   The merge is sound here: with left `Inexact(100 MB)` and right `Exact(2000 
MB)`, `supports_swap_join_order` is `100MB > 2000MB` = false, so the post-swap 
type equals the declared type in every case.
   
   Add a join type to DataFusion and neither current list fails. The merged 
form does. It also subsumes 
`to_actual_join_stages_the_build_side_for_the_q8_shape` above. ~54 lines.



##########
ballista/scheduler/src/state/aqe/test/mod.rs:
##########
@@ -46,6 +48,16 @@ use datafusion::prelude::{SessionConfig, SessionContext};
 use std::sync::Arc;
 
 pub(crate) fn mock_partitions_with_statistics() -> Vec<Vec<PartitionLocation>> 
{
+    mock_partitions_with_size(42, 10)
+}
+
+/// Shuffle output reporting `num_rows` rows over `num_bytes` bytes in a single
+/// partition. Tests that turn on the *value* of a measured size, rather than
+/// just its presence, pick their own figures.
+pub(crate) fn mock_partitions_with_size(
+    num_rows: u64,
+    num_bytes: u64,
+) -> Vec<Vec<PartitionLocation>> {

Review Comment:
   Good extraction, but it stopped one function short. 
`mock_partitions_with_statistics_no_data` right below (L85-108) is still a full 
24-line copy of the same `PartitionLocation` literal, differing only in 
`PartitionStats::new(Some(0), None, Some(0))`:
   
   ```rust
   pub(crate) fn mock_partitions_with_statistics_no_data() -> 
Vec<Vec<PartitionLocation>> {
       mock_partitions_with_size(0, 0)
   }
   ```
   
   ~21 lines. An extraction that covers two of three cases is the state most 
likely to rot.



##########
ballista/scheduler/src/state/aqe/test/stats_table.rs:
##########
@@ -224,6 +224,26 @@ pub(crate) fn sizeless_statistics(schema: &Schema, 
num_rows: usize) -> Statistic
     }
 }
 
+/// Statistics for a table whose size is an *estimate* rather than a
+/// measurement, as a filtered scan reports it before the filter has run.
+///
+/// This is the shape that makes staging a join's build side worthwhile: the
+/// planner falls back to `default_filter_selectivity` and produces a figure
+/// that is over the broadcast budget but carries no evidence behind it.
+pub(crate) fn estimated_statistics(
+    schema: &Schema,
+    num_rows: usize,
+    total_byte_size: usize,
+) -> Statistics {
+    use datafusion::common::{ColumnStatistics, stats::Precision};
+
+    Statistics {
+        num_rows: Precision::Inexact(num_rows),
+        total_byte_size: Precision::Inexact(total_byte_size),
+        column_statistics: vec![ColumnStatistics::new_unknown(); 
schema.fields().len()],
+    }
+}
+

Review Comment:
   This is `sized_statistics` with `Exact` swapped for `Inexact`, and 
DataFusion ships the conversion:
   
   ```rust
   pub(crate) fn estimated_statistics(schema: &Schema, num_rows: usize, 
total_byte_size: usize) -> Statistics {
       sized_statistics(schema, num_rows, total_byte_size).to_inexact()
   }
   ```
   
   (`Statistics::to_inexact`, `datafusion-common-55.0.0/src/stats.rs:485`.)
   
   Counterpoint: three sibling constructors is this file's existing shape, so 
low value either way.



##########
ballista/scheduler/src/state/aqe/test/stage_build_side.rs:
##########
@@ -0,0 +1,229 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Staging a join's build side, across the two passes it takes to pay off.
+//!
+//! The unit tests in `dynamic_join` pin the *decision* in isolation: given
+//! these statistics, does the resolver choose to stage? What they cannot show
+//! is the mechanism the feature exists for, which only appears across two
+//! passes — shuffle the build side alone, read back what that shuffle actually
+//! measured, then decide the join for real against a number instead of a 
guess.
+//! These tests drive `AdaptivePlanner` through both passes and assert on the
+//! plan that comes out of each.
+//!
+//! Sizes are declared ([`StatsTable`]) rather than materialised, because the
+//! shape staging turns on is a build side whose size is an *estimate* sitting
+//! over the broadcast budget. No table a test is willing to build produces 
that.
+
+use crate::state::aqe::planner::AdaptivePlanner;
+use crate::state::aqe::test::stats_table::{
+    StatsTable, estimated_statistics, sized_statistics,
+};
+use crate::state::aqe::test::{
+    mock_partitions_with_size, mock_partitions_with_statistics,
+};
+use ballista_core::assert_plan;
+use ballista_core::extension::SessionConfigExt;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::common::Statistics;
+use datafusion::execution::{
+    SessionStateBuilder, config::SessionConfig, context::SessionContext,
+};
+use std::sync::Arc;
+
+const MB: usize = 1024 * 1024;
+
+/// A build side whose size is only a guess, sitting over the shipped 128 MB
+/// broadcast budget but well inside the 32x window past which no measurement
+/// could bring it back under. This is the TPC-H q8 `part` shape: a filtered
+/// dimension scan the planner sizes with `default_filter_selectivity`.
+const GUESSED_BUILD_BYTES: usize = 200 * MB;
+
+/// A fact-table probe side, far past the 10x the ratio requires.
+const FACT_PROBE_BYTES: usize = 40 * 1024 * MB;
+
+/// A measured build size that still cannot be broadcast, for the fallthrough.
+const MEASURED_TOO_LARGE_BYTES: u64 = (200 * MB) as u64;
+
+fn join_schema() -> Arc<Schema> {
+    Arc::new(Schema::new(vec![
+        Field::new("id", DataType::Int32, false),
+        Field::new("val", DataType::Int32, false),
+    ]))
+}
+
+/// A context carrying Ballista's shipped configuration, so these tests run
+/// against the thresholds a deployment uses — including
+/// `ballista.optimizer.stage_build_side`, which defaults to on.
+fn ballista_ctx() -> SessionContext {
+    let config = SessionConfig::new_with_ballista()
+        .with_target_partitions(4)
+        .with_round_robin_repartition(false);
+    let state = SessionStateBuilder::new_with_default_features()
+        .with_config(config)
+        .build();
+    SessionContext::new_with_state(state)
+}
+
+fn register(ctx: &SessionContext, name: &str, stats: Statistics) {
+    ctx.register_table(name, Arc::new(StatsTable::new(join_schema(), stats, 
4)))
+        .unwrap();
+}

Review Comment:
   I diffed these against `broadcast_thresholds.rs`:
   
   - `ballista_ctx()` is **byte-identical** to `broadcast_thresholds.rs:80-88`
   - `join_schema()` is **byte-identical** to its `narrow_schema()` (L71-76)
   - `register()` is its L90-93 with the `schema` param dropped, so strictly 
less capable
   - `const MB` is a third copy (L47 there, and `dynamic_join.rs:911`)
   
   `test/mod.rs` is already the home for cross-file AQE fixtures 
(`mock_context`, `mock_schema`, `mock_partitions_*`). Hoisting these there and 
importing in both files nets ~25 lines, and more importantly stops "the shipped 
session config" from meaning two different things.



##########
docs/source/contributors-guide/benchmarking.md:
##########
@@ -24,10 +24,22 @@ Current TPC-H **SF1000** results for Ballista, compared 
against a vanilla
 
 ## Versions under test
 
-| Engine   | Version                                                           
                                                                                
                        |
-| -------- | 
-------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 |
-| Ballista | 
[`67b3a19b`](https://github.com/apache/datafusion-ballista/commit/67b3a19bb442db879c7056722d14696cc33341b9)
 (`main`, 2026-09-03), Cargo pkg `54.0.0`, DataFusion `55.0.0` |
-| Spark    | 3.4 (vanilla, no acceleration plugin)                             
                                                                                
                        |
+| Engine   | Version                                                           
                                                                                
                           |
+| -------- | 
----------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 |
+| Ballista | 
[`da92ba8d`](https://github.com/apache/datafusion-ballista/pull/2434/commits/da92ba8d1bce6e2fd76d01aa86a8c971fa34ac2f)
 (2026-09-08), Cargo pkg `54.0.0`, DataFusion `55.0.0` |
+| Spark    | 3.4 (vanilla, no acceleration plugin)                             
                                                                                
                           |
+
+```{note}
+`da92ba8d` is the scheduler commit of
+[#2434](https://github.com/apache/datafusion-ballista/pull/2434) — `main`
+as of 2026-09-08 plus that PR's change, which stages a join's build side
+and reads its measured statistics before `DynamicJoinSelectionExec`
+commits to a strategy. It accounts for almost all of the improvement over
+the previous (`67b3a19b`) result set — Q8 96.02 → 40.68, Q9 107.89 →
+51.94, Q21 95.63 → 65.69 — so a run from `main` before #2434 lands will
+not reproduce these numbers. Re-pin this row to the merge commit once it
+does.

Review Comment:
   Checked-in docs pinned to an unmerged PR commit, carrying an unowned TODO 
("Re-pin this row to the merge commit once it does"). This is also the only 
`{note}` admonition in the whole `docs/source` tree, so it is not an 
established pattern here.
   
   Leaving the row on `67b3a19b` and landing a separate docs commit after merge 
avoids the self-invalidating state.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to