From 31f8b6e43d2775700a17acbd65017755ca14020d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 17 Jul 2026 09:49:51 -0600 Subject: [PATCH 01/22] 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 --- ballista/core/src/config.rs | 4 +- .../state/aqe/execution_plan/dynamic_join.rs | 107 ++++++++++++++++-- 2 files changed, 98 insertions(+), 13 deletions(-) diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index 406d1b9e82..142b69f0e3 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -236,7 +236,9 @@ static CONFIG_ENTRIES: LazyLock> = 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_CLIENT_PULL.to_string(), 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 c7d22ae8aa..4d388c9cc7 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,35 @@ impl DynamicJoinSelectionExec { config: &ConfigOptions, ) -> Result { let prefer_hash_join = config.optimizer.prefer_hash_join; - let threshold_collect_left_join_bytes = - config.optimizer.hash_join_single_partition_threshold; + // The byte threshold comes from Ballista's own broadcast config so that a + // single key governs broadcast selection under both the static planner + // (`maybe_promote_to_broadcast`) and AQE. A value of 0 disables broadcast + // promotion entirely, matching the static planner. The row threshold is + // only a fallback used when byte statistics are absent, so it stays on + // DataFusion's setting. + let bc = config + .extensions + .get::() + .cloned() + .unwrap_or_default(); + let threshold_collect_left_join_bytes = bc.broadcast_join_threshold_bytes(); 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, - ); + // 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 @@ -445,6 +461,7 @@ mod tests { use super::*; use crate::state::aqe::execution_plan::ExchangeExec; use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::config::ExtensionOptions; use datafusion::common::{ColumnStatistics, Statistics, stats::Precision}; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::expressions::Column; @@ -472,6 +489,27 @@ mod tests { 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 { let left = stats_exec(left_rows); let right = stats_exec(right_rows); @@ -492,8 +530,14 @@ 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(); + config.extensions.insert(bc); dj.to_actual_join(&config).unwrap() } @@ -553,6 +597,45 @@ 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" + ); + } + + // 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`. From aade4aace656e493823d74c89359799b0570f8e7 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 17 Jul 2026 09:57:09 -0600 Subject: [PATCH 02/22] 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. --- ballista/core/src/config.rs | 23 +++ ballista/core/src/extension.rs | 43 +++++- .../state/aqe/execution_plan/dynamic_join.rs | 138 +++++++++++++----- docs/source/user-guide/tuning-guide.md | 11 +- 4 files changed, 170 insertions(+), 45 deletions(-) diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index 142b69f0e3..f1128fa4af 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. @@ -241,6 +248,14 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| 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((128 * 1024).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, @@ -571,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 aaa68c6735..1777db80a5 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::() + .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::().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 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 4d388c9cc7..9391709ea6 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -215,20 +215,19 @@ impl DynamicJoinSelectionExec { config: &ConfigOptions, ) -> Result { let prefer_hash_join = config.optimizer.prefer_hash_join; - // The byte threshold comes from Ballista's own broadcast config so that a - // single key governs broadcast selection under both the static planner - // (`maybe_promote_to_broadcast`) and AQE. A value of 0 disables broadcast - // promotion entirely, matching the static planner. The row threshold is - // only a fallback used when byte statistics are absent, so it stays on - // DataFusion's setting. + // 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::() .cloned() .unwrap_or_default(); let threshold_collect_left_join_bytes = bc.broadcast_join_threshold_bytes(); - let threshold_collect_left_join_rows = - config.optimizer.hash_join_single_partition_threshold_rows; + 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 @@ -461,8 +460,8 @@ mod tests { use super::*; use crate::state::aqe::execution_plan::ExchangeExec; use datafusion::arrow::datatypes::{DataType, Field, Schema}; - use datafusion::config::ExtensionOptions; 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; @@ -482,37 +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( - 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, - ) + /// 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 { + 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)]), + )) } - /// 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( + /// Core driver: runs the join-strategy decision for the given children with + /// explicit Ballista broadcast thresholds. + fn run_to_actual_join( + left: Arc, + right: Arc, 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 { @@ -530,17 +521,62 @@ mod tests { }; let mut config = ConfigOptions::new(); config.optimizer.prefer_hash_join = prefer_hash_join; - 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, @@ -626,6 +662,38 @@ mod tests { ); } + // 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] diff --git a/docs/source/user-guide/tuning-guide.md b/docs/source/user-guide/tuning-guide.md index c8845ba225..9f989b4208 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 | 131072 | 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. From 2729768000e3a898b4984b098546d2266f7d4dbc Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 17 Jul 2026 10:10:51 -0600 Subject: [PATCH 03/22] 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. --- ballista/core/src/config.rs | 2 +- docs/source/user-guide/tuning-guide.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index f1128fa4af..bebcdc0e7f 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -255,7 +255,7 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| Applies to adaptive query planning (AQE). Set to 0 to disable \ promotion via the row-count path.".to_string(), DataType::UInt64, - Some((128 * 1024).to_string())), + 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, diff --git a/docs/source/user-guide/tuning-guide.md b/docs/source/user-guide/tuning-guide.md index 9f989b4208..f8889c409b 100644 --- a/docs/source/user-guide/tuning-guide.md +++ b/docs/source/user-guide/tuning-guide.md @@ -200,7 +200,7 @@ Adaptive Query Planning is EXPERIMENTAL, should be used for testing purposes onl | ------------------------------------------------ | ------- | -------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | 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 | 131072 | 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. | +| 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 From bd940a372e76321390d5c574ca3816f9a388db1a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 17 Jul 2026 10:20:00 -0600 Subject: [PATCH 04/22] 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. --- ballista/core/src/extension.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index 1777db80a5..b0a58289b9 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -774,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 @@ -798,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 From a42245a8be8bdc12b7082a99ed669dbd4d0f474e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 17 Jul 2026 11:03:05 -0600 Subject: [PATCH 05/22] 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. --- ballista/scheduler/src/planner.rs | 176 ++++++++++++++++++++++++++---- 1 file changed, 155 insertions(+), 21 deletions(-) diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index cf713ce83d..0128540c5b 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, config: &ConfigOptions, ) -> Result> { + let threshold_bytes = config + .extensions + .get::() + .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::() && *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::() - .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, 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::>())), + Arc::new(Int32Array::from((0..10_000).collect::>())), + ], + ) + .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> = + vec![stage.clone() as Arc]; + while let Some(node) = walker.pop() { + if let Some(hj) = node.downcast_ref::() { + 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 From 6809fa5e3aa2c8c2a2aa5ed6fe76b0ec097bb15d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 17 Jul 2026 11:15:40 -0600 Subject: [PATCH 06/22] docs: apply prettier formatting to AQE tuning-guide table --- docs/source/user-guide/tuning-guide.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/source/user-guide/tuning-guide.md b/docs/source/user-guide/tuning-guide.md index f8889c409b..c2daa312b2 100644 --- a/docs/source/user-guide/tuning-guide.md +++ b/docs/source/user-guide/tuning-guide.md @@ -196,11 +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. | +| 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 From bf185ce970f529ee5676b89e17368522a633b16d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 17 Jul 2026 14:18:46 -0600 Subject: [PATCH 07/22] 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/scheduler/src/state/aqe/test/join_selection.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/ballista/scheduler/src/state/aqe/test/join_selection.rs b/ballista/scheduler/src/state/aqe/test/join_selection.rs index b652ca7d30..aa4aac860c 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() From 806fcb78d9e02d08858b384f9231df3e212208bc Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 17 Jul 2026 02:32:56 -0600 Subject: [PATCH 08/22] 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. --- .../state/aqe/execution_plan/dynamic_join.rs | 230 +++++++++++++++++- 1 file changed, 224 insertions(+), 6 deletions(-) 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 9391709ea6..3eb4fce2ca 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, @@ -365,12 +366,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,6 +473,65 @@ impl DynamicJoinSelectionExec { } } +/// 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 { + 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::*; @@ -466,10 +543,151 @@ mod tests { 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; + fn test_schema() -> Arc { 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, + ) -> Arc { + 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) -> 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 = 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 = 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 = 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 = 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 { Arc::new(StatisticsExec::new( Statistics { From f33bef64f5bdaa48399bf0a6794beaa1fc6915b1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 17 Jul 2026 03:12:42 -0600 Subject: [PATCH 09/22] 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. --- .../state/aqe/test/broadcast_thresholds.rs | 299 ++++++++++++++++++ ballista/scheduler/src/state/aqe/test/mod.rs | 4 + .../src/state/aqe/test/stats_table.rs | 232 ++++++++++++++ 3 files changed, 535 insertions(+) create mode 100644 ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs create mode 100644 ballista/scheduler/src/state/aqe/test/stats_table.rs 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 0000000000..dfa0fac75e --- /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 { + 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 { + 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, 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 a8640d3e30..e4bffbbb7c 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 0000000000..1e4174b9ff --- /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>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + 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, +} + +impl StatsExec { + fn try_new( + schema: SchemaRef, + statistics: Statistics, + partitions: usize, + projection: Option<&Vec>, + ) -> Result { + 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 { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + internal_err!("StatsExec holds no data and cannot be executed") + } + + fn partition_statistics(&self, partition: Option) -> Result> { + // 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()], + } +} From affeec0b9c3f7bf2488e46da364892301724e730 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 18 Jul 2026 11:16:53 -0600 Subject: [PATCH 10/22] 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. --- .../execution_plans/sort_shuffle/writer.rs | 43 +++++++++++++------ .../state/aqe/execution_plan/dynamic_join.rs | 42 ++++++++++++------ 2 files changed, 59 insertions(+), 26 deletions(-) diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index 8612acdb33..a08a512223 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -65,7 +65,7 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, Statistics, displayable, }; use futures::{StreamExt, TryFutureExt, TryStreamExt}; -use log::debug; +use log::{debug, info, warn}; use crate::serde::scheduler::PartitionStats; @@ -350,18 +350,35 @@ impl SortShuffleWriterExec { // Reservation drops naturally; nothing left to free. drop(reservation); - debug!( - "Sort shuffle write for partition {} completed in {} seconds. \ - Output: {:?}, Index: {:?}, Spill events: {}, Spill batches: {}, \ - Spill bytes: {}", - input_partition, - now.elapsed().as_secs(), - data_path, - index_path, - spill_events, - total_spilled_batches, - total_bytes_spilled - ); + let elapsed_secs = now.elapsed().as_secs(); + if total_bytes_spilled > 0 { + warn!( + "Sort shuffle spill: job={} stage={} input_partition={} \ + spilled {} bytes in {} batches ({} events) under memory \ + pressure; write completed in {} seconds", + job_id, + stage_id, + input_partition, + total_bytes_spilled, + total_spilled_batches, + spill_events, + elapsed_secs, + ); + } else { + info!( + "Sort shuffle write for partition {} completed in {} seconds. \ + Output: {:?}, Index: {:?}, Rows: {}, Spill events: {}, \ + Spill batches: {}, Spill bytes: {}", + input_partition, + elapsed_secs, + data_path, + index_path, + total_rows, + spill_events, + total_spilled_batches, + total_bytes_spilled + ); + } let mut results = Vec::new(); for (part_id, num_batches, num_rows, num_bytes) in partition_stats { 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 3eb4fce2ca..4674fc515d 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -31,7 +31,7 @@ use datafusion::{ }, }, }; -use log::debug; +use log::info; use std::sync::Arc; use crate::physical_optimizer::join_selection::collect_left_broadcast_safe; @@ -271,17 +271,7 @@ 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), @@ -316,7 +306,33 @@ 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)", + }; + + info!( + "AQE join decision plan_id={} action={} partition_mode={:?} \ + under_threshold={} left=(rows={:?}, bytes={:?}) \ + right=(rows={:?}, bytes={:?}) byte_threshold={} row_threshold={}", + self.plan_id, + action_label, + partition_mode, + under_threshold, + 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( From b53c1e35f688cb32a4df6c107656bdbbc023abd7 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 18 Jul 2026 12:18:09 -0600 Subject: [PATCH 11/22] 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. --- ballista/core/proto/ballista.proto | 3 ++ ballista/core/src/config.rs | 9 ++-- .../execution_plans/sort_shuffle/config.rs | 15 +++--- .../execution_plans/sort_shuffle/writer.rs | 9 ++-- ballista/core/src/serde/generated/ballista.rs | 4 ++ ballista/core/src/serde/mod.rs | 46 ++++++++++++++++++- 6 files changed, 72 insertions(+), 14 deletions(-) diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index 59d844abc9..5cf355a595 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -77,6 +77,9 @@ message SortShuffleWriterExecNode { datafusion.PhysicalHashRepartition output_partitioning = 4; // Target batch size in rows when materializing buffered shuffle data. uint64 batch_size = 8; + // Per-task buffered-bytes budget at which the writer spills to disk. 0 means + // no cap (spill is driven solely by memory-pool pressure). + uint64 memory_limit_per_task_bytes = 9; } message UnresolvedShuffleExecNode { diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index bebcdc0e7f..82eb604a56 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -235,11 +235,12 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| Some(DEFAULT_SHUFFLE_CHANNEL_CAPACITY.to_string())), ConfigEntry::new(BALLISTA_SHUFFLE_SORT_BASED_MEMORY_LIMIT_PER_TASK_BYTES.to_string(), "Per-task buffered-bytes budget at which the sort shuffle writer spills its \ - in-memory batches to disk. Counted independently of the runtime memory pool, so \ - spilling kicks in even when the pool is unbounded. Total worst-case sort shuffle \ - memory per executor is approximately vcores * this value.".to_string(), + in-memory batches to disk, counted independently of the runtime memory pool. A \ + value of 0 (the default) disables this cap so spilling is driven solely by \ + memory-pool pressure. When non-zero, total worst-case sort shuffle memory per \ + executor is approximately vcores * this value.".to_string(), DataType::UInt64, - Some((256 * 1024 * 1024).to_string())), + Some("0".to_string())), 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. \ diff --git a/ballista/core/src/execution_plans/sort_shuffle/config.rs b/ballista/core/src/execution_plans/sort_shuffle/config.rs index 5dd4b15095..461cc2e63a 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/config.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/config.rs @@ -26,8 +26,10 @@ pub struct SortShuffleConfig { /// `interleave_record_batch` (default: 8192). pub batch_size: usize, /// Per-task buffered-bytes budget at which the writer spills its in-memory - /// batches to disk. Counted independently of the runtime `MemoryPool`, so - /// spilling kicks in even when the pool is unbounded. + /// batches to disk, counted independently of the runtime `MemoryPool`. A + /// value of 0 (the default) disables this cap so spilling is driven solely + /// by memory-pool pressure; a non-zero value adds a second spill trigger + /// that bounds buffered memory even when the pool is unbounded. pub memory_limit_per_task_bytes: usize, } @@ -36,13 +38,14 @@ impl Default for SortShuffleConfig { Self { enabled: false, batch_size: 8192, - memory_limit_per_task_bytes: 256 * 1024 * 1024, + memory_limit_per_task_bytes: 0, } } } impl SortShuffleConfig { - /// Creates a new configuration with the default per-task memory limit. + /// Creates a new configuration with the default (uncapped) per-task memory + /// limit. pub fn new(enabled: bool, batch_size: usize) -> Self { Self { enabled, @@ -67,7 +70,7 @@ mod tests { let config = SortShuffleConfig::default(); assert!(!config.enabled); assert_eq!(config.batch_size, 8192); - assert_eq!(config.memory_limit_per_task_bytes, 256 * 1024 * 1024); + assert_eq!(config.memory_limit_per_task_bytes, 0); } #[test] @@ -75,7 +78,7 @@ mod tests { let config = SortShuffleConfig::new(true, 4096); assert!(config.enabled); assert_eq!(config.batch_size, 4096); - assert_eq!(config.memory_limit_per_task_bytes, 256 * 1024 * 1024); + assert_eq!(config.memory_limit_per_task_bytes, 0); } #[test] diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index a08a512223..5ed2ab08fc 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -246,8 +246,9 @@ impl SortShuffleWriterExec { let mut hash_buffer: Vec = 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; let memory_limit = config.memory_limit_per_task_bytes; @@ -282,7 +283,9 @@ impl SortShuffleWriterExec { let pool_rejected_growth = reservation.try_grow(growth).is_err(); buffered_bytes = buffered_bytes.saturating_add(growth); - if pool_rejected_growth || buffered_bytes >= memory_limit { + if pool_rejected_growth + || (memory_limit > 0 && buffered_bytes >= memory_limit) + { let spill_timer = metrics.spill_time.timer(); let (event_batches, event_bytes) = spill_all_partitions( &mut buffered, diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index 7d08dca420..64e22d2f53 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -93,6 +93,10 @@ pub struct SortShuffleWriterExecNode { /// Target batch size in rows when materializing buffered shuffle data. #[prost(uint64, tag = "8")] pub batch_size: u64, + /// Per-task buffered-bytes budget at which the writer spills to disk. 0 means + /// no cap (spill is driven solely by memory-pool pressure). + #[prost(uint64, tag = "9")] + pub memory_limit_per_task_bytes: u64, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct UnresolvedShuffleExecNode { diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 045b303f01..ae25f78f96 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -423,7 +423,10 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { } else { 8192 }; - let config = SortShuffleConfig::new(true, batch_size); + let config = SortShuffleConfig::new(true, batch_size) + .with_memory_limit_per_task_bytes( + sort_shuffle_writer.memory_limit_per_task_bytes as usize, + ); Ok(Arc::new(SortShuffleWriterExec::try_new( sort_shuffle_writer.job_id.clone().into(), @@ -622,6 +625,8 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { input: None, output_partitioning, batch_size: config.batch_size as u64, + memory_limit_per_task_bytes: config.memory_limit_per_task_bytes + as u64, }, )), }; @@ -1009,6 +1014,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 = 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 = 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::() + .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(); From c6a474a9abfbde5aed97d4becc931073b2ab2730 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 18 Jul 2026 22:20:30 -0600 Subject: [PATCH 12/22] feat: add hash_join_max_build_partition_bytes config --- ballista/core/src/config.rs | 26 ++++++++++++++++++++++++++ ballista/core/src/extension.rs | 14 ++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index 82eb604a56..b0f2e1e254 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -110,6 +110,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. @@ -257,6 +264,12 @@ static CONFIG_ENTRIES: LazyLock> = 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, @@ -595,6 +608,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) @@ -859,4 +877,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/extension.rs b/ballista/core/src/extension.rs index b0a58289b9..28fd77e650 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::() + .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 From 99e66ee4d84cc64aa6b825dfa600fe9445f13d12 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 18 Jul 2026 22:29:09 -0600 Subject: [PATCH 13/22] 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). --- .../state/aqe/execution_plan/dynamic_join.rs | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) 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 4674fc515d..f92dbddc2a 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -489,6 +489,43 @@ impl DynamicJoinSelectionExec { } } +/// 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 +/// no partition reports a byte size. +#[allow(dead_code)] // not wired into `to_actual_join` yet; a follow-up consumes this +fn max_per_partition_build_bytes(build: &Arc) -> Option { + let exchange = build.downcast_ref::()?; + let partitions = exchange.shuffle_partitions()?; + partitions + .iter() + .map(|locations| { + locations + .iter() + .filter_map(|location| location.partition_stats.num_bytes()) + .sum::() + }) + .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`. @@ -552,15 +589,21 @@ fn estimated_value_width(data_type: &DataType) -> usize { 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::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 { Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)])) @@ -1061,4 +1104,63 @@ 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 { + let schema = test_schema(); + let input = Arc::new(EmptyExec::new(schema)) as Arc; + 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 + } + + // 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)); + } } From 51746083a367e93d5a1a22728e7b3f9987716f90 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 18 Jul 2026 22:36:46 -0600 Subject: [PATCH 14/22] feat: fall back to SMJ when hash-join build exceeds per-slot budget --- .../state/aqe/execution_plan/dynamic_join.rs | 129 +++++++++++++++++- 1 file changed, 127 insertions(+), 2 deletions(-) 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 f92dbddc2a..4c86afb536 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -229,6 +229,7 @@ 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(); // The `!= 0` guard sits here rather than inside // `supports_collect_by_thresholds`: that helper falls back to the row @@ -276,7 +277,7 @@ impl DynamicJoinSelectionExec { .to_hash_join(PartitionMode::CollectLeft) .map(JoinSelectionAction::CollectLeft), (JoinInputState::Repartitioned, PartitionMode::Partitioned) - if prefer_hash_join => + if prefer_hash_join && self.hash_build_fits(max_build_bytes) => { self.to_hash_join(PartitionMode::Partitioned) .map(JoinSelectionAction::Hash) @@ -335,6 +336,21 @@ impl DynamicJoinSelectionExec { Ok(action) } + /// Whether the build (left) side of a `Partitioned` hash join fits the + /// per-slot memory budget: `max_build_bytes == 0` means the check is + /// disabled (always fits, matching current behavior before this check + /// existed), and an unknown build size is treated as fitting too, since + /// there is no evidence to force a fallback. + fn hash_build_fits(&self, max_build_bytes: usize) -> bool { + if max_build_bytes == 0 { + return true; + } + match max_per_partition_build_bytes(&self.left) { + Some(bytes) => bytes <= max_build_bytes, + None => true, // unknown size → don't force a fallback + } + } + pub(crate) fn to_hash_join( &self, partition_mode: PartitionMode, @@ -510,7 +526,6 @@ impl DynamicJoinSelectionExec { /// 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 /// no partition reports a byte size. -#[allow(dead_code)] // not wired into `to_actual_join` yet; a follow-up consumes this fn max_per_partition_build_bytes(build: &Arc) -> Option { let exchange = build.downcast_ref::()?; let partitions = exchange.shuffle_partitions()?; @@ -597,6 +612,7 @@ mod tests { 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; @@ -1163,4 +1179,113 @@ mod tests { 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 { + match action { + JoinSelectionAction::Repartition(exec) => { + Arc::clone(exec) as Arc + } + JoinSelectionAction::CollectLeft(exec) => { + Arc::clone(exec) as Arc + } + JoinSelectionAction::LateCollectLeft(exec) => { + Arc::clone(exec) as Arc + } + JoinSelectionAction::Hash(exec) => Arc::clone(exec) as Arc, + JoinSelectionAction::Sort(exec) => Arc::clone(exec) as Arc, + } + } + + // 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}"); + } } From 45d9844acbf1457e395e8ab44909af9bdcb41061 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 18 Jul 2026 22:42:40 -0600 Subject: [PATCH 15/22] feat: log hash-join build-fit decision at INFO --- .../src/state/aqe/execution_plan/dynamic_join.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) 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 4c86afb536..685daef7af 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -318,13 +318,16 @@ impl DynamicJoinSelectionExec { }; info!( - "AQE join decision plan_id={} action={} partition_mode={:?} \ - under_threshold={} left=(rows={:?}, bytes={:?}) \ - right=(rows={:?}, bytes={:?}) byte_threshold={} row_threshold={}", + "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, + max_per_partition_build_bytes(&self.left), + max_build_bytes, stats_left.num_rows, stats_left.total_byte_size, stats_right.num_rows, From afef9afcfb5d2e5188c0794c0baed7f98013ed37 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 18 Jul 2026 22:49:01 -0600 Subject: [PATCH 16/22] refactor: hoist build-size call and tidy hash-join fit-check docs --- .../state/aqe/execution_plan/dynamic_join.rs | 44 +++++++++++-------- 1 file changed, 26 insertions(+), 18 deletions(-) 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 685daef7af..568e2f330c 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -230,6 +230,7 @@ impl DynamicJoinSelectionExec { 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 @@ -277,7 +278,8 @@ impl DynamicJoinSelectionExec { .to_hash_join(PartitionMode::CollectLeft) .map(JoinSelectionAction::CollectLeft), (JoinInputState::Repartitioned, PartitionMode::Partitioned) - if prefer_hash_join && self.hash_build_fits(max_build_bytes) => + if prefer_hash_join + && hash_build_fits(max_build_bytes, build_max_partition_bytes) => { self.to_hash_join(PartitionMode::Partitioned) .map(JoinSelectionAction::Hash) @@ -326,7 +328,7 @@ impl DynamicJoinSelectionExec { action_label, partition_mode, under_threshold, - max_per_partition_build_bytes(&self.left), + build_max_partition_bytes, max_build_bytes, stats_left.num_rows, stats_left.total_byte_size, @@ -339,21 +341,6 @@ impl DynamicJoinSelectionExec { Ok(action) } - /// Whether the build (left) side of a `Partitioned` hash join fits the - /// per-slot memory budget: `max_build_bytes == 0` means the check is - /// disabled (always fits, matching current behavior before this check - /// existed), and an unknown build size is treated as fitting too, since - /// there is no evidence to force a fallback. - fn hash_build_fits(&self, max_build_bytes: usize) -> bool { - if max_build_bytes == 0 { - return true; - } - match max_per_partition_build_bytes(&self.left) { - Some(bytes) => bytes <= max_build_bytes, - None => true, // unknown size → don't force a fallback - } - } - pub(crate) fn to_hash_join( &self, partition_mode: PartitionMode, @@ -508,6 +495,24 @@ 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, +) -> 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 @@ -528,7 +533,10 @@ impl DynamicJoinSelectionExec { /// /// 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 -/// no partition reports a byte size. +/// 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) -> Option { let exchange = build.downcast_ref::()?; let partitions = exchange.shuffle_partitions()?; From 3fe1771918678ace8b2ae00446909f1811280530 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 19 Jul 2026 10:08:49 -0600 Subject: [PATCH 17/22] 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. --- .../source/contributors-guide/benchmarking.md | 115 ++++++++++++++---- 1 file changed, 88 insertions(+), 27 deletions(-) diff --git a/docs/source/contributors-guide/benchmarking.md b/docs/source/contributors-guide/benchmarking.md index a066f0aa05..9ce84f4e73 100644 --- a/docs/source/contributors-guide/benchmarking.md +++ b/docs/source/contributors-guide/benchmarking.md @@ -276,35 +276,39 @@ it, and only the exact commit makes that visible. Every figure is from a **full 22-query suite run** — one query after another in a single session, on freshly started executors. -| Query | Spark | Comet | Ballista (AQE off) | Ballista (AQE on) | Rows | -| --------: | ---------: | ---------: | -----------------: | ----------------: | -----: | -| 1 | 444.8 | 49.3 | 70.8 | 36.2 | 4 | -| 2 | 74.3 | 37.3 | 155.5 | 63.6 | 100 | -| 3 | 158.4 | 99.1 | 206.2 | 255.0 | 10 | -| 4 | 104.1 | 42.3 | 76.8 | 53.9 | 5 | -| 5 | 364.9 | 234.6 | 542.7 | 700.5 | 5 | -| 6 | 22.0 | 15.3 | 18.3 | 12.7 | 1 | -| 7 | 196.1 | 141.8 | 575.5 | 531.4 | 4 | -| 8 | 412.7 | 291.6 | 657.1 | 801.2 | 2 | -| 9 | 570.2 | 392.1 | 891.8 | 1025.2 | 175 | -| 10 | 147.3 | 112.5 | 240.3 | 166.7 | 20 | -| 11 | 58.3 | 48.7 | 105.1 | 76.2 | 0 [1] | -| 12 | 75.9 | 52.9 | 79.4 | 85.2 | 2 | -| 13 | 114.1 | 71.9 | 96.7 | 95.6 | 30 | -| 14 | 44.6 | 29.0 | 38.0 | 34.3 | 1 | -| 15 | 108.9 | 63.8 | 45.9 | 52.7 | \* [2] | -| 16 | 33.9 | 18.7 | 21.9 | 23.0 | 27840 | -| 17 | 519.5 | 308.2 | 697.8 | 400.4 | 1 | -| 18 | 492.8 | 234.2 | 620.7 | 820.9 | 100 | -| 19 | 53.7 | 39.4 | 39.6 | 42.1 | 1 | -| 20 | 108.1 | 74.6 | 80.5 | 82.9 | 110759 | -| 21 | 536.4 | 351.4 | 1083.0 | 926.6 | 100 | -| 22 | 47.0 | 29.8 | 34.9 | 34.6 | 7 | -| **Total** | **4687.9** | **2738.5** | **6378.2** | **6320.7** | | - -**All 22 queries completed on all four configurations**; there are no failures to +| Query | Spark | Comet | Ballista (AQE on) | Rows | +| --------: | ---------: | ---------: | ----------------: | -----: | +| 1 | 444.8 | 49.3 | 36.2 | 4 | +| 2 | 74.3 | 37.3 | 63.6 | 100 | +| 3 | 158.4 | 99.1 | 255.0 | 10 | +| 4 | 104.1 | 42.3 | 53.9 | 5 | +| 5 | 364.9 | 234.6 | 700.5 | 5 | +| 6 | 22.0 | 15.3 | 12.7 | 1 | +| 7 | 196.1 | 141.8 | 531.4 | 4 | +| 8 | 412.7 | 291.6 | 801.2 | 2 | +| 9 | 570.2 | 392.1 | 1025.2 | 175 | +| 10 | 147.3 | 112.5 | 166.7 | 20 | +| 11 | 58.3 | 48.7 | 76.2 | 0 [1] | +| 12 | 75.9 | 52.9 | 85.2 | 2 | +| 13 | 114.1 | 71.9 | 95.6 | 30 | +| 14 | 44.6 | 29.0 | 34.3 | 1 | +| 15 | 108.9 | 63.8 | 52.7 | \* [2] | +| 16 | 33.9 | 18.7 | 23.0 | 27840 | +| 17 | 519.5 | 308.2 | 400.4 | 1 | +| 18 | 492.8 | 234.2 | 820.9 | 100 | +| 19 | 53.7 | 39.4 | 42.1 | 1 | +| 20 | 108.1 | 74.6 | 82.9 | 110759 | +| 21 | 536.4 | 351.4 | 926.6 | 100 | +| 22 | 47.0 | 29.8 | 34.6 | 7 | +| **Total** | **4687.9** | **2738.5** | **6320.7** | | + +**All 22 queries completed on all three configurations**; there are no failures to report. Row counts agreed across every engine except Q15. +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. @@ -325,6 +329,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. From ad16347b4078c02251bdfc3815f7096194cc3987 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 20:03:24 -0600 Subject: [PATCH 18/22] 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). --- .../source/contributors-guide/benchmarking.md | 46 +++++++++---------- 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/docs/source/contributors-guide/benchmarking.md b/docs/source/contributors-guide/benchmarking.md index 7889b7c5ce..512b7d7ac5 100644 --- a/docs/source/contributors-guide/benchmarking.md +++ b/docs/source/contributors-guide/benchmarking.md @@ -272,7 +272,7 @@ Versions under test: | Engine | Version | | -------- | ----------------------------------------------- | -| Ballista | `main` @ `49d1fec8` | +| Ballista | `#2084` @ `becb3767` | | Spark | 3.5.3 (vanilla, Comet disabled) | | Comet | `main` @ `b0165552` (1.0.0-SNAPSHOT, Spark 3.5) | @@ -287,29 +287,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 From fac1fa220bcea7c846d162dbae4f7f645f8242ba Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 20:26:46 -0600 Subject: [PATCH 19/22] 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. --- ballista/core/src/config.rs | 10 +++++----- .../src/execution_plans/sort_shuffle/config.rs | 15 +++++++-------- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index 0125bed93e..1ec5cadbd6 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -249,12 +249,12 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| "Per-task buffered-bytes budget at which the sort shuffle writer spills its \ in-memory batches to disk. Counted independently of the runtime memory pool, so \ spilling kicks in even when the pool is unbounded. Total worst-case sort shuffle \ - memory per executor is approximately vcores * this value. The default of 0 disables \ - the per-task budget and relies solely on runtime memory-pool pressure to trigger \ - spilling; this is safe only with a bounded memory pool, otherwise the writer never \ - spills and may run out of memory.".to_string(), + memory per executor is approximately vcores * this value. Set to 0 to disable the \ + per-task budget and rely solely on runtime memory-pool pressure to trigger spilling; \ + this is safe only with a bounded memory pool, otherwise the writer never spills and \ + may run out of memory.".to_string(), DataType::UInt64, - Some("0".to_string())), + Some((256 * 1024 * 1024).to_string())), 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. \ diff --git a/ballista/core/src/execution_plans/sort_shuffle/config.rs b/ballista/core/src/execution_plans/sort_shuffle/config.rs index 80b10280b1..9d40958f96 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/config.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/config.rs @@ -27,9 +27,9 @@ pub struct SortShuffleConfig { pub batch_size: usize, /// Per-task buffered-bytes budget at which the writer spills its in-memory /// batches to disk. Counted independently of the runtime `MemoryPool`, so - /// spilling kicks in even when the pool is unbounded. A value of 0 (the - /// default) disables the per-task budget, leaving the runtime `MemoryPool` - /// as the sole spill trigger. + /// spilling kicks in even when the pool is unbounded. A value of 0 disables + /// the per-task budget, leaving the runtime `MemoryPool` as the sole spill + /// trigger. pub memory_limit_per_task_bytes: usize, } @@ -38,14 +38,13 @@ impl Default for SortShuffleConfig { Self { enabled: false, batch_size: 8192, - memory_limit_per_task_bytes: 0, + memory_limit_per_task_bytes: 256 * 1024 * 1024, } } } impl SortShuffleConfig { - /// Creates a new configuration with the default (uncapped) per-task memory - /// limit. + /// Creates a new configuration with the default per-task memory limit. pub fn new(enabled: bool, batch_size: usize) -> Self { Self { enabled, @@ -70,7 +69,7 @@ mod tests { let config = SortShuffleConfig::default(); assert!(!config.enabled); assert_eq!(config.batch_size, 8192); - assert_eq!(config.memory_limit_per_task_bytes, 0); + assert_eq!(config.memory_limit_per_task_bytes, 256 * 1024 * 1024); } #[test] @@ -78,7 +77,7 @@ mod tests { let config = SortShuffleConfig::new(true, 4096); assert!(config.enabled); assert_eq!(config.batch_size, 4096); - assert_eq!(config.memory_limit_per_task_bytes, 0); + assert_eq!(config.memory_limit_per_task_bytes, 256 * 1024 * 1024); } #[test] From 45db50c3d39b596965d3f3f9e1e2704894eefc1f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 20 Jul 2026 20:29:36 -0600 Subject: [PATCH 20/22] 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. --- docs/source/contributors-guide/benchmarking.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/source/contributors-guide/benchmarking.md b/docs/source/contributors-guide/benchmarking.md index 512b7d7ac5..88352e70d2 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 | `#2084` @ `becb3767` | +| Ballista | `#2084` @ `fac1fa22` | | Spark | 3.5.3 (vanilla, Comet disabled) | | Comet | `main` @ `b0165552` (1.0.0-SNAPSHOT, Spark 3.5) | From 2939ecefb478caf0628095c1badf3d4fb43be04f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 06:31:50 -0600 Subject: [PATCH 21/22] Update ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs Co-authored-by: Phillip LeBlanc --- ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 568e2f330c..dc89aa2893 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -319,7 +319,7 @@ impl DynamicJoinSelectionExec { JoinSelectionAction::Sort(_) => "SortMerge(Partitioned)", }; - info!( + 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={:?}) \ From b563844b87d795b230c68e685cf0209567b049b0 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 21 Jul 2026 06:43:43 -0600 Subject: [PATCH 22/22] 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. --- ballista/core/src/execution_plans/sort_shuffle/writer.rs | 4 ++-- .../scheduler/src/state/aqe/execution_plan/dynamic_join.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index ede8a1d4da..c71d51ace9 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -62,7 +62,7 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, Statistics, displayable, }; use futures::{StreamExt, TryStreamExt}; -use log::{debug, info, warn}; +use log::{debug, warn}; /// Result of finalizing shuffle output: (data_path, index_path, partition_write_stats) /// where partition_write_stats is (partition_id, num_batches, num_rows, num_bytes) @@ -568,7 +568,7 @@ impl SortShuffleWriterExec { write_time, ); } else { - info!( + debug!( "Sort shuffle write for partition {} completed. \ Output: {:?}, Index: {:?}, Rows: {}, \ repart_time={:?} spill_time={:?} write_time={:?}, \ 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 dc89aa2893..fbc09c4510 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -31,7 +31,7 @@ use datafusion::{ }, }, }; -use log::info; +use log::debug; use std::sync::Arc; use crate::physical_optimizer::join_selection::collect_left_broadcast_safe;