Refactor stage struct

# Conflicts:
#	src/execution_plans/distributed.rs

Author: gabotechs

benchmarks/src/run.rs

@@ -304,7 +304,7 @@ impl RunOpt {
         let mut n_tasks = 0;
         physical_plan.clone().transform_down(|node| {
             if let Some(node) = node.as_network_boundary() {
-                n_tasks += node.input_stage().tasks.len()
+                n_tasks += node.input_stage().task_count()
             }
             Ok(Transformed::no(node))
         })?;

console/examples/tpcds_runner.rs

@@ -189,7 +189,8 @@ async fn run_single_query(
     let batches = stream.try_collect::<Vec<_>>().await?;
     if explain_analyze {
         let output =
-            datafusion_distributed::explain_analyze(plan, DistributedMetricsFormat::Aggregated)?;
+            datafusion_distributed::explain_analyze(plan, DistributedMetricsFormat::Aggregated)
+                .await?;
         println!("{output}");
     }
     Ok(batches)

src/execution_plans/benchmarks/shuffle_bench.rs

@@ -2,9 +2,10 @@ use super::fixture::{
     InMemoryChannelsResolver, benchmark_schema, make_input_partitions, rows_for_producer,
 };
 use crate::common::task_ctx_with_extension;
+use crate::stage::RemoteStage;
 use crate::worker::WorkerConnectionPool;
 use crate::worker::test_utils::worker_handles::MemoryWorkerHandle;
-use crate::{DistributedExt, DistributedTaskContext, ExecutionTask, NetworkShuffleExec, Stage};
+use crate::{DistributedExt, DistributedTaskContext, NetworkShuffleExec, Stage};
 use arrow::datatypes::Schema;
 use arrow_ipc::CompressionType;
 use datafusion::common::Result;
@@ -165,16 +166,14 @@ impl ShuffleBench {
             .build()
             .task_ctx();
         let input_stage_tasks = (0..bench.producer_tasks)
-            .map(|i| ExecutionTask {
-                url: Some(Url::parse(&format!("http://localhost:{i}")).unwrap()),
-            })
+            .map(|i| Url::parse(&format!("http://localhost:{i}")).unwrap())
             .collect();
 
         Ok(ShuffleFixture {
             bench,
             schema,
             task_ctx,
-            input_stage_tasks,
+            input_stage_workers: input_stage_tasks,
             workers,
         })
     }
@@ -190,7 +189,7 @@ pub struct ShuffleFixture {
     bench: ShuffleBench,
     schema: Arc<Schema>,
     task_ctx: Arc<datafusion::execution::TaskContext>,
-    input_stage_tasks: Vec<ExecutionTask>,
+    input_stage_workers: Vec<Url>,
     workers: Vec<MemoryWorkerHandle>,
 }
 
@@ -213,12 +212,11 @@ impl ShuffleFixture {
                 .await?;
         }
 
-        let input_stage = Stage {
+        let input_stage = Stage::Remote(RemoteStage {
             query_id,
             num: 0,
-            plan: None,
-            tasks: self.input_stage_tasks.clone(),
-        };
+            workers: self.input_stage_workers.clone(),
+        });
 
         let mut join_set = JoinSet::default();
         for task_index in 0..self.bench.consumer_tasks {

src/execution_plans/benchmarks/transport_bench.rs

@@ -2,10 +2,10 @@ use super::fixture::{
     InMemoryChannelsResolver, benchmark_schema, make_input_partitions, rows_for_producer,
 };
 use crate::common::task_ctx_with_extension;
+use crate::stage::RemoteStage;
 use crate::worker::test_utils::worker_handles::{MemoryWorkerHandle, TcpWorkerHandle};
 use crate::{
-    DefaultChannelResolver, DistributedExt, DistributedTaskContext, ExecutionTask,
-    NetworkShuffleExec, Stage,
+    DefaultChannelResolver, DistributedExt, DistributedTaskContext, NetworkShuffleExec, Stage,
 };
 use arrow::datatypes::Schema;
 use arrow_ipc::CompressionType;
@@ -202,9 +202,7 @@ impl TransportBench {
                 .build()
                 .task_ctx(),
             input_stage_tasks: (0..self.producer_tasks)
-                .map(|i| ExecutionTask {
-                    url: Some(Url::parse(&format!("http://localhost:{i}")).unwrap()),
-                })
+                .map(|i| Url::parse(&format!("http://localhost:{i}")).unwrap())
                 .collect(),
             workers: PreparedTransportWorkers::InMemory(workers),
         })
@@ -239,12 +237,7 @@ impl TransportBench {
                 .with_distributed_compression(self.compression)?
                 .build()
                 .task_ctx(),
-            input_stage_tasks: workers
-                .iter()
-                .map(|worker| ExecutionTask {
-                    url: Some(worker.url.clone()),
-                })
-                .collect(),
+            input_stage_tasks: workers.iter().map(|worker| worker.url.clone()).collect(),
             workers: PreparedTransportWorkers::Tcp(workers),
         })
     }
@@ -260,7 +253,7 @@ pub struct TransportFixture {
     bench: TransportBench,
     schema: Arc<Schema>,
     task_ctx: Arc<datafusion::execution::TaskContext>,
-    input_stage_tasks: Vec<ExecutionTask>,
+    input_stage_tasks: Vec<Url>,
     workers: PreparedTransportWorkers,
 }
 
@@ -269,12 +262,11 @@ impl TransportFixture {
         let query_id = Uuid::new_v4();
         self.workers.register_plans(query_id).await?;
 
-        let input_stage = Stage {
+        let input_stage = Stage::Remote(RemoteStage {
             query_id,
             num: 0,
-            plan: None,
-            tasks: self.input_stage_tasks.clone(),
-        };
+            workers: self.input_stage_tasks.clone(),
+        });
 
         let mut join_set = JoinSet::default();
         for task_index in 0..self.bench.consumer_tasks {

src/execution_plans/distributed.rs

@@ -5,7 +5,7 @@ use crate::execution_plans::ChildrenIsolatorUnionExec;
 use crate::networking::get_distributed_worker_resolver;
 use crate::passthrough_headers::get_passthrough_headers;
 use crate::protobuf::{DistributedCodec, tonic_status_to_datafusion_error};
-use crate::stage::{ExecutionTask, Stage};
+use crate::stage::{LocalStage, RemoteStage, Stage};
 use crate::worker::generated::worker as pb;
 use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration;
 use crate::worker::generated::worker::{
@@ -21,7 +21,7 @@ use datafusion::common::HashMap;
 use datafusion::common::instant::Instant;
 use datafusion::common::runtime::JoinSet;
 use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
-use datafusion::common::{Result, exec_err, internal_err};
+use datafusion::common::{Result, exec_err};
 use datafusion::common::{exec_datafusion_err, internal_datafusion_err};
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
@@ -126,10 +126,10 @@ impl DistributedExec {
         let _ = self.plan.apply(|plan| {
             if let Some(boundary) = plan.as_network_boundary() {
                 let stage = boundary.input_stage();
-                for i in 0..stage.tasks.len() {
+                for i in 0..stage.task_count() {
                     expected_keys.push(TaskKey {
-                        query_id: stage.query_id.as_bytes().to_vec(),
-                        stage_id: stage.num as u64,
+                        query_id: serialize_uuid(&stage.query_id()),
+                        stage_id: stage.num() as u64,
                         task_number: i as u64,
                     });
                 }
@@ -193,61 +193,63 @@ impl DistributedExec {
                 return Ok(Transformed::no(plan));
             };
 
-            let task_estimator = get_distributed_task_estimator(ctx.session_config())?;
+            let Stage::Local(stage) = plan.input_stage() else {
+                return exec_err!("Input stage from network boundary was not in Local state");
+            };
 
-            let stage = plan.input_stage();
+            let task_estimator = get_distributed_task_estimator(ctx.session_config())?;
 
-            let mut spawner =
-                CoordinatorToWorkerTaskSpawner::new(stage, &metrics, &codec, &mut join_set)?;
+            let mut spawner = CoordinatorToWorkerTaskSpawner::new(
+                stage,
+                &metrics,
+                &self.task_metrics,
+                &codec,
+                &mut join_set,
+            )?;
 
             let routed_urls = match task_estimator.route_tasks(&TaskRoutingContext {
                 task_ctx: Arc::clone(ctx),
-                plan: stage.plan.as_ref().ok_or_else(|| {
-                    internal_datafusion_err!(
-                        "network boundary stage missing input plan during routing"
-                    )
-                })?,
-                task_count: stage.tasks.len(),
+                plan: &stage.plan,
+                task_count: stage.tasks,
                 available_urls: &available_urls,
             }) {
                 Ok(Some(routed_urls)) => routed_urls,
                 // If the user has not defined custom routing with a `route_tasks` implementation, we
                 // default to round-robin task assignation from a randomized starting point.
                 Ok(None) => {
                     let start_idx = rand::rng().random_range(0..available_urls.len());
-                    (0..stage.tasks.len())
+                    (0..stage.tasks)
                         .map(|i| available_urls[(start_idx + i) % available_urls.len()].clone())
                         .collect()
                 }
                 Err(e) => return Err(exec_datafusion_err!("error routing tasks to workers: {e}")),
             };
 
-            if routed_urls.len() != stage.tasks.len() {
+            if routed_urls.len() != stage.tasks {
                 return Err(exec_datafusion_err!(
                     "number of tasks ({}) was not equal to number of urls ({}) at execution time",
-                    stage.tasks.len(),
+                    stage.tasks,
                     routed_urls.len()
                 ));
             }
 
-            let mut tasks = Vec::with_capacity(stage.tasks.len());
+            let mut workers = Vec::with_capacity(stage.tasks);
             for (i, routed_url) in routed_urls.into_iter().enumerate() {
-                tasks.push(ExecutionTask {
-                    url: Some(routed_url.clone()),
-                });
+                workers.push(routed_url.clone());
                 // Spawn a task that sends the subplan to the chosen URL.
                 // There will be as many spawned tasks as workers.
                 let (tx, worker_rx) = spawner.send_plan_task(Arc::clone(ctx), i, routed_url)?;
-                spawner.metrics_collection_task(i, worker_rx, Arc::clone(&self.task_metrics));
+                spawner.metrics_collection_task(i, worker_rx);
                 spawner.work_unit_feed_task(Arc::clone(ctx), i, tx)?;
             }
 
-            Ok(Transformed::yes(plan.with_input_stage(Stage {
-                query_id: stage.query_id,
-                num: stage.num,
-                plan: None,
-                tasks,
-            })?))
+            Ok(Transformed::yes(plan.with_input_stage(Stage::Remote(
+                RemoteStage {
+                    query_id: stage.query_id,
+                    num: stage.num,
+                    workers,
+                },
+            ))?))
         })?;
         Ok(PreparedPlan {
             plan: prepared.data,
@@ -367,32 +369,31 @@ struct CoordinatorToWorkerTaskSpawner<'a> {
     stage_id: usize,
     task_count: usize,
     metrics: &'a CoordinatorToWorkerMetrics,
+    task_metrics: &'a Arc<MetricsStore>,
     join_set: &'a mut JoinSet<Result<()>>,
 }
 
 impl<'a> CoordinatorToWorkerTaskSpawner<'a> {
     /// Builds a new [CoordinatorToWorkerTaskSpawner] based on the [Stage] that needs to be
     /// fanned out to multiple workers.
     fn new(
-        stage: &'a Stage,
+        stage: &'a LocalStage,
         metrics: &'a CoordinatorToWorkerMetrics,
+        task_metrics: &'a Arc<MetricsStore>,
         codec: &'a dyn PhysicalExtensionCodec,
         join_set: &'a mut JoinSet<Result<()>>,
     ) -> Result<Self> {
-        let Some(plan) = &stage.plan else {
-            return internal_err!("Plan is not set for stage {}", stage.num);
-        };
-
-        let plan_proto =
-            PhysicalPlanNode::try_from_physical_plan(Arc::clone(plan), codec)?.encode_to_vec();
+        let plan_proto = PhysicalPlanNode::try_from_physical_plan(Arc::clone(&stage.plan), codec)?
+            .encode_to_vec();
 
         Ok(Self {
-            plan,
+            plan: &stage.plan,
             plan_proto,
             query_id: stage.query_id,
             stage_id: stage.num,
-            task_count: stage.tasks.len(),
+            task_count: stage.tasks,
             metrics,
+            task_metrics,
             join_set,
         })
     }
@@ -520,13 +521,13 @@ impl<'a> CoordinatorToWorkerTaskSpawner<'a> {
         &mut self,
         task_i: usize,
         mut worker_to_coordinator_rx: WorkerResponseRx,
-        task_metrics_collection: Arc<MetricsStore>,
     ) {
         let task_key = TaskKey {
             query_id: serialize_uuid(&self.query_id),
             stage_id: self.stage_id as u64,
             task_number: task_i as u64,
         };
+        let task_metrics_collection = Arc::clone(self.task_metrics);
         tokio::spawn(async move {
             while let Some(Ok(msg)) = worker_to_coordinator_rx.recv().await {
                 let Some(worker_to_coordinator_msg::Inner::TaskMetrics(pre_order_metrics)) =