feat: Flotilla stage redesign (#4306)

## Changes Made

Add tentative redesign (and comment for rationale) for stages, subject
to change.

## 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)

---------

Co-authored-by: Sammy Sidhu <[email protected]>

Author: colin-ho

Cargo.lock

@@ -4,8 +4,10 @@ common-daft-config = {path = "../common/daft-config", default-features = false}
 common-error = {path = "../common/error", default-features = false}
 common-partitioning = {path = "../common/partitioning", default-features = false}
 common-treenode = {path = "../common/treenode", default-features = false}
+daft-dsl = {path = "../daft-dsl", default-features = false}
 daft-local-plan = {path = "../daft-local-plan", default-features = false}
 daft-logical-plan = {path = "../daft-logical-plan", default-features = false}
+daft-schema = {path = "../daft-schema", default-features = false}
 futures = {workspace = true}
 pyo3 = {workspace = true, optional = true}
 pyo3-async-runtimes = {workspace = true, optional = true}

src/daft-distributed/Cargo.toml

@@ -11,13 +11,15 @@ use crate::{
     stage::StageContext,
 };
 
+#[allow(dead_code)]
 pub(crate) struct CollectNode {
     local_physical_plans: Vec<LocalPhysicalPlanRef>,
     children: Vec<Box<dyn DistributedPipelineNode>>,
     input_psets: HashMap<String, Vec<PartitionRef>>,
 }
 
 impl CollectNode {
+    #[allow(dead_code)]
     pub fn new(
         local_physical_plans: Vec<LocalPhysicalPlanRef>,
         children: Vec<Box<dyn DistributedPipelineNode>>,
@@ -37,14 +39,15 @@ impl CollectNode {
         }
     }
 
+    #[allow(dead_code)]
     async fn execution_loop(
         _task_dispatcher_handle: TaskDispatcherHandle,
         _local_physical_plans: Vec<LocalPhysicalPlanRef>,
         _psets: HashMap<String, Vec<PartitionRef>>,
         _input_node: Option<RunningPipelineNode>,
         _result_tx: Sender<PipelineOutput>,
     ) -> DaftResult<()> {
-        todo!("Implement collect execution sloop");
+        todo!("FLOTILLA_MS1: Implement collect execution sloop");
     }
 }
 

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

@@ -20,6 +20,7 @@ pub(crate) struct LimitNode {
 }
 
 impl LimitNode {
+    #[allow(dead_code)]
     pub fn new(
         limit: usize,
         local_physical_plans: Vec<LocalPhysicalPlanRef>,
@@ -41,14 +42,15 @@ impl LimitNode {
         }
     }
 
+    #[allow(dead_code)]
     async fn execution_loop(
         _task_dispatcher_handle: TaskDispatcherHandle,
         _local_physical_plans: Vec<LocalPhysicalPlanRef>,
         _input_node: Option<RunningPipelineNode>,
         _input_psets: HashMap<String, Vec<PartitionRef>>,
         _result_tx: Sender<PipelineOutput>,
     ) -> DaftResult<()> {
-        todo!("Implement limit execution loop");
+        todo!("FLOTILLA_MS1: Implement limit execution loop");
     }
 }
 

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

@@ -30,6 +30,7 @@ pub(crate) trait DistributedPipelineNode: Send + Sync {
     fn name(&self) -> &'static str;
     #[allow(dead_code)]
     fn children(&self) -> Vec<&dyn DistributedPipelineNode>;
+    #[allow(dead_code)]
     fn start(&mut self, stage_context: &mut StageContext) -> RunningPipelineNode;
 }
 
@@ -39,6 +40,7 @@ pub(crate) struct RunningPipelineNode {
 }
 
 impl RunningPipelineNode {
+    #[allow(dead_code)]
     fn new(result_receiver: Receiver<PipelineOutput>) -> Self {
         Self { result_receiver }
     }
@@ -53,7 +55,7 @@ impl Stream for RunningPipelineNode {
     type Item = DaftResult<PipelineOutput>;
 
     fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        todo!("Implement stream for running pipeline node");
+        todo!("FLOTILLA_MS1: Implement stream for running pipeline node");
     }
 }
 
@@ -64,6 +66,7 @@ pub(crate) enum PipelineOutput {
     Running(Box<dyn SwordfishTaskResultHandle>),
 }
 
+#[allow(dead_code)]
 pub(crate) fn logical_plan_to_pipeline_node(
     plan: LogicalPlanRef,
     config: Arc<DaftExecutionConfig>,
@@ -100,7 +103,7 @@ pub(crate) fn logical_plan_to_pipeline_node(
                         std::mem::take(&mut self.psets),
                     ))];
                     // Here we will have to return a placeholder, essentially cutting off the plan
-                    todo!("Implement pipeline node boundary splitter for limit");
+                    todo!("FLOTILLA_MS1: Implement pipeline node boundary splitter for limit");
                 }
                 _ if is_root => {
                     let input_nodes = std::mem::take(&mut self.current_nodes);

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

@@ -8,5 +8,5 @@ pub(crate) fn translate_pipeline_plan_to_local_physical_plans(
     _logical_plan: LogicalPlanRef,
     _execution_config: &DaftExecutionConfig,
 ) -> DaftResult<Vec<LocalPhysicalPlanRef>> {
-    todo!("Implement translate pipeline plan to local physical plans");
+    todo!("FLOTILLA_MS1: Implement translate pipeline plan to local physical plans");
 }

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

@@ -6,20 +6,21 @@ use std::{
 };
 
 use common_daft_config::DaftExecutionConfig;
-use common_error::{DaftError, DaftResult};
+use common_error::DaftResult;
 use common_partitioning::PartitionRef;
 use daft_logical_plan::{LogicalPlanBuilder, LogicalPlanRef};
-use futures::{Stream, StreamExt};
+use futures::Stream;
 
 use crate::{
-    channel::{create_channel, Receiver, Sender},
+    channel::{create_channel, Receiver},
     runtime::{get_or_init_runtime, JoinHandle},
     scheduling::worker::WorkerManagerFactory,
-    stage::split_at_stage_boundary,
+    stage::StagePlan,
 };
 
 pub struct DistributedPhysicalPlan {
-    remaining_logical_plan: Option<LogicalPlanRef>,
+    #[allow(dead_code)]
+    logical_plan: LogicalPlanRef,
     config: Arc<DaftExecutionConfig>,
 }
 
@@ -29,74 +30,40 @@ impl DistributedPhysicalPlan {
         config: Arc<DaftExecutionConfig>,
     ) -> DaftResult<Self> {
         let plan = builder.build();
-        if !can_translate_logical_plan(&plan) {
-            return Err(DaftError::InternalError(
-                "Cannot run this physical plan on distributed swordfish yet".to_string(),
-            ));
-        }
 
         Ok(Self {
-            remaining_logical_plan: Some(plan),
+            logical_plan: plan,
             config,
         })
     }
 
-    async fn run_plan_loop(
-        logical_plan: LogicalPlanRef,
-        config: Arc<DaftExecutionConfig>,
-        worker_manager_factory: Box<dyn WorkerManagerFactory>,
-        psets: HashMap<String, Vec<PartitionRef>>,
-        result_sender: Sender<PartitionRef>,
+    async fn execute_stages(
+        _stage_plan: StagePlan,
+        _psets: HashMap<String, Vec<PartitionRef>>,
+        _worker_manager_factory: Box<dyn WorkerManagerFactory>,
     ) -> DaftResult<()> {
-        let (stage, _remaining_plan) = split_at_stage_boundary(&logical_plan, &config)?;
-        let mut running_stage = stage.run_stage(psets, worker_manager_factory)?;
-        while let Some(result) = running_stage.next().await {
-            if result_sender.send(result?).await.is_err() {
-                break;
-            }
-        }
-        todo!("Implement stage running loop");
-    }
-
-    #[allow(dead_code)]
-    fn update_plan(
-        _plan: LogicalPlanRef,
-        _results: Vec<PartitionRef>,
-    ) -> DaftResult<LogicalPlanRef> {
-        // Update the logical plan with the results of the previous stage.
-        // This is where the AQE magic happens.
-        todo!("Implement plan updating and AQE");
+        todo!("FLOTILLA_MS1: Implement execute stages");
     }
 
     pub fn run_plan(
         &self,
         psets: HashMap<String, Vec<PartitionRef>>,
         worker_manager_factory: Box<dyn WorkerManagerFactory>,
-    ) -> PlanResult {
-        let (result_sender, result_receiver) = create_channel(1);
+    ) -> DaftResult<PlanResult> {
+        let (_result_sender, result_receiver) = create_channel(1);
         let runtime = get_or_init_runtime();
-        let handle = runtime.spawn(Self::run_plan_loop(
-            self.remaining_logical_plan
-                .as_ref()
-                .expect("Expected remaining logical plan")
-                .clone(),
-            self.config.clone(),
-            worker_manager_factory,
-            psets,
-            result_sender,
-        ));
-        PlanResult::new(handle, result_receiver)
+        let stage_plan = StagePlan::from_logical_plan(self.logical_plan.clone())?;
+        let handle = runtime.spawn(async move {
+            Self::execute_stages(stage_plan, psets, worker_manager_factory).await
+        });
+        Ok(PlanResult::new(handle, result_receiver))
     }
 
     pub fn execution_config(&self) -> &Arc<DaftExecutionConfig> {
         &self.config
     }
 }
 
-fn can_translate_logical_plan(_plan: &LogicalPlanRef) -> bool {
-    todo!("Implement logical plan translation check");
-}
-
 // This is the output of a plan, a receiver to receive the results of the plan.
 // And the join handle to the task that runs the plan.
 pub struct PlanResult {
@@ -117,6 +84,6 @@ impl Stream for PlanResult {
     type Item = DaftResult<PartitionRef>;
 
     fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        todo!("Implement stream for plan result");
+        todo!("FLOTILLA_MS1: Implement stream for plan result");
     }
 }

src/daft-distributed/src/plan/mod.rs

@@ -95,7 +95,7 @@ impl PyDistributedPhysicalPlan {
         );
         let part_stream = self
             .planner
-            .run_plan(psets, Box::new(worker_manager_factory));
+            .run_plan(psets, Box::new(worker_manager_factory))?;
         let part_stream = PythonPartitionRefStream {
             inner: Arc::new(Mutex::new(part_stream)),
         };

src/daft-distributed/src/python/mod.rs

@@ -17,6 +17,8 @@ pub(crate) struct RayWorkerManager {
     task_locals: pyo3_async_runtimes::TaskLocals,
 }
 
+// TODO(FLOTILLA_MS1): Make Ray worker manager live for the duration of the program
+// so that we don't have to recreate it on every stage.
 impl RayWorkerManager {
     pub fn try_new(
         daft_execution_config: Arc<DaftExecutionConfig>,

src/daft-distributed/src/python/ray/worker_manager.rs

@@ -43,7 +43,7 @@ impl TaskDispatcher {
         _dispatcher: Self,
         _task_rx: Receiver<DispatchedTask>,
     ) -> DaftResult<()> {
-        todo!("Implement run dispatch loop");
+        todo!("FLOTILLA_MS1: Implement run dispatch loop");
     }
 }