[WIP][Experimental] feat: use DataFusion's isolated_partitions option for broadcast-safe join selection

Cargo.toml

@@ -149,3 +149,14 @@ lto = false
 debug = false
 debug-assertions = false
 incremental = false
+
+# Experimental: DataFusion 55.2.0 with the `isolated_partitions` backport from
+# apache/datafusion (see the linked DataFusion PR). Remove once Ballista upgrades
+# to a DataFusion release that contains it.
+[patch.crates-io]
+datafusion = { git = "https://github.com/andygrove/datafusion.git", rev = "448ed538be3bafdc5d0e5a5a1c11fdcc1431e7c7" }
+datafusion-cli = { git = "https://github.com/andygrove/datafusion.git", rev = "448ed538be3bafdc5d0e5a5a1c11fdcc1431e7c7" }
+datafusion-functions-aggregate-common = { git = "https://github.com/andygrove/datafusion.git", rev = "448ed538be3bafdc5d0e5a5a1c11fdcc1431e7c7" }
+datafusion-proto = { git = "https://github.com/andygrove/datafusion.git", rev = "448ed538be3bafdc5d0e5a5a1c11fdcc1431e7c7" }
+datafusion-proto-common = { git = "https://github.com/andygrove/datafusion.git", rev = "448ed538be3bafdc5d0e5a5a1c11fdcc1431e7c7" }
+datafusion-substrait = { git = "https://github.com/andygrove/datafusion.git", rev = "448ed538be3bafdc5d0e5a5a1c11fdcc1431e7c7" }

ballista/core/src/extension.rs

@@ -755,6 +755,9 @@ impl SessionConfigHelperExt for SessionConfig {
             //
             // See https://github.com/apache/datafusion-ballista/issues/1375
             .set_bool("datafusion.optimizer.enable_dynamic_filter_pushdown", false)
+            // Ballista runs each partition as an isolated task, so DataFusion
+            // must not broadcast a build side whose rows the join emits.
+            .set_bool("datafusion.execution.isolated_partitions", true)
     }
 }
 

ballista/scheduler/src/physical_optimizer/join_selection.rs

@@ -113,31 +113,6 @@ 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)]
@@ -650,7 +625,6 @@ 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},
@@ -664,40 +638,6 @@ mod test {
         },
     };
 
-    // 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

ballista/scheduler/src/planner.rs

@@ -53,7 +53,7 @@ use datafusion::physical_plan::{
 use log::debug;
 
 use crate::physical_optimizer::join_selection::{
-    collect_left_broadcast_safe, should_swap_join_order,
+    should_swap_join_order,
 };
 use crate::state::task_builder::restrict_plan_to_partitions;
 
@@ -466,7 +466,7 @@ impl DefaultDistributedPlanner {
             // Broadcasting is only correct for probe-driven join types. If the
             // join type is not broadcast-safe, demote it back to a partitioned
             // (shuffle) join. Correctness guard, independent of the threshold.
-            if !collect_left_broadcast_safe(*hash_join.join_type()) {
+            if hash_join.join_type().emits_build_side_rows() {
                 debug!(
                     "broadcast check: demoting DataFusion-promoted CollectLeft join with unsafe join_type={:?} to Partitioned",
                     hash_join.join_type(),
@@ -550,7 +550,7 @@ impl DefaultDistributedPlanner {
         } else {
             *hash_join.join_type()
         };
-        if !collect_left_broadcast_safe(promoted_join_type) {
+        if promoted_join_type.emits_build_side_rows() {
             debug!(
                 "broadcast check: join type {promoted_join_type:?} is not broadcast-safe, skipping promotion"
             );
@@ -1596,7 +1596,6 @@ order by
     #[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 = [
@@ -1631,7 +1630,7 @@ order by
                         && *hj.partition_mode() == PartitionMode::CollectLeft
                     {
                         assert!(
-                            collect_left_broadcast_safe(*hj.join_type()),
+                            !hj.join_type().emits_build_side_rows(),
                             "unsafe CollectLeft join_type={:?} promoted for `{sql}`",
                             hj.join_type()
                         );

ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs

@@ -41,7 +41,6 @@ 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;
 
@@ -353,7 +352,7 @@ impl DynamicJoinSelectionExec {
         }
 
         let partition_mode = if self.null_aware
-            || (under_threshold && collect_left_broadcast_safe(build_side_join_type))
+            || (under_threshold && !build_side_join_type.emits_build_side_rows())
         {
             PartitionMode::CollectLeft
         } else {
@@ -664,7 +663,7 @@ fn requires_build_staging(
     // Measuring can only change the decision if `CollectLeft` is reachable at
     // all, and that is a property of the post-swap join type alone: no size
     // makes a `Full`, or a left-sided semi/anti/outer, join broadcastable.
-    if !collect_left_broadcast_safe(build_side_join_type) {
+    if build_side_join_type.emits_build_side_rows() {
         return false;
     }
 
@@ -1136,10 +1135,10 @@ mod tests {
                 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.emits_build_side_rows(),
                     "join_type {join_type:?} (prefer_hash_join={prefer_hash_join}): collected={}, expected safe={}",
                     is_collected(&action),
-                    collect_left_broadcast_safe(join_type),
+                    !join_type.emits_build_side_rows(),
                 );
             }
         }
@@ -1417,7 +1416,7 @@ mod tests {
     // Staging pays off only where `CollectLeft` is reachable, which is a
     // property of the post-swap join type alone. The left side is the smaller
     // one here, so `supports_swap_join_order` is false and the post-swap type
-    // equals the declared one, making `collect_left_broadcast_safe` the oracle.
+    // equals the declared one, making `JoinType::emits_build_side_rows` the oracle.
     // Looping all ten means a join type added to DataFusion cannot slip through
     // untested — which two hand-partitioned lists would have let it do.
     #[test]
@@ -1437,7 +1436,7 @@ mod tests {
             let action = q8_shaped_action(join_type, true);
             assert_eq!(
                 matches!(action, JoinSelectionAction::StageBuildSide { .. }),
-                collect_left_broadcast_safe(join_type),
+                !join_type.emits_build_side_rows(),
                 "join_type {join_type:?}: staging must follow broadcast safety",
             );
         }

ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs

@@ -17,7 +17,6 @@
 
 pub mod chaos_exec;
 pub mod coalesce_partitions;
-pub mod demote_broadcast_join;
 pub mod distributed_exchange;
 pub mod join_selection;
 pub mod normalize_interleave;
@@ -26,7 +25,6 @@ pub mod prefix_window;
 pub mod propagate_empty;
 
 pub use coalesce_partitions::*;
-pub use demote_broadcast_join::*;
 pub use distributed_exchange::*;
 pub use join_selection::*;
 pub use normalize_interleave::*;

ballista/scheduler/src/state/aqe/planner.rs

@@ -22,7 +22,7 @@ use crate::state::aqe::execution_plan::{
 };
 use crate::state::aqe::optimizer_rule::chaos_exec::ChaosCreatingRule;
 use crate::state::aqe::optimizer_rule::{
-    CoalescePartitionsRule, DelayJoinSelectionRule, DemoteUnsafeBroadcastJoinRule,
+    CoalescePartitionsRule, DelayJoinSelectionRule,
     DistributedExchangeRule, NormalizeInterleaveRule, ParallelWindowRule,
     PrefixWindowRule, PropagateEmptyExecRule, SelectJoinRule,
 };
@@ -606,14 +606,9 @@ impl AdaptivePlanner {
                 // `r_name = EUROPE` in q2 — and every copy is evaluated per row
                 // and again in the row-group pruning predicate.
                 "FilterPushdown" => vec![],
-                // `join_selection` promotes a small build side to `CollectLeft`
-                // without restricting by join type -- safe in one process, but
-                // Ballista runs one task per probe partition. Demote the unsafe
-                // ones straight after, while `EnsureRequirements` can still add
-                // the repartitions a `Partitioned` join needs.
-                "join_selection" => {
-                    vec![r, Arc::new(DemoteUnsafeBroadcastJoinRule::default())]
-                }
+                // `join_selection` only broadcasts join types that are safe for
+                // isolated tasks, because Ballista sets
+                // `datafusion.execution.isolated_partitions`.
                 _ => vec![r],
             })
             .collect()