From 59396c9941d5edbe5cde5aaf4eabaa912b0d9e7c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 26 Jun 2026 10:51:34 -0600 Subject: [PATCH 1/3] feat(scheduler): broadcast small build side of SortMergeJoinExec Add ballista.optimizer.broadcast_sort_merge_join_enabled (default false). When enabled, the static distributed planner converts a SortMergeJoinExec whose smaller side is under broadcast_join_threshold_bytes into a CollectLeft hash join and lowers the build side as a broadcast stage, reusing the existing hash-join broadcast path. Redundant input sorts are dropped during conversion. TPC-H SF10 (AQE off): ~16% faster across the suite (q8 1.6x, q11 1.5x, q2 1.4x), identical row counts. --- ballista/core/src/config.rs | 25 +++ ballista/scheduler/src/planner.rs | 250 ++++++++++++++++++++++++++++-- 2 files changed, 261 insertions(+), 14 deletions(-) diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index fdad8d7f8b..c930257428 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. Disabled 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(false.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_disabled_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 ed396719f4..76fee31f5d 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -40,9 +40,12 @@ use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_optimizer::enforce_sorting::EnforceSorting; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder}; -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, @@ -295,9 +298,31 @@ impl DefaultDistributedPlanner { debug!("broadcast check: threshold is 0, broadcast disabled"); return Ok(plan); } - let Some(hash_join) = plan.as_any().downcast_ref::() else { - 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.as_any().downcast_ref::() { + Some(smj) if smj_broadcast_enabled => { + Some(convert_sort_merge_to_hash_join(smj)?) + } + _ => None, + }; + let hash_join: &HashJoinExec = + match (plan.as_any().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(), @@ -397,10 +422,25 @@ 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.as_any().is::() { + (promoted.clone(), None) + } else if let Some(proj) = promoted.as_any().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 .as_any() .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() @@ -413,13 +453,46 @@ 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), + } } } +/// 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.as_any().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)) +} + /// Workaround until DF 54 migration: /// datafusion-proto 53.1.0 encodes both `Some([])` (zero cols) and `None` (all cols) /// as an empty list, decoding back to `None` and shifting column indices (#1838). @@ -995,7 +1068,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, false)?; let df = ctx .sql("select count(*) from big join small on big.k = small.k") @@ -1039,12 +1112,153 @@ order by Ok(()) } + #[tokio::test] + async fn distributed_broadcast_sort_merge_join_plan() -> Result<(), BallistaError> { + use datafusion::physical_plan::joins::PartitionMode; + use datafusion::physical_plan::joins::SortMergeJoinExec; + use datafusion::physical_plan::sorts::sort::SortExec; + + // prefer_hash_join=false -> DataFusion plans a SortMergeJoinExec. + // broadcast_sort_merge_join_enabled=true -> the small side should be + // converted to a broadcast CollectLeft hash join. + 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)?; + for (i, stage) in stages.iter().enumerate() { + println!("Stage {i}:\n{}", displayable(stage.as_ref()).indent(false)); + } + + let mut found_broadcast_join = false; + for stage in &stages { + let mut walker: Vec> = + vec![stage.clone() as Arc]; + while let Some(node) = walker.pop() { + assert!( + node.as_any().downcast_ref::().is_none(), + "no SortMergeJoinExec should remain after broadcast conversion" + ); + assert!( + node.as_any().downcast_ref::().is_none(), + "redundant SortExec should be stripped during conversion" + ); + if let Some(hj) = node.as_any().downcast_ref::() { + assert_eq!(*hj.partition_mode(), PartitionMode::CollectLeft); + let left = hj.children()[0].clone(); + let unresolved = left + .as_any() + .downcast_ref::() + .expect("left input should be UnresolvedShuffleExec"); + assert!(unresolved.broadcast, "left input should be broadcast"); + found_broadcast_join = true; + } + walker.extend(node.children().iter().map(|c| (*c).clone())); + } + } + assert!( + found_broadcast_join, + "expected a broadcast HashJoinExec converted from SortMergeJoinExec" + ); + + Ok(()) + } + + #[tokio::test] + async fn distributed_sort_merge_join_unchanged_when_flag_disabled() + -> Result<(), BallistaError> { + use datafusion::physical_plan::joins::SortMergeJoinExec; + + // SMJ planned (prefer_hash_join=false), flag OFF -> no conversion. + 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)?; + + let mut found_smj = false; + for stage in &stages { + let mut walker: Vec> = + vec![stage.clone() as Arc]; + while let Some(node) = walker.pop() { + if node.as_any().downcast_ref::().is_some() { + found_smj = true; + } + if let Some(unresolved) = + node.as_any().downcast_ref::() + { + assert!( + !unresolved.broadcast, + "no broadcast expected when SMJ flag is disabled" + ); + } + walker.extend(node.children().iter().map(|c| (*c).clone())); + } + } + assert!(found_smj, "SortMergeJoinExec should be left unchanged"); + + Ok(()) + } + + #[tokio::test] + async fn distributed_sort_merge_join_unchanged_when_sides_too_large() + -> Result<(), BallistaError> { + use datafusion::physical_plan::joins::SortMergeJoinExec; + + // Flag ON but threshold tiny (1 byte) -> neither side qualifies. + 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)?; + + let mut found_smj = false; + for stage in &stages { + let mut walker: Vec> = + vec![stage.clone() as Arc]; + while let Some(node) = walker.pop() { + if node.as_any().downcast_ref::().is_some() { + found_smj = true; + } + walker.extend(node.children().iter().map(|c| (*c).clone())); + } + } + assert!( + found_smj, + "SortMergeJoinExec should be left unchanged when no side is under threshold" + ); + + 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)?; + 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") @@ -1091,7 +1305,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)?; + 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") @@ -1386,6 +1601,8 @@ order by fn make_broadcast_test_ctx( threshold_bytes: usize, small_partitions: usize, + prefer_hash_join: bool, + smj_broadcast_enabled: bool, ) -> Result< ( datafusion::prelude::SessionContext, @@ -1433,11 +1650,16 @@ order by let session_config = SessionConfig::new() .with_target_partitions(2) + .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)?; From d56882ca3fa935eed9de0eb4bd5f4764a068076b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 08:12:35 -0600 Subject: [PATCH 2/3] feat(config): default broadcast_sort_merge_join_enabled to true Enable the SortMergeJoinExec broadcast conversion by default so small build sides are broadcast without opt-in (~16% on TPC-H SF10, AQE off). Disable the flag explicitly in the sort-merge-join client test so it still exercises the plain SortMergeJoinExec execution path. --- ballista/client/tests/context_checks.rs | 8 ++++++++ ballista/core/src/config.rs | 8 ++++---- 2 files changed, 12 insertions(+), 4 deletions(-) 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 c930257428..9a9bc6d7f4 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -88,7 +88,7 @@ pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES: &str = /// 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. Disabled by default. +/// static distributed planner. Enabled by default. pub const BALLISTA_BROADCAST_SORT_MERGE_JOIN_ENABLED: &str = "ballista.optimizer.broadcast_sort_merge_join_enabled"; @@ -210,7 +210,7 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| 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(false.to_string())), + 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, @@ -762,8 +762,8 @@ mod tests { } #[test] - fn broadcast_sort_merge_join_disabled_by_default() { + fn broadcast_sort_merge_join_enabled_by_default() { let config = BallistaConfig::default(); - assert!(!config.broadcast_sort_merge_join_enabled()); + assert!(config.broadcast_sort_merge_join_enabled()); } } From 7f868fca903a34fee8317096cbd157c0d81f5be9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 2 Jul 2026 08:22:03 -0600 Subject: [PATCH 3/3] test: assert SMJ broadcast plans via string snapshots Rewrite the three SortMergeJoinExec broadcast tests to assert the stage plan as a string via the assert_plan! macro instead of walking the plan tree with manual downcasts, making the expected plan explicit and the tests easier to read. --- ballista/scheduler/src/planner.rs | 121 +++++++++++------------------- 1 file changed, 44 insertions(+), 77 deletions(-) diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index b0a44b9016..d5602f57be 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -721,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; @@ -1077,13 +1078,11 @@ order by #[tokio::test] async fn distributed_broadcast_sort_merge_join_plan() -> Result<(), BallistaError> { - use datafusion::physical_plan::joins::PartitionMode; - use datafusion::physical_plan::joins::SortMergeJoinExec; - use datafusion::physical_plan::sorts::sort::SortExec; - // prefer_hash_join=false -> DataFusion plans a SortMergeJoinExec. - // broadcast_sort_merge_join_enabled=true -> the small side should be - // converted to a broadcast CollectLeft hash join. + // 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 @@ -1097,39 +1096,18 @@ order by let job_uuid = Uuid::new_v4(); let stages = planner.plan_query_stages(&job_uuid.to_string().into(), plan, &options)?; - for (i, stage) in stages.iter().enumerate() { - println!("Stage {i}:\n{}", displayable(stage.as_ref()).indent(false)); - } - let mut found_broadcast_join = false; - for stage in &stages { - let mut walker: Vec> = - vec![stage.clone() as Arc]; - while let Some(node) = walker.pop() { - assert!( - node.downcast_ref::().is_none(), - "no SortMergeJoinExec should remain after broadcast conversion" - ); - assert!( - node.downcast_ref::().is_none(), - "redundant SortExec should be stripped during conversion" - ); - if let Some(hj) = node.downcast_ref::() { - assert_eq!(*hj.partition_mode(), PartitionMode::CollectLeft); - let left = hj.children()[0].clone(); - let unresolved = left - .downcast_ref::() - .expect("left input should be UnresolvedShuffleExec"); - assert!(unresolved.broadcast, "left input should be broadcast"); - found_broadcast_join = true; - } - walker.extend(node.children().iter().map(|c| (*c).clone())); - } - } - assert!( - found_broadcast_join, - "expected a broadcast HashJoinExec converted from SortMergeJoinExec" - ); + // 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(()) } @@ -1137,9 +1115,8 @@ order by #[tokio::test] async fn distributed_sort_merge_join_unchanged_when_flag_disabled() -> Result<(), BallistaError> { - use datafusion::physical_plan::joins::SortMergeJoinExec; - - // SMJ planned (prefer_hash_join=false), flag OFF -> no conversion. + // 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 @@ -1153,24 +1130,18 @@ order by let stages = planner.plan_query_stages(&job_uuid.to_string().into(), plan, &options)?; - let mut found_smj = false; - for stage in &stages { - let mut walker: Vec> = - vec![stage.clone() as Arc]; - while let Some(node) = walker.pop() { - if node.downcast_ref::().is_some() { - found_smj = true; - } - if let Some(unresolved) = node.downcast_ref::() { - assert!( - !unresolved.broadcast, - "no broadcast expected when SMJ flag is disabled" - ); - } - walker.extend(node.children().iter().map(|c| (*c).clone())); - } - } - assert!(found_smj, "SortMergeJoinExec should be left unchanged"); + // 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(()) } @@ -1178,9 +1149,8 @@ order by #[tokio::test] async fn distributed_sort_merge_join_unchanged_when_sides_too_large() -> Result<(), BallistaError> { - use datafusion::physical_plan::joins::SortMergeJoinExec; - - // Flag ON but threshold tiny (1 byte) -> neither side qualifies. + // 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 @@ -1194,21 +1164,18 @@ order by let stages = planner.plan_query_stages(&job_uuid.to_string().into(), plan, &options)?; - let mut found_smj = false; - for stage in &stages { - let mut walker: Vec> = - vec![stage.clone() as Arc]; - while let Some(node) = walker.pop() { - if node.downcast_ref::().is_some() { - found_smj = true; - } - walker.extend(node.children().iter().map(|c| (*c).clone())); - } - } - assert!( - found_smj, - "SortMergeJoinExec should be left unchanged when no side is under threshold" - ); + // 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(()) }