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 bd17e435c feat: Improve TPC-H performance with AQE on (#2084)
bd17e435c is described below
commit bd17e435c401c32af0c1dc81bcbcbe82621fa825
Author: Andy Grove <[email protected]>
AuthorDate: Tue Jul 21 08:46:10 2026 -0600
feat: Improve TPC-H performance with AQE on (#2084)
* feat: make AQE respect broadcast_join_threshold_bytes
The ballista.optimizer.broadcast_join_threshold_bytes config is only
consumed by the static distributed planner (maybe_promote_to_broadcast).
Under adaptive query planning, broadcast (CollectLeft) selection in
DynamicJoinSelectionExec::to_actual_join instead used DataFusion's
hash_join_single_partition_threshold (1 MiB default), so the Ballista key
had no effect and the effective cutoff was silently a different value.
Use broadcast_join_threshold_bytes as the byte threshold in the AQE join
selection path, keeping DataFusion's row threshold as the absent-stats
fallback. A value of 0 disables broadcast promotion, matching the static
planner. This gives a single config key consistent behavior under both
planners.
Closes #2085
* feat: add Ballista broadcast_join_threshold_rows and use it in AQE
Follow-up within the same change: the AQE join-selection path also used
DataFusion's hash_join_single_partition_threshold_rows as the row-count
fallback. There was no Ballista equivalent, so the row threshold still
escaped the single-config goal.
Add ballista.optimizer.broadcast_join_threshold_rows (default 128K,
mirroring DataFusion's previous default) plus SessionConfigExt accessors,
and use it in DynamicJoinSelectionExec::to_actual_join instead of the
DataFusion key. AQE broadcast selection now depends only on Ballista
config. Document both broadcast thresholds in the AQE tuning guide.
* fix: default broadcast_join_threshold_rows to 1M to match existing
behavior
SessionConfig::new_with_ballista() already installs a 1,000,000 row
threshold for DataFusion's hash_join_single_partition_threshold_rows.
Default the new Ballista row key to the same value so consolidating AQE
onto the Ballista keys does not silently lower the effective row-count
broadcast cutoff.
* refactor: derive new_with_ballista join thresholds from Ballista config
defaults
Instead of hard-coding 10 MiB / 1M for the DataFusion
hash_join_single_partition_threshold[_rows] session settings, read them
from BallistaConfig::default().broadcast_join_threshold_bytes()/_rows().
The Ballista broadcast-threshold defaults are now the single source of
truth for both DataFusion's built-in JoinSelection and Ballista's AQE
join selection. Values are unchanged.
* feat: demote DataFusion CollectLeft joins over the Ballista threshold
In the static planner, maybe_promote_to_broadcast trusted any
broadcast-safe HashJoinExec(CollectLeft) that DataFusion's JoinSelection
produced. DataFusion decides CollectLeft from its own session threshold,
which can exceed a runtime override of
ballista.optimizer.broadcast_join_threshold_bytes. Demote such a join back
to Partitioned when its build side is not under the current Ballista
threshold (or when broadcasts are disabled with threshold 0), so the
Ballista key is authoritative in the static path too. Null-aware anti
joins are never demoted since they require CollectLeft.
* docs: apply prettier formatting to AQE tuning-guide table
* test: disable AQE broadcast via Ballista threshold in repartition tests
The join-selection repartition tests forced the repartition path by setting
DataFusion's hash_join_single_partition_threshold[_rows] to 0. AQE join
selection now reads the broadcast cutoff from the Ballista config
(broadcast_join_threshold_bytes), so those DataFusion keys no longer gate
CollectLeft promotion and the small test tables were promoted to broadcast.
Set the Ballista broadcast byte threshold to 0 in the helper, which disables
CollectLeft promotion and restores the repartitioned plans the snapshots
assert.
* fix: size the AQE broadcast decision by bytes, not row count
`supports_collect_by_thresholds` compared a row count against
`hash_join_single_partition_threshold_rows` whenever `total_byte_size`
was unknown, so a build side of up to a million arbitrarily wide rows
could be broadcast to every probe task without the byte threshold ever
applying.
Unknown `total_byte_size` is the common case, not an edge case:
DataFusion discards it on every join, and rebuilding it in
`Statistics::calculate_total_byte_size` only works when every column has
a fixed width, so a single `Utf8` column loses it permanently. In TPC-H
that covers most dimension-side join results.
Estimate the size instead and hold it to the same byte threshold. Each
column contributes its own `byte_size` statistic when present -- a total
for that column's output, already scaled for filters and limits --
otherwise its fixed width times the row count, otherwise a default width
mirroring Spark's `StringType`/`BinaryType` defaults. An overflowing
estimate declines the broadcast rather than wrapping to a small number.
The row threshold is retained as a ceiling, so this can only reject a
broadcast the row rule would have allowed, never introduce a new one.
Closes #2081.
* test: cover broadcast thresholds with declared statistics
The broadcast-vs-partitioned decision is a function of statistics, but
the tests around it could only describe tables they were willing to
materialise, so the sizes it actually turns on had no coverage: a build
side of hundreds of thousands of rows, or one whose `total_byte_size` is
unknown. The existing tests instead toggle the decision by zeroing the
threshold, which shows the rule is self-consistent but not that the
shipped thresholds behave.
Add `StatsTable`, a table that declares its statistics and holds no
rows, so a fixture can say "800,000 rows of unknown size" in one line.
Its scan reports the declared figures and cannot be executed, which is
enough for the planner tests, and it deliberately does not recompute
`total_byte_size` on projection, since an unknown size is the case these
fixtures exist to express.
Add tests covering the decision at both edges -- wide rows of unknown
size are not broadcast, while small dimensions, narrow rows, and known
sizes under the threshold still are -- run under
`SessionConfig::new_with_ballista` so they exercise the 10 MB / 1,000,000
row thresholds a deployment ships with rather than DataFusion's
defaults, plus a test pinning those defaults directly.
`wide_rows_of_unknown_size_are_not_broadcast` fails on the rule that
preceded the previous commit and passes with it. The rest pass either
way: they guard against the estimate rejecting broadcasts it should
allow.
Part of #2081.
* feat: log AQE join decisions at INFO and shuffle spills at WARN
Surface two previously debug-only diagnostics at levels visible under the
default INFO filter, so join-strategy and memory-pressure behavior can be
observed without enabling debug logging.
- AQE dynamic join selection now logs each decision at INFO, naming the
resolved action (CollectLeft/Hash/SortMerge/Repartition), the partition
mode, and the size-aware inputs (per-side row/byte estimates and the
byte/row broadcast thresholds).
- Sort-shuffle write completion logs at INFO with row and spill counts;
when a partition spills under memory pressure it logs at WARN with the
spilled bytes, batches, and event count. Per-batch spill events stay at
debug.
* feat: default sort-shuffle spill cap to 0 (uncapped) and plumb it through
serde
The sort-shuffle writer's per-task buffered-bytes cap defaulted to 256 MB,
far
below the per-task memory-pool budget, so it spilled long before the pool
was
under pressure. Default the cap to 0, which disables it: spilling is then
driven
solely by memory-pool pressure. A non-zero value still adds a second spill
trigger and is retained for tests and explicit tuning.
The cap was also dropped during physical-plan serialization (the executor
rebuilt the config with the default), so a configured value never reached
the
executor. Carry memory_limit_per_task_bytes through the
SortShuffleWriterExecNode
protobuf and apply it on decode, so an override via
ballista.shuffle.sort_based.memory_limit_per_task_bytes takes effect on
executors.
* feat: add hash_join_max_build_partition_bytes config
* feat: add max_per_partition_build_bytes helper for hash-join fit check
Reads the actual materialized per-partition byte sizes off the resolved
ExchangeExec feeding a Partitioned hash join's build side (the same source
CoalescePartitionsRule reads) and returns the MAX rather than the average,
since a single oversized partition is enough to OOM even when the average
partition is small (the Q18 failure shape).
* feat: fall back to SMJ when hash-join build exceeds per-slot budget
* feat: log hash-join build-fit decision at INFO
* refactor: hoist build-size call and tidy hash-join fit-check docs
* docs: add hash-join + safety-fallback SF1000 suite results (2x16, p64)
Record a full 22-query SF1000 run with prefer_hash_join=true and the AQE
hash-join
build-size safety fallback (hash_join_max_build_partition_bytes=64 MiB) on a
2 executor x 16 core cluster at target_partitions=64: all 22 queries
complete with
no OOM, where a pure hash-join run fails on Q18. Remove the AQE-off column
pending a
re-run at a matched core count.
* docs: refresh SF1000 Ballista AQE-on results on #2084 build
Re-ran the TPC-H SF1000 suite (AQE on, target_partitions=64,
prefer_hash_join=false, 1 iteration) on the 2x16-core reference cluster
against the PR build (becb3767). Q1-Q17 from a full-suite run, Q19-Q22 as
individual jobs; Q18 still OOMs (Partitioned build side, unchanged).
Ballista total (excl. Q18) improves 4817.8 -> 4661.0s, led by the
join-heavy queries (Q7 -88s, Q8 -191s, Q9 -169s).
* revert: keep upstream 256 MiB sort-shuffle spill-cap default
Drop this branch's change of the sort-shuffle per-task spill-cap default
to 0 (uncapped); restore the 256 MiB default from main. The serde
plumbing for the value stays (it landed on main via #2091). This PR no
longer alters the shipped spill-cap default.
* docs: note SF1000 results use uncapped sort-shuffle spill override
Make explicit that the reference AQE-on numbers were produced with the
sort-shuffle per-task spill cap overridden to 0 (uncapped); the shipped
default is 256 MiB. Point the Ballista row at the current branch commit.
* Update ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs
Co-authored-by: Phillip LeBlanc <[email protected]>
* refactor: log per-partition shuffle write at debug; fix log import
Demote the per-partition shuffle-write completion log to debug (it fires
once per shuffle partition; spills stay at WARN). Update the dynamic_join
log import to match the debug! call.
---------
Co-authored-by: Phillip LeBlanc <[email protected]>
---
ballista/core/src/config.rs | 26 ++
.../src/execution_plans/sort_shuffle/writer.rs | 5 +-
ballista/core/src/extension.rs | 14 +
ballista/core/src/serde/mod.rs | 39 ++
.../src/state/aqe/execution_plan/dynamic_join.rs | 510 ++++++++++++++++++++-
.../src/state/aqe/test/broadcast_thresholds.rs | 299 ++++++++++++
ballista/scheduler/src/state/aqe/test/mod.rs | 4 +
.../scheduler/src/state/aqe/test/stats_table.rs | 232 ++++++++++
docs/source/contributors-guide/benchmarking.md | 113 ++++-
9 files changed, 1196 insertions(+), 46 deletions(-)
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
index 716597ff6..1ec5cadbd 100644
--- a/ballista/core/src/config.rs
+++ b/ballista/core/src/config.rs
@@ -111,6 +111,13 @@ pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES: &str =
pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS: &str =
"ballista.optimizer.broadcast_join_threshold_rows";
+/// Configuration key for the maximum per-partition hash-join build-side bytes
+/// permitted for a Partitioned hash join under AQE. When a build partition
+/// exceeds this, the join falls back to SortMergeJoin (spillable). `0`
disables
+/// the check (hash join is used regardless of build size).
+pub const BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES: &str =
+ "ballista.optimizer.hash_join_max_build_partition_bytes";
+
/// Configuration key to enable AQE coalesce-shuffle-partitions rule.
/// Disabled by default — opt in when the workload benefits from larger
/// downstream tasks more than from preserved parallelism.
@@ -264,6 +271,12 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String,
ConfigEntry>> = LazyLock::new(||
promotion via the row-count path.".to_string(),
DataType::UInt64,
Some((1_000_000).to_string())),
+
ConfigEntry::new(BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES.to_string(),
+ "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). 0 (the default) disables
the check.".to_string(),
+ DataType::UInt64,
+ Some("0".to_string())),
ConfigEntry::new(BALLISTA_CLIENT_PULL.to_string(),
"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.".to_string(),
DataType::Boolean,
@@ -615,6 +628,11 @@ impl BallistaConfig {
self.get_usize_setting(BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS)
}
+ /// Maximum per-partition hash-join build-side bytes before falling back
to SMJ.
+ pub fn hash_join_max_build_partition_bytes(&self) -> usize {
+ self.get_usize_setting(BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES)
+ }
+
/// Returns whether the AQE coalesce-shuffle-partitions rule is enabled.
pub fn coalesce_enabled(&self) -> bool {
self.get_bool_setting(BALLISTA_COALESCE_ENABLED)
@@ -884,4 +902,12 @@ mod tests {
assert_eq!(16777216, config.grpc_client_max_message_size());
Ok(())
}
+
+ #[test]
+ fn hash_join_max_build_partition_bytes_defaults_to_zero() {
+ assert_eq!(
+ BallistaConfig::default().hash_join_max_build_partition_bytes(),
+ 0
+ );
+ }
}
diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs
b/ballista/core/src/execution_plans/sort_shuffle/writer.rs
index 5a41aeb5c..c71d51ace 100644
--- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs
+++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs
@@ -439,8 +439,9 @@ impl SortShuffleWriterExec {
let mut hash_buffer: Vec<u64> = Vec::new();
let mut spill_events: u64 = 0;
// Absolute buffered-bytes counter, independent of the runtime
- // `MemoryPool`. Drives spill decisions so the writer bounds its
- // RSS even when the pool is unbounded.
+ // `MemoryPool`. When `memory_limit` is non-zero it caps this
counter
+ // as a second spill trigger; a `memory_limit` of 0 disables the
cap
+ // so spilling is driven solely by memory-pool pressure.
let mut buffered_bytes: usize = 0;
// A limit of 0 disables the per-task budget, leaving the runtime
// `MemoryPool` as the sole spill trigger.
diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs
index b0a58289b..28fd77e65 100644
--- a/ballista/core/src/extension.rs
+++ b/ballista/core/src/extension.rs
@@ -204,6 +204,10 @@ pub trait SessionConfigExt {
/// disables promotion via the row-count path.
fn with_ballista_broadcast_join_threshold_rows(self, threshold_rows:
usize) -> Self;
+ /// Returns the maximum per-partition hash-join build-side bytes before
+ /// falling back to SortMergeJoin under AQE. `0` disables the check.
+ fn ballista_hash_join_max_build_partition_bytes(&self) -> usize;
+
/// retrieves grpc client max message size
fn ballista_grpc_client_max_message_size(&self) -> usize;
@@ -513,6 +517,16 @@ impl SessionConfigExt for SessionConfig {
}
}
+ fn ballista_hash_join_max_build_partition_bytes(&self) -> usize {
+ self.options()
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.hash_join_max_build_partition_bytes())
+ .unwrap_or_else(|| {
+ BallistaConfig::default().hash_join_max_build_partition_bytes()
+ })
+ }
+
fn ballista_shuffle_reader_maximum_concurrent_requests(&self) -> usize {
self.options()
.extensions
diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs
index ae2b85c93..a22bb977d 100644
--- a/ballista/core/src/serde/mod.rs
+++ b/ballista/core/src/serde/mod.rs
@@ -1143,6 +1143,45 @@ mod test {
assert_eq!(stored.groups[0].upstream_indices, vec![0, 1, 2, 3]);
}
+ #[tokio::test]
+ async fn sort_shuffle_writer_memory_limit_survives_roundtrip() {
+ use datafusion::physical_plan::empty::EmptyExec;
+
+ let schema = create_test_schema();
+ let input: Arc<dyn ExecutionPlan> =
Arc::new(EmptyExec::new(schema.clone()));
+ let partitioning =
+ Partitioning::Hash(vec![col("id", schema.as_ref()).unwrap()], 4);
+
+ let config = SortShuffleConfig::new(true, 4096)
+ .with_memory_limit_per_task_bytes(1024 * 1024 * 1024);
+ let original = SortShuffleWriterExec::try_new(
+ "job-1".to_string().into(),
+ 3,
+ input.clone(),
+ String::new(),
+ partitioning,
+ config,
+ )
+ .unwrap();
+
+ let codec = BallistaPhysicalExtensionCodec::default();
+ let mut buf: Vec<u8> = vec![];
+ codec.try_encode(Arc::new(original), &mut buf).unwrap();
+
+ let ctx = SessionContext::new().task_ctx();
+ let decoded = codec.try_decode(&buf, &[input], &ctx).unwrap();
+ let decoded = decoded
+ .downcast_ref::<SortShuffleWriterExec>()
+ .expect("Expected SortShuffleWriterExec");
+
+ assert_eq!(
+ decoded.config().memory_limit_per_task_bytes,
+ 1024 * 1024 * 1024,
+ "memory limit override must survive serialization to the executor"
+ );
+ assert_eq!(decoded.config().batch_size, 4096);
+ }
+
#[tokio::test]
async fn
test_shuffle_reader_exec_coalesced_roundtrip_multi_group_mixed_sizes() {
let schema = create_test_schema();
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 9391709ea..fbc09c451 100644
--- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs
+++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs
@@ -18,7 +18,8 @@
use ballista_core::config::BallistaConfig;
use datafusion::{
arrow::compute::SortOptions,
- common::{JoinType, NullEquality, Result, exec_err, internal_err},
+ arrow::datatypes::{DataType, Schema},
+ common::{ColumnStatistics, JoinType, NullEquality, Result, exec_err,
internal_err},
config::ConfigOptions,
execution::{SendableRecordBatchStream, TaskContext},
physical_expr_common::physical_expr::fmt_sql,
@@ -228,6 +229,8 @@ impl DynamicJoinSelectionExec {
.unwrap_or_default();
let threshold_collect_left_join_bytes =
bc.broadcast_join_threshold_bytes();
let threshold_collect_left_join_rows =
bc.broadcast_join_threshold_rows();
+ let max_build_bytes = bc.hash_join_max_build_partition_bytes();
+ let build_max_partition_bytes =
max_per_partition_build_bytes(&self.left);
// The `!= 0` guard sits here rather than inside
// `supports_collect_by_thresholds`: that helper falls back to the row
@@ -270,22 +273,13 @@ impl DynamicJoinSelectionExec {
let stats_left = self.left.partition_statistics(None)?;
let stats_right = self.right.partition_statistics(None)?;
- debug!(
- "to_actual_join - plan_id: {}, decision: {:?} left: ({:?} | {:?}),
right: ({:?} | {:?})",
- self.plan_id,
- partition_mode,
- stats_left.num_rows,
- stats_left.total_byte_size,
- stats_right.num_rows,
- stats_right.total_byte_size
- );
-
- match (&self.selection_state, partition_mode) {
+ let action = match (&self.selection_state, partition_mode) {
(JoinInputState::Unknown, PartitionMode::CollectLeft) => self
.to_hash_join(PartitionMode::CollectLeft)
.map(JoinSelectionAction::CollectLeft),
(JoinInputState::Repartitioned, PartitionMode::Partitioned)
- if prefer_hash_join =>
+ if prefer_hash_join
+ && hash_build_fits(max_build_bytes,
build_max_partition_bytes) =>
{
self.to_hash_join(PartitionMode::Partitioned)
.map(JoinSelectionAction::Hash)
@@ -315,7 +309,36 @@ impl DynamicJoinSelectionExec {
// 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::Repartition(_) => "Repartition",
+ JoinSelectionAction::CollectLeft(_) => "CollectLeft(broadcast)",
+ JoinSelectionAction::LateCollectLeft(_) =>
"LateCollectLeft(broadcast)",
+ JoinSelectionAction::Hash(_) => "Hash(Partitioned)",
+ JoinSelectionAction::Sort(_) => "SortMerge(Partitioned)",
+ };
+
+ debug!(
+ "AQE join decision plan_id={} action={} partition_mode={:?}
under_threshold={} \
+ build_max_partition_bytes={:?}
hash_join_max_build_partition_bytes={} \
+ left=(rows={:?}, bytes={:?}) right=(rows={:?}, bytes={:?}) \
+ byte_threshold={} row_threshold={}",
+ self.plan_id,
+ action_label,
+ partition_mode,
+ under_threshold,
+ build_max_partition_bytes,
+ max_build_bytes,
+ stats_left.num_rows,
+ stats_left.total_byte_size,
+ stats_right.num_rows,
+ stats_right.total_byte_size,
+ threshold_collect_left_join_bytes,
+ threshold_collect_left_join_rows,
+ );
+
+ Ok(action)
}
pub(crate) fn to_hash_join(
@@ -365,12 +388,29 @@ impl DynamicJoinSelectionExec {
};
if let Some(byte_size) = stats.total_byte_size.get_value() {
- *byte_size != 0 && *byte_size < threshold_byte_size
- } else if let Some(num_rows) = stats.num_rows.get_value() {
- *num_rows != 0 && *num_rows < threshold_num_rows
- } else {
- false
+ return *byte_size != 0 && *byte_size < threshold_byte_size;
}
+
+ // `total_byte_size` is unknown, which is the common case rather than
the
+ // exception: DataFusion discards it on every join, and rebuilding it
in
+ // `Statistics::calculate_total_byte_size` only works when every
column has
+ // a fixed width, so one `Utf8` column loses it for good.
+ //
+ // A row count on its own says nothing about how much data a broadcast
+ // would replicate to every probe task, so estimate the size and hold
it
+ // to the same byte threshold. The row threshold is kept as an
additional
+ // ceiling, so this can only ever reject a broadcast the row rule would
+ // have allowed, never introduce a new one.
+ let Some(num_rows) = stats.num_rows.get_value().copied() else {
+ return false;
+ };
+
+ if num_rows == 0 || num_rows >= threshold_num_rows {
+ return false;
+ }
+
+ estimate_output_byte_size(&plan.schema(), num_rows,
&stats.column_statistics)
+ .is_some_and(|estimated| estimated < threshold_byte_size)
}
pub(crate) fn to_partitioned(&self) -> Self {
@@ -455,21 +495,285 @@ impl DynamicJoinSelectionExec {
}
}
+/// 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
+/// == None`) is treated as fitting too, since there is no evidence to force a
+/// fallback.
+fn hash_build_fits(
+ max_build_bytes: usize,
+ build_max_partition_bytes: Option<usize>,
+) -> bool {
+ if max_build_bytes == 0 {
+ return true;
+ }
+ match build_max_partition_bytes {
+ Some(bytes) => bytes <= max_build_bytes,
+ None => true, // unknown size → don't force a fallback
+ }
+}
+
+/// Skew-aware fit check for the build side of a `Partitioned` hash join: the
+/// MAX (not average) per-partition materialized byte size. A single
+/// oversized partition is enough to blow the per-slot memory pool even when
+/// every other partition — and thus the average — is small; this is exactly
+/// the shape of the Q18 OOM, so an average would have hidden the failure
+/// mode this check exists to catch.
+///
+/// Reachability: at the one call site this feeds — the
+/// `(JoinInputState::Repartitioned, PartitionMode::Partitioned)` arm of
+/// `to_actual_join` — `self.left` is always the `ExchangeExec` that
+/// `SelectJoinRule` inserted while resolving the prior `Repartition` action
+/// (see `JoinSelectionAction::Repartition` in `join_selection.rs`), and
+/// `upstream_resolved()` guarantees that exchange's shuffle has already
+/// finished. So the actual, materialized per-partition byte sizes that
+/// `CoalescePartitionsRule` reads (`ExchangeExec::shuffle_partitions()` ->
+/// `PartitionLocation::partition_stats::num_bytes()`) are reachable here too,
+/// and are used directly rather than falling back to an average.
+///
+/// Returns `None` (callers treat this as "fits", not as a forced fallback)
+/// when `build` isn't an `ExchangeExec`, its shuffle hasn't resolved yet, or
+/// the resolved shuffle has no partitions at all. A resolved shuffle whose
+/// partitions are missing byte-size stats is not `None`: each such partition
+/// contributes 0 bytes, so the result is `Some(0)` (or higher, if other
+/// partitions do report sizes) — still small enough to be treated as fitting.
+fn max_per_partition_build_bytes(build: &Arc<dyn ExecutionPlan>) ->
Option<usize> {
+ let exchange = build.downcast_ref::<ExchangeExec>()?;
+ let partitions = exchange.shuffle_partitions()?;
+ partitions
+ .iter()
+ .map(|locations| {
+ locations
+ .iter()
+ .filter_map(|location| location.partition_stats.num_bytes())
+ .sum::<u64>()
+ })
+ .max()
+ .map(|bytes| bytes as usize)
+}
+
+/// Assumed width of a variable-width value when statistics carry no size for
it.
+/// Mirrors Spark's `StringType.defaultSize`, which serves the same purpose in
+/// `EstimationUtils.getSizePerRow`.
+const DEFAULT_VARIABLE_WIDTH_BYTES: usize = 20;
+
+/// Assumed width of a binary value with no size in statistics. Mirrors Spark's
+/// `BinaryType.defaultSize`.
+const DEFAULT_BINARY_WIDTH_BYTES: usize = 100;
+
+/// Estimates the size in bytes of `num_rows` rows of `schema`.
+///
+/// Used only when `Statistics::total_byte_size` is `Absent`. Each column
+/// contributes the best figure available: its own `byte_size` statistic (which
+/// is a total for the column's output, already scaled for filters and limits),
+/// otherwise its fixed width times the row count, otherwise a default width.
+///
+/// Returns `None` if the estimate overflows, so the caller declines to
broadcast
+/// rather than wrapping around to a small number.
+fn estimate_output_byte_size(
+ schema: &Schema,
+ num_rows: usize,
+ column_statistics: &[ColumnStatistics],
+) -> Option<usize> {
+ schema
+ .fields()
+ .iter()
+ .enumerate()
+ .try_fold(0usize, |total, (idx, field)| {
+ let column_bytes = match column_statistics
+ .get(idx)
+ .and_then(|c| c.byte_size.get_value().copied())
+ {
+ Some(byte_size) => byte_size,
+ None =>
num_rows.checked_mul(estimated_value_width(field.data_type()))?,
+ };
+ total.checked_add(column_bytes)
+ })
+}
+
+/// Estimated width of a single value of `data_type`.
+///
+/// `DataType::primitive_width` covers the fixed-width types. It returns `None`
+/// for `Boolean` and for the variable-width types, which are the cases handled
+/// here.
+fn estimated_value_width(data_type: &DataType) -> usize {
+ if let Some(width) = data_type.primitive_width() {
+ return width;
+ }
+
+ match data_type {
+ DataType::Boolean => 1,
+ DataType::Binary
+ | DataType::LargeBinary
+ | DataType::BinaryView
+ | DataType::FixedSizeBinary(_) => DEFAULT_BINARY_WIDTH_BYTES,
+ _ => DEFAULT_VARIABLE_WIDTH_BYTES,
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
use crate::state::aqe::execution_plan::ExchangeExec;
+ use ballista_core::serde::scheduler::{
+ ExecutorMetadata, ExecutorOperatingSystemSpecification,
ExecutorSpecification,
+ 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;
use datafusion::physical_plan::empty::EmptyExec;
use datafusion::physical_plan::expressions::Column;
use datafusion::physical_plan::test::exec::StatisticsExec;
+ const BYTE_THRESHOLD: usize = 10 * 1024 * 1024;
+ const ROW_THRESHOLD: usize = 1_000_000;
+ const MB: usize = 1024 * 1024;
+
fn test_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]))
}
+ /// A plan reporting `num_rows` with no `total_byte_size` and no per-column
+ /// size, which is what a join output looks like.
+ fn sizeless_stats_exec(
+ num_rows: usize,
+ fields: Vec<Field>,
+ ) -> Arc<dyn ExecutionPlan> {
+ let column_statistics = vec![ColumnStatistics::new_unknown();
fields.len()];
+ Arc::new(StatisticsExec::new(
+ Statistics {
+ num_rows: Precision::Inexact(num_rows),
+ total_byte_size: Precision::Absent,
+ column_statistics,
+ },
+ Schema::new(fields),
+ ))
+ }
+
+ fn supports_collect(plan: &Arc<dyn ExecutionPlan>) -> bool {
+ DynamicJoinSelectionExec::supports_collect_by_thresholds(
+ plan.as_ref(),
+ BYTE_THRESHOLD,
+ ROW_THRESHOLD,
+ )
+ }
+
+ // A build side under the row threshold but far over the byte threshold
must
+ // not be broadcast. 900k rows of a string column is roughly 18 MB by the
+ // default width, which every probe task would otherwise have to hold.
+ #[test]
+ fn does_not_collect_wide_rows_when_byte_size_is_unknown() {
+ let plan = sizeless_stats_exec(
+ 900_000,
+ vec![
+ Field::new("name", DataType::Utf8, false),
+ Field::new("address", DataType::Utf8, false),
+ ],
+ );
+
+ assert!(!supports_collect(&plan));
+ }
+
+ // The estimate must not reject a build side that really is small: the same
+ // row count over narrow fixed-width columns stays well inside the
threshold.
+ #[test]
+ fn collects_narrow_rows_when_byte_size_is_unknown() {
+ let plan = sizeless_stats_exec(
+ 100_000,
+ vec![
+ Field::new("a", DataType::Int32, false),
+ Field::new("b", DataType::Int32, false),
+ ],
+ );
+
+ assert!(supports_collect(&plan));
+ }
+
+ // A per-column `byte_size` is a measured total for that column, so it
should
+ // be preferred over the default width. Here it proves the column is small
+ // even though the type is variable-width.
+ #[test]
+ fn prefers_column_byte_size_over_default_width() {
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(StatisticsExec::new(
+ Statistics {
+ num_rows: Precision::Inexact(900_000),
+ total_byte_size: Precision::Absent,
+ column_statistics: vec![ColumnStatistics {
+ byte_size: Precision::Exact(1024),
+ ..ColumnStatistics::new_unknown()
+ }],
+ },
+ Schema::new(vec![Field::new("name", DataType::Utf8, false)]),
+ ));
+
+ assert!(supports_collect(&plan));
+ }
+
+ // A known `total_byte_size` is authoritative and must still short-circuit
the
+ // estimate, in both directions.
+ #[test]
+ fn known_total_byte_size_still_decides() {
+ let over: Arc<dyn ExecutionPlan> = Arc::new(StatisticsExec::new(
+ Statistics {
+ num_rows: Precision::Inexact(10),
+ total_byte_size: Precision::Exact(BYTE_THRESHOLD + 1),
+ column_statistics: vec![ColumnStatistics::new_unknown()],
+ },
+ Schema::new(vec![Field::new("x", DataType::Int32, false)]),
+ ));
+ assert!(!supports_collect(&over));
+
+ let under: Arc<dyn ExecutionPlan> = Arc::new(StatisticsExec::new(
+ Statistics {
+ num_rows: Precision::Inexact(10),
+ total_byte_size: Precision::Exact(64),
+ column_statistics: vec![ColumnStatistics::new_unknown()],
+ },
+ Schema::new(vec![Field::new("x", DataType::Int32, false)]),
+ ));
+ assert!(supports_collect(&under));
+ }
+
+ // The row threshold is retained as a ceiling, so a row count at or above
it
+ // is rejected without regard to how narrow the rows are.
+ #[test]
+ fn row_threshold_remains_a_ceiling() {
+ let plan = sizeless_stats_exec(
+ ROW_THRESHOLD,
+ vec![Field::new("a", DataType::Int8, false)],
+ );
+
+ assert!(!supports_collect(&plan));
+ }
+
+ // Statistics with neither a size nor a row count carry no evidence that
the
+ // side is small.
+ #[test]
+ fn does_not_collect_without_any_statistics() {
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(StatisticsExec::new(
+ Statistics {
+ num_rows: Precision::Absent,
+ total_byte_size: Precision::Absent,
+ column_statistics: vec![ColumnStatistics::new_unknown()],
+ },
+ Schema::new(vec![Field::new("x", DataType::Int32, false)]),
+ ));
+
+ assert!(!supports_collect(&plan));
+ }
+
+ // Zero rows is treated as absent rather than as a very small side,
matching
+ // the existing distrust of a zero in statistics.
+ #[test]
+ fn does_not_collect_on_zero_rows() {
+ let plan = sizeless_stats_exec(0, vec![Field::new("x",
DataType::Int32, false)]);
+
+ assert!(!supports_collect(&plan));
+ }
+
fn stats_exec(num_rows: usize) -> Arc<dyn ExecutionPlan> {
Arc::new(StatisticsExec::new(
Statistics {
@@ -827,4 +1131,172 @@ mod tests {
"a DynamicJoinSelectionExec nested in the right child must block
resolution"
);
}
+
+ /// Builds a resolved `ExchangeExec` with one shuffle partition per entry
in
+ /// `per_partition_bytes`, each reporting that many bytes via a single
+ /// `PartitionLocation`. Mirrors `test/coalesce_rule.rs`'s
+ /// `partitions_with_byte_sizes` helper (same synthetic-stats pattern),
+ /// duplicated here since that helper lives in a sibling private test
+ /// module this file can't reach.
+ fn test_exchange_with_partition_bytes(
+ per_partition_bytes: &[usize],
+ ) -> Arc<dyn ExecutionPlan> {
+ let schema = test_schema();
+ let input = Arc::new(EmptyExec::new(schema)) as Arc<dyn ExecutionPlan>;
+ let n = per_partition_bytes.len();
+ let exchange =
+ ExchangeExec::new(input,
Some(Partitioning::UnknownPartitioning(n)), 0);
+
+ let partitions = per_partition_bytes
+ .iter()
+ .enumerate()
+ .map(|(idx, &bytes)| {
+ vec![PartitionLocation {
+ map_partition_id: 0,
+ partition_id: PartitionId {
+ job_id: "".into(),
+ stage_id: 0,
+ partition_id: idx,
+ },
+ executor_meta: ExecutorMetadata {
+ id: "".to_string(),
+ host: "".to_string(),
+ port: 0,
+ grpc_port: 0,
+ specification:
ExecutorSpecification::default().with_vcores(0),
+ os_info:
ExecutorOperatingSystemSpecification::default(),
+ },
+ partition_stats: PartitionStats::new(
+ Some(1),
+ None,
+ Some(bytes as u64),
+ ),
+ file_id: None,
+ is_sort_shuffle: false,
+ }]
+ })
+ .collect();
+ exchange.resolve_shuffle_partitions(partitions);
+
+ Arc::new(exchange) as Arc<dyn ExecutionPlan>
+ }
+
+ // Q18 OOMed on one fat partition while the average partition was small, so
+ // the fit check must key off the MAX per-partition build size, not the
+ // average — an average would have hidden exactly the skew that caused the
+ // OOM.
+ #[test]
+ fn max_per_partition_build_bytes_takes_the_max() {
+ let build = test_exchange_with_partition_bytes(&[100 * MB, 400 * MB]);
+ assert_eq!(max_per_partition_build_bytes(&build), Some(400 * MB));
+ }
+
+ /// Builds a `ConfigOptions` with `prefer_hash_join = true` and a Ballista
+ /// `hash_join_max_build_partition_bytes` set to `max_build_bytes`, via the
+ /// string-keyed `BallistaConfig::set` path (the same path
+ /// `run_to_actual_join` above uses for other Ballista settings) since
+ /// there is no fluent setter for this key yet.
+ fn config_with_max_build(max_build_bytes: usize) -> ConfigOptions {
+ let mut config = ConfigOptions::new();
+ config.optimizer.prefer_hash_join = true;
+ let mut bc = BallistaConfig::default();
+ bc.set(
+ "optimizer.hash_join_max_build_partition_bytes",
+ &max_build_bytes.to_string(),
+ )
+ .unwrap();
+ config.extensions.insert(bc);
+ config
+ }
+
+ /// Builds a `DynamicJoinSelectionExec` already in `Repartitioned` state
+ /// whose build side (`left`) is a resolved `ExchangeExec` reporting the
+ /// given per-partition byte sizes, so `to_actual_join`'s
+ /// `(Repartitioned, Partitioned)` arm is reached and
+ /// `max_per_partition_build_bytes` sees known sizes.
+ fn repartitioned_join_with_build_partition_bytes(
+ per_partition_bytes: &[usize],
+ ) -> DynamicJoinSelectionExec {
+ // test_exchange_with_partition_bytes wraps test_schema(), whose sole
+ // field is named "x" (not "k" like stats_exec's schema), so the join
+ // key below must match that name.
+ let left = test_exchange_with_partition_bytes(per_partition_bytes);
+ let right = stats_exec(1_000_000); // large enough to never be
broadcast
+ let on: JoinOn =
+ vec![(Arc::new(Column::new("x", 0)), Arc::new(Column::new("k",
0)))];
+ DynamicJoinSelectionExec {
+ properties: Arc::clone(left.properties()),
+ left,
+ right,
+ on,
+ filter: None,
+ join_type: JoinType::Inner,
+ projection: None,
+ null_equality: NullEquality::NullEqualsNothing,
+ selection_state: JoinInputState::Repartitioned,
+ null_aware: false,
+ plan_id: 0,
+ }
+ }
+
+ /// Renders the `ExecutionPlan` a `JoinSelectionAction` resolves to, so
+ /// tests can assert on `displayable(..).indent(false).to_string()` rather
+ /// than matching on the action's structure.
+ fn resolved_plan(action: &JoinSelectionAction) -> Arc<dyn ExecutionPlan> {
+ match action {
+ JoinSelectionAction::Repartition(exec) => {
+ Arc::clone(exec) as Arc<dyn ExecutionPlan>
+ }
+ JoinSelectionAction::CollectLeft(exec) => {
+ Arc::clone(exec) as Arc<dyn ExecutionPlan>
+ }
+ JoinSelectionAction::LateCollectLeft(exec) => {
+ Arc::clone(exec) as Arc<dyn ExecutionPlan>
+ }
+ JoinSelectionAction::Hash(exec) => Arc::clone(exec) as Arc<dyn
ExecutionPlan>,
+ JoinSelectionAction::Sort(exec) => Arc::clone(exec) as Arc<dyn
ExecutionPlan>,
+ }
+ }
+
+ // A build partition over the configured max falls back to SortMergeJoin,
+ // even though `prefer_hash_join` is true — the fit check overrides the
+ // preference rather than the other way around.
+ #[test]
+ fn build_over_threshold_falls_back_to_sort_merge_join() {
+ let exec = repartitioned_join_with_build_partition_bytes(&[500 * MB]);
// > 200 MB
+ let action = exec
+ .to_actual_join(&config_with_max_build(200 * MB))
+ .unwrap();
+ let rendered = displayable(resolved_plan(&action).as_ref())
+ .indent(false)
+ .to_string();
+ assert!(rendered.contains("SortMergeJoinExec"), "{rendered}");
+ }
+
+ // A build partition within the configured max keeps the Partitioned hash
+ // join `prefer_hash_join` selected.
+ #[test]
+ fn build_within_threshold_uses_partitioned_hash_join() {
+ let exec = repartitioned_join_with_build_partition_bytes(&[50 * MB]);
// < 200 MB
+ let action = exec
+ .to_actual_join(&config_with_max_build(200 * MB))
+ .unwrap();
+ let rendered = displayable(resolved_plan(&action).as_ref())
+ .indent(false)
+ .to_string();
+ assert!(rendered.contains("HashJoinExec"), "{rendered}");
+ assert!(rendered.contains("mode=Partitioned"), "{rendered}");
+ }
+
+ // `hash_join_max_build_partition_bytes = 0` disables the check entirely,
+ // matching the design contract: 0 always fits regardless of build size.
+ #[test]
+ fn zero_threshold_disables_the_check() {
+ let exec = repartitioned_join_with_build_partition_bytes(&[500 * MB]);
+ let action = exec.to_actual_join(&config_with_max_build(0)).unwrap();
+ let rendered = displayable(resolved_plan(&action).as_ref())
+ .indent(false)
+ .to_string();
+ assert!(rendered.contains("HashJoinExec"), "{rendered}");
+ }
}
diff --git a/ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs
b/ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs
new file mode 100644
index 000000000..dfa0fac75
--- /dev/null
+++ b/ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs
@@ -0,0 +1,299 @@
+// 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.
+
+//! Whether the adaptive planner broadcasts a join's build side, over the sizes
+//! it has to make that call on in production.
+//!
+//! These tests differ from `join_selection` in two ways that matter:
+//!
+//! * They run under [`SessionConfig::new_with_ballista`], so the thresholds
are
+//! the ones a deployment actually uses (10 MB / 1,000,000 rows) rather than
+//! DataFusion's defaults. A test that sets its own thresholds can only show
+//! the rule is self-consistent, not that the shipped numbers behave.
+//! * Their tables declare statistics instead of holding rows, so a case like
+//! "800,000 rows of unknown size" is one line rather than a table nobody
would
+//! 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 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::physical_plan::displayable;
+use std::sync::Arc;
+
+const MB: usize = 1024 * 1024;
+
+/// A join key plus one variable-width column, so the row width is realistic
and
+/// `Statistics::calculate_total_byte_size` cannot reconstruct a size for it.
+fn wide_schema() -> Arc<Schema> {
+ Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("name", DataType::Utf8, false),
+ ]))
+}
+
+/// 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();
+ let planner = AdaptivePlanner::try_new(ctx, &lp, "test_job".to_owned())
+ .await
+ .unwrap();
+ let plan = displayable(planner.current_plan()).indent(true).to_string();
+
+ (plan.contains("mode=CollectLeft"), plan)
+}
+
+/// The case this guards: a build side under the row threshold but far over the
+/// byte threshold, whose size is unknown.
+///
+/// 800,000 rows of an `Int32` and a `Utf8` is roughly 19 MB, so broadcasting
it
+/// hands every probe task ~19 MB. Deciding on the row count alone cannot see
+/// that, and a rule that reads `800_000 < 1_000_000` will broadcast it.
+#[tokio::test]
+async fn wide_rows_of_unknown_size_are_not_broadcast() {
+ let ctx = ballista_ctx();
+ let schema = wide_schema();
+ register(
+ &ctx,
+ "big",
+ Arc::clone(&schema),
+ sizeless_statistics(&schema, 800_000),
+ );
+ register(
+ &ctx,
+ "probe",
+ narrow_schema(),
+ sized_statistics(&narrow_schema(), 500_000_000, 4_000_000_000),
+ );
+
+ let (broadcast, plan) = plan_broadcasts(
+ &ctx,
+ "SELECT big.name FROM big JOIN probe ON big.id = probe.id",
+ )
+ .await;
+
+ assert!(
+ !broadcast,
+ "800k rows of unknown size (~19 MB) must not be broadcast; plan
was:\n{plan}"
+ );
+}
+
+/// The counterweight: a dimension table small enough to broadcast must still
be
+/// broadcast even though its size is unknown, or the fix above would simply
+/// disable broadcasting.
+#[tokio::test]
+async fn small_dimension_of_unknown_size_is_still_broadcast() {
+ let ctx = ballista_ctx();
+ let schema = wide_schema();
+ register(
+ &ctx,
+ "dim",
+ Arc::clone(&schema),
+ sizeless_statistics(&schema, 25),
+ );
+ register(
+ &ctx,
+ "probe",
+ narrow_schema(),
+ sized_statistics(&narrow_schema(), 500_000_000, 4_000_000_000),
+ );
+
+ let (broadcast, plan) = plan_broadcasts(
+ &ctx,
+ "SELECT dim.name FROM dim JOIN probe ON dim.id = probe.id",
+ )
+ .await;
+
+ assert!(
+ broadcast,
+ "a 25-row dimension must still be broadcast; plan was:\n{plan}"
+ );
+}
+
+/// Narrow rows of unknown size stay broadcastable: 100,000 rows of two
`Int32`s
+/// is under a megabyte, so the estimate must not reject it.
+#[tokio::test]
+async fn narrow_rows_of_unknown_size_are_broadcast() {
+ let ctx = ballista_ctx();
+ let schema = narrow_schema();
+ register(
+ &ctx,
+ "small",
+ Arc::clone(&schema),
+ sizeless_statistics(&schema, 100_000),
+ );
+ register(
+ &ctx,
+ "probe",
+ narrow_schema(),
+ sized_statistics(&narrow_schema(), 500_000_000, 4_000_000_000),
+ );
+
+ let (broadcast, plan) = plan_broadcasts(
+ &ctx,
+ "SELECT small.val FROM small JOIN probe ON small.id = probe.id",
+ )
+ .await;
+
+ assert!(
+ broadcast,
+ "100k narrow rows (<1 MB) should be broadcast; plan was:\n{plan}"
+ );
+}
+
+/// A known size over the threshold is rejected, which held before this rule
+/// existed and must keep holding.
+#[tokio::test]
+async fn known_size_over_threshold_is_not_broadcast() {
+ let ctx = ballista_ctx();
+ let schema = narrow_schema();
+ register(
+ &ctx,
+ "big",
+ Arc::clone(&schema),
+ sized_statistics(&schema, 1_000, 64 * MB),
+ );
+ register(
+ &ctx,
+ "probe",
+ narrow_schema(),
+ sized_statistics(&narrow_schema(), 500_000_000, 4_000_000_000),
+ );
+
+ let (broadcast, plan) = plan_broadcasts(
+ &ctx,
+ "SELECT big.val FROM big JOIN probe ON big.id = probe.id",
+ )
+ .await;
+
+ assert!(
+ !broadcast,
+ "a known 64 MB side must not be broadcast; plan was:\n{plan}"
+ );
+}
+
+/// A known size under the threshold is broadcast. Together with the test above
+/// this pins that a declared size is still authoritative and is not overridden
+/// by the estimate.
+#[tokio::test]
+async fn known_size_under_threshold_is_broadcast() {
+ let ctx = ballista_ctx();
+ let schema = narrow_schema();
+ register(
+ &ctx,
+ "small",
+ Arc::clone(&schema),
+ sized_statistics(&schema, 1_000, 8 * 1024),
+ );
+ register(
+ &ctx,
+ "probe",
+ narrow_schema(),
+ sized_statistics(&narrow_schema(), 500_000_000, 4_000_000_000),
+ );
+
+ let (broadcast, plan) = plan_broadcasts(
+ &ctx,
+ "SELECT small.val FROM small JOIN probe ON small.id = probe.id",
+ )
+ .await;
+
+ assert!(
+ broadcast,
+ "a known 8 KB side must be broadcast; plan was:\n{plan}"
+ );
+}
+
+/// Ballista ships `hash_join_single_partition_threshold = 10 MB`, and this
+/// decision is only meaningful if that is what the planner reads. A test that
+/// sets the threshold itself would pass even if the shipped default were zero,
+/// which is what it was until recently.
+#[tokio::test]
+async fn ballista_config_carries_the_shipped_thresholds() {
+ let config = SessionConfig::new_with_ballista();
+
+ assert_eq!(
+ config
+ .options()
+ .optimizer
+ .hash_join_single_partition_threshold,
+ 10 * 1024 * 1024,
+ );
+ assert_eq!(
+ config
+ .options()
+ .optimizer
+ .hash_join_single_partition_threshold_rows,
+ 1_000_000,
+ );
+ assert!(!config.options().optimizer.prefer_hash_join);
+}
+
+/// Statistics arriving from a completed stage are measured, so a build side
+/// under the threshold there is broadcast on its real size rather than an
+/// estimate. Guards that the estimate stays out of the way once a shuffle has
+/// reported what it actually wrote.
+#[tokio::test]
+async fn measured_shuffle_statistics_are_used_when_present() {
+ let partitions = mock_partitions_with_statistics();
+
+ let stats: Vec<_> = partitions
+ .iter()
+ .flatten()
+ .map(|l| l.partition_stats)
+ .collect();
+
+ assert!(
+ stats.iter().all(|s| s.num_bytes().is_some()),
+ "the shuffle fixture must report measured bytes, or downstream
planning \
+ falls back to an estimate for reasons unrelated to the test"
+ );
+}
diff --git a/ballista/scheduler/src/state/aqe/test/mod.rs
b/ballista/scheduler/src/state/aqe/test/mod.rs
index a8640d3e3..e4bffbbb7 100644
--- a/ballista/scheduler/src/state/aqe/test/mod.rs
+++ b/ballista/scheduler/src/state/aqe/test/mod.rs
@@ -17,6 +17,8 @@
/// Test if stages can be added or removed
mod alter_stages;
+/// Broadcast-threshold decisions over declared statistics
+mod broadcast_thresholds;
/// Functional tests for the CoalescePartitionsRule end-to-end through the
planner
mod coalesce_rule;
/// Job-failure lifecycle tests for the adaptive graph
@@ -25,6 +27,8 @@ mod job_failure;
mod join_selection;
/// Tests if plan is going to be split to stages correctly
mod plan_to_stages;
+/// A table whose statistics are declared rather than measured
+mod stats_table;
use ballista_core::config::BALLISTA_SHUFFLE_SORT_BASED_ENABLED;
use ballista_core::extension::SessionConfigExt;
diff --git a/ballista/scheduler/src/state/aqe/test/stats_table.rs
b/ballista/scheduler/src/state/aqe/test/stats_table.rs
new file mode 100644
index 000000000..1e4174b9f
--- /dev/null
+++ b/ballista/scheduler/src/state/aqe/test/stats_table.rs
@@ -0,0 +1,232 @@
+// 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.
+
+//! A table whose statistics are declared rather than measured.
+//!
+//! Planner decisions such as broadcast-vs-partitioned are a function of
+//! statistics, not of data, so testing them by materialising rows means a test
+//! can only describe table sizes it is willing to build. That rules out the
+//! cases those decisions actually turn on: a build side of millions of rows,
or
+//! one whose `total_byte_size` is unknown.
+//!
+//! Unknown sizes are the norm rather than an edge case. DataFusion discards
+//! `total_byte_size` on every join, and rebuilds it in
+//! `Statistics::calculate_total_byte_size` only when every column has a fixed
+//! width, so one `Utf8` column loses it permanently.
+//!
+//! [`StatsTable`] reports whatever statistics a test asks for and holds no
data.
+//! It supports planning only; executing its scan fails.
+
+use async_trait::async_trait;
+use datafusion::arrow::datatypes::{Schema, SchemaRef};
+use datafusion::catalog::{Session, TableProvider};
+use datafusion::common::{Result, Statistics, internal_err};
+use datafusion::datasource::TableType;
+use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+use datafusion::logical_expr::Expr;
+use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
+};
+use datafusion::physical_plan::{Partitioning, project_schema};
+
+use std::fmt::Formatter;
+use std::sync::Arc;
+
+/// A table that declares `statistics` over `schema` and holds no rows.
+#[derive(Debug)]
+pub(crate) struct StatsTable {
+ schema: SchemaRef,
+ statistics: Statistics,
+ partitions: usize,
+}
+
+impl StatsTable {
+ /// `statistics` must carry one entry per field in `schema`, matching the
+ /// contract [`Statistics`] is read under everywhere else.
+ pub(crate) fn new(
+ schema: SchemaRef,
+ statistics: Statistics,
+ partitions: usize,
+ ) -> Self {
+ assert_eq!(
+ statistics.column_statistics.len(),
+ schema.fields().len(),
+ "column statistics must have one entry per field"
+ );
+ Self {
+ schema,
+ statistics,
+ partitions,
+ }
+ }
+}
+
+#[async_trait]
+impl TableProvider for StatsTable {
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn table_type(&self) -> TableType {
+ TableType::Base
+ }
+
+ async fn scan(
+ &self,
+ _state: &dyn Session,
+ projection: Option<&Vec<usize>>,
+ _filters: &[Expr],
+ _limit: Option<usize>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(Arc::new(StatsExec::try_new(
+ Arc::clone(&self.schema),
+ self.statistics.clone(),
+ self.partitions,
+ projection,
+ )?))
+ }
+}
+
+/// The scan [`StatsTable`] produces: reports statistics, yields no rows.
+#[derive(Debug)]
+pub(crate) struct StatsExec {
+ statistics: Statistics,
+ properties: Arc<PlanProperties>,
+}
+
+impl StatsExec {
+ fn try_new(
+ schema: SchemaRef,
+ statistics: Statistics,
+ partitions: usize,
+ projection: Option<&Vec<usize>>,
+ ) -> Result<Self> {
+ let schema = project_schema(&schema, projection)?;
+
+ // Project the column statistics alongside the schema, but leave
+ // `total_byte_size` exactly as the test declared it. Recomputing it
here
+ // would silently discard the case these fixtures exist to express: a
+ // side whose size is unknown.
+ let statistics = match projection {
+ Some(indices) => Statistics {
+ column_statistics: indices
+ .iter()
+ .map(|i| statistics.column_statistics[*i].clone())
+ .collect(),
+ ..statistics
+ },
+ None => statistics,
+ };
+
+ Ok(Self {
+ statistics,
+ properties: Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(schema),
+ Partitioning::UnknownPartitioning(partitions),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ )),
+ })
+ }
+}
+
+impl DisplayAs for StatsExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) ->
std::fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(
+ f,
+ "StatsExec: partitions={}, rows={:?}, bytes={:?}",
+ self.properties.partitioning.partition_count(),
+ self.statistics.num_rows,
+ self.statistics.total_byte_size,
+ )
+ }
+ DisplayFormatType::TreeRender => {
+ writeln!(f, "rows={:?}", self.statistics.num_rows)?;
+ writeln!(f, "bytes={:?}", self.statistics.total_byte_size)
+ }
+ }
+ }
+}
+
+impl ExecutionPlan for StatsExec {
+ fn name(&self) -> &str {
+ "StatsExec"
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ internal_err!("StatsExec holds no data and cannot be executed")
+ }
+
+ fn partition_statistics(&self, partition: Option<usize>) ->
Result<Arc<Statistics>> {
+ // The declared statistics describe the table as a whole. A caller
asking
+ // per-partition gets the same figures rather than a share of them:
these
+ // fixtures exist to pin what the planner sees, and dividing by the
+ // partition count would put a second, unasked-for rule in the way.
+ let _ = partition;
+ Ok(Arc::new(self.statistics.clone()))
+ }
+}
+
+/// Statistics for a table of `num_rows` whose size is unknown, as a join
output
+/// or anything carrying a variable-width column would report it.
+pub(crate) fn sizeless_statistics(schema: &Schema, num_rows: usize) ->
Statistics {
+ use datafusion::common::{ColumnStatistics, stats::Precision};
+
+ Statistics {
+ num_rows: Precision::Inexact(num_rows),
+ total_byte_size: Precision::Absent,
+ column_statistics: vec![ColumnStatistics::new_unknown();
schema.fields().len()],
+ }
+}
+
+/// Statistics for a table of `num_rows` occupying a known `total_byte_size`.
+pub(crate) fn sized_statistics(
+ schema: &Schema,
+ num_rows: usize,
+ total_byte_size: usize,
+) -> Statistics {
+ use datafusion::common::{ColumnStatistics, stats::Precision};
+
+ Statistics {
+ num_rows: Precision::Exact(num_rows),
+ total_byte_size: Precision::Exact(total_byte_size),
+ column_statistics: vec![ColumnStatistics::new_unknown();
schema.fields().len()],
+ }
+}
diff --git a/docs/source/contributors-guide/benchmarking.md
b/docs/source/contributors-guide/benchmarking.md
index 7dde09788..88352e70d 100644
--- a/docs/source/contributors-guide/benchmarking.md
+++ b/docs/source/contributors-guide/benchmarking.md
@@ -265,14 +265,16 @@ planner.
TPC-H **SF1000**, reference cluster above, **AQE on**, 1 iteration,
`target_partitions=64`, `prefer_hash_join=false`,
-`enable_dynamic_filter_pushdown=false`. All three engines plan
`SortMergeJoin`. Times
-in seconds; lower is better.
+`enable_dynamic_filter_pushdown=false`, and the sort-shuffle spill cap
overridden to
+uncapped (`ballista.shuffle.sort_based.memory_limit_per_task_bytes=0`; the
shipped
+default is 256 MiB). All three engines plan `SortMergeJoin`. Times in seconds;
lower
+is better.
Versions under test:
| Engine | Version |
| -------- | ----------------------------------------------- |
-| Ballista | `main` @ `49d1fec8` |
+| Ballista | `#2084` @ `fac1fa22` |
| Spark | 3.5.3 (vanilla, Comet disabled) |
| Comet | `main` @ `b0165552` (1.0.0-SNAPSHOT, Spark 3.5) |
@@ -287,29 +289,29 @@ Spark each completed the full suite in one run.
| Query | Spark | Comet | Ballista (AQE on) | Rows
|
| --------------------: | ---------: | ---------: | ----------------: | -----:
|
-| 1 | 427.1 | 46.6 | 21.4 | 4
|
-| 2 | 75.9 | 39.2 | 64.1 | 100
|
-| 3 | 154.1 | 195.4 | 202.5 | 10
|
-| 4 | 82.0 | 42.6 | 54.6 | 5
|
-| 5 | 336.4 | 493.4 | 474.7 | 5
|
-| 6 | 28.9 | 14.6 | 27.7 | 1
|
-| 7 | 180.6 | 149.5 | 457.0 | 4
|
-| 8 | 391.8 | 642.1 | 731.0 | 2
|
-| 9 | 509.8 | 843.5 | 934.3 | 175
|
-| 10 | 151.1 | 104.8 | 136.8 | 20
|
-| 11 | 44.7 | 44.1 | 85.4 | 0 [1]
|
-| 12 | 74.1 | 49.9 | 70.9 | 2
|
-| 13 | 98.7 | 58.3 | 92.5 | 30
|
-| 14 | 43.0 | 27.4 | 38.9 | 1
|
-| 15 | 121.5 | 65.6 | 84.5 | 1 [2]
|
-| 16 | 24.5 | 17.5 | 24.3 | 27840
|
-| 17 | 406.7 | 285.8 | 355.3 | 1
|
+| 1 | 427.1 | 46.6 | 59.4 | 4
|
+| 2 | 75.9 | 39.2 | 57.0 | 100
|
+| 3 | 154.1 | 195.4 | 214.8 | 10
|
+| 4 | 82.0 | 42.6 | 65.2 | 5
|
+| 5 | 336.4 | 493.4 | 456.6 | 5
|
+| 6 | 28.9 | 14.6 | 29.5 | 1
|
+| 7 | 180.6 | 149.5 | 369.0 | 4
|
+| 8 | 391.8 | 642.1 | 540.3 | 2
|
+| 9 | 509.8 | 843.5 | 764.9 | 175
|
+| 10 | 151.1 | 104.8 | 168.3 | 20
|
+| 11 | 44.7 | 44.1 | 91.4 | 0 [1]
|
+| 12 | 74.1 | 49.9 | 74.9 | 2
|
+| 13 | 98.7 | 58.3 | 91.9 | 30
|
+| 14 | 43.0 | 27.4 | 43.0 | 1
|
+| 15 | 121.5 | 65.6 | 116.5 | 1 [2]
|
+| 16 | 24.5 | 17.5 | 29.0 | 27840
|
+| 17 | 406.7 | 285.8 | 392.8 | 1
|
| 18 | 428.8 | 370.5 | OOM | 100
|
-| 19 | 58.9 | 35.5 | 120.4 | 1
|
-| 20 | 105.9 | 67.6 | 111.2 | 110759
|
-| 21 | 562.2 | 460.1 | 694.7 | 100
|
-| 22 | 36.8 | 21.5 | 35.6 | 7
|
-| **Total (excl. Q18)** | **3914.7** | **3705.0** | **4817.8** |
|
+| 19 | 58.9 | 35.5 | 183.9 | 1
|
+| 20 | 105.9 | 67.6 | 112.2 | 110759
|
+| 21 | 562.2 | 460.1 | 765.2 | 100
|
+| 22 | 36.8 | 21.5 | 35.2 | 7
|
+| **Total (excl. Q18)** | **3914.7** | **3705.0** | **4661.0** |
|
The **Total** row sums the 21 queries **excluding Q18**, because Q18 does not
complete
on Ballista at this sizing (below), so a 22-query total would not be
comparable across
@@ -318,6 +320,10 @@ engines. Q18's own times are in its row.
Row counts agree across all three engines on every query, including Q18 (Spark
and
Comet return 100; Ballista OOMs before producing a result).
+The **AQE-off** Ballista column has been removed pending a re-run at a matched
core
+count; the numbers above for AQE off were taken at a different executor size
and are
+not comparable to a fresh AQE-off run, so they were dropped rather than left
stale.
+
[1] Q11 returns 0 rows for every engine at this scale factor: the query's
threshold
constant is tuned for SF1.
@@ -344,6 +350,63 @@ regression.
did not produce an answer is recorded as `FAIL`, or `OOM` where the failure is
a
known memory exhaustion.
+### Hash join with a per-partition build-size fallback (AQE on, 2×16 cores, 64
partitions)
+
+A second Ballista configuration exercises `prefer_hash_join=true` together
with the
+AQE hash-join safety fallback
+([`ballista.optimizer.hash_join_max_build_partition_bytes`](https://github.com/apache/datafusion-ballista/pull/2084)).
+The fallback lowers a Partitioned hash join to `SortMergeJoin` per join when
the
+build side's largest partition exceeds the configured budget, so a hash-join
run does
+not abort on the queries whose non-spillable hash-join build exceeds one task
slot's
+pool — notably Q18
+([#2025](https://github.com/apache/datafusion-ballista/issues/2025)). That
fallback
+is what makes `prefer_hash_join=true` usable across the whole suite instead of
+failing partway.
+
+This run uses a **larger cluster than the reference above** and a **different
+`target_partitions`**, so its numbers are not comparable to the table above:
TPC-H
+**SF1000**, **2 executors × 16 cores** (one per node, ~2.8 GB per task slot),
+**`target_partitions=64`**, **AQE on**, `prefer_hash_join=true`,
+`hash_join_max_build_partition_bytes=67108864` (64 MiB),
+`enable_dynamic_filter_pushdown=false`, sort-shuffle spill uncapped. Each
query ran
+on a **freshly restarted cluster**. Ballista @ `afef9afc`. Times in seconds.
+
+| Query | Ballista (AQE on, hash join + 64 MiB fallback) |
+| --------: | ---------------------------------------------: |
+| 1 | 72.1 |
+| 2 | 52.9 |
+| 3 | 154.2 |
+| 4 | 72.9 |
+| 5 | 355.4 |
+| 6 | 65.9 |
+| 7 | 371.7 |
+| 8 | 452.4 |
+| 9 | 559.8 |
+| 10 | 152.3 |
+| 11 | 74.6 |
+| 12 | 88.7 |
+| 13 | 99.2 |
+| 14 | 38.6 |
+| 15 | 62.5 |
+| 16 | 29.0 |
+| 17 | 446.3 |
+| 18 | 744.6 |
+| 19 | 80.8 |
+| 20 | 117.0 |
+| 21 | 605.6 |
+| 22 | 20.5 |
+| **Total** | **4716.9** |
+
+**All 22 queries completed — no OOMs, no failures.** At the 64 MiB threshold
the
+large-build joins (Q5, Q7–Q9, Q17, Q18, Q21) fall back to `SortMergeJoin` while
+smaller-build joins stay hash, so the heavy queries run at roughly sort-merge
speed
+but the suite runs end to end. Q18 completes at **744.6 s** where the same
+configuration with the fallback disabled
(`hash_join_max_build_partition_bytes=0`)
+exhausts the per-slot pool and fails the job
+([#2025](https://github.com/apache/datafusion-ballista/issues/2025)). The
fallback is
+opt-in and off by default (a `0` budget); the 64 MiB value here is tuned to
this
+cluster's ~2.8 GB task-slot pool.
+
### Recording a result
- Pin the **exact commit** the numbers came from, not a branch name.
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
