Gene.bordegaray/2025/12/add broadcast exec

benchmarks/cdk/bin/worker.rs

@@ -38,6 +38,10 @@ struct Cmd {
     /// The bucket name.
     #[structopt(long, default_value = "datafusion-distributed-benchmarks")]
     bucket: String,
+
+    // Turns broadcast joins on.
+    #[structopt(long)]
+    broadcast_joins: bool,
 }
 
 #[tokio::main]
@@ -67,12 +71,15 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let runtime_env = Arc::new(RuntimeEnv::default());
     runtime_env.register_object_store(&s3_url, s3);
 
-    let state = SessionStateBuilder::new()
+    let mut state = SessionStateBuilder::new()
         .with_default_features()
         .with_runtime_env(Arc::clone(&runtime_env))
         .with_distributed_worker_resolver(Ec2WorkerResolver::new())
         .with_physical_optimizer_rule(Arc::new(DistributedPhysicalOptimizerRule))
         .build();
+    if cmd.broadcast_joins {
+        state = state.with_distributed_broadcast_joins_enabled(true)?;
+    }
     let ctx = SessionContext::from(state);
 
     let worker = Worker::default().with_runtime_env(runtime_env);

benchmarks/src/run.rs

@@ -93,6 +93,10 @@ pub struct RunOpt {
     #[structopt(long)]
     children_isolator_unions: bool,
 
+    /// Turns on broadcast joins.
+    #[structopt(long)]
+    enable_broadcast_joins: bool,
+
     /// Collects metrics across network boundaries
     #[structopt(long)]
     collect_metrics: bool,
@@ -210,6 +214,7 @@ impl RunOpt {
                 self.cardinality_task_sf.unwrap_or(1.0),
             )?
             .with_distributed_children_isolator_unions(self.children_isolator_unions)?
+            .with_distributed_broadcast_joins_enabled(self.enable_broadcast_joins)?
             .with_distributed_metrics_collection(self.collect_metrics)?
             .build();
         let ctx = SessionContext::new_with_state(state);

src/distributed_ext.rs

@@ -449,6 +449,23 @@ pub trait DistributedExt: Sized {
         &mut self,
         enabled: bool,
     ) -> Result<(), DataFusionError>;
+
+    /// Enables broadcast joins for CollectLeft hash joins. When enabled, the build side of
+    /// a CollectLeft join is broadcast to all consumer tasks instead of being coalesced
+    /// into a single partition.
+    ///
+    /// Note: This option is disabled by default until the implementation is smarter about when to
+    /// broadcast.
+    fn with_distributed_broadcast_joins_enabled(
+        self,
+        enabled: bool,
+    ) -> Result<Self, DataFusionError>;
+
+    /// Same as [DistributedExt::with_distributed_broadcast_joins_enabled] but with an in-place mutation.
+    fn set_distributed_broadcast_joins_enabled(
+        &mut self,
+        enabled: bool,
+    ) -> Result<(), DataFusionError>;
 }
 
 impl DistributedExt for SessionConfig {
@@ -529,6 +546,15 @@ impl DistributedExt for SessionConfig {
         Ok(())
     }
 
+    fn set_distributed_broadcast_joins_enabled(
+        &mut self,
+        enabled: bool,
+    ) -> Result<(), DataFusionError> {
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        d_cfg.broadcast_joins_enabled = enabled;
+        Ok(())
+    }
+
     delegate! {
         to self {
             #[call(set_distributed_option_extension)]
@@ -574,6 +600,10 @@ impl DistributedExt for SessionConfig {
             #[call(set_distributed_children_isolator_unions)]
             #[expr($?;Ok(self))]
             fn with_distributed_children_isolator_unions(mut self, enabled: bool) -> Result<Self, DataFusionError>;
+
+            #[call(set_distributed_broadcast_joins_enabled)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_broadcast_joins_enabled(mut self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -635,6 +665,11 @@ impl DistributedExt for SessionStateBuilder {
             #[call(set_distributed_children_isolator_unions)]
             #[expr($?;Ok(self))]
             fn with_distributed_children_isolator_unions(mut self, enabled: bool) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_broadcast_joins_enabled(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_broadcast_joins_enabled)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_broadcast_joins_enabled(mut self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -696,6 +731,11 @@ impl DistributedExt for SessionState {
             #[call(set_distributed_children_isolator_unions)]
             #[expr($?;Ok(self))]
             fn with_distributed_children_isolator_unions(mut self, enabled: bool) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_broadcast_joins_enabled(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_broadcast_joins_enabled)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_broadcast_joins_enabled(mut self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -757,6 +797,11 @@ impl DistributedExt for SessionContext {
             #[call(set_distributed_children_isolator_unions)]
             #[expr($?;Ok(self))]
             fn with_distributed_children_isolator_unions(self, enabled: bool) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_broadcast_joins_enabled(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_broadcast_joins_enabled)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_broadcast_joins_enabled(self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }

src/distributed_planner/distributed_config.rs

@@ -39,6 +39,12 @@ extensions_options! {
         /// Propagate collected metrics from all nodes in the plan across network boundaries
         /// so that they can be reconstructed on the head node of the plan.
         pub collect_metrics: bool, default = true
+        /// Enable broadcast joins for CollectLeft hash joins. When enabled, the build side of
+        /// a CollectLeft join is broadcast to all consumer tasks.
+        /// TODO: This option exists temporarily until we become smarter about when to actually
+        /// use broadcasting like checking build side size.
+        /// For now, broadcasting all CollectLeft joins is not always beneficial.
+        pub broadcast_joins_enabled: bool, default = false
         /// Collection of [TaskEstimator]s that will be applied to leaf nodes in order to
         /// estimate how many tasks should be spawned for the [Stage] containing the leaf node.
         pub(crate) __private_task_estimator: CombinedTaskEstimator, default = CombinedTaskEstimator::default()