feat: Flotilla broadcast join (#4867)

## Changes Made

Add broadcast joins to flotilla. The translation logic for determining
join strategy + broadcast join threshold is copied from the old ray
runner's physical plan translator.

The broadcast join works by simply materializing the partition refs for
the broadcast side, then copying them to the each of the receiver sides
task (aka broadcast), and then adding a hash join instruction.

## Related Issues

<!-- Link to related GitHub issues, e.g., "Closes #123" -->

## Checklist

- [ ] Documented in API Docs (if applicable)
- [ ] Documented in User Guide (if applicable)
- [ ] If adding a new documentation page, doc is added to
`docs/mkdocs.yml` navigation
- [ ] Documentation builds and is formatted properly (tag @/ccmao1130
for docs review)

Author: colin-ho

src/daft-distributed/src/pipeline_node/gather.rs

@@ -9,7 +9,7 @@ use futures::TryStreamExt;
 use super::{DistributedPipelineNode, SubmittableTaskStream};
 use crate::{
     pipeline_node::{
-        make_in_memory_scan_from_materialized_outputs, NodeID, NodeName, PipelineNodeConfig,
+        make_in_memory_task_from_materialized_outputs, NodeID, NodeName, PipelineNodeConfig,
         PipelineNodeContext,
     },
     scheduling::{
@@ -80,7 +80,7 @@ impl GatherNode {
             .await?;
 
         let self_clone = self.clone();
-        let task = make_in_memory_scan_from_materialized_outputs(
+        let task = make_in_memory_task_from_materialized_outputs(
             TaskContext::from((&self_clone.context, task_id_counter.next())),
             materialized,
             &(self_clone as Arc<dyn DistributedPipelineNode>),

src/daft-distributed/src/pipeline_node/join/broadcast_join.rs

@@ -0,0 +1,234 @@
+use std::{collections::HashMap, sync::Arc};
+
+use common_display::{tree::TreeDisplay, DisplayLevel};
+use common_error::DaftResult;
+use daft_dsl::expr::bound_expr::BoundExpr;
+use daft_local_plan::LocalPhysicalPlan;
+use daft_logical_plan::{stats::StatsState, JoinType};
+use daft_schema::schema::SchemaRef;
+use futures::{StreamExt, TryStreamExt};
+
+use crate::{
+    pipeline_node::{
+        make_in_memory_scan_from_materialized_outputs, DistributedPipelineNode, NodeID, NodeName,
+        PipelineNodeConfig, PipelineNodeContext, SubmittableTaskStream,
+    },
+    scheduling::{
+        scheduler::{SchedulerHandle, SubmittableTask},
+        task::{SchedulingStrategy, SwordfishTask, TaskContext},
+    },
+    stage::{StageConfig, StageExecutionContext, TaskIDCounter},
+    utils::channel::{create_channel, Sender},
+};
+
+pub(crate) struct BroadcastJoinNode {
+    config: PipelineNodeConfig,
+    context: PipelineNodeContext,
+
+    // Join properties
+    left_on: Vec<BoundExpr>,
+    right_on: Vec<BoundExpr>,
+    null_equals_nulls: Option<Vec<bool>>,
+    join_type: JoinType,
+    is_swapped: bool,
+
+    broadcaster: Arc<dyn DistributedPipelineNode>,
+    broadcaster_schema: SchemaRef,
+    receiver: Arc<dyn DistributedPipelineNode>,
+}
+
+impl BroadcastJoinNode {
+    const NODE_NAME: NodeName = "BroadcastJoin";
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn new(
+        node_id: NodeID,
+        logical_node_id: Option<NodeID>,
+        stage_config: &StageConfig,
+        left_on: Vec<BoundExpr>,
+        right_on: Vec<BoundExpr>,
+        null_equals_nulls: Option<Vec<bool>>,
+        join_type: JoinType,
+        is_swapped: bool,
+        broadcaster: Arc<dyn DistributedPipelineNode>,
+        receiver: Arc<dyn DistributedPipelineNode>,
+        output_schema: SchemaRef,
+    ) -> Self {
+        let context = PipelineNodeContext::new(
+            stage_config,
+            node_id,
+            Self::NODE_NAME,
+            vec![broadcaster.node_id(), receiver.node_id()],
+            vec![broadcaster.name(), receiver.name()],
+            logical_node_id,
+        );
+
+        // For broadcast joins, we use the receiver's clustering spec since the broadcaster
+        // will be gathered to all partitions
+        let config = PipelineNodeConfig::new(
+            output_schema,
+            stage_config.config.clone(),
+            receiver.config().clustering_spec.clone(),
+        );
+
+        let broadcaster_schema = broadcaster.config().schema.clone();
+        Self {
+            config,
+            context,
+            left_on,
+            right_on,
+            null_equals_nulls,
+            join_type,
+            is_swapped,
+            broadcaster,
+            broadcaster_schema,
+            receiver,
+        }
+    }
+
+    pub fn arced(self) -> Arc<dyn DistributedPipelineNode> {
+        Arc::new(self)
+    }
+
+    fn multiline_display(&self) -> Vec<String> {
+        use itertools::Itertools;
+        let mut res = vec!["Broadcast Join".to_string()];
+        res.push(format!(
+            "Left on: {}",
+            self.left_on.iter().map(|e| e.to_string()).join(", ")
+        ));
+        res.push(format!(
+            "Right on: {}",
+            self.right_on.iter().map(|e| e.to_string()).join(", ")
+        ));
+        res.push(format!("Join type: {}", self.join_type));
+        res.push(format!("Is swapped: {}", self.is_swapped));
+        if let Some(null_equals_nulls) = &self.null_equals_nulls {
+            res.push(format!(
+                "Null equals nulls: [{}]",
+                null_equals_nulls.iter().map(|b| b.to_string()).join(", ")
+            ));
+        }
+        res
+    }
+
+    async fn execution_loop(
+        self: Arc<Self>,
+        broadcaster_input: SubmittableTaskStream,
+        mut receiver_input: SubmittableTaskStream,
+        task_id_counter: TaskIDCounter,
+        result_tx: Sender<SubmittableTask<SwordfishTask>>,
+        scheduler_handle: SchedulerHandle<SwordfishTask>,
+    ) -> DaftResult<()> {
+        let materialized_broadcast_data = broadcaster_input
+            .materialize(scheduler_handle.clone())
+            .try_collect::<Vec<_>>()
+            .await?;
+        let materialized_broadcast_data_plan = make_in_memory_scan_from_materialized_outputs(
+            &materialized_broadcast_data,
+            self.broadcaster_schema.clone(),
+            self.node_id(),
+        )?;
+        let broadcast_psets = HashMap::from([(
+            self.node_id().to_string(),
+            materialized_broadcast_data
+                .into_iter()
+                .flat_map(|output| output.into_inner().0)
+                .collect::<Vec<_>>(),
+        )]);
+        while let Some(task) = receiver_input.next().await {
+            let input_plan = task.task().plan();
+            let (left_plan, right_plan) = if self.is_swapped {
+                (input_plan, materialized_broadcast_data_plan.clone())
+            } else {
+                (materialized_broadcast_data_plan.clone(), input_plan)
+            };
+            let join_plan = LocalPhysicalPlan::hash_join(
+                left_plan,
+                right_plan,
+                self.left_on.clone(),
+                self.right_on.clone(),
+                self.null_equals_nulls.clone(),
+                self.join_type,
+                self.config.schema.clone(),
+                StatsState::NotMaterialized,
+            );
+
+            let mut psets = task.task().psets().clone();
+            psets.extend(broadcast_psets.clone());
+
+            let config = task.task().config().clone();
+
+            let task = task.with_new_task(SwordfishTask::new(
+                TaskContext::from((self.context(), task_id_counter.next())),
+                join_plan,
+                config,
+                psets,
+                SchedulingStrategy::Spread,
+                self.context().to_hashmap(),
+            ));
+            if result_tx.send(task).await.is_err() {
+                break;
+            }
+        }
+        Ok(())
+    }
+}
+
+impl TreeDisplay for BroadcastJoinNode {
+    fn display_as(&self, level: DisplayLevel) -> String {
+        match level {
+            DisplayLevel::Compact => self.get_name(),
+            _ => self.multiline_display().join("\n"),
+        }
+    }
+
+    fn get_children(&self) -> Vec<&dyn TreeDisplay> {
+        vec![
+            self.broadcaster.as_tree_display(),
+            self.receiver.as_tree_display(),
+        ]
+    }
+
+    fn get_name(&self) -> String {
+        Self::NODE_NAME.to_string()
+    }
+}
+
+impl DistributedPipelineNode for BroadcastJoinNode {
+    fn context(&self) -> &PipelineNodeContext {
+        &self.context
+    }
+
+    fn config(&self) -> &PipelineNodeConfig {
+        &self.config
+    }
+
+    fn children(&self) -> Vec<Arc<dyn DistributedPipelineNode>> {
+        vec![self.broadcaster.clone(), self.receiver.clone()]
+    }
+
+    fn produce_tasks(
+        self: Arc<Self>,
+        stage_context: &mut StageExecutionContext,
+    ) -> SubmittableTaskStream {
+        let broadcaster_input = self.broadcaster.clone().produce_tasks(stage_context);
+        let receiver_input = self.receiver.clone().produce_tasks(stage_context);
+
+        let (result_tx, result_rx) = create_channel(1);
+        let execution_loop = self.execution_loop(
+            broadcaster_input,
+            receiver_input,
+            stage_context.task_id_counter(),
+            result_tx,
+            stage_context.scheduler_handle(),
+        );
+        stage_context.spawn(execution_loop);
+
+        SubmittableTaskStream::from(result_rx)
+    }
+
+    fn as_tree_display(&self) -> &dyn TreeDisplay {
+        self
+    }
+}

src/daft-distributed/src/pipeline_node/hash_join.rs renamed to src/daft-distributed/src/pipeline_node/join/hash_join.rs

@@ -2,21 +2,18 @@ use std::{cmp::max, sync::Arc};
 
 use common_daft_config::DaftExecutionConfig;
 use common_display::{tree::TreeDisplay, DisplayLevel};
-use common_error::DaftResult;
-use daft_dsl::{expr::bound_expr::BoundExpr, join::normalize_join_keys};
+use daft_dsl::expr::bound_expr::BoundExpr;
 use daft_local_plan::LocalPhysicalPlan;
 use daft_logical_plan::{
-    ops::join::JoinPredicate, partitioning::HashClusteringConfig, stats::StatsState,
-    ClusteringSpec, JoinType,
+    partitioning::HashClusteringConfig, stats::StatsState, ClusteringSpec, JoinType,
 };
 use daft_schema::schema::SchemaRef;
 use futures::StreamExt;
 
-use super::{DistributedPipelineNode, SubmittableTaskStream};
 use crate::{
     pipeline_node::{
-        repartition::RepartitionNode, translate::LogicalPlanToPipelineNodeTranslator, NodeID,
-        NodeName, PipelineNodeConfig, PipelineNodeContext,
+        DistributedPipelineNode, NodeID, NodeName, PipelineNodeConfig, PipelineNodeContext,
+        SubmittableTaskStream,
     },
     scheduling::{
         scheduler::SubmittableTask,
@@ -201,7 +198,7 @@ impl DistributedPipelineNode for HashJoinNode {
     }
 }
 
-fn gen_num_partitions(
+pub(crate) fn gen_num_partitions(
     left_spec: &ClusteringSpec,
     right_spec: &ClusteringSpec,
     cfg: &DaftExecutionConfig,
@@ -228,79 +225,3 @@ fn gen_num_partitions(
         (_, _, a, b) => max(a, b),
     }
 }
-
-impl LogicalPlanToPipelineNodeTranslator {
-    pub(crate) fn gen_hash_join_nodes(
-        &mut self,
-        logical_node_id: Option<NodeID>,
-        join_on: JoinPredicate,
-
-        left: Arc<dyn DistributedPipelineNode>,
-        right: Arc<dyn DistributedPipelineNode>,
-
-        join_type: JoinType,
-        output_schema: SchemaRef,
-    ) -> DaftResult<Arc<dyn DistributedPipelineNode>> {
-        let (_, left_on, right_on, null_equals_nulls) = join_on.split_eq_preds();
-
-        let (left_on, right_on) = normalize_join_keys(
-            left_on,
-            right_on,
-            left.config().schema.clone(),
-            right.config().schema.clone(),
-        )?;
-        let left_on = BoundExpr::bind_all(&left_on, &left.config().schema)?;
-        let right_on = BoundExpr::bind_all(&right_on, &right.config().schema)?;
-
-        let num_partitions = gen_num_partitions(
-            left.config().clustering_spec.as_ref(),
-            right.config().clustering_spec.as_ref(),
-            self.stage_config.config.as_ref(),
-        );
-
-        let left = RepartitionNode::new(
-            self.get_next_pipeline_node_id(),
-            logical_node_id,
-            &self.stage_config,
-            daft_logical_plan::partitioning::RepartitionSpec::Hash(
-                daft_logical_plan::partitioning::HashRepartitionConfig::new(
-                    Some(num_partitions),
-                    left_on.clone().into_iter().map(|e| e.into()).collect(),
-                ),
-            ),
-            left.config().schema.clone(),
-            left,
-        )
-        .arced();
-
-        let right = RepartitionNode::new(
-            self.get_next_pipeline_node_id(),
-            logical_node_id,
-            &self.stage_config,
-            daft_logical_plan::partitioning::RepartitionSpec::Hash(
-                daft_logical_plan::partitioning::HashRepartitionConfig::new(
-                    Some(num_partitions),
-                    right_on.clone().into_iter().map(|e| e.into()).collect(),
-                ),
-            ),
-            right.config().schema.clone(),
-            right,
-        )
-        .arced();
-
-        Ok(HashJoinNode::new(
-            self.get_next_pipeline_node_id(),
-            logical_node_id,
-            &self.stage_config,
-            left_on,
-            right_on,
-            Some(null_equals_nulls),
-            join_type,
-            num_partitions,
-            left,
-            right,
-            output_schema,
-        )
-        .arced())
-    }
-}

src/daft-distributed/src/pipeline_node/join/mod.rs

@@ -0,0 +1,6 @@
+mod broadcast_join;
+pub(crate) mod hash_join;
+pub(crate) mod translate_join;
+
+pub(crate) use broadcast_join::BroadcastJoinNode;
+pub(crate) use hash_join::HashJoinNode;