diff --git a/ballista/client/tests/context_checks.rs b/ballista/client/tests/context_checks.rs index 2638942886..75d17f37ff 100644 --- a/ballista/client/tests/context_checks.rs +++ b/ballista/client/tests/context_checks.rs @@ -520,17 +520,14 @@ 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())] #[tokio::test] - async fn should_disable_collect_left( + async fn should_set_collect_left_thresholds( #[future(awt)] #[case] ctx: SessionContext, @@ -542,12 +539,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 dd5ec26926..aaa68c6735 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 @@ -805,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) } } @@ -1046,7 +1056,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() @@ -1056,7 +1068,7 @@ mod test { .options_mut() .set( "datafusion.optimizer.hash_join_single_partition_threshold", - "10485760", + "5242880", ) .unwrap(); @@ -1068,7 +1080,7 @@ mod test { .options() .optimizer .hash_join_single_partition_threshold, - 10485760 + 5242880 ); } @@ -1084,7 +1096,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 ); } diff --git a/ballista/scheduler/src/physical_optimizer/join_selection.rs b/ballista/scheduler/src/physical_optimizer/join_selection.rs index 7410ec0270..20d4649283 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)] @@ -617,6 +642,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}, @@ -629,6 +655,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 95fcc6c9b0..6538cb45bb 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -45,7 +45,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>); @@ -270,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::() @@ -281,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(), @@ -361,6 +383,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})" ); @@ -403,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( @@ -888,7 +956,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") @@ -936,7 +1004,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") @@ -974,6 +1042,243 @@ 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.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(()) + } + + // 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 + // 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.downcast_ref::() { + assert_ne!( + *hj.partition_mode(), + PartitionMode::CollectLeft, + "LEFT join with broadcast build side must not be CollectLeft" + ); + } + if let Some(unresolved) = node.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.downcast_ref::() { + assert_ne!( + *hj.partition_mode(), + PartitionMode::CollectLeft, + "no CollectLeft expected when joins are sort-merge" + ); + } + if let Some(unresolved) = node.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> { @@ -981,7 +1286,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") @@ -1273,6 +1578,7 @@ order by fn make_broadcast_test_ctx( threshold_bytes: usize, small_partitions: usize, + prefer_hash_join: bool, ) -> Result< ( datafusion::prelude::SessionContext, @@ -1320,6 +1626,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 2bac00644b..c7d22ae8aa 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -32,7 +32,9 @@ 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; /// has children of this join been /// repartitioned @@ -217,7 +219,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, @@ -225,12 +227,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)?; @@ -424,12 +445,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 7dba88e81f..b652ca7d30 100644 --- a/ballista/scheduler/src/state/aqe/test/join_selection.rs +++ b/ballista/scheduler/src/state/aqe/test/join_selection.rs @@ -141,6 +141,38 @@ 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(()) +} + #[tokio::test] async fn test_hash_join_two_tables_repartition() -> datafusion::common::Result<()> { let ctx = make_ctx_without_collect_left(true);