diff --git a/ballista/client/tests/context_checks.rs b/ballista/client/tests/context_checks.rs index 75d17f37ff..dc0576d9d6 100644 --- a/ballista/client/tests/context_checks.rs +++ b/ballista/client/tests/context_checks.rs @@ -960,6 +960,14 @@ mod supported { ctx: SessionContext, test_data: String, ) -> datafusion::error::Result<()> { + // Exercise the plain sort-merge-join execution path: disable the + // default SortMergeJoinExec broadcast conversion so the join stays a + // SortMergeJoinExec in the distributed plan. + ctx.sql("SET ballista.optimizer.broadcast_sort_merge_join_enabled = false") + .await? + .collect() + .await?; + ctx.register_parquet( "t0", &format!("{test_data}/alltypes_plain.parquet"), diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index fdad8d7f8b..9a9bc6d7f4 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -86,6 +86,12 @@ 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 to enable broadcasting a small build side of a +/// `SortMergeJoinExec` by converting it to a `CollectLeft` hash join in the +/// static distributed planner. Enabled by default. +pub const BALLISTA_BROADCAST_SORT_MERGE_JOIN_ENABLED: &str = + "ballista.optimizer.broadcast_sort_merge_join_enabled"; + /// 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. @@ -199,6 +205,12 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| Set to 0 to disable promotion.".to_string(), DataType::UInt64, Some((10 * 1024 * 1024).to_string())), + ConfigEntry::new(BALLISTA_BROADCAST_SORT_MERGE_JOIN_ENABLED.to_string(), + "Broadcast a small build side of a SortMergeJoinExec by converting it \ + to a CollectLeft hash join in the static distributed planner. \ + The build side must also fit under broadcast_join_threshold_bytes.".to_string(), + DataType::Boolean, + Some(true.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, @@ -496,6 +508,13 @@ impl BallistaConfig { self.get_usize_setting(BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES) } + /// Returns whether broadcasting a small build side of a `SortMergeJoinExec` + /// (by converting it to a `CollectLeft` hash join) is enabled in the static + /// distributed planner. + pub fn broadcast_sort_merge_join_enabled(&self) -> bool { + self.get_bool_setting(BALLISTA_BROADCAST_SORT_MERGE_JOIN_ENABLED) + } + /// Returns whether the AQE coalesce-shuffle-partitions rule is enabled. pub fn coalesce_enabled(&self) -> bool { self.get_bool_setting(BALLISTA_COALESCE_ENABLED) @@ -741,4 +760,10 @@ mod tests { assert_eq!(16777216, config.grpc_client_max_message_size()); Ok(()) } + + #[test] + fn broadcast_sort_merge_join_enabled_by_default() { + let config = BallistaConfig::default(); + assert!(config.broadcast_sort_merge_join_enabled()); + } } diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index 6538cb45bb..d5602f57be 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -37,8 +37,12 @@ use datafusion::config::ConfigOptions; use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_optimizer::enforce_sorting::EnforceSorting; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; +use datafusion::physical_plan::joins::{ + HashJoinExec, HashJoinExecBuilder, PartitionMode, SortMergeJoinExec, +}; +use datafusion::physical_plan::projection::ProjectionExec; use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::{ ExecutionPlan, Partitioning, with_new_children_if_necessary, @@ -272,10 +276,6 @@ 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 @@ -283,8 +283,11 @@ impl DefaultDistributedPlanner { // 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 { + // regardless of the Ballista broadcast threshold below. A + // `SortMergeJoinExec` falls through to the SMJ-broadcast path below. + if let Some(hash_join) = plan.downcast_ref::() + && *hash_join.partition_mode() == PartitionMode::CollectLeft + { if collect_left_broadcast_safe(*hash_join.join_type()) { return Ok(plan); } @@ -306,6 +309,31 @@ impl DefaultDistributedPlanner { debug!("broadcast check: threshold is 0, broadcast disabled"); return Ok(plan); } + + // The candidate is a `HashJoinExec(Partitioned)`: either the plan itself, + // or a `SortMergeJoinExec` converted to one when SMJ broadcast is enabled. + // `converted` owns the converted join so `hash_join` can borrow from it, + // and the original `plan` is returned unchanged when no promotion happens. + let smj_broadcast_enabled = config + .extensions + .get::() + .map(|c| c.broadcast_sort_merge_join_enabled()) + .unwrap_or_else(|| { + BallistaConfig::default().broadcast_sort_merge_join_enabled() + }); + let converted: Option> = + match plan.downcast_ref::() { + Some(smj) if smj_broadcast_enabled => { + Some(convert_sort_merge_to_hash_join(smj)?) + } + _ => None, + }; + let hash_join: &HashJoinExec = + match (plan.downcast_ref::(), &converted) { + (Some(hj), _) => hj, + (None, Some(hj)) => hj.as_ref(), + (None, None) => return Ok(plan), + }; debug!( "broadcast check: evaluating HashJoinExec mode={:?} join_type={:?} threshold={threshold_bytes}", hash_join.partition_mode(), @@ -421,9 +449,24 @@ impl DefaultDistributedPlanner { ) }; - let promoted_join = promoted + // `swap_inputs` may wrap the join in a `ProjectionExec` to restore the + // original column order. Locate the `HashJoinExec` (bare or under that + // projection) so the build side can be coalesced, then re-wrap if needed. + let (join_node, projection): ( + Arc, + Option>, + ) = if promoted.downcast_ref::().is_some() { + (promoted.clone(), None) + } else if let Some(proj) = promoted.downcast_ref::() { + (proj.input().clone(), Some(promoted.clone())) + } else { + debug!("broadcast check: unexpected promoted plan shape, skipping"); + return Ok(plan); + }; + + let promoted_join = join_node .downcast_ref::() - .expect("promoted plan must still be a HashJoinExec"); + .expect("promoted join node must be a HashJoinExec"); let new_left: Arc = if promoted_join .left() .properties() @@ -436,10 +479,16 @@ impl DefaultDistributedPlanner { promoted_join.left().clone() }; let new_right = promoted_join.right().clone(); - Ok(with_new_children_if_necessary( - promoted, - vec![new_left, new_right], - )?) + let rebuilt_join = + with_new_children_if_necessary(join_node, vec![new_left, new_right])?; + + // Re-wrap in the projection if `swap_inputs` added one. The recursive + // lowering descends into the projection and broadcasts the build side of + // the inner `HashJoinExec(CollectLeft)`. + match projection { + Some(proj) => Ok(with_new_children_if_necessary(proj, vec![rebuilt_join])?), + None => Ok(rebuilt_join), + } } /// Rewrites a `HashJoinExec(CollectLeft)` as an equivalent @@ -473,6 +522,33 @@ impl DefaultDistributedPlanner { } } +/// Strips a top-level `SortExec`, returning its input. A `SortMergeJoinExec` +/// requires sorted inputs, but the broadcast `CollectLeft` hash join converted +/// from it does not, so the sort is dropped during conversion. +fn strip_sort(plan: Arc) -> Arc { + if let Some(sort) = plan.downcast_ref::() { + sort.input().clone() + } else { + plan + } +} + +/// Converts a `SortMergeJoinExec` into an equivalent `HashJoinExec(Partitioned)`, +/// dropping the now-redundant input sorts. The result is then evaluated by the +/// normal hash-join broadcast path, which promotes it to `CollectLeft` when a +/// side fits under the threshold. +fn convert_sort_merge_to_hash_join(smj: &SortMergeJoinExec) -> Result> { + let left = strip_sort(smj.left().clone()); + let right = strip_sort(smj.right().clone()); + let hash_join = + HashJoinExecBuilder::new(left, right, smj.on().to_vec(), smj.join_type()) + .with_filter(smj.filter().clone()) + .with_partition_mode(PartitionMode::Partitioned) + .with_null_equality(smj.null_equality) + .build()?; + Ok(Arc::new(hash_join)) +} + fn create_unresolved_shuffle( shuffle_writer: &dyn ShuffleWriter, ) -> Arc { @@ -645,6 +721,7 @@ pub(crate) fn create_shuffle_writer_with_config( #[cfg(test)] mod test { + use crate::assert_plan; use crate::planner::{DefaultDistributedPlanner, DistributedPlanner}; use crate::test_utils::datafusion_test_context; use ballista_core::error::BallistaError; @@ -956,7 +1033,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, true)?; + let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, true, false)?; let df = ctx .sql("select count(*) from big join small on big.k = small.k") @@ -999,12 +1076,116 @@ order by Ok(()) } + #[tokio::test] + async fn distributed_broadcast_sort_merge_join_plan() -> Result<(), BallistaError> { + // prefer_hash_join=false -> DataFusion plans a SortMergeJoinExec. + // broadcast_sort_merge_join_enabled=true -> the small build side is + // converted to a broadcast CollectLeft hash join: the SortMergeJoinExec + // and its input SortExecs are gone, and the build input is an + // UnresolvedShuffleExec with broadcast=true. + let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, false, true)?; + + 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)?; + + // Stage 1 holds the join: a broadcast CollectLeft hash join, no + // SortMergeJoinExec and no SortExec. + assert_plan!(stages[1].as_ref(), @r" + ShuffleWriterExec: partitioning: None + AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + ProjectionExec: expr=[] + ProjectionExec: expr=[k@1 as k, k@0 as k] + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(k@0, k@0)] + UnresolvedShuffleExec: broadcast=true, upstream_partitions: 1 + DataSourceExec: partitions=1, partition_sizes=[1] + "); + + Ok(()) + } + + #[tokio::test] + async fn distributed_sort_merge_join_unchanged_when_flag_disabled() + -> Result<(), BallistaError> { + // SMJ planned (prefer_hash_join=false), flag OFF -> no conversion: the + // join stays a SortMergeJoinExec over sorted inputs. + let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, false, 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)?; + + // Stage 0 holds the join, still a SortMergeJoinExec over sorted inputs + // (no HashJoinExec, no broadcast UnresolvedShuffleExec). + assert_plan!(stages[0].as_ref(), @r" + ShuffleWriterExec: partitioning: None + AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + ProjectionExec: expr=[] + SortMergeJoinExec: join_type=Inner, on=[(k@0, k@0)] + SortExec: expr=[k@0 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + SortExec: expr=[k@0 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + "); + + Ok(()) + } + + #[tokio::test] + async fn distributed_sort_merge_join_unchanged_when_sides_too_large() + -> Result<(), BallistaError> { + // Flag ON but threshold tiny (1 byte) -> neither side qualifies, so the + // join stays a SortMergeJoinExec over sorted inputs. + let (ctx, options) = make_broadcast_test_ctx(1, 1, false, true)?; + + 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)?; + + // Stage 0 holds the join, still a SortMergeJoinExec over sorted inputs + // (no HashJoinExec, no broadcast UnresolvedShuffleExec). + assert_plan!(stages[0].as_ref(), @r" + ShuffleWriterExec: partitioning: None + AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + ProjectionExec: expr=[] + SortMergeJoinExec: join_type=Inner, on=[(k@0, k@0)] + SortExec: expr=[k@0 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + SortExec: expr=[k@0 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + "); + + Ok(()) + } + #[tokio::test] async fn distributed_join_plan_no_broadcast_when_threshold_zero() -> Result<(), BallistaError> { use datafusion::physical_plan::joins::PartitionMode; - let (ctx, options) = make_broadcast_test_ctx(0, 1, true)?; + let (ctx, options) = make_broadcast_test_ctx(0, 1, true, false)?; let df = ctx .sql("select count(*) from big join small on big.k = small.k") @@ -1065,7 +1246,8 @@ order by ]; for sql in sqls { - let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, true)?; + let (ctx, options) = + make_broadcast_test_ctx(10 * 1024 * 1024, 1, true, false)?; let plan = ctx.sql(sql).await?.into_optimized_plan()?; let plan = ctx.state().create_physical_plan(&plan).await?; @@ -1197,7 +1379,7 @@ order by { use datafusion::physical_plan::joins::PartitionMode; - let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, true)?; + let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, true, false)?; let df = ctx .sql("select * from small left join big on small.k = big.k") @@ -1242,7 +1424,7 @@ order by { use datafusion::physical_plan::joins::PartitionMode; - let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, false)?; + let (ctx, options) = make_broadcast_test_ctx(10 * 1024 * 1024, 1, false, false)?; let df = ctx .sql("select count(*) from big join small on big.k = small.k") @@ -1286,7 +1468,8 @@ 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, true)?; + let (ctx, options) = + make_broadcast_test_ctx(10 * 1024 * 1024 * 1024, 3, true, false)?; let df = ctx .sql("select count(*) from big join small on big.k = small.k") @@ -1579,6 +1762,7 @@ order by threshold_bytes: usize, small_partitions: usize, prefer_hash_join: bool, + smj_broadcast_enabled: bool, ) -> Result< ( datafusion::prelude::SessionContext, @@ -1626,12 +1810,16 @@ order by let session_config = SessionConfig::new() .with_target_partitions(2) - .set_bool("datafusion.optimizer.prefer_hash_join", prefer_hash_join) + .with_ballista_broadcast_join_threshold_bytes(threshold_bytes) .set_usize( "datafusion.optimizer.hash_join_single_partition_threshold", 0, ) - .with_ballista_broadcast_join_threshold_bytes(threshold_bytes); + .set_bool("datafusion.optimizer.prefer_hash_join", prefer_hash_join) + .set_bool( + "ballista.optimizer.broadcast_sort_merge_join_enabled", + smj_broadcast_enabled, + ); let ctx = datafusion::prelude::SessionContext::new_with_config(session_config); ctx.register_batch("big", big_batch)?;