This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 2c1357f0d feat(scheduler): stage a join's build side before deciding
the join (#2434)
2c1357f0d is described below
commit 2c1357f0d61770e007eebd608391b10969ab36de
Author: Andy Grove <[email protected]>
AuthorDate: Sat Sep 12 15:08:20 2026 -0600
feat(scheduler): stage a join's build side before deciding the join (#2434)
* feat(scheduler): stage a join's build side before deciding the join
AQE's `Repartition` action shuffles both join inputs at once. When the
build side's size is only an estimate, that estimate is the sole reason
the join is not a broadcast, and the probe side is a fact table, this
commits to the single most expensive stage in the query before any
measurement exists. TPC-H q8 at SF1000 shuffles all 6B rows of lineitem
to join against 1.33M rows of filtered `part`, and AQE can never walk it
back: once `Repartition` fires, `selection_state` is `Repartitioned` and
only Hash and SortMerge remain reachable.
Add `JoinSelectionAction::StageBuildSide`, which shuffles the smaller
side alone on the join key and leaves the probe side untouched. The
resolver re-runs once that stage completes, reads the exchange's
measured statistics, and usually promotes the join to `CollectLeft`, so
the probe side is never shuffled. It fires at most once per join,
guarded on neither child already being an exchange, and `Repartition`
now reuses an exchange already on the join key rather than nesting a
second stage boundary inside it.
Also stop applying `broadcast_join_threshold_rows` as a ceiling over the
schema-derived byte estimate. The config documents it as the fallback
for when byte statistics are unavailable, but the code applied both, so
a narrow build side well inside the byte budget was rejected on row
count alone: 1M `i64` keys is 8 MB against a 128 MiB default budget.
The row rule now decides only when no byte estimate can be formed, and
`0` still disables that path.
* docs(bench): refresh TPC-H SF1000 results at da92ba8d
Q8 96.02 -> 40.68, Q9 107.89 -> 51.94 and Q21 95.63 -> 65.69 as the
build side is staged and measured before the join strategy is chosen.
Suite total 817.07 -> 696.09 against an unchanged Spark 3.4 baseline
of 845.21 on the same cluster shape.
Pin the version row to this PR's head and note that the numbers do not
reproduce from main until it lands.
* docs(bench): correct commit reference in provenance note
* refactor(scheduler): gate build-side staging on broadcast-safe join types
Staging a join's build side can only change the outcome when CollectLeft
is reachable, and that is a property of the post-swap join type alone:
partition_mode is CollectLeft only when `under_threshold &&
collect_left_broadcast_safe(..)`. Without that term in the staging guard
a Full join (Full.swap() == Full), or a left-sided semi/anti/outer join
whose build side is already the left input, would stage, fail the
exchange guard on the second pass, and fall through to Repartition
having bought a serialised stage boundary and no decision.
The swap cannot rescue those cases either: supports_swap_join_order is
left_bytes > right_bytes, so a build side that is already the left input
is the smaller one, and measuring it only shrinks it further.
Also from review:
- Replace `build_is_left: bool` with DataFusion's existing `JoinSide`
rather than a new enum, and reject `JoinSide::None` explicitly.
- Rename `should_stage_build_side` to `requires_build_staging`.
- Promote the two gating figures to config keys, read through a new
BuildStagingLimits struct, so they can be tuned without rebuilding:
stage_build_side_min_probe_ratio and
stage_build_side_max_estimate_multiple.
- Restore the hash_build_fits doc comment, which the previous commit
orphaned onto a constant.
- Add state/aqe/test/stage_build_side.rs, covering the two-pass
mechanism end to end: stage then broadcast against the measured size,
the too-large fallthrough reusing the staged exchange rather than
nesting a second one, and a Left join never staging.
The broadcast_join_threshold_rows semantics change moves to #2444.
* refactor(scheduler): simplify build-side staging after review
Quality pass following comphead's reuse and simplification review, which
found no correctness issues. Behaviour preserving: the only snapshots
that changed are the four in stage_build_side.rs, and their plan shapes
and plan_ids are identical to before.
Production code:
- Drop the `join` payload from `JoinSelectionAction::StageBuildSide`. It
was a field-identical copy of the node `SelectJoinRule` is already
visiting, so the arm now rebuilds from that node.
- Replace `build_side: JoinSide` with `build_is_right: bool`. `JoinSide`
is three variants for a two-state fact, which forced an unreachable
`JoinSide::None` arm; the producer already computes the bool as
`swap_inputs`.
- Collapse the two `(Unknown, Partitioned)` arms into one. `!null_aware`
was implied by matching `Partitioned`, and the join-type condition
moves inside `requires_build_staging`.
- Drop `with_selection_state`, a no-op at its only new call site, and
restore `to_partitioned`'s inline body.
- Replace `BuildStagingLimits` with plain params, matching the three
sibling predicates in the same file.
- Hoist `probe_stats` beside `build_stats` instead of a third swap
ternary.
- `join_key_partitioning` returns a tuple by value over the owned Vec,
dropping a dead clone. `exchange_on` becomes
`SelectJoinRule::ensure_exchange`, keeping its laziness through the
receiver and documenting why it uses structural equality rather than
`Partitioning::satisfaction`.
Tests:
- Fold six single-assert staging tests into two behind a `stages(..)`
predicate helper.
- Merge the two hand-partitioned join-type lists into one loop over all
ten types, with `collect_left_broadcast_safe` as the oracle, so a join
type added to DataFusion cannot slip through untested.
- Un-fork `q8_shaped_action*` from `run_to_actual_join` by giving the
latter a `stage_build_side` param.
- Parameterise the fourth `StatisticsExec` builder; the others wrap it.
- Hoist `ballista_ctx`, `narrow_schema`, `register_stats_table` and `MB`
into `test/mod.rs` so "the shipped session config" means one thing.
- `mock_partitions_with_statistics_no_data` delegates to
`mock_partitions_with_size`; `estimated_statistics` uses DataFusion's
`Statistics::to_inexact`.
- Drop `does_not_stage_a_left_join` and
`exchange_on_wraps_a_non_exchange_child`, both covered elsewhere.
Config entries now state that staging applies only to estimated
statistics.
Reverts the benchmarking.md update: it pinned a row to an unmerged PR
commit behind an unowned TODO, and those numbers predate the join-type
guard. They will be re-measured and landed separately.
---
ballista/core/src/config.rs | 72 +++++
.../src/state/aqe/execution_plan/dynamic_join.rs | 354 +++++++++++++++++++--
.../src/state/aqe/optimizer_rule/join_selection.rs | 175 ++++++++--
.../src/state/aqe/test/broadcast_thresholds.rs | 43 +--
ballista/scheduler/src/state/aqe/test/mod.rs | 74 +++--
.../src/state/aqe/test/stage_build_side.rs | 167 ++++++++++
.../scheduler/src/state/aqe/test/stats_table.rs | 14 +
docs/source/user-guide/configs.md | 61 ++--
8 files changed, 827 insertions(+), 133 deletions(-)
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
index 526dd967a..d2f5475d5 100644
--- a/ballista/core/src/config.rs
+++ b/ballista/core/src/config.rs
@@ -119,6 +119,31 @@ pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS: &str =
pub const BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES: &str =
"ballista.optimizer.hash_join_max_build_partition_bytes";
+/// Configuration key controlling whether AQE stages a join's prospective build
+/// side as its own shuffle before deciding how to run the join. When a build
+/// side's size is only an estimate and the probe side is far larger, shuffling
+/// both sides commits to the expensive shuffle before any measurement exists.
+/// With this enabled the planner shuffles the smaller side alone, reads back
+/// its measured size, and only then chooses broadcast or repartition. Enabled
+/// by default; set to `false` to always shuffle both sides at once.
+pub const BALLISTA_STAGE_BUILD_SIDE: &str =
"ballista.optimizer.stage_build_side";
+
+/// Configuration key for how many times larger the probe side must be before
+/// AQE stages a join's build side on its own. Staging serialises two shuffles
+/// that would otherwise run concurrently, so the deferral costs at most the
+/// smaller side's runtime; this ratio is what bounds the cost of being wrong.
+/// Only consulted when `ballista.optimizer.stage_build_side` is enabled.
+pub const BALLISTA_STAGE_BUILD_SIDE_MIN_PROBE_RATIO: &str =
+ "ballista.optimizer.stage_build_side_min_probe_ratio";
+
+/// Configuration key for how far over `broadcast_join_threshold_bytes` an
+/// *estimated* build side may sit and still be worth measuring. Past this
+/// multiple the side is large on any reading and no measurement brings it back
+/// under budget, so it is shuffled without the extra round trip. Only
consulted
+/// when `ballista.optimizer.stage_build_side` is enabled.
+pub const BALLISTA_STAGE_BUILD_SIDE_MAX_ESTIMATE_MULTIPLE: &str =
+ "ballista.optimizer.stage_build_side_max_estimate_multiple";
+
/// Configuration key controlling the logical rewrite of uncorrelated
/// `NOT IN (subquery)` filter predicates into a plain anti join plus a
/// one-row count aggregate. The rewrite avoids DataFusion's null-aware hash
@@ -294,6 +319,35 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String,
ConfigEntry>> = LazyLock::new(||
which makes AQE use a hash join regardless of build
size.".to_string(),
DataType::UInt64,
Some((64 * 1024 * 1024).to_string())),
+ ConfigEntry::new(BALLISTA_STAGE_BUILD_SIDE.to_string(),
+ "Stages a join's prospective build side as its own
shuffle before \
+ choosing the join strategy. Applies only when that
side's size is \
+ an estimate rather than a measurement, and the probe
side is far \
+ larger: an exact size over the broadcast budget is a
fact that no \
+ further measurement can change. Lets AQE measure the
build side and \
+ pick a broadcast join instead of committing to a
full shuffle of \
+ the probe side up front. Set to false to always
shuffle both sides \
+ at once.".to_string(),
+ DataType::Boolean,
+ Some(true.to_string())),
+ ConfigEntry::new(BALLISTA_STAGE_BUILD_SIDE_MIN_PROBE_RATIO.to_string(),
+ "How many times larger the probe side must be before
AQE stages a \
+ join's build side on its own. Staging serialises two
shuffles that \
+ would otherwise run concurrently, so the deferral
costs at most the \
+ smaller side's runtime; this ratio bounds the cost
of being wrong. \
+ Only consulted when stage_build_side is enabled, and
only for build \
+ sides whose size is an estimate rather than a
measurement.".to_string(),
+ DataType::UInt64,
+ Some((10).to_string())),
+
ConfigEntry::new(BALLISTA_STAGE_BUILD_SIDE_MAX_ESTIMATE_MULTIPLE.to_string(),
+ "How far over broadcast_join_threshold_bytes an
estimated build side \
+ may sit and still be worth measuring. Past this
multiple the side is \
+ large on any reading and no measurement brings it
back under budget, \
+ so it is shuffled without the extra round trip. Only
consulted when \
+ stage_build_side is enabled, and only for build
sides whose size is \
+ an estimate rather than a measurement.".to_string(),
+ DataType::UInt64,
+ Some((32).to_string())),
ConfigEntry::new(BALLISTA_NOT_IN_SUBQUERY_REWRITE.to_string(),
"Rewrites uncorrelated NOT IN (subquery) filter
predicates into a \
plain anti join plus a one-row count aggregate during
logical \
@@ -709,6 +763,24 @@ impl BallistaConfig {
self.get_usize_setting(BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES)
}
+ /// Returns whether AQE stages a join's prospective build side as its own
+ /// shuffle before choosing the join strategy.
+ pub fn stage_build_side_enabled(&self) -> bool {
+ self.get_bool_setting(BALLISTA_STAGE_BUILD_SIDE)
+ }
+
+ /// How many times larger the probe side must be before AQE stages a join's
+ /// build side on its own.
+ pub fn stage_build_side_min_probe_ratio(&self) -> usize {
+ self.get_usize_setting(BALLISTA_STAGE_BUILD_SIDE_MIN_PROBE_RATIO)
+ }
+
+ /// How far over `broadcast_join_threshold_bytes` an estimated build side
may
+ /// sit and still be worth measuring.
+ pub fn stage_build_side_max_estimate_multiple(&self) -> usize {
+ self.get_usize_setting(BALLISTA_STAGE_BUILD_SIDE_MAX_ESTIMATE_MULTIPLE)
+ }
+
/// Whether uncorrelated `NOT IN (subquery)` filter predicates are
rewritten
/// into a plain anti join plus a one-row count aggregate during logical
/// optimization, avoiding the single-task null-aware hash join.
diff --git a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs
b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs
index 9d8577fcc..fb7e1942a 100644
--- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs
+++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs
@@ -22,8 +22,8 @@ use datafusion::{
arrow::compute::SortOptions,
arrow::datatypes::{DataType, Schema},
common::{
- ColumnStatistics, JoinType, NullEquality, Result, exec_err,
internal_err,
- plan_err, tree_node::TreeNodeRecursion,
+ ColumnStatistics, JoinType, NullEquality, Result, Statistics, exec_err,
+ internal_err, plan_err, stats::Precision, tree_node::TreeNodeRecursion,
},
config::ConfigOptions,
execution::{SendableRecordBatchStream, TaskContext},
@@ -212,6 +212,17 @@ 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`].
+ ///
+ /// The join itself is not carried: `SelectJoinRule` is already visiting
the
+ /// node this was resolved from, and staging changes none of its fields.
+ StageBuildSide {
+ /// Whether the build side is the join's right input rather than its
+ /// left, which is exactly whether the resolver swapped the inputs.
+ build_is_right: bool,
+ },
Repartition(Arc<DynamicJoinSelectionExec>),
CollectLeft(Arc<HashJoinExec>),
LateCollectLeft(Arc<HashJoinExec>),
@@ -315,10 +326,10 @@ impl DynamicJoinSelectionExec {
.compute(self.left.as_ref(), &StatisticsArgs::new())?;
let stats_right = StatisticsContext::new()
.compute(self.right.as_ref(), &StatisticsArgs::new())?;
- let build_stats = if swap_inputs {
- &stats_right
+ let (build_stats, probe_stats) = if swap_inputs {
+ (&stats_right, &stats_left)
} else {
- &stats_left
+ (&stats_left, &stats_right)
};
let build_size_known = match build_stats.total_byte_size.get_value() {
Some(bytes) => *bytes != 0,
@@ -385,15 +396,37 @@ impl DynamicJoinSelectionExec {
JoinInputState::Repartitioned,
PartitionMode::Partitioned | PartitionMode::CollectLeft,
) => self.to_sort_merge_join().map(JoinSelectionAction::Sort),
- (JoinInputState::Unknown, PartitionMode::Partitioned) => Ok(
-
JoinSelectionAction::Repartition(Arc::new(self.to_partitioned())),
- ),
+ (JoinInputState::Unknown, PartitionMode::Partitioned) => {
+ // Staging fires at most once per join: once it does one child
is
+ // an `ExchangeExec`, and this guard never passes again.
+ if bc.stage_build_side_enabled()
+ && !self.left.is::<ExchangeExec>()
+ && !self.right.is::<ExchangeExec>()
+ && requires_build_staging(
+ build_stats,
+ probe_stats,
+ build_side_join_type,
+ threshold_collect_left_join_bytes,
+ bc.stage_build_side_min_probe_ratio(),
+ bc.stage_build_side_max_estimate_multiple(),
+ )
+ {
+ Ok(JoinSelectionAction::StageBuildSide {
+ build_is_right: swap_inputs,
+ })
+ } else {
+ Ok(JoinSelectionAction::Repartition(Arc::new(
+ self.to_partitioned(),
+ )))
+ }
+ }
// this method calculates partition mode, and at the moment it
// can't calculate it as PartitionMode::Auto
(_, PartitionMode::Auto) => internal_err!("this case should not be
possible"),
}?;
let action_label = match &action {
+ JoinSelectionAction::StageBuildSide { .. } => "StageBuildSide",
JoinSelectionAction::Repartition(_) => "Repartition",
JoinSelectionAction::CollectLeft(_) => "CollectLeft(broadcast)",
JoinSelectionAction::LateCollectLeft(_) =>
"LateCollectLeft(broadcast)",
@@ -601,6 +634,65 @@ impl DynamicJoinSelectionExec {
}
}
+/// 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.
+///
+/// Staging serialises two shuffles that would otherwise run concurrently, so
the
+/// deferral costs at most the *smaller* side's runtime; `min_probe_ratio` is
+/// what bounds the cost of being wrong.
+///
+/// `max_estimate_multiple` is measured evidence rather than a round number:
+/// 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.
+fn requires_build_staging(
+ build_stats: &Statistics,
+ probe_stats: &Statistics,
+ build_side_join_type: JoinType,
+ broadcast_threshold_bytes: usize,
+ min_probe_ratio: usize,
+ max_estimate_multiple: usize,
+) -> bool {
+ // Measuring can only change the decision if `CollectLeft` is reachable at
+ // all, and that is a property of the post-swap join type alone: no size
+ // makes a `Full`, or a left-sided semi/anti/outer, join broadcastable.
+ if !collect_left_broadcast_safe(build_side_join_type) {
+ return false;
+ }
+
+ if broadcast_threshold_bytes == 0 {
+ return false;
+ }
+
+ let Precision::Inexact(build_bytes) = build_stats.total_byte_size else {
+ return false;
+ };
+
+ if build_bytes < broadcast_threshold_bytes {
+ return false;
+ }
+
+ if build_bytes >
broadcast_threshold_bytes.saturating_mul(max_estimate_multiple) {
+ return false;
+ }
+
+ // A missing probe-side size declines: with nothing to compare against
there
+ // is no evidence the round trip pays off.
+ let Some(probe_bytes) = probe_stats.total_byte_size.get_value().copied()
else {
+ return false;
+ };
+
+ probe_bytes >= build_bytes.saturating_mul(min_probe_ratio)
+}
+
/// Whether the build (left) side of a `Partitioned` hash join fits the
/// per-slot memory budget. A `max_build_bytes` of 0 disables the check (the
/// join always uses hash join); an unknown build size
(`build_max_partition_bytes
@@ -741,7 +833,6 @@ mod tests {
PartitionId, PartitionLocation, PartitionStats,
};
use datafusion::arrow::datatypes::{DataType, Field, Schema};
- use datafusion::common::{ColumnStatistics, Statistics, stats::Precision};
use datafusion::config::ExtensionOptions;
use datafusion::physical_plan::Partitioning;
use datafusion::physical_plan::displayable;
@@ -895,32 +986,38 @@ mod tests {
assert!(!supports_collect(&plan));
}
- fn stats_exec(num_rows: usize) -> Arc<dyn ExecutionPlan> {
+ /// A single-column `Int32` source reporting exactly the statistics given.
+ /// The wrappers around it differ only in how `total_byte_size` is filled.
+ fn stats_exec_with(
+ num_rows: Precision<usize>,
+ total_byte_size: Precision<usize>,
+ ) -> Arc<dyn ExecutionPlan> {
Arc::new(StatisticsExec::new(
Statistics {
- num_rows: Precision::Inexact(num_rows),
- total_byte_size: Precision::Inexact(num_rows * 16),
+ num_rows,
+ total_byte_size,
column_statistics: vec![ColumnStatistics::new_unknown()],
},
Schema::new(vec![Field::new("k", DataType::Int32, false)]),
))
}
+ /// A source reporting `num_rows` rows at 16 bytes each.
+ fn stats_exec(num_rows: usize) -> Arc<dyn ExecutionPlan> {
+ stats_exec_with(
+ Precision::Inexact(num_rows),
+ Precision::Inexact(num_rows * 16),
+ )
+ }
+
/// A source that reports a row count but no byte-size statistic, so the
/// resolver must fall back to the row-count broadcast threshold.
fn stats_exec_rows_only(num_rows: usize) -> Arc<dyn ExecutionPlan> {
- Arc::new(StatisticsExec::new(
- Statistics {
- num_rows: Precision::Inexact(num_rows),
- total_byte_size: Precision::Absent,
- column_statistics: vec![ColumnStatistics::new_unknown()],
- },
- Schema::new(vec![Field::new("k", DataType::Int32, false)]),
- ))
+ stats_exec_with(Precision::Inexact(num_rows), Precision::Absent)
}
/// Core driver: runs the join-strategy decision for the given children
with
- /// explicit Ballista broadcast thresholds.
+ /// explicit Ballista broadcast thresholds and build-side staging setting.
fn run_to_actual_join(
left: Arc<dyn ExecutionPlan>,
right: Arc<dyn ExecutionPlan>,
@@ -928,6 +1025,7 @@ mod tests {
prefer_hash_join: bool,
broadcast_threshold_bytes: usize,
broadcast_threshold_rows: usize,
+ stage_build_side: bool,
) -> JoinSelectionAction {
let on: JoinOn =
vec![(Arc::new(Column::new("k", 0)), Arc::new(Column::new("k",
0)))];
@@ -957,6 +1055,8 @@ mod tests {
&broadcast_threshold_rows.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()
}
@@ -999,6 +1099,9 @@ mod tests {
// Large row threshold: byte stats are present here, so the row
path
// is never consulted; keep it out of the way of the byte decision.
1_000_000,
+ // The shipped default. These sides are far too small to stage, so
+ // it only matters that the setting matches production.
+ true,
)
}
@@ -1103,6 +1206,7 @@ mod tests {
true,
10 * 1024 * 1024,
200,
+ true,
)),
"smaller side (100 rows) is under a 200-row threshold and must
collect"
);
@@ -1114,6 +1218,7 @@ mod tests {
true,
10 * 1024 * 1024,
50,
+ true,
)),
"neither side is under a 50-row threshold, so the join must
repartition"
);
@@ -1147,6 +1252,7 @@ mod tests {
10_000,
// Row ceiling well clear of both sides, so bytes decide.
1_000_000,
+ true,
);
assert!(
@@ -1166,6 +1272,209 @@ 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> {
+ stats_exec_with(Precision::Inexact(1_000), total_byte_size)
+ }
+
+ /// 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;
+
+ /// Statistics carrying only a byte size, which is all
+ /// [`requires_build_staging`] reads off each side.
+ fn byte_stats(total_byte_size: Precision<usize>) -> Statistics {
+ Statistics {
+ num_rows: Precision::Inexact(1_000),
+ total_byte_size,
+ column_statistics: vec![],
+ }
+ }
+
+ /// [`requires_build_staging`] for an `Inner` join under 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 stages(
+ build: Precision<usize>,
+ probe: Precision<usize>,
+ broadcast_threshold_bytes: usize,
+ ) -> bool {
+ let bc = BallistaConfig::default();
+ requires_build_staging(
+ &byte_stats(build),
+ &byte_stats(probe),
+ JoinType::Inner,
+ broadcast_threshold_bytes,
+ bc.stage_build_side_min_probe_ratio(),
+ bc.stage_build_side_max_estimate_multiple(),
+ )
+ }
+
+ // Staging turns on a build side whose size is a *guess* over the budget,
+ // against a probe side big enough that deferring its shuffle repays the
+ // round trip. Each negative below removes exactly one of those properties.
+ #[test]
+ fn stages_only_a_guessed_build_side_against_a_much_larger_probe() {
+ let ratio =
BallistaConfig::default().stage_build_side_min_probe_ratio();
+
+ assert!(
+ stages(
+ Precision::Inexact(GUESSED_BUILD_BYTES),
+ Precision::Exact(FACT_PROBE_BYTES),
+ STAGE_THRESHOLD
+ ),
+ "a guessed-large build side against a fact-table probe must stage"
+ );
+ assert!(
+ !stages(
+ Precision::Exact(GUESSED_BUILD_BYTES),
+ Precision::Exact(FACT_PROBE_BYTES),
+ STAGE_THRESHOLD
+ ),
+ "an exact size over budget is a fact, with nothing left to measure"
+ );
+ assert!(
+ !stages(
+ Precision::Inexact(GUESSED_BUILD_BYTES),
+ Precision::Exact(GUESSED_BUILD_BYTES * (ratio - 1)),
+ STAGE_THRESHOLD
+ ),
+ "a comparable probe side does not repay serialising the shuffles"
+ );
+ assert!(
+ !stages(
+ Precision::Inexact(GUESSED_BUILD_BYTES),
+ Precision::Absent,
+ STAGE_THRESHOLD
+ ),
+ "a missing probe size is no evidence the round trip pays off"
+ );
+ }
+
+ // The broadcast budget bounds staging at both ends, and `0` disables it
+ // along with the promotion it exists to reach.
+ #[test]
+ fn staging_is_bounded_by_the_broadcast_budget() {
+ let multiple =
BallistaConfig::default().stage_build_side_max_estimate_multiple();
+ let far_over = STAGE_THRESHOLD * (multiple + 1);
+
+ assert!(
+ !stages(
+ Precision::Inexact(STAGE_THRESHOLD - 1),
+ Precision::Exact(FACT_PROBE_BYTES),
+ STAGE_THRESHOLD
+ ),
+ "a build side already under budget is not what forced Partitioned"
+ );
+ assert!(
+ !stages(
+ Precision::Inexact(far_over),
+ Precision::Exact(far_over * 1_000),
+ STAGE_THRESHOLD
+ ),
+ "past the multiple no measurement brings the side under budget"
+ );
+ assert!(
+ !stages(
+ Precision::Inexact(GUESSED_BUILD_BYTES),
+ Precision::Exact(FACT_PROBE_BYTES),
+ 0
+ ),
+ "broadcast promotion disabled leaves no decision to revisit"
+ );
+ }
+
+ // 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 for `join_type`, staging on or off.
+ fn q8_shaped_action(
+ join_type: JoinType,
+ stage_build_side: bool,
+ ) -> JoinSelectionAction {
+ run_to_actual_join(
+ sized_stats_exec(Precision::Inexact(GUESSED_BUILD_BYTES)),
+ sized_stats_exec(Precision::Exact(FACT_PROBE_BYTES)),
+ join_type,
+ true,
+ STAGE_THRESHOLD,
+ ROW_THRESHOLD,
+ stage_build_side,
+ )
+ }
+
+ // Staging pays off only where `CollectLeft` is reachable, which is a
+ // property of the post-swap join type alone. The left side is the smaller
+ // one here, so `supports_swap_join_order` is false and the post-swap type
+ // equals the declared one, making `collect_left_broadcast_safe` the
oracle.
+ // Looping all ten means a join type added to DataFusion cannot slip
through
+ // untested — which two hand-partitioned lists would have let it do.
+ #[test]
+ fn stages_exactly_the_join_types_that_can_broadcast() {
+ for join_type in [
+ JoinType::Inner,
+ JoinType::Left,
+ JoinType::Right,
+ JoinType::Full,
+ JoinType::LeftSemi,
+ JoinType::RightSemi,
+ JoinType::LeftAnti,
+ JoinType::RightAnti,
+ JoinType::LeftMark,
+ JoinType::RightMark,
+ ] {
+ let action = q8_shaped_action(join_type, true);
+ assert_eq!(
+ matches!(action, JoinSelectionAction::StageBuildSide { .. }),
+ collect_left_broadcast_safe(join_type),
+ "join_type {join_type:?}: staging must follow broadcast
safety",
+ );
+ }
+ }
+
+ #[test]
+ fn stage_build_side_config_turns_the_behaviour_off() {
+ assert!(matches!(
+ q8_shaped_action(JoinType::Inner, false),
+ JoinSelectionAction::Repartition(_)
+ ));
+ }
+
+ // Staging must fire at most once. Once a child is an exchange the resolver
+ // falls through to `Repartition`, which shuffles the remaining side and
+ // reuses the exchange already in place.
+ #[test]
+ fn does_not_stage_again_once_a_child_is_an_exchange() {
+ let staged: Arc<dyn ExecutionPlan> = Arc::new(ExchangeExec::new(
+ sized_stats_exec(Precision::Inexact(GUESSED_BUILD_BYTES)),
+ None,
+ 0,
+ ));
+
+ let action = run_to_actual_join(
+ staged,
+ sized_stats_exec(Precision::Exact(FACT_PROBE_BYTES)),
+ JoinType::Inner,
+ true,
+ STAGE_THRESHOLD,
+ ROW_THRESHOLD,
+ true,
+ );
+
+ assert!(matches!(action, JoinSelectionAction::Repartition(_)));
+ }
+
/// Constructs a minimal `DynamicJoinSelectionExec` around the given
children.
/// Properties are borrowed from the left child — correct enough for unit
tests
/// that only exercise `children_resolved`.
@@ -1382,6 +1691,9 @@ mod tests {
/// than matching on the action's structure.
fn resolved_plan(action: &JoinSelectionAction) -> Arc<dyn ExecutionPlan> {
match action {
+ JoinSelectionAction::StageBuildSide { .. } => {
+ unreachable!("staging carries no resolved plan to render")
+ }
JoinSelectionAction::Repartition(exec) => {
Arc::clone(exec) as Arc<dyn ExecutionPlan>
}
diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs
b/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs
index cc804ba0b..a48636e47 100644
--- a/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs
@@ -24,7 +24,7 @@ use datafusion::{
error::DataFusionError,
physical_optimizer::PhysicalOptimizerRule,
physical_plan::{
- ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions,
+ ChildrenPropertiesMode, ExecutionPlan, Partitioning,
ReplaceChildrenOptions,
joins::{HashJoinExec, PartitionMode, SortMergeJoinExec},
},
};
@@ -134,6 +134,51 @@ impl SelectJoinRule {
}
}
+/// The shuffle partitioning a join's `(left, right)` inputs need to be
+/// co-partitioned on its keys.
+fn join_key_partitioning(
+ join: &DynamicJoinSelectionExec,
+ partition_count: usize,
+) -> (Partitioning, Partitioning) {
+ let mut required = join
+ ._required_input_distribution()
+ .into_iter()
+ .map(|d| d.create_partitioning(partition_count));
+ let left = required.next().expect("a join requires two distributions");
+ let right = required.next().expect("a join requires two distributions");
+ (left, right)
+}
+
+impl SelectJoinRule {
+ /// 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.
+ ///
+ /// The match is structural equality, deliberately narrower than
+ /// [`DynamicJoinSelectionExec::inputs_already_partitioned`], which asks
the
+ /// same question through the equivalence-aware
`Partitioning::satisfaction`.
+ /// The only exchange worth reusing here is one this rule itself just
placed
+ /// on these join keys, so the plainer comparison is the honest test and
the
+ /// two should not be unified.
+ fn ensure_exchange(
+ &self,
+ child: Arc<dyn ExecutionPlan>,
+ partitioning: Partitioning,
+ ) -> 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), self.plan_id()))
+ }
+}
+
/// Whether giving the build side its own broadcast stage is worth it.
fn broadcast_build_side_pays_off(
build: &dyn ExecutionPlan,
@@ -294,30 +339,59 @@ impl PhysicalOptimizerRule for SelectJoinRule {
Ok(Transformed::yes(exec.data))
}
+ JoinSelectionAction::StageBuildSide {
build_is_right } => {
+ // 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 (left_key, right_key) =
join_key_partitioning(
+ dynamic_join,
+ config.execution.target_partitions,
+ );
+
+ let children = if build_is_right {
+ vec![
+ dynamic_join.left.clone(),
+ Arc::new(ExchangeExec::new(
+ dynamic_join.right.clone(),
+ Some(right_key),
+ self.plan_id(),
+ ))
+ as Arc<dyn ExecutionPlan>,
+ ]
+ } else {
+ vec![
+ Arc::new(ExchangeExec::new(
+ dynamic_join.left.clone(),
+ Some(left_key),
+ self.plan_id(),
+ ))
+ as Arc<dyn ExecutionPlan>,
+ dynamic_join.right.clone(),
+ ]
+ };
+
+ let staged =
Arc::clone(&node).replace_children(
+ children,
+ ReplaceChildrenOptions::new(
+ ChildrenPropertiesMode::Recompute,
+ ),
+ )?;
+
+ Ok(Transformed::yes(staged))
+ }
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(),
- ));
-
- let right = Arc::new(ExchangeExec::new(
- right,
- Some(partitioning[1].clone()),
- self.plan_id(),
- ));
+ let (left_key, right_key) =
join_key_partitioning(
+ &dynamic_join,
+ config.execution.target_partitions,
+ );
+
+ let left = self
+
.ensure_exchange(dynamic_join.left.clone(), left_key);
+ let right = self.ensure_exchange(
+ dynamic_join.right.clone(),
+ right_key,
+ );
let dynamic_join =
dynamic_join.replace_children(
vec![left, right],
@@ -391,6 +465,59 @@ 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 =
+ SelectJoinRule::default().ensure_exchange(Arc::clone(&staged),
partitioning);
+
+ 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 = SelectJoinRule::default()
+ .ensure_exchange(staged, key_partitioning("k", &schema));
+ let outer = wrapped.downcast_ref::<ExchangeExec>().unwrap();
+
+ assert!(outer.input().is::<ExchangeExec>());
+ }
+
fn make_table(schema: Arc<Schema>) -> Arc<MemTable> {
let batch = RecordBatch::try_new(
Arc::clone(&schema),
diff --git a/ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs
b/ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs
index 907df743a..47644e90b 100644
--- a/ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs
+++ b/ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs
@@ -29,23 +29,19 @@
//! build. Sizes are the input to this decision, so they need to be the thing
//! the test varies.
-use crate::state::aqe::test::stats_table::{
- StatsTable, sized_statistics, sizeless_statistics,
-};
-use crate::state::aqe::{
- planner::AdaptivePlanner, test::mock_partitions_with_statistics,
+use crate::state::aqe::planner::AdaptivePlanner;
+use crate::state::aqe::test::stats_table::{sized_statistics,
sizeless_statistics};
+use crate::state::aqe::test::{
+ MB, ballista_ctx, mock_partitions_with_statistics, narrow_schema,
+ register_stats_table as register,
};
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 datafusion::execution::config::SessionConfig;
+use datafusion::execution::context::SessionContext;
use datafusion::physical_plan::displayable;
use std::sync::Arc;
-const MB: usize = 1024 * 1024;
-
/// A join key plus two variable-width columns, so the row width is realistic
and
/// `Statistics::calculate_total_byte_size` cannot reconstruct a size for it.
///
@@ -67,31 +63,6 @@ fn wide_schema() -> Arc<Schema> {
]))
}
-/// A join key plus one fixed-width column.
-fn narrow_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
exercise
-/// the thresholds a deployment runs with.
-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, schema: Arc<Schema>, stats:
Statistics) {
- ctx.register_table(name, Arc::new(StatsTable::new(schema, stats, 4)))
- .unwrap();
-}
-
/// Resolves the join and reports whether any side ended up broadcast.
async fn plan_broadcasts(ctx: &SessionContext, sql: &str) -> (bool, String) {
let lp = ctx.sql(sql).await.unwrap().into_optimized_plan().unwrap();
diff --git a/ballista/scheduler/src/state/aqe/test/mod.rs
b/ballista/scheduler/src/state/aqe/test/mod.rs
index 5d3e2cca1..9dbbbf32c 100644
--- a/ballista/scheduler/src/state/aqe/test/mod.rs
+++ b/ballista/scheduler/src/state/aqe/test/mod.rs
@@ -30,9 +30,12 @@ mod plan_to_stages;
/// Regression tests for range-repartition planning end-to-end
/// through `AdaptivePlanner` (DER → routing park → filter injection).
mod range_repartition;
+/// Multi-pass coverage for staging a join's build side before deciding the
join
+mod stage_build_side;
/// A table whose statistics are declared rather than measured
mod stats_table;
+use ballista_core::extension::SessionConfigExt;
use ballista_core::serde::scheduler::{
ExecutorMetadata, ExecutorOperatingSystemSpecification,
ExecutorSpecification,
PartitionId, PartitionLocation, PartitionStats,
@@ -40,12 +43,58 @@ use ballista_core::serde::scheduler::{
use datafusion::arrow::array::{Int32Array, RecordBatch};
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion::catalog::TableProvider;
+use datafusion::common::Statistics;
use datafusion::datasource::MemTable;
use datafusion::execution::SessionStateBuilder;
use datafusion::prelude::{SessionConfig, SessionContext};
+use stats_table::StatsTable;
use std::sync::Arc;
+/// One mebibyte, for declaring fixture sizes.
+pub(crate) const MB: usize = 1024 * 1024;
+
+/// A join key plus one fixed-width column.
+pub(crate) fn narrow_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 tests exercise the
+/// thresholds a deployment runs with rather than DataFusion's defaults.
+pub(crate) 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)
+}
+
+/// Registers a [`StatsTable`] declaring `stats` over `schema`.
+pub(crate) fn register_stats_table(
+ ctx: &SessionContext,
+ name: &str,
+ schema: Arc<Schema>,
+ stats: Statistics,
+) {
+ ctx.register_table(name, Arc::new(StatsTable::new(schema, stats, 4)))
+ .unwrap();
+}
+
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>> {
let location = PartitionLocation {
// next few properties are generic values
map_partition_id: 0,
@@ -63,7 +112,7 @@ pub(crate) fn mock_partitions_with_statistics() ->
Vec<Vec<PartitionLocation>> {
os_info: ExecutorOperatingSystemSpecification::default(),
},
// next few properties are needed
- partition_stats: PartitionStats::new(Some(42), None, Some(10)),
+ partition_stats: PartitionStats::new(Some(num_rows), None,
Some(num_bytes)),
file_id: None,
is_sort_shuffle: false,
};
@@ -71,28 +120,7 @@ pub(crate) fn mock_partitions_with_statistics() ->
Vec<Vec<PartitionLocation>> {
}
pub(crate) fn mock_partitions_with_statistics_no_data() ->
Vec<Vec<PartitionLocation>> {
- let location = PartitionLocation {
- // next few properties are generic values
- map_partition_id: 0,
- partition_id: PartitionId {
- job_id: "".into(),
- stage_id: 0,
- partition_id: 0,
- },
- executor_meta: ExecutorMetadata {
- id: "".to_string(),
- host: "".to_string(),
- port: 0,
- grpc_port: 0,
- specification: ExecutorSpecification::default().with_vcores(0),
- os_info: ExecutorOperatingSystemSpecification::default(),
- },
- // next few properties are needed
- partition_stats: PartitionStats::new(Some(0), None, Some(0)),
- file_id: None,
- is_sort_shuffle: false,
- };
- vec![vec![location]]
+ mock_partitions_with_size(0, 0)
}
/// Returns schema with three columns (a,b,c) all of [DataType::Int32] type
diff --git a/ballista/scheduler/src/state/aqe/test/stage_build_side.rs
b/ballista/scheduler/src/state/aqe/test/stage_build_side.rs
new file mode 100644
index 000000000..428fcb5bf
--- /dev/null
+++ b/ballista/scheduler/src/state/aqe/test/stage_build_side.rs
@@ -0,0 +1,167 @@
+// 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 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::{estimated_statistics,
sized_statistics};
+use crate::state::aqe::test::{
+ MB, ballista_ctx, mock_partitions_with_size,
mock_partitions_with_statistics,
+ narrow_schema, register_stats_table,
+};
+use ballista_core::assert_plan;
+use datafusion::execution::context::SessionContext;
+
+/// 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;
+
+const Q8_SHAPED_JOIN: &str = "SELECT dim.val FROM dim JOIN fact ON dim.id =
fact.id";
+
+/// 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 = narrow_schema();
+ register_stats_table(
+ &ctx,
+ "dim",
+ schema.clone(),
+ estimated_statistics(&schema, 1_000_000, GUESSED_BUILD_BYTES),
+ );
+ register_stats_table(
+ &ctx,
+ "fact",
+ schema.clone(),
+ 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()
+}
+
+/// 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)
+ ");
+}
diff --git a/ballista/scheduler/src/state/aqe/test/stats_table.rs
b/ballista/scheduler/src/state/aqe/test/stats_table.rs
index 2f0aa80de..a41a93faf 100644
--- a/ballista/scheduler/src/state/aqe/test/stats_table.rs
+++ b/ballista/scheduler/src/state/aqe/test/stats_table.rs
@@ -224,6 +224,20 @@ 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 {
+ sized_statistics(schema, num_rows, total_byte_size).to_inexact()
+}
+
/// Statistics for a table of `num_rows` occupying a known `total_byte_size`.
pub(crate) fn sized_statistics(
schema: &Schema,
diff --git a/docs/source/user-guide/configs.md
b/docs/source/user-guide/configs.md
index 91b068356..72dafc54e 100644
--- a/docs/source/user-guide/configs.md
+++ b/docs/source/user-guide/configs.md
@@ -104,35 +104,38 @@ standard DataFusion settings.
<!-- BEGIN GENERATED CONFIG REFERENCE
prefix=ballista.job.,ballista.standalone.,ballista.cache.,ballista.client.,ballista.optimizer.,ballista.planner.,ballista.scheduler.
-->
<!-- prettier-ignore -->
-| key | type | default
| description
[...]
-| ------------------------------------------------------ | ------- |
----------------------------- |
-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
[...]
-| ballista.cache.noop | Boolean | true
| Disable default cache node extension
[...]
-| ballista.client.grpc_connect_timeout_seconds | UInt64 | 20
| Connection timeout for gRPC client in seconds
[...]
-| ballista.client.grpc_http2_keepalive_interval_seconds | UInt64 | 300
| HTTP/2 keep-alive interval for gRPC client in seconds
[...]
-| ballista.client.grpc_max_message_size | UInt64 | 16777216
| Configuration for max message size in gRPC clients
[...]
-| ballista.client.grpc_tcp_keepalive_seconds | UInt64 | 3600
| TCP keep-alive interval for gRPC client in seconds
[...]
-| ballista.client.grpc_timeout_seconds | UInt64 | 20
| Request timeout for gRPC client in seconds
[...]
-| ballista.client.initial_connection_window_size | UInt64 | 67108864
| HTTP/2 initial connection-level flow-control window for
gRPC data-plane clients, in bytes. Should be >= the shuffle governor byte
budget so the governor, not the transport window, is the binding backpressure.
0 leaves the tonic default.
[...]
-| ballista.client.initial_stream_window_size | UInt64 | 16777216
| HTTP/2 initial stream-level flow-control window for gRPC
data-plane clients, in bytes. 0 leaves the tonic default.
[...]
-| ballista.client.io_retries_times | UInt16 | 3
| Number of retries for IO operations in the Ballista
client.
[...]
-| ballista.client.io_retry_wait_time_ms | UInt64 | 3000
| Wait time in milliseconds between IO retries in the
Ballista client.
[...]
-| ballista.client.pull | Boolean | false
| Should client employ pull or push job tracking. In pull
mode client will make a request to server in the loop, until job finishes. Pull
mode is kept for legacy clients.
[...]
-| ballista.client.use_tls | Boolean | false
| Should connection between client, scheduler, and
executors use TLS.
[...]
-| ballista.job.name | Utf8 | (none)
| Sets the job name that will appear in the web user
interface for any submitted jobs
[...]
-| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 134217728
| Byte-size threshold below which a hash join's smaller
side is promoted to CollectLeft and lowered via the broadcast pattern. Governs
broadcast selection under both the static distributed planner and adaptive
query planning (AQE). It also caps null-aware anti joins with a known build
size because they require single-task CollectLeft execution. Set to 0 to
disable promotion and reject null- [...]
-| ballista.optimizer.broadcast_join_threshold_rows | UInt64 | 1000000
| Row-count threshold below which a hash join's smaller
side is promoted to CollectLeft and lowered via the broadcast pattern, used as
a fallback when byte-size statistics are unavailable. Applies to adaptive query
planning (AQE). Set to 0 to disable promotion via the row-count path.
[...]
-| ballista.optimizer.hash_join_max_build_partition_bytes | UInt64 | 67108864
| Maximum per-partition hash-join build-side bytes for a
Partitioned hash join under AQE. A build partition larger than this falls back
to SortMergeJoin (spillable). Defaults to 64 MiB; 0 disables the check, which
makes AQE use a hash join regardless of build size.
[...]
-| ballista.optimizer.not_in_subquery_rewrite | Boolean | true
| Rewrites uncorrelated NOT IN (subquery) filter predicates
into a plain anti join plus a one-row count aggregate during logical
optimization. The rewrite avoids DataFusion's null-aware hash join, which
Ballista must otherwise execute in a single task. Set to false to keep the
null-aware join and its single-task lowering.
[...]
-| ballista.planner.adaptive.enabled | Boolean | true
| Enables Adaptive Query Planning: joins and partition
counts are chosen from measured runtime statistics instead of planning-time
estimates. Set to false to use the static distributed planner.
[...]
-| ballista.planner.adaptive_join.enabled | Boolean | true
| Enables the AQE dynamic join-selection rule
(SelectJoinRule). When true (default), DynamicJoinSelectionExec nodes are
resolved to concrete HashJoin or CollectLeft join implementations at runtime.
Disable only for debugging.
[...]
-| ballista.planner.coalesce.enabled | Boolean | false
| Enables the AQE coalesce-shuffle-partitions rule.
Disabled by default — opt in when fewer/larger downstream tasks matter more
than parallelism.
[...]
-| ballista.planner.coalesce.merged_partition_factor | Float64 | 1.2
| Two adjacent partitions are merged when their combined
size is below target_partition_bytes times this factor. Mirrors Spark's legacy
coalesce semantics.
[...]
-| ballista.planner.coalesce.small_partition_factor | Float64 | 0.2
| A coalesced partition smaller than target_partition_bytes
times this factor counts as small and is merged into its neighbour. Mirrors
Spark's legacy coalesce semantics.
[...]
-| ballista.planner.coalesce.target_partition_bytes | UInt64 | 67108864
| Target post-coalesce partition size in bytes. Mirrors
Spark's advisoryPartitionSizeInBytes.
[...]
-| ballista.planner.parallel_window.enabled | Boolean | false
| Enables the AQE parallel-window rule
(ParallelWindowRule), which rewrites bounded-RANGE-frame windows into a
distributed range-shuffle so BoundedWindowAggExec's single-partition constraint
is not a serial bottleneck. Disabled by default — opt in when the workload
contains matching window shapes.
[...]
-| ballista.planner.propagate_empty.enabled | Boolean | true
| Enables the AQE propagate-empty-relation rule. Injects
EmptyExec into the plan where an input is known to be empty, such as one side
of a join, allowing downstream work to be skipped.
[...]
-| ballista.scheduler.max_partitions_per_task | UInt64 | 0
| Upper bound on the number of input partitions packed into
a single task's `partition_slice`. `0` (default) means unbounded — the
scheduler fills each task up to the executor's free vcore count. `1` means one
task per input partition. Does not apply to collapse stages, which must pack
their full pending queue into a single task for correctness.
[...]
-| ballista.standalone.parallelism | UInt16 | number of
available CPU cores | Number of concurrent tasks a standalone in-process
executor will run.
[...]
+| key | type |
default | description
[...]
+| --------------------------------------------------------- | ------- |
----------------------------- |
----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
[...]
+| ballista.cache.noop | Boolean | true
| Disable default cache node extension
[...]
+| ballista.client.grpc_connect_timeout_seconds | UInt64 | 20
| Connection timeout for gRPC client in seconds
[...]
+| ballista.client.grpc_http2_keepalive_interval_seconds | UInt64 | 300
| HTTP/2 keep-alive interval for gRPC client in seconds
[...]
+| ballista.client.grpc_max_message_size | UInt64 |
16777216 | Configuration for max message size in gRPC
clients
[...]
+| ballista.client.grpc_tcp_keepalive_seconds | UInt64 | 3600
| TCP keep-alive interval for gRPC client in seconds
[...]
+| ballista.client.grpc_timeout_seconds | UInt64 | 20
| Request timeout for gRPC client in seconds
[...]
+| ballista.client.initial_connection_window_size | UInt64 |
67108864 | HTTP/2 initial connection-level flow-control
window for gRPC data-plane clients, in bytes. Should be >= the shuffle governor
byte budget so the governor, not the transport window, is the binding
backpressure. 0 leaves the tonic default.
[...]
+| ballista.client.initial_stream_window_size | UInt64 |
16777216 | HTTP/2 initial stream-level flow-control window
for gRPC data-plane clients, in bytes. 0 leaves the tonic default.
[...]
+| ballista.client.io_retries_times | UInt16 | 3
| Number of retries for IO operations in the Ballista
client.
[...]
+| ballista.client.io_retry_wait_time_ms | UInt64 | 3000
| Wait time in milliseconds between IO retries in the
Ballista client.
[...]
+| ballista.client.pull | Boolean | false
| Should client employ pull or push job tracking. In
pull mode client will make a request to server in the loop, until job finishes.
Pull mode is kept for legacy clients.
[...]
+| ballista.client.use_tls | Boolean | false
| Should connection between client, scheduler, and
executors use TLS.
[...]
+| ballista.job.name | Utf8 | (none)
| Sets the job name that will appear in the web user
interface for any submitted jobs
[...]
+| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 |
134217728 | Byte-size threshold below which a hash join's
smaller side is promoted to CollectLeft and lowered via the broadcast pattern.
Governs broadcast selection under both the static distributed planner and
adaptive query planning (AQE). It also caps null-aware anti joins with a known
build size because they require single-task CollectLeft execution. Set to 0 to
disable promotion and reject nu [...]
+| ballista.optimizer.broadcast_join_threshold_rows | UInt64 |
1000000 | Row-count threshold below which a hash join's
smaller side is promoted to CollectLeft and lowered via the broadcast pattern,
used as a fallback when byte-size statistics are unavailable. Applies to
adaptive query planning (AQE). Set to 0 to disable promotion via the row-count
path.
[...]
+| ballista.optimizer.hash_join_max_build_partition_bytes | UInt64 |
67108864 | Maximum per-partition hash-join build-side
bytes for a Partitioned hash join under AQE. A build partition larger than this
falls back to SortMergeJoin (spillable). Defaults to 64 MiB; 0 disables the
check, which makes AQE use a hash join regardless of build size.
[...]
+| ballista.optimizer.not_in_subquery_rewrite | Boolean | true
| Rewrites uncorrelated NOT IN (subquery) filter
predicates into a plain anti join plus a one-row count aggregate during logical
optimization. The rewrite avoids DataFusion's null-aware hash join, which
Ballista must otherwise execute in a single task. Set to false to keep the
null-aware join and its single-task lowering.
[...]
+| ballista.optimizer.stage_build_side | Boolean | true
| Stages a join's prospective build side as its own
shuffle before choosing the join strategy. Applies only when that side's size
is an estimate rather than a measurement, and the probe side is far larger: an
exact size over the broadcast budget is a fact that no further measurement can
change. Lets AQE measure the build side and pick a broadcast join instead of
committing to a full shuf [...]
+| ballista.optimizer.stage_build_side_max_estimate_multiple | UInt64 | 32
| How far over broadcast_join_threshold_bytes an
estimated build side may sit and still be worth measuring. Past this multiple
the side is large on any reading and no measurement brings it back under
budget, so it is shuffled without the extra round trip. Only consulted when
stage_build_side is enabled, and only for build sides whose size is an estimate
rather than a measurement. [...]
+| ballista.optimizer.stage_build_side_min_probe_ratio | UInt64 | 10
| How many times larger the probe side must be before
AQE stages a join's build side on its own. Staging serialises two shuffles that
would otherwise run concurrently, so the deferral costs at most the smaller
side's runtime; this ratio bounds the cost of being wrong. Only consulted when
stage_build_side is enabled, and only for build sides whose size is an estimate
rather than a measure [...]
+| ballista.planner.adaptive.enabled | Boolean | true
| Enables Adaptive Query Planning: joins and partition
counts are chosen from measured runtime statistics instead of planning-time
estimates. Set to false to use the static distributed planner.
[...]
+| ballista.planner.adaptive_join.enabled | Boolean | true
| Enables the AQE dynamic join-selection rule
(SelectJoinRule). When true (default), DynamicJoinSelectionExec nodes are
resolved to concrete HashJoin or CollectLeft join implementations at runtime.
Disable only for debugging.
[...]
+| ballista.planner.coalesce.enabled | Boolean | false
| Enables the AQE coalesce-shuffle-partitions rule.
Disabled by default — opt in when fewer/larger downstream tasks matter more
than parallelism.
[...]
+| ballista.planner.coalesce.merged_partition_factor | Float64 | 1.2
| Two adjacent partitions are merged when their combined
size is below target_partition_bytes times this factor. Mirrors Spark's legacy
coalesce semantics.
[...]
+| ballista.planner.coalesce.small_partition_factor | Float64 | 0.2
| A coalesced partition smaller than
target_partition_bytes times this factor counts as small and is merged into its
neighbour. Mirrors Spark's legacy coalesce semantics.
[...]
+| ballista.planner.coalesce.target_partition_bytes | UInt64 |
67108864 | Target post-coalesce partition size in bytes.
Mirrors Spark's advisoryPartitionSizeInBytes.
[...]
+| ballista.planner.parallel_window.enabled | Boolean | false
| Enables the AQE parallel-window rule
(ParallelWindowRule), which rewrites bounded-RANGE-frame windows into a
distributed range-shuffle so BoundedWindowAggExec's single-partition constraint
is not a serial bottleneck. Disabled by default — opt in when the workload
contains matching window shapes.
[...]
+| ballista.planner.propagate_empty.enabled | Boolean | true
| Enables the AQE propagate-empty-relation rule. Injects
EmptyExec into the plan where an input is known to be empty, such as one side
of a join, allowing downstream work to be skipped.
[...]
+| ballista.scheduler.max_partitions_per_task | UInt64 | 0
| Upper bound on the number of input partitions packed
into a single task's `partition_slice`. `0` (default) means unbounded — the
scheduler fills each task up to the executor's free vcore count. `1` means one
task per input partition. Does not apply to collapse stages, which must pack
their full pending queue into a single task for correctness.
[...]
+| ballista.standalone.parallelism | UInt16 | number
of available CPU cores | Number of concurrent tasks a standalone in-process
executor will run.
[...]
<!-- END GENERATED CONFIG REFERENCE -->
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
