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 fd5aa65d6 feat: make broadcast_join_threshold_bytes/rows authoritative
under AQE and the static planner (#2086)
fd5aa65d6 is described below
commit fd5aa65d6e332d532cd8220b524bac2daeca15f9
Author: Andy Grove <[email protected]>
AuthorDate: Sat Jul 18 08:54:00 2026 -0600
feat: make broadcast_join_threshold_bytes/rows authoritative under AQE and
the static planner (#2086)
* 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.
---
ballista/core/src/config.rs | 27 ++-
ballista/core/src/extension.rs | 52 +++++-
ballista/scheduler/src/planner.rs | 176 +++++++++++++++---
.../src/state/aqe/execution_plan/dynamic_join.rs | 197 ++++++++++++++++++---
.../scheduler/src/state/aqe/test/join_selection.rs | 5 +
docs/source/user-guide/tuning-guide.md | 11 +-
6 files changed, 411 insertions(+), 57 deletions(-)
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
index 406d1b9e8..bebcdc0e7 100644
--- a/ballista/core/src/config.rs
+++ b/ballista/core/src/config.rs
@@ -103,6 +103,13 @@ pub const
BALLISTA_SHUFFLE_SORT_BASED_MEMORY_LIMIT_PER_TASK_BYTES: &str =
pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES: &str =
"ballista.optimizer.broadcast_join_threshold_bytes";
+/// Configuration key for the 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.
+/// Set to `0` to disable promotion via the row-count path.
+pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS: &str =
+ "ballista.optimizer.broadcast_join_threshold_rows";
+
/// 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.
@@ -236,9 +243,19 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String,
ConfigEntry>> = LazyLock::new(||
ConfigEntry::new(BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES.to_string(),
"Byte-size threshold below which a hash join's
smaller side is \
promoted to CollectLeft and lowered via the
broadcast pattern. \
- Set to 0 to disable promotion.".to_string(),
+ Governs broadcast selection under both the static
distributed \
+ planner and adaptive query planning (AQE). Set to 0
to disable \
+ promotion.".to_string(),
DataType::UInt64,
Some((10 * 1024 * 1024).to_string())),
+ ConfigEntry::new(BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS.to_string(),
+ "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.".to_string(),
+ DataType::UInt64,
+ Some((1_000_000).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,
@@ -569,6 +586,14 @@ impl BallistaConfig {
self.get_usize_setting(BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES)
}
+ /// Returns the 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. `0`
+ /// disables promotion via the row-count path.
+ pub fn broadcast_join_threshold_rows(&self) -> usize {
+ self.get_usize_setting(BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS)
+ }
+
/// Returns whether the AQE coalesce-shuffle-partitions rule is enabled.
pub fn coalesce_enabled(&self) -> bool {
self.get_bool_setting(BALLISTA_COALESCE_ENABLED)
diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs
index aaa68c673..b0a58289b 100644
--- a/ballista/core/src/extension.rs
+++ b/ballista/core/src/extension.rs
@@ -16,13 +16,13 @@
// under the License.
use crate::config::{
- BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES,
BALLISTA_CLIENT_GRPC_MAX_MESSAGE_SIZE,
- BALLISTA_CLIENT_USE_TLS, BALLISTA_COALESCE_ENABLED,
- BALLISTA_COALESCE_MERGED_PARTITION_FACTOR,
BALLISTA_COALESCE_SMALL_PARTITION_FACTOR,
- BALLISTA_COALESCE_TARGET_PARTITION_BYTES, BALLISTA_JOB_NAME,
- BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ,
BALLISTA_SHUFFLE_READER_MAX_REQUESTS,
- BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT,
BALLISTA_STANDALONE_PARALLELISM,
- BallistaConfig,
+ BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES,
BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS,
+ BALLISTA_CLIENT_GRPC_MAX_MESSAGE_SIZE, BALLISTA_CLIENT_USE_TLS,
+ BALLISTA_COALESCE_ENABLED, BALLISTA_COALESCE_MERGED_PARTITION_FACTOR,
+ BALLISTA_COALESCE_SMALL_PARTITION_FACTOR,
BALLISTA_COALESCE_TARGET_PARTITION_BYTES,
+ BALLISTA_JOB_NAME, BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ,
+ BALLISTA_SHUFFLE_READER_MAX_REQUESTS,
BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT,
+ BALLISTA_STANDALONE_PARALLELISM, BallistaConfig,
};
use crate::planner::BallistaQueryPlanner;
use crate::serde::protobuf::KeyValuePair;
@@ -192,6 +192,18 @@ pub trait SessionConfigExt {
fn with_ballista_broadcast_join_threshold_bytes(self, threshold_bytes:
usize)
-> Self;
+ /// retrieves the 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. `0` disables
+ /// promotion via the row-count path.
+ fn ballista_broadcast_join_threshold_rows(&self) -> usize;
+
+ /// Sets the 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. Setting `0`
+ /// disables promotion via the row-count path.
+ fn with_ballista_broadcast_join_threshold_rows(self, threshold_rows:
usize) -> Self;
+
/// retrieves grpc client max message size
fn ballista_grpc_client_max_message_size(&self) -> usize;
@@ -484,6 +496,23 @@ impl SessionConfigExt for SessionConfig {
}
}
+ fn ballista_broadcast_join_threshold_rows(&self) -> usize {
+ self.options()
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.broadcast_join_threshold_rows())
+ .unwrap_or_else(||
BallistaConfig::default().broadcast_join_threshold_rows())
+ }
+
+ fn with_ballista_broadcast_join_threshold_rows(self, threshold_rows:
usize) -> Self {
+ if self.options().extensions.get::<BallistaConfig>().is_some() {
+ self.set_usize(BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS,
threshold_rows)
+ } else {
+ self.with_option_extension(BallistaConfig::default())
+ .set_usize(BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS,
threshold_rows)
+ }
+ }
+
fn ballista_shuffle_reader_maximum_concurrent_requests(&self) -> usize {
self.options()
.extensions
@@ -745,6 +774,7 @@ impl SessionConfigHelperExt for SessionConfig {
}
fn ballista_restricted_configuration(self) -> Self {
+ let ballista_defaults = BallistaConfig::default();
self
// round robbin repartition does not work well with ballista.
// this setting it will also be enforced by the scheduler
@@ -769,13 +799,17 @@ impl SessionConfigHelperExt for SessionConfig {
//
// A build side smaller than these thresholds is collected into a
// CollectLeft (broadcast) hash join rather than being
repartitioned.
+ // The values mirror Ballista's own broadcast thresholds so a
single
+ // set of `ballista.optimizer.broadcast_join_threshold_*` defaults
+ // drives both DataFusion's built-in JoinSelection (static planner)
+ // and Ballista's AQE join selection.
.set_u64(
"datafusion.optimizer.hash_join_single_partition_threshold",
- 10 * 1024 * 1024,
+ ballista_defaults.broadcast_join_threshold_bytes() as u64,
)
.set_u64(
"datafusion.optimizer.hash_join_single_partition_threshold_rows",
- 1_000_000,
+ ballista_defaults.broadcast_join_threshold_rows() as u64,
)
//
// DataFusion's hash join has no spill support, so each parallel
diff --git a/ballista/scheduler/src/planner.rs
b/ballista/scheduler/src/planner.rs
index cf713ce83..0128540c5 100644
--- a/ballista/scheduler/src/planner.rs
+++ b/ballista/scheduler/src/planner.rs
@@ -266,43 +266,74 @@ impl DefaultDistributedPlanner {
self.next_stage_id
}
- /// If `plan` is a `HashJoinExec(Partitioned)` whose smaller side fits
- /// under the broadcast threshold, returns a rewritten
- /// `HashJoinExec(CollectLeft)` (with a swap if the small side was on
- /// the right) wrapped so the build subtree is a single-partition input.
+ /// Reconciles a join's partition mode with the Ballista broadcast
+ /// threshold (`ballista.optimizer.broadcast_join_threshold_bytes`).
+ ///
+ /// - A `HashJoinExec(Partitioned)` whose smaller side fits under the
+ /// threshold is rewritten as a `HashJoinExec(CollectLeft)` (with a swap
+ /// if the small side was on the right) wrapped so the build subtree is a
+ /// single-partition input.
+ /// - A `HashJoinExec(CollectLeft)` that DataFusion's `JoinSelection` chose
+ /// using its own session threshold is demoted back to `Partitioned` when
+ /// its join type is not broadcast-safe, or when the build side is not
+ /// under the Ballista threshold (including a threshold of `0`, which
+ /// disables broadcast joins). This makes the Ballista key authoritative
+ /// even when it is overridden at runtime below the DataFusion session
+ /// value. Null-aware anti joins are never demoted (they require
+ /// `CollectLeft`).
+ ///
/// Otherwise returns the input unchanged.
fn maybe_promote_to_broadcast(
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
+ let threshold_bytes = config
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.broadcast_join_threshold_bytes())
+ .unwrap_or_else(|| {
+ BallistaConfig::default().broadcast_join_threshold_bytes()
+ });
+
// DataFusion's own `JoinSelection` may have already stamped
// `CollectLeft` on this join (it does so for any build side under
// `datafusion.optimizer.hash_join_single_partition_threshold`, without
// restricting by join type). A `CollectLeft` join replicates the build
- // side to every probe task, which is only correct for probe-driven
join
- // types. If the join type is not broadcast-safe, demote it back to a
- // partitioned (shuffle) join. This is a correctness guard, so it runs
- // regardless of the Ballista broadcast threshold below.
+ // side to every probe task.
if let Some(hash_join) = plan.downcast_ref::<HashJoinExec>()
&& *hash_join.partition_mode() == PartitionMode::CollectLeft
{
- if collect_left_broadcast_safe(*hash_join.join_type()) {
+ // Broadcasting is only correct for probe-driven join types. If the
+ // join type is not broadcast-safe, demote it back to a partitioned
+ // (shuffle) join. Correctness guard, independent of the threshold.
+ if !collect_left_broadcast_safe(*hash_join.join_type()) {
+ debug!(
+ "broadcast check: demoting DataFusion-promoted CollectLeft
join with unsafe join_type={:?} to Partitioned",
+ hash_join.join_type(),
+ );
+ return Self::demote_collect_left_to_partitioned(hash_join,
config);
+ }
+ // Null-aware anti joins must stay `CollectLeft`: they track
+ // probe-side state that a partitioned join cannot reconstruct, so
+ // never demote them regardless of the threshold.
+ if hash_join.null_aware {
return Ok(plan);
}
- debug!(
- "broadcast check: demoting DataFusion-promoted CollectLeft
join with unsafe join_type={:?} to Partitioned",
- hash_join.join_type(),
- );
- return Self::demote_collect_left_to_partitioned(hash_join, config);
+ // Safe join type: honor the Ballista broadcast threshold.
DataFusion
+ // decided `CollectLeft` using its own session threshold, which can
+ // exceed a user's runtime `broadcast_join_threshold_bytes`
override.
+ // If broadcasts are disabled (0) or the build (left) side is not
+ // under the Ballista threshold, demote so the Ballista key is
+ // authoritative in the static planner path too.
+ if threshold_bytes == 0 || !under(&**hash_join.left(),
threshold_bytes) {
+ debug!(
+ "broadcast check: demoting CollectLeft join to
Partitioned; build side not under Ballista threshold={threshold_bytes} (or
broadcasts disabled)",
+ );
+ return Self::demote_collect_left_to_partitioned(hash_join,
config);
+ }
+ return Ok(plan);
}
- let threshold_bytes = config
- .extensions
- .get::<BallistaConfig>()
- .map(|c| c.broadcast_join_threshold_bytes())
- .unwrap_or_else(|| {
- BallistaConfig::default().broadcast_join_threshold_bytes()
- });
if threshold_bytes == 0 {
debug!("broadcast check: threshold is 0, broadcast disabled");
return Ok(plan);
@@ -1324,6 +1355,109 @@ order by
Ok(())
}
+ // Plans `big join small` with DataFusion's own hash-join threshold high
+ // enough for its `JoinSelection` to promote to `CollectLeft`, and the
+ // Ballista broadcast threshold set to `ballista_threshold_bytes`. Returns
+ // the partition modes of every `HashJoinExec` across the resulting stages.
+ async fn collect_left_modes_for_ballista_threshold(
+ ballista_threshold_bytes: usize,
+ ) -> Result<Vec<datafusion::physical_plan::joins::PartitionMode>,
BallistaError> {
+ use ballista_core::extension::SessionConfigExt;
+ use datafusion::arrow::array::Int32Array;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::arrow::record_batch::RecordBatch;
+ use datafusion::prelude::SessionConfig;
+
+ let big_schema = Arc::new(Schema::new(vec![
+ Field::new("k", DataType::Int32, false),
+ Field::new("v", DataType::Int32, false),
+ ]));
+ let small_schema = Arc::new(Schema::new(vec![
+ Field::new("k", DataType::Int32, false),
+ Field::new("w", DataType::Int32, false),
+ ]));
+ let big_batch = RecordBatch::try_new(
+ big_schema,
+ vec![
+ Arc::new(Int32Array::from((0..10_000).collect::<Vec<_>>())),
+ Arc::new(Int32Array::from((0..10_000).collect::<Vec<_>>())),
+ ],
+ )
+ .map_err(|e| BallistaError::General(e.to_string()))?;
+ let small_batch = RecordBatch::try_new(
+ small_schema,
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2])),
+ Arc::new(Int32Array::from(vec![10, 20])),
+ ],
+ )
+ .map_err(|e| BallistaError::General(e.to_string()))?;
+
+ // DF threshold high so DataFusion's JoinSelection promotes to
+ // CollectLeft on its own; the Ballista threshold is what we vary.
+ let session_config = SessionConfig::new()
+ .with_target_partitions(2)
+ .set_bool("datafusion.optimizer.prefer_hash_join", true)
+ .set_usize(
+ "datafusion.optimizer.hash_join_single_partition_threshold",
+ 10 * 1024 * 1024,
+ )
+
.with_ballista_broadcast_join_threshold_bytes(ballista_threshold_bytes);
+ let ctx =
datafusion::prelude::SessionContext::new_with_config(session_config);
+ ctx.register_batch("big", big_batch)?;
+ ctx.register_batch("small", small_batch)?;
+ let options = ctx.state().config().options().clone();
+
+ let plan = ctx
+ .sql("select * from big join small on big.k = small.k")
+ .await?
+ .into_optimized_plan()?;
+ let plan = ctx.state().create_physical_plan(&plan).await?;
+
+ let mut planner = DefaultDistributedPlanner::new();
+ let job_uuid = Uuid::new_v4();
+ let stages =
+ planner.plan_query_stages(&job_uuid.to_string().into(), plan,
&options)?;
+
+ let mut modes = vec![];
+ for stage in &stages {
+ let mut walker: Vec<Arc<dyn ExecutionPlan>> =
+ vec![stage.clone() as Arc<dyn ExecutionPlan>];
+ while let Some(node) = walker.pop() {
+ if let Some(hj) = node.downcast_ref::<HashJoinExec>() {
+ modes.push(*hj.partition_mode());
+ }
+ walker.extend(node.children().iter().map(|c| (*c).clone()));
+ }
+ }
+ Ok(modes)
+ }
+
+ // DataFusion's `JoinSelection` promotes a small build side to
`CollectLeft`
+ // using its own session threshold. When the Ballista broadcast threshold
is
+ // lowered below that build size at runtime, the distributed planner must
+ // demote the `CollectLeft` join back to `Partitioned` so the Ballista key
is
+ // authoritative; a high Ballista threshold leaves the broadcast in place.
+ #[tokio::test]
+ async fn df_collect_left_demoted_when_over_ballista_threshold()
+ -> Result<(), BallistaError> {
+ use datafusion::physical_plan::joins::PartitionMode;
+
+ let kept = collect_left_modes_for_ballista_threshold(10 * 1024 *
1024).await?;
+ assert!(
+ kept.contains(&PartitionMode::CollectLeft),
+ "a high Ballista threshold must keep the DataFusion-promoted
CollectLeft join, got {kept:?}"
+ );
+
+ let demoted = collect_left_modes_for_ballista_threshold(8).await?;
+ assert!(
+ !demoted.contains(&PartitionMode::CollectLeft),
+ "a Ballista threshold below the build side must demote the
CollectLeft join to Partitioned, got {demoted:?}"
+ );
+
+ Ok(())
+ }
+
// A LEFT join with the small side on the left builds (broadcasts) the left
// side and emits a null-padded row for every unmatched left row; each
probe
// task would emit those independently. The guard must keep it
repartitioned
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 c7d22ae8a..9391709ea 100644
--- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs
+++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+use ballista_core::config::BallistaConfig;
use datafusion::{
arrow::compute::SortOptions,
common::{JoinType, NullEquality, Result, exec_err, internal_err},
@@ -214,20 +215,34 @@ impl DynamicJoinSelectionExec {
config: &ConfigOptions,
) -> Result<JoinSelectionAction> {
let prefer_hash_join = config.optimizer.prefer_hash_join;
- let threshold_collect_left_join_bytes =
- config.optimizer.hash_join_single_partition_threshold;
- let threshold_collect_left_join_rows =
- config.optimizer.hash_join_single_partition_threshold_rows;
-
- let under_threshold = Self::supports_collect_by_thresholds(
- self.left.as_ref(),
- threshold_collect_left_join_bytes,
- threshold_collect_left_join_rows,
- ) || Self::supports_collect_by_thresholds(
- self.right.as_ref(),
- threshold_collect_left_join_bytes,
- threshold_collect_left_join_rows,
- );
+ // Both broadcast thresholds come from Ballista's own config so that a
+ // single set of keys governs broadcast selection under both the static
+ // planner (`maybe_promote_to_broadcast`) and AQE. A byte threshold of 0
+ // disables broadcast promotion entirely, matching the static planner.
+ // The row threshold is only consulted as a fallback when byte-size
+ // statistics are absent.
+ let bc = config
+ .extensions
+ .get::<BallistaConfig>()
+ .cloned()
+ .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();
+
+ // The `!= 0` guard sits here rather than inside
+ // `supports_collect_by_thresholds`: that helper falls back to the row
+ // check when the byte threshold is 0/absent, so gating at this level
is
+ // what makes a 0 threshold disable both the byte and row broadcast
paths.
+ let under_threshold = threshold_collect_left_join_bytes != 0
+ && (Self::supports_collect_by_thresholds(
+ self.left.as_ref(),
+ threshold_collect_left_join_bytes,
+ threshold_collect_left_join_rows,
+ ) || Self::supports_collect_by_thresholds(
+ self.right.as_ref(),
+ threshold_collect_left_join_bytes,
+ threshold_collect_left_join_rows,
+ ));
// The resolver collects the smaller side onto the build (left) input,
// swapping the join type when the right side is the smaller one (the
@@ -446,6 +461,7 @@ mod tests {
use crate::state::aqe::execution_plan::ExchangeExec;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::common::{ColumnStatistics, Statistics, stats::Precision};
+ use datafusion::config::ExtensionOptions;
use datafusion::physical_plan::empty::EmptyExec;
use datafusion::physical_plan::expressions::Column;
use datafusion::physical_plan::test::exec::StatisticsExec;
@@ -465,16 +481,29 @@ mod tests {
))
}
- /// Run the resolver's join-strategy decision for `join_type` with
`left_rows`
- /// and `right_rows` on the two sides, both small enough to be collectable.
- fn actual_join_for(
+ /// 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)]),
+ ))
+ }
+
+ /// Core driver: runs the join-strategy decision for the given children
with
+ /// explicit Ballista broadcast thresholds.
+ fn run_to_actual_join(
+ left: Arc<dyn ExecutionPlan>,
+ right: Arc<dyn ExecutionPlan>,
join_type: JoinType,
- left_rows: usize,
- right_rows: usize,
prefer_hash_join: bool,
+ broadcast_threshold_bytes: usize,
+ broadcast_threshold_rows: usize,
) -> JoinSelectionAction {
- let left = stats_exec(left_rows);
- let right = stats_exec(right_rows);
let on: JoinOn =
vec![(Arc::new(Column::new("k", 0)), Arc::new(Column::new("k",
0)))];
let dj = DynamicJoinSelectionExec {
@@ -492,11 +521,62 @@ mod tests {
};
let mut config = ConfigOptions::new();
config.optimizer.prefer_hash_join = prefer_hash_join;
- config.optimizer.hash_join_single_partition_threshold = 10 * 1024 *
1024;
- config.optimizer.hash_join_single_partition_threshold_rows = 1_000_000;
+ let mut bc = BallistaConfig::default();
+ bc.set(
+ "optimizer.broadcast_join_threshold_bytes",
+ &broadcast_threshold_bytes.to_string(),
+ )
+ .unwrap();
+ bc.set(
+ "optimizer.broadcast_join_threshold_rows",
+ &broadcast_threshold_rows.to_string(),
+ )
+ .unwrap();
+ config.extensions.insert(bc);
dj.to_actual_join(&config).unwrap()
}
+ /// Run the resolver's join-strategy decision for `join_type` with
`left_rows`
+ /// and `right_rows` on the two sides, both small enough to be collectable.
+ fn actual_join_for(
+ join_type: JoinType,
+ left_rows: usize,
+ right_rows: usize,
+ prefer_hash_join: bool,
+ ) -> JoinSelectionAction {
+ // 10 MiB is large enough for the tiny StatisticsExec sides below to be
+ // collectable, matching the default broadcast threshold.
+ actual_join_with_threshold(
+ join_type,
+ left_rows,
+ right_rows,
+ prefer_hash_join,
+ 10 * 1024 * 1024,
+ )
+ }
+
+ /// Same as [`actual_join_for`] but with an explicit
+ /// `broadcast_join_threshold_bytes`, so tests can prove the Ballista
config
+ /// drives the byte cutoff (0 disables broadcast promotion).
+ fn actual_join_with_threshold(
+ join_type: JoinType,
+ left_rows: usize,
+ right_rows: usize,
+ prefer_hash_join: bool,
+ broadcast_threshold_bytes: usize,
+ ) -> JoinSelectionAction {
+ run_to_actual_join(
+ stats_exec(left_rows),
+ stats_exec(right_rows),
+ join_type,
+ prefer_hash_join,
+ broadcast_threshold_bytes,
+ // 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,
+ )
+ }
+
fn is_collected(action: &JoinSelectionAction) -> bool {
matches!(
action,
@@ -553,6 +633,77 @@ mod tests {
);
}
+ // The byte cutoff comes from
`ballista.optimizer.broadcast_join_threshold_bytes`,
+ // not DataFusion's `hash_join_single_partition_threshold`. `stats_exec(n)`
+ // reports `n * 16` bytes, so the smaller side here is 1600 bytes. A
threshold
+ // above it collects; a threshold below it repartitions.
+ #[test]
+ fn broadcast_threshold_bytes_drives_collect_decision() {
+ // smaller (left) side = 100 * 16 = 1600 bytes, larger = 16000 bytes
+ assert!(
+ is_collected(&actual_join_with_threshold(
+ JoinType::Inner,
+ 100,
+ 1000,
+ true,
+ 2000,
+ )),
+ "smaller side (1600 B) is under a 2000 B threshold and must
collect"
+ );
+ assert!(
+ !is_collected(&actual_join_with_threshold(
+ JoinType::Inner,
+ 100,
+ 1000,
+ true,
+ 1000,
+ )),
+ "neither side is under a 1000 B threshold, so the join must
repartition"
+ );
+ }
+
+ // When byte-size statistics are absent, the decision falls back to the
+ // Ballista row-count threshold (`broadcast_join_threshold_rows`), not
+ // DataFusion's `hash_join_single_partition_threshold_rows`.
+ #[test]
+ fn broadcast_threshold_rows_drives_fallback_decision() {
+ // Byte-size is Absent on both sides, so only the row threshold
applies.
+ // Smaller side = 100 rows. A generous byte threshold keeps the byte
+ // guard open so the row fallback is what decides.
+ assert!(
+ is_collected(&run_to_actual_join(
+ stats_exec_rows_only(100),
+ stats_exec_rows_only(1000),
+ JoinType::Inner,
+ true,
+ 10 * 1024 * 1024,
+ 200,
+ )),
+ "smaller side (100 rows) is under a 200-row threshold and must
collect"
+ );
+ assert!(
+ !is_collected(&run_to_actual_join(
+ stats_exec_rows_only(100),
+ stats_exec_rows_only(1000),
+ JoinType::Inner,
+ true,
+ 10 * 1024 * 1024,
+ 50,
+ )),
+ "neither side is under a 50-row threshold, so the join must
repartition"
+ );
+ }
+
+ // A threshold of 0 disables broadcast promotion entirely, matching the
static
+ // planner, even for a side small enough to collect under any positive
cutoff.
+ #[test]
+ fn zero_broadcast_threshold_disables_collect() {
+ assert!(
+ !is_collected(&actual_join_with_threshold(JoinType::Inner, 1, 1,
true, 0,)),
+ "broadcast_join_threshold_bytes=0 must disable CollectLeft
promotion"
+ );
+ }
+
/// 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`.
diff --git a/ballista/scheduler/src/state/aqe/test/join_selection.rs
b/ballista/scheduler/src/state/aqe/test/join_selection.rs
index b652ca7d3..aa4aac860 100644
--- a/ballista/scheduler/src/state/aqe/test/join_selection.rs
+++ b/ballista/scheduler/src/state/aqe/test/join_selection.rs
@@ -19,6 +19,7 @@ use crate::{
assert_plan,
state::aqe::{planner::AdaptivePlanner,
test::mock_partitions_with_statistics},
};
+use ballista_core::extension::SessionConfigExt;
use datafusion::{
arrow::{
array::{Int32Array, RecordBatch},
@@ -85,6 +86,10 @@ fn make_ctx_without_collect_left(prefer_hash_join: bool) ->
SessionContext {
"datafusion.optimizer.hash_join_single_partition_threshold_rows",
0,
)
+ // AQE join selection reads the broadcast cutoff from Ballista's own
+ // config; a byte threshold of 0 disables CollectLeft promotion,
keeping
+ // these joins repartitioned.
+ .with_ballista_broadcast_join_threshold_bytes(0)
.with_target_partitions(4)
.with_round_robin_repartition(false);
let state = SessionStateBuilder::new_with_default_features()
diff --git a/docs/source/user-guide/tuning-guide.md
b/docs/source/user-guide/tuning-guide.md
index c8845ba22..c2daa312b 100644
--- a/docs/source/user-guide/tuning-guide.md
+++ b/docs/source/user-guide/tuning-guide.md
@@ -196,9 +196,11 @@ Adaptive Query Planning is EXPERIMENTAL, should be used
for testing purposes onl
### Configuration
-| key | type | default | description
|
-| --------------------------------- | ------- | ------- |
------------------------------------------- |
-| ballista.planner.adaptive.enabled | Boolean | false | Enables the adaptive
planner. Experimental. |
+| key | type | default |
description
|
+| ------------------------------------------------- | ------- | -------- |
----------------------------------------------------------------------------------------------------------------------------------------------------------------------
|
+| ballista.planner.adaptive.enabled | Boolean | false |
Enables the adaptive planner. Experimental.
|
+| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 10485760 |
Byte-size threshold below which a hash join's smaller side is broadcast
(`CollectLeft`). Governs both the static planner and AQE. Set to 0 to disable
broadcast joins. |
+| ballista.optimizer.broadcast_join_threshold_rows | UInt64 | 1000000 |
Row-count fallback threshold used when byte-size statistics are unavailable.
Applies to AQE. Set to 0 to disable promotion via the row-count path.
|
### What AQE does today
@@ -209,6 +211,9 @@ implemented:
- **Join reordering.** Uses runtime row counts from completed stages so the
smaller side drives the join.
+- **Broadcast join selection.** When a join input's runtime size falls under
+ `ballista.optimizer.broadcast_join_threshold_bytes` (or the row-count
+ fallback), the smaller side is broadcast (`CollectLeft`) instead of shuffled.
- **Empty stage elimination.** When a completed stage produces zero rows, its
downstream exchange is replaced with an empty execution node, and emptiness
is propagated up the plan so downstream stages are skipped entirely.
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
