From fc6ec0dc8643ea282f0bb61d15c80bf75f6352b9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 26 Jun 2026 09:05:56 -0600 Subject: [PATCH 1/7] feat(core): broadcast small build sides via non-zero CollectLeft thresholds Set hash_join_single_partition_threshold to 10MB and hash_join_single_partition_threshold_rows to 1M so the AQE join resolver collects small build sides into a broadcast (CollectLeft) hash join instead of repartitioning them. Update the client settings test to assert the new default values. --- ballista/client/tests/context_checks.rs | 14 +++++++------- ballista/core/src/extension.rs | 11 ++++------- 2 files changed, 11 insertions(+), 14 deletions(-) diff --git a/ballista/client/tests/context_checks.rs b/ballista/client/tests/context_checks.rs index 2638942886..71d5d71b31 100644 --- a/ballista/client/tests/context_checks.rs +++ b/ballista/client/tests/context_checks.rs @@ -530,7 +530,7 @@ mod supported { #[case::standalone(standalone_context())] #[case::remote(remote_context())] #[tokio::test] - async fn should_disable_collect_left( + async fn should_set_collect_left_thresholds( #[future(awt)] #[case] ctx: SessionContext, @@ -542,12 +542,12 @@ mod supported { .await?; let expected = [ - "+----------------------------------------------------------------+-------+", - "| name | value |", - "+----------------------------------------------------------------+-------+", - "| datafusion.optimizer.hash_join_single_partition_threshold | 0 |", - "| datafusion.optimizer.hash_join_single_partition_threshold_rows | 0 |", - "+----------------------------------------------------------------+-------+", + "+----------------------------------------------------------------+----------+", + "| name | value |", + "+----------------------------------------------------------------+----------+", + "| datafusion.optimizer.hash_join_single_partition_threshold | 10485760 |", + "| datafusion.optimizer.hash_join_single_partition_threshold_rows | 1000000 |", + "+----------------------------------------------------------------+----------+", ]; assert_batches_eq!(expected, &result); diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index d53307339c..0b0c751c09 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -767,18 +767,15 @@ impl SessionConfigHelperExt for SessionConfig { // same like previous comment .set_bool("datafusion.sql_parser.map_string_types_to_utf8view", false) // - // As mentioned in https://github.com/apache/datafusion-ballista/issues/1055 - // "Left/full outer join incorrect for CollectLeft / broadcast" - // - // In order to make correct results (decreasing performance) CollectLeft - // has been disabled until fixed + // A build side smaller than these thresholds is collected into a + // CollectLeft (broadcast) hash join rather than being repartitioned. .set_u64( "datafusion.optimizer.hash_join_single_partition_threshold", - 0, + 10 * 1024 * 1024, ) .set_u64( "datafusion.optimizer.hash_join_single_partition_threshold_rows", - 0, + 1_000_000, ) // // DataFusion's hash join has no spill support, so each parallel From 79651d7d8ea6de810805452c670eb72ab23d9743 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 28 Jun 2026 17:17:09 -0600 Subject: [PATCH 2/7] fix(scheduler): restrict CollectLeft broadcast to safe join types The adaptive join resolver chose CollectLeft purely on size thresholds. A CollectLeft join broadcasts the build (left) side to every probe task, where each task sees only its slice of the probe side, so join types that emit rows on behalf of the build side (Left, Full, LeftSemi, LeftAnti, LeftMark) would have every task emit them independently and produce duplicate or spurious rows. Restrict CollectLeft to join types whose output is driven entirely by the probe side (Inner, Right, RightSemi, RightAnti, RightMark), evaluated on the post-swap join type using the same swap decision as the resolver, and keep unsafe join types repartitioned. Add resolver tests covering left/left-anti (repartitioned) and right (still collected) joins. --- ballista/client/tests/context_checks.rs | 9 +- .../state/aqe/execution_plan/dynamic_join.rs | 50 ++++++++++- .../src/state/aqe/test/join_selection.rs | 87 +++++++++++++++++++ 3 files changed, 136 insertions(+), 10 deletions(-) diff --git a/ballista/client/tests/context_checks.rs b/ballista/client/tests/context_checks.rs index 71d5d71b31..75d17f37ff 100644 --- a/ballista/client/tests/context_checks.rs +++ b/ballista/client/tests/context_checks.rs @@ -520,12 +520,9 @@ mod supported { Ok(()) } - // As mentioned in https://github.com/apache/datafusion-ballista/issues/1055 - // "Left/full outer join incorrect for CollectLeft / broadcast" - // - // In order to make correct results (decreasing performance) CollectLeft - // has been disabled until fixed - + // Ballista raises these DataFusion thresholds above zero so the adaptive + // join resolver can collect a small build side into a broadcast + // (CollectLeft) hash join instead of repartitioning it. #[rstest] #[case::standalone(standalone_context())] #[case::remote(remote_context())] 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 3c7853c26c..d4729c8b25 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -33,6 +33,29 @@ use log::debug; use std::sync::Arc; use crate::state::aqe::execution_plan::ExchangeExec; +use crate::state::aqe::optimizer_rule::join_selection::SelectJoinRule; + +/// Whether a `CollectLeft` (broadcast) hash join is correct for `join_type` in +/// Ballista's distributed execution. +/// +/// `CollectLeft` replicates the build (left) side to every probe task, and each +/// task processes only its own slice of the probe (right) side. Output rows that +/// are determined by a single probe-side row are therefore produced exactly +/// once. Join types that emit rows on behalf of the build side — unmatched outer +/// rows, or semi/anti/mark rows that depend on a build row's global match status +/// — would instead be emitted independently by every task, producing duplicate +/// or spurious rows. Only join types whose output is driven entirely by the +/// probe side are safe to broadcast. +pub(crate) fn collect_left_broadcast_safe(join_type: JoinType) -> bool { + matches!( + join_type, + JoinType::Inner + | JoinType::Right + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::RightMark + ) +} /// has children of this join been /// repartitioned @@ -221,7 +244,7 @@ impl DynamicJoinSelectionExec { let threshold_collect_left_join_rows = config.optimizer.hash_join_single_partition_threshold_rows; - let partition_mode = if Self::supports_collect_by_thresholds( + let under_threshold = Self::supports_collect_by_thresholds( self.left.as_ref(), threshold_collect_left_join_bytes, threshold_collect_left_join_rows, @@ -229,12 +252,31 @@ impl DynamicJoinSelectionExec { self.right.as_ref(), threshold_collect_left_join_bytes, threshold_collect_left_join_rows, - ) { - PartitionMode::CollectLeft + ); + + // The resolver collects the smaller side onto the build (left) input, + // swapping the join type when the right side is the smaller one (the + // `swap_inputs` calls in `SelectJoinRule`). A `CollectLeft` join then + // broadcasts that build side to every probe task, which is only correct + // for join types that never emit rows on behalf of the build side. Use + // the same swap decision as the resolver to determine the resulting join + // type and skip the broadcast when it would be unsafe. + let build_side_join_type = if SelectJoinRule::supports_swap_join_order( + self.left.as_ref(), + self.right.as_ref(), + )? { + self.join_type.swap() } else { - PartitionMode::Partitioned + self.join_type }; + let partition_mode = + if under_threshold && collect_left_broadcast_safe(build_side_join_type) { + PartitionMode::CollectLeft + } else { + PartitionMode::Partitioned + }; + let stats_left = self.left.partition_statistics(None)?; let stats_right = self.right.partition_statistics(None)?; diff --git a/ballista/scheduler/src/state/aqe/test/join_selection.rs b/ballista/scheduler/src/state/aqe/test/join_selection.rs index 7dba88e81f..394120775f 100644 --- a/ballista/scheduler/src/state/aqe/test/join_selection.rs +++ b/ballista/scheduler/src/state/aqe/test/join_selection.rs @@ -141,6 +141,93 @@ async fn test_hash_join_two_tables_coalesce() -> datafusion::common::Result<()> Ok(()) } +// A `CollectLeft` join broadcasts (replicates) the build/left side to every +// probe task, each of which sees only its slice of the probe/right side. For a +// LEFT join that emits a null-padded row for every unmatched left row, each +// task would emit those rows independently, producing duplicate/spurious rows +// (apache/datafusion-ballista#1055). The resolver must therefore keep an +// unsafe-to-broadcast join type repartitioned instead of collecting it. +#[tokio::test] +async fn test_left_join_not_collected_left() -> datafusion::common::Result<()> { + let ctx = make_ctx(true); + register_2tables(&ctx); + + let lp = ctx + .sql("SELECT t1.id, t2.val FROM t1 LEFT JOIN t2 ON t1.id = t2.id") + .await + .unwrap() + .into_optimized_plan() + .unwrap(); + + let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?; + let plan = planner.current_plan(); + assert_plan!(plan, @ r" + AdaptiveDatafusionExec: is_final=false, plan_id=3, stage_id=pending, stage_resolved=false + ProjectionExec: expr=[id@0 as id, val@2 as val] + DynamicJoinSelectionExec: plan_id=0, join_type=Left, on=[(id@0, id@0)] repartitioned=true + ExchangeExec: partitioning=Hash([id@0], 4), plan_id=1, stage_id=pending, stage_resolved=false + DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1] + ExchangeExec: partitioning=Hash([id@0], 4), plan_id=2, stage_id=pending, stage_resolved=false + DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1] + "); + Ok(()) +} + +// A LEFT ANTI join emits a left row when it has no match anywhere on the right. +// With the left side broadcast, a task only sees its slice of the right side, so +// it would emit left rows that actually match in another task's slice. Must stay +// repartitioned. +#[tokio::test] +async fn test_left_anti_join_not_collected_left() -> datafusion::common::Result<()> { + let ctx = make_ctx(true); + register_2tables(&ctx); + + let lp = ctx + .sql("SELECT t1.id FROM t1 WHERE t1.id NOT IN (SELECT t2.id FROM t2)") + .await + .unwrap() + .into_optimized_plan() + .unwrap(); + + let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?; + let plan = planner.current_plan(); + let rendered = datafusion::physical_plan::displayable(plan) + .indent(true) + .to_string(); + assert!( + !rendered.contains("mode=CollectLeft"), + "LEFT ANTI join must not be broadcast as CollectLeft, got:\n{rendered}" + ); + Ok(()) +} + +// A RIGHT join emits all right/probe rows; the probe side stays partitioned so +// each row is emitted by exactly one task. Broadcasting the left/build side is +// safe, so the guard must still allow CollectLeft here. +#[tokio::test] +async fn test_right_join_still_collected_left() -> datafusion::common::Result<()> { + let ctx = make_ctx(true); + register_2tables(&ctx); + + let lp = ctx + .sql("SELECT t1.id, t2.val FROM t1 RIGHT JOIN t2 ON t1.id = t2.id") + .await + .unwrap() + .into_optimized_plan() + .unwrap(); + + let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?; + let plan = planner.current_plan(); + let rendered = datafusion::physical_plan::displayable(plan) + .indent(true) + .to_string(); + assert!( + rendered.contains("mode=CollectLeft"), + "RIGHT join is broadcast-safe and should be collected left, got:\n{rendered}" + ); + Ok(()) +} + #[tokio::test] async fn test_hash_join_two_tables_repartition() -> datafusion::common::Result<()> { let ctx = make_ctx_without_collect_left(true); From 3ec52c3dc7eb20e4b4b82b93356df452d9672cc5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 28 Jun 2026 17:29:53 -0600 Subject: [PATCH 3/7] fix(scheduler): guard static planner broadcast and add join-type test matrix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The static DefaultDistributedPlanner promoted small-side hash joins to CollectLeft on size alone, with no join-type check, so a LEFT/FULL/semi/anti/ mark join whose build side fit under broadcast_join_threshold_bytes was broadcast and produced duplicate or spurious rows — the same hazard already guarded on the adaptive path. Move the broadcast-safety predicate into physical_optimizer::join_selection so both planners share one definition, and apply it in maybe_promote_to_broadcast using the post-swap join type. Add a test matrix covering both planners, both build-side orientations, every join type, prefer_hash_join on/off, and the post-swap decision, plus a unit test that locks down the broadcast-safe join-type set. --- .../src/physical_optimizer/join_selection.rs | 61 ++++++ ballista/scheduler/src/planner.rs | 179 +++++++++++++++++- .../state/aqe/execution_plan/dynamic_join.rs | 125 +++++++++--- .../src/state/aqe/test/join_selection.rs | 55 ------ 4 files changed, 339 insertions(+), 81 deletions(-) diff --git a/ballista/scheduler/src/physical_optimizer/join_selection.rs b/ballista/scheduler/src/physical_optimizer/join_selection.rs index b50b0f3c21..0bcb8542d8 100644 --- a/ballista/scheduler/src/physical_optimizer/join_selection.rs +++ b/ballista/scheduler/src/physical_optimizer/join_selection.rs @@ -111,6 +111,31 @@ fn supports_collect_by_thresholds( } } +/// Whether a `CollectLeft` (broadcast) hash join is correct for `join_type` in +/// Ballista's distributed execution. +/// +/// `CollectLeft` replicates the build (left) side to every probe task, and each +/// task processes only its own slice of the probe (right) side. Output rows that +/// are determined by a single probe-side row are therefore produced exactly +/// once. Join types that emit rows on behalf of the build side — unmatched outer +/// rows, or semi/anti/mark rows that depend on a build row's global match status +/// — would instead be emitted independently by every task, producing duplicate +/// or spurious rows. Only join types whose output is driven entirely by the +/// probe side are safe to broadcast. +/// +/// `join_type` must be the join type *after* any build-side swap, since the +/// build side is always the left input of the resulting `HashJoinExec`. +pub(crate) fn collect_left_broadcast_safe(join_type: JoinType) -> bool { + matches!( + join_type, + JoinType::Inner + | JoinType::Right + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::RightMark + ) +} + /// Predicate that checks whether the given join type supports input swapping. #[deprecated(since = "45.0.0", note = "use JoinType::supports_swap instead")] #[allow(dead_code)] @@ -618,6 +643,7 @@ fn apply_subrules( mod test { use std::sync::Arc; + use super::collect_left_broadcast_safe; use datafusion::{ arrow::datatypes::{DataType, Field, Schema}, common::{ColumnStatistics, JoinSide, JoinType, Statistics, stats::Precision}, @@ -630,6 +656,41 @@ mod test { test::exec::StatisticsExec, }, }; + + // A CollectLeft broadcast replicates the build (left) side to every probe + // task. Only join types whose output is driven entirely by the probe (right) + // side may be broadcast; any type that emits rows on behalf of the build side + // would duplicate them across tasks. This locks down the full set so a new or + // reclassified `JoinType` can't silently become broadcastable. + #[test] + fn collect_left_broadcast_safe_matches_probe_driven_join_types() { + let safe = [ + JoinType::Inner, + JoinType::Right, + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::RightMark, + ]; + let unsafe_types = [ + JoinType::Left, + JoinType::Full, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + ]; + for join_type in safe { + assert!( + collect_left_broadcast_safe(join_type), + "{join_type:?} should be safe to broadcast" + ); + } + for join_type in unsafe_types { + assert!( + !collect_left_broadcast_safe(join_type), + "{join_type:?} must not be broadcast" + ); + } + } // // join selection should not change order of joins for // nested loop join as we're not able to insert new CoalescePartitions diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index ed396719f4..29e974142c 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -49,7 +49,9 @@ use datafusion::physical_plan::{ }; use log::debug; -use crate::physical_optimizer::join_selection::should_swap_join_order; +use crate::physical_optimizer::join_selection::{ + collect_left_broadcast_safe, should_swap_join_order, +}; type PartialQueryStageResult = (Arc, Vec>); @@ -375,6 +377,22 @@ impl DefaultDistributedPlanner { right_under }; + // A CollectLeft join broadcasts the build (left) side to every probe + // task, which is only correct for join types that never emit rows on + // behalf of the build side. The build side is the left input of the + // resulting join, so check the join type after the swap is applied. + let promoted_join_type = if swap { + hash_join.join_type().swap() + } else { + *hash_join.join_type() + }; + if !collect_left_broadcast_safe(promoted_join_type) { + debug!( + "broadcast check: join type {promoted_join_type:?} is not broadcast-safe, skipping promotion" + ); + return Ok(plan); + } + debug!( "broadcast check: promoting to CollectLeft (left_under={left_under}, right_under={right_under}, swap={swap})" ); @@ -995,7 +1013,7 @@ order by async fn distributed_broadcast_join_plan() -> Result<(), BallistaError> { use datafusion::physical_plan::joins::PartitionMode; - let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1)?; + let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, true)?; let df = ctx .sql("select count(*) from big join small on big.k = small.k") @@ -1044,7 +1062,7 @@ order by -> Result<(), BallistaError> { use datafusion::physical_plan::joins::PartitionMode; - let (ctx, options) = make_broadcast_test_ctx(0, 1)?; + let (ctx, options) = make_broadcast_test_ctx(0, 1, true)?; let df = ctx .sql("select count(*) from big join small on big.k = small.k") @@ -1084,6 +1102,157 @@ order by Ok(()) } + // Across every expressible join type and both build-side orientations, any + // CollectLeft join the static planner promotes must be broadcast-safe: a + // CollectLeft join replicates the build (left) side to every probe task, so + // its (post-swap) join type must be driven by the probe side. `small` is + // under the threshold and `big` is not. + #[tokio::test] + async fn distributed_broadcast_only_promotes_safe_join_types() + -> Result<(), BallistaError> { + use crate::physical_optimizer::join_selection::collect_left_broadcast_safe; + use datafusion::physical_plan::joins::PartitionMode; + + let sqls = [ + "select * from big join small on big.k = small.k", + "select * from small join big on small.k = big.k", + "select * from big left join small on big.k = small.k", + "select * from small left join big on small.k = big.k", + "select * from big right join small on big.k = small.k", + "select * from small right join big on small.k = big.k", + "select * from big full join small on big.k = small.k", + "select * from small full join big on small.k = big.k", + ]; + + for sql in sqls { + let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, true)?; + let plan = ctx.sql(sql).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, + )?; + + for stage in &stages { + let mut walker: Vec> = + vec![stage.clone() as Arc]; + while let Some(node) = walker.pop() { + if let Some(hj) = node.as_any().downcast_ref::() + && *hj.partition_mode() == PartitionMode::CollectLeft + { + assert!( + collect_left_broadcast_safe(*hj.join_type()), + "unsafe CollectLeft join_type={:?} promoted for `{sql}`", + hj.join_type() + ); + } + walker.extend(node.children().iter().map(|c| (*c).clone())); + } + } + } + + 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 + // even though the small side is under the broadcast threshold. + #[tokio::test] + async fn distributed_left_join_small_build_not_broadcast() -> Result<(), BallistaError> + { + use datafusion::physical_plan::joins::PartitionMode; + + let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, true)?; + + let df = ctx + .sql("select * from small left join big on small.k = big.k") + .await?; + let plan = df.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)?; + + for stage in &stages { + let mut walker: Vec> = + vec![stage.clone() as Arc]; + while let Some(node) = walker.pop() { + if let Some(hj) = node.as_any().downcast_ref::() { + assert_ne!( + *hj.partition_mode(), + PartitionMode::CollectLeft, + "LEFT join with broadcast build side must not be CollectLeft" + ); + } + if let Some(unresolved) = + node.as_any().downcast_ref::() + { + assert!( + !unresolved.broadcast, + "no broadcast reader expected for unsafe LEFT join" + ); + } + walker.extend(node.children().iter().map(|c| (*c).clone())); + } + } + + Ok(()) + } + + // The static planner only promotes `HashJoinExec`. With prefer_hash_join + // disabled the join plans as a `SortMergeJoinExec`, so no broadcast happens + // even though the build side is under the threshold. + #[tokio::test] + async fn distributed_no_broadcast_when_sort_merge_join() -> Result<(), BallistaError> + { + use datafusion::physical_plan::joins::PartitionMode; + + let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, false)?; + + let df = ctx + .sql("select count(*) from big join small on big.k = small.k") + .await?; + let plan = df.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)?; + + for stage in &stages { + let mut walker: Vec> = + vec![stage.clone() as Arc]; + while let Some(node) = walker.pop() { + if let Some(hj) = node.as_any().downcast_ref::() { + assert_ne!( + *hj.partition_mode(), + PartitionMode::CollectLeft, + "no CollectLeft expected when joins are sort-merge" + ); + } + if let Some(unresolved) = + node.as_any().downcast_ref::() + { + assert!( + !unresolved.broadcast, + "no broadcast reader expected for sort-merge join" + ); + } + walker.extend(node.children().iter().map(|c| (*c).clone())); + } + } + + Ok(()) + } + #[tokio::test] async fn distributed_broadcast_join_plan_multi_partition_build() -> Result<(), BallistaError> { @@ -1091,7 +1260,7 @@ order by // broadcast build stage has 3 input partitions and writes 3 shuffle // files. The broadcast UnresolvedShuffleExec must report // upstream_partition_count = 3. - let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024 * 1024, 3)?; + let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024 * 1024, 3, true)?; let df = ctx .sql("select count(*) from big join small on big.k = small.k") @@ -1386,6 +1555,7 @@ order by fn make_broadcast_test_ctx( threshold_bytes: usize, small_partitions: usize, + prefer_hash_join: bool, ) -> Result< ( datafusion::prelude::SessionContext, @@ -1433,6 +1603,7 @@ order by let session_config = SessionConfig::new() .with_target_partitions(2) + .set_bool("datafusion.optimizer.prefer_hash_join", prefer_hash_join) .set_usize( "datafusion.optimizer.hash_join_single_partition_threshold", 0, 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 d4729c8b25..cb4153c9ea 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -32,31 +32,10 @@ use datafusion::{ use log::debug; use std::sync::Arc; +use crate::physical_optimizer::join_selection::collect_left_broadcast_safe; use crate::state::aqe::execution_plan::ExchangeExec; use crate::state::aqe::optimizer_rule::join_selection::SelectJoinRule; -/// Whether a `CollectLeft` (broadcast) hash join is correct for `join_type` in -/// Ballista's distributed execution. -/// -/// `CollectLeft` replicates the build (left) side to every probe task, and each -/// task processes only its own slice of the probe (right) side. Output rows that -/// are determined by a single probe-side row are therefore produced exactly -/// once. Join types that emit rows on behalf of the build side — unmatched outer -/// rows, or semi/anti/mark rows that depend on a build row's global match status -/// — would instead be emitted independently by every task, producing duplicate -/// or spurious rows. Only join types whose output is driven entirely by the -/// probe side are safe to broadcast. -pub(crate) fn collect_left_broadcast_safe(join_type: JoinType) -> bool { - matches!( - join_type, - JoinType::Inner - | JoinType::Right - | JoinType::RightSemi - | JoinType::RightAnti - | JoinType::RightMark - ) -} - /// has children of this join been /// repartitioned #[derive(Debug, Clone, PartialEq, Eq)] @@ -472,12 +451,114 @@ mod tests { use super::*; use crate::state::aqe::execution_plan::ExchangeExec; use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::{ColumnStatistics, Statistics, stats::Precision}; use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::expressions::Column; + use datafusion::physical_plan::test::exec::StatisticsExec; fn test_schema() -> Arc { Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)])) } + fn stats_exec(num_rows: usize) -> Arc { + Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(num_rows), + total_byte_size: Precision::Inexact(num_rows * 16), + column_statistics: vec![ColumnStatistics::new_unknown()], + }, + Schema::new(vec![Field::new("k", DataType::Int32, false)]), + )) + } + + /// 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 { + 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 { + properties: Arc::clone(left.properties()), + left, + right, + on, + filter: None, + join_type, + projection: None, + null_equality: NullEquality::NullEqualsNothing, + selection_state: JoinInputState::Unknown, + null_aware: false, + plan_id: 0, + }; + let mut config = ConfigOptions::new(); + 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; + dj.to_actual_join(&config).unwrap() + } + + fn is_collected(action: &JoinSelectionAction) -> bool { + matches!( + action, + JoinSelectionAction::CollectLeft(_) | JoinSelectionAction::LateCollectLeft(_) + ) + } + + // With equal-sized sides (no swap) the resolver must collect a small build + // side into a CollectLeft broadcast only for join types whose output is + // driven by the probe side; everything that emits build-side rows must be + // repartitioned instead. `prefer_hash_join` only selects the repartitioned + // fallback (hash vs sort-merge), so the broadcast-safety decision must be + // identical under both settings. + #[test] + fn to_actual_join_collects_only_broadcast_safe_join_types() { + for prefer_hash_join in [true, false] { + for join_type in [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::RightSemi, + JoinType::LeftAnti, + JoinType::RightAnti, + JoinType::LeftMark, + JoinType::RightMark, + ] { + let action = actual_join_for(join_type, 100, 100, prefer_hash_join); + assert_eq!( + is_collected(&action), + collect_left_broadcast_safe(join_type), + "join_type {join_type:?} (prefer_hash_join={prefer_hash_join}): collected={}, expected safe={}", + is_collected(&action), + collect_left_broadcast_safe(join_type), + ); + } + } + } + + // Safety is evaluated on the join type *after* the build-side swap. A LEFT + // join with the smaller side on the right swaps to a (safe) RIGHT join and is + // collected; a RIGHT join with the smaller side on the right swaps to an + // (unsafe) LEFT join and must be repartitioned. + #[test] + fn to_actual_join_uses_post_swap_join_type() { + assert!( + is_collected(&actual_join_for(JoinType::Left, 1000, 10, true)), + "LEFT with small right swaps to a safe RIGHT join and should collect" + ); + assert!( + !is_collected(&actual_join_for(JoinType::Right, 1000, 10, true)), + "RIGHT with small right swaps to an unsafe LEFT join and must repartition" + ); + } + /// Constructs a minimal `DynamicJoinSelectionExec` around the given children. /// Properties are borrowed from the left child — correct enough for unit tests /// that only exercise `children_resolved`. diff --git a/ballista/scheduler/src/state/aqe/test/join_selection.rs b/ballista/scheduler/src/state/aqe/test/join_selection.rs index 394120775f..b652ca7d30 100644 --- a/ballista/scheduler/src/state/aqe/test/join_selection.rs +++ b/ballista/scheduler/src/state/aqe/test/join_selection.rs @@ -173,61 +173,6 @@ async fn test_left_join_not_collected_left() -> datafusion::common::Result<()> { Ok(()) } -// A LEFT ANTI join emits a left row when it has no match anywhere on the right. -// With the left side broadcast, a task only sees its slice of the right side, so -// it would emit left rows that actually match in another task's slice. Must stay -// repartitioned. -#[tokio::test] -async fn test_left_anti_join_not_collected_left() -> datafusion::common::Result<()> { - let ctx = make_ctx(true); - register_2tables(&ctx); - - let lp = ctx - .sql("SELECT t1.id FROM t1 WHERE t1.id NOT IN (SELECT t2.id FROM t2)") - .await - .unwrap() - .into_optimized_plan() - .unwrap(); - - let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?; - let plan = planner.current_plan(); - let rendered = datafusion::physical_plan::displayable(plan) - .indent(true) - .to_string(); - assert!( - !rendered.contains("mode=CollectLeft"), - "LEFT ANTI join must not be broadcast as CollectLeft, got:\n{rendered}" - ); - Ok(()) -} - -// A RIGHT join emits all right/probe rows; the probe side stays partitioned so -// each row is emitted by exactly one task. Broadcasting the left/build side is -// safe, so the guard must still allow CollectLeft here. -#[tokio::test] -async fn test_right_join_still_collected_left() -> datafusion::common::Result<()> { - let ctx = make_ctx(true); - register_2tables(&ctx); - - let lp = ctx - .sql("SELECT t1.id, t2.val FROM t1 RIGHT JOIN t2 ON t1.id = t2.id") - .await - .unwrap() - .into_optimized_plan() - .unwrap(); - - let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?; - let plan = planner.current_plan(); - let rendered = datafusion::physical_plan::displayable(plan) - .indent(true) - .to_string(); - assert!( - rendered.contains("mode=CollectLeft"), - "RIGHT join is broadcast-safe and should be collected left, got:\n{rendered}" - ); - Ok(()) -} - #[tokio::test] async fn test_hash_join_two_tables_repartition() -> datafusion::common::Result<()> { let ctx = make_ctx_without_collect_left(true); From 58f38675d8300fe672c3f82b9347b4c5d8ce9d4c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 28 Jun 2026 17:42:23 -0600 Subject: [PATCH 4/7] test(core): update upgrade defaults test for non-zero CollectLeft threshold #1902 added tests asserting the Ballista default hash_join_single_partition_threshold is 0; this PR raises that default to 10 MB (and the row threshold to 1M), so should_apply_defaults_when_upgrading_plain_config now asserts the new values. Give should_preserve_user_overrides_on_upgrade an override distinct from the default so it still proves the user value survives. --- ballista/core/src/extension.rs | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index 0b0c751c09..18854e5321 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -1027,7 +1027,9 @@ mod test { // re-applying the defaults would discard them. See #1901. #[test] fn should_preserve_user_overrides_on_upgrade() { - // Ballista defaults these to prefer_hash_join=false and threshold=0. + // Ballista defaults these to prefer_hash_join=false and the threshold to + // 10 MB. The overrides below differ from those defaults so the assertions + // prove the user's values survived `upgrade_for_ballista`. let mut config = SessionConfig::new_with_ballista(); config .options_mut() @@ -1037,7 +1039,7 @@ mod test { .options_mut() .set( "datafusion.optimizer.hash_join_single_partition_threshold", - "10485760", + "5242880", ) .unwrap(); @@ -1049,7 +1051,7 @@ mod test { .options() .optimizer .hash_join_single_partition_threshold, - 10485760 + 5242880 ); } @@ -1065,7 +1067,14 @@ mod test { .options() .optimizer .hash_join_single_partition_threshold, - 0 + 10 * 1024 * 1024 + ); + assert_eq!( + config + .options() + .optimizer + .hash_join_single_partition_threshold_rows, + 1_000_000 ); } From fbd116ce14292d34d1e7921ec252a228e3dc4a8e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 30 Jun 2026 08:04:26 -0600 Subject: [PATCH 5/7] fix(scheduler): use dyn ExecutionPlan::downcast_ref in broadcast tests DataFusion 54 removed ExecutionPlan::as_any in favor of the inherent downcast_ref on dyn ExecutionPlan, which auto-derefs through Arc. The new broadcast-safety tests still used as_any().downcast_ref(), which no longer compiles. Drop the as_any() hop to match the rest of the file. --- ballista/scheduler/src/planner.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index d66ad35d7a..f6d29cb79f 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -1031,7 +1031,7 @@ order by let mut walker: Vec> = vec![stage.clone() as Arc]; while let Some(node) = walker.pop() { - if let Some(hj) = node.as_any().downcast_ref::() + if let Some(hj) = node.downcast_ref::() && *hj.partition_mode() == PartitionMode::CollectLeft { assert!( @@ -1074,7 +1074,7 @@ order by let mut walker: Vec> = vec![stage.clone() as Arc]; while let Some(node) = walker.pop() { - if let Some(hj) = node.as_any().downcast_ref::() { + if let Some(hj) = node.downcast_ref::() { assert_ne!( *hj.partition_mode(), PartitionMode::CollectLeft, @@ -1082,7 +1082,7 @@ order by ); } if let Some(unresolved) = - node.as_any().downcast_ref::() + node.downcast_ref::() { assert!( !unresolved.broadcast, @@ -1121,7 +1121,7 @@ order by let mut walker: Vec> = vec![stage.clone() as Arc]; while let Some(node) = walker.pop() { - if let Some(hj) = node.as_any().downcast_ref::() { + if let Some(hj) = node.downcast_ref::() { assert_ne!( *hj.partition_mode(), PartitionMode::CollectLeft, @@ -1129,7 +1129,7 @@ order by ); } if let Some(unresolved) = - node.as_any().downcast_ref::() + node.downcast_ref::() { assert!( !unresolved.broadcast, From b40a6701e9b5c234ea33deba765d1766dbc4a23f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 30 Jun 2026 11:01:54 -0600 Subject: [PATCH 6/7] fix: demote DataFusion-promoted unsafe CollectLeft joins to partitioned DataFusion's JoinSelection runs during create_physical_plan and, with the non-zero hash_join_single_partition_threshold default this PR sets, can stamp CollectLeft on a join without restricting by join type. The distributed planner's broadcast-safety guard only covered the Partitioned -> CollectLeft promotion path, so a pre-promoted unsafe join (e.g. a LEFT join with the small side already on the left) reached the broadcast-lowering branch and replicated the outer side across every probe task, producing duplicate/spurious rows. In maybe_promote_to_broadcast, demote any CollectLeft hash join whose join type is not broadcast-safe back to a Partitioned (shuffle) join, hash-partitioning both inputs on the join keys. This is a correctness guard, so it runs regardless of the Ballista broadcast threshold. Add a test that mirrors the production session config so DataFusion's own JoinSelection performs the promotion. --- ballista/scheduler/src/planner.rs | 154 ++++++++++++++++++++++++++++-- 1 file changed, 145 insertions(+), 9 deletions(-) diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index f6d29cb79f..6538cb45bb 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -272,6 +272,29 @@ impl DefaultDistributedPlanner { plan: Arc, config: &ConfigOptions, ) -> Result> { + let Some(hash_join) = plan.downcast_ref::() else { + return Ok(plan); + }; + + // 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. + if *hash_join.partition_mode() == PartitionMode::CollectLeft { + if collect_left_broadcast_safe(*hash_join.join_type()) { + 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); + } + let threshold_bytes = config .extensions .get::() @@ -283,9 +306,6 @@ impl DefaultDistributedPlanner { debug!("broadcast check: threshold is 0, broadcast disabled"); return Ok(plan); } - let Some(hash_join) = plan.downcast_ref::() else { - return Ok(plan); - }; debug!( "broadcast check: evaluating HashJoinExec mode={:?} join_type={:?} threshold={threshold_bytes}", hash_join.partition_mode(), @@ -421,6 +441,36 @@ impl DefaultDistributedPlanner { vec![new_left, new_right], )?) } + + /// Rewrites a `HashJoinExec(CollectLeft)` as an equivalent + /// `HashJoinExec(Partitioned)`, hash-partitioning both inputs on the join + /// keys so the join is executed as a shuffle join rather than by + /// broadcasting the build side. Used to undo an unsafe `CollectLeft` + /// promotion that DataFusion's `JoinSelection` applied before the + /// distributed planner ran. + fn demote_collect_left_to_partitioned( + hash_join: &HashJoinExec, + config: &ConfigOptions, + ) -> Result> { + let partitions = config.execution.target_partitions; + let left_keys: Vec<_> = + hash_join.on().iter().map(|(l, _)| Arc::clone(l)).collect(); + let right_keys: Vec<_> = + hash_join.on().iter().map(|(_, r)| Arc::clone(r)).collect(); + let new_left: Arc = Arc::new(RepartitionExec::try_new( + hash_join.left().clone(), + Partitioning::Hash(left_keys, partitions), + )?); + let new_right: Arc = Arc::new(RepartitionExec::try_new( + hash_join.right().clone(), + Partitioning::Hash(right_keys, partitions), + )?); + Ok(hash_join + .builder() + .with_partition_mode(PartitionMode::Partitioned) + .with_new_children(vec![new_left, new_right])? + .build_exec()?) + } } fn create_unresolved_shuffle( @@ -1048,6 +1098,96 @@ order by Ok(()) } + // When DataFusion's threshold is non-zero, DataFusion's own `JoinSelection` + // runs during `create_physical_plan` and can stamp `CollectLeft` on a LEFT + // join whose small side is already on the left (no swap, no join-type + // check). The distributed planner must demote such an unsafe `CollectLeft` + // join back to a partitioned join rather than broadcasting the outer side. + #[tokio::test] + async fn df_promoted_left_join_must_not_broadcast() -> 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::common::JoinType; + use datafusion::physical_plan::joins::PartitionMode; + 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()))?; + + // Match the PR's production session config: DF and Ballista thresholds + // both at 10 MB so DF's `JoinSelection` can promote on its own. + 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(10 * 1024 * 1024); + 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 small left join big on small.k = big.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)?; + + 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::() { + assert!( + !(matches!(hj.join_type(), JoinType::Left) + && *hj.partition_mode() == PartitionMode::CollectLeft), + "LEFT join must not be CollectLeft (would broadcast the \ + outer side across probe tasks)" + ); + } + if let Some(unresolved) = node.downcast_ref::() { + assert!( + !unresolved.broadcast, + "no broadcast reader expected for an unsafe LEFT join" + ); + } + walker.extend(node.children().iter().map(|c| (*c).clone())); + } + } + 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 @@ -1081,9 +1221,7 @@ order by "LEFT join with broadcast build side must not be CollectLeft" ); } - if let Some(unresolved) = - node.downcast_ref::() - { + if let Some(unresolved) = node.downcast_ref::() { assert!( !unresolved.broadcast, "no broadcast reader expected for unsafe LEFT join" @@ -1128,9 +1266,7 @@ order by "no CollectLeft expected when joins are sort-merge" ); } - if let Some(unresolved) = - node.downcast_ref::() - { + if let Some(unresolved) = node.downcast_ref::() { assert!( !unresolved.broadcast, "no broadcast reader expected for sort-merge join" From 382b8e1ec4461df8e374716726e5f4e2a221bae6 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 30 Jun 2026 13:34:20 -0600 Subject: [PATCH 7/7] fix: disable DataFusion dynamic filter pushdown in session defaults MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DataFusion 54 populates dynamic filters at runtime from an upstream operator (a hash join build side, a TopK heap, a partial aggregate) and reads them in a downstream scan within the same plan. Ballista splits a plan into stages at shuffle and broadcast boundaries that execute as independent tasks, so when the producing operator and the consuming scan land in different stages the filter is never populated across the boundary and the consuming scan blocks forever. This deadlocks every multi-join TPC-H query under the adaptive (AQE) planner, where broadcast (CollectLeft) hash joins create the cross-stage producer/ consumer split — the build stages complete but the probe stages hang. The static planner avoids it only because it plans sort-merge joins, which carry no join dynamic filter. Pin datafusion.optimizer.enable_dynamic_filter_pushdown to false in the Ballista session defaults, alongside the other DataFusion options Ballista already disables because they rely on in-process state that cannot cross stage boundaries. Tracked by #1375. --- ballista/core/src/extension.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index ef96be2007..aaa68c6735 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -802,6 +802,19 @@ impl SessionConfigHelperExt for SessionConfig { "datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery", false, ) + // + // DataFusion's dynamic filters are populated at runtime by an + // upstream operator (a hash join build side, a TopK heap, a partial + // aggregate) and read by a downstream scan within the same plan. + // Ballista splits a plan into stages at shuffle and broadcast + // boundaries that run as independent tasks, so when the producing + // operator and the consuming scan land in different stages the + // filter is never populated across the boundary and the scan blocks + // forever. Disable dynamic filter pushdown until Ballista can carry + // dynamic filters across stage boundaries. + // + // See https://github.com/apache/datafusion-ballista/issues/1375 + .set_bool("datafusion.optimizer.enable_dynamic_filter_pushdown", false) } }