WIP commit for execution functionality

Author: robtandy

Cargo.lock

@@ -8,6 +8,7 @@ datafusion = { version = "48.0.0" }
 datafusion-proto = { version = "48.0.0" }
 arrow-flight = "55.2.0"
 async-trait = "0.1.88"
+itertools = "0.14.0"
 tokio = { version = "1.46.1", features = ["full"] }
 # Fixed to 0.12.3 because of arrow-flight
 tonic = { version = "0.12.3", features = ["transport"] }
@@ -21,4 +22,5 @@ dashmap = "6.1.0"
 prost = "0.13.5"
 
 [dev-dependencies]
-insta = { version = "1.43.1" , features = ["filters"]}
+insta = { version = "1.43.1", features = ["filters"] }
+

Cargo.toml

@@ -4,13 +4,14 @@ mod exec;
 //mod flight_service;
 mod plan;
 mod remote;
-mod stage;
-pub mod task;
+#[cfg(test)]
 mod test_utils;
 
 pub use channel_manager::{
     ArrowFlightChannel, BoxCloneSyncChannel, ChannelManager, ChannelResolver,
 };
 
+pub mod physical_optimizer;
+pub mod stage;
+pub mod task;
 pub use plan::ArrowFlightReadExec;
-pub use stage::{display_stage, display_stage_tree, ExecutionStage, StagePlanner};

src/lib.rs

@@ -10,22 +10,97 @@ use datafusion::{
     config::ConfigOptions,
     datasource::physical_plan::{FileScanConfig, FileSource},
     error::Result,
+    physical_optimizer::PhysicalOptimizerRule,
     physical_plan::{
         displayable, execution_plan::need_data_exchange, ExecutionPlan, ExecutionPlanProperties,
     },
 };
+use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 
 use crate::ArrowFlightReadExec;
 
 use super::stage::ExecutionStage;
 
+#[derive(Debug, Default)]
+pub struct DistributedPhysicalOptimizerRule {
+    /// Optional codec to assist in serializing and deserializing any custom
+    /// ExecutionPlan nodes
+    codec: Option<Arc<dyn PhysicalExtensionCodec>>,
+    /// maximum number of partitions per task. This is used to determine how many
+    /// tasks to create for each stage
+    partitions_per_task: Option<usize>,
+}
+
+impl DistributedPhysicalOptimizerRule {
+    pub fn new() -> Self {
+        DistributedPhysicalOptimizerRule {
+            codec: None,
+            partitions_per_task: None,
+        }
+    }
+
+    /// Set a codec to use to assist in serializing and deserializing
+    /// custom ExecutionPlan nodes.
+    pub fn with_codec(mut self, codec: Arc<dyn PhysicalExtensionCodec>) -> Self {
+        self.codec = Some(codec);
+        self
+    }
+
+    /// Set the maximum number of partitions per task. This is used to determine how many
+    /// tasks to create for each stage.
+    ///
+    /// If a stage holds a plan with 10 partitions, and this is set to 3,
+    /// then the stage will be split into 4 tasks:
+    /// - Task 1: partitions 0, 1, 2
+    /// - Task 2: partitions 3, 4, 5
+    /// - Task 3: partitions 6, 7, 8
+    /// - Task 4: partitions 9
+    ///
+    /// Each task will be executed on a separate host
+    pub fn with_maximum_partitions_per_task(mut self, partitions_per_task: usize) -> Self {
+        self.partitions_per_task = Some(partitions_per_task);
+        self
+    }
+}
+
+impl PhysicalOptimizerRule for DistributedPhysicalOptimizerRule {
+    fn optimize(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        _config: &ConfigOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        // We can only optimize plans that are not already distributed
+        if plan.as_any().is::<ExecutionStage>() {
+            return Ok(plan);
+        }
+        println!(
+            "DistributedPhysicalOptimizerRule: optimizing plan: {}",
+            displayable(plan.as_ref()).indent(false)
+        );
+
+        let mut planner = StagePlanner::new(self.codec.clone(), self.partitions_per_task);
+        plan.rewrite(&mut planner)?;
+        planner
+            .finish()
+            .map(|stage| stage as Arc<dyn ExecutionPlan>)
+    }
+
+    fn name(&self) -> &str {
+        "DistributedPhysicalOptimizer"
+    }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
+}
+
 /// StagePlanner is a TreeNodeRewriter that walks the plan tree and creates
 /// a tree of ExecutionStage nodes that represent discrete stages of execution
 /// can are separated by a data shuffle.
 ///
 /// See https://howqueryengineswork.com/13-distributed-query.html for more information
 /// about distributed execution.
-pub struct StagePlanner {
+struct StagePlanner {
     /// used to keep track of the current plan head
     plan_head: Option<Arc<dyn ExecutionPlan>>,
     /// Current depth in the plan tree, as we walk the tree
@@ -38,48 +113,39 @@ pub struct StagePlanner {
     input_stages: Vec<(usize, Arc<ExecutionStage>)>,
     /// current stage number
     stage_counter: usize,
-}
-
-/// Create an ExecutionStage from a plan.   
-///
-/// The resulting ExecutionStage cannot be executed directly, but is used as input,
-/// together with a vec of Worker Addresses, to further break up the stages into [`ExecutionTask`]s
-/// Which are ultimately what facilitates distributed execution.
-impl TryFrom<Arc<dyn ExecutionPlan>> for ExecutionStage {
-    type Error = DataFusionError;
-    fn try_from(plan: Arc<dyn ExecutionPlan>) -> Result<Self> {
-        let mut planner = StagePlanner {
-            plan_head: Some(plan.clone()),
-            depth: 0,
-            input_stages: vec![],
-            stage_counter: 1,
-        };
-
-        plan.rewrite(&mut planner)?;
-        planner.finish().and_then(|arc_stage| {
-            Arc::into_inner(arc_stage).ok_or(internal_datafusion_err!(
-                "Failed to convert Arc<ExecutionStage> to ExecutionStage"
-            ))
-        })
-    }
+    /// Optional codec to assist in serializing and deserializing any custom
+    codec: Option<Arc<dyn PhysicalExtensionCodec>>,
+    /// partitions_per_task is used to determine how many tasks to create for each stage
+    partitions_per_task: Option<usize>,
 }
 
 impl StagePlanner {
-    pub fn new() -> Self {
+    fn new(
+        codec: Option<Arc<dyn PhysicalExtensionCodec>>,
+        partitions_per_task: Option<usize>,
+    ) -> Self {
         StagePlanner {
             plan_head: None,
             depth: 0,
             input_stages: vec![],
-            stage_counter: 0,
+            stage_counter: 1,
+            codec,
+            partitions_per_task,
         }
     }
 
-    pub fn finish(mut self) -> Result<Arc<ExecutionStage>> {
+    fn finish(mut self) -> Result<Arc<ExecutionStage>> {
         if self.input_stages.is_empty() {
-            return internal_err!("No input stages found, did you forget to call rewrite()?");
-        }
-
-        if self.depth < self.input_stages[0].0 {
+            Ok(Arc::new(ExecutionStage::new(
+                self.stage_counter,
+                self.plan_head
+                    .take()
+                    .ok_or_else(|| internal_datafusion_err!("No plan head set"))?,
+                vec![],
+            )))
+        } else if self.depth < self.input_stages[0].0 {
+            // There is more plan above the last stage we created, so we need to
+            // create a new stage that includes the last plan head
             Ok(Arc::new(ExecutionStage::new(
                 self.stage_counter,
                 self.plan_head
@@ -88,10 +154,12 @@ impl StagePlanner {
                 self.input_stages
                     .iter()
                     .map(|(_, stage)| stage.clone())
-                    .collect::<Vec<_>>(),
+                    .collect(),
             )))
         } else {
-            Ok(self.input_stages.remove(0).1)
+            // We have a plan head, and we are at the same depth as the last stage we created,
+            // so we can just return the last stage
+            Ok(self.input_stages.last().unwrap().1.clone())
         }
     }
 }
@@ -142,13 +210,16 @@ impl TreeNodeRewriter for StagePlanner {
 
             self.input_stages.retain(|(depth, _)| *depth <= self.depth);
 
-            let stage = Arc::new(ExecutionStage::new(
-                self.stage_counter,
-                plan.clone(),
-                child_stages,
-            ));
+            let mut stage = ExecutionStage::new(self.stage_counter, plan.clone(), child_stages);
+
+            if let Some(partitions_per_task) = self.partitions_per_task {
+                stage = stage.with_maximum_partitions_per_task(partitions_per_task);
+            }
+            if let Some(codec) = self.codec.as_ref() {
+                stage = stage.with_codec(codec.clone());
+            }
 
-            self.input_stages.push((self.depth, stage));
+            self.input_stages.push((self.depth, Arc::new(stage)));
 
             // As we are walking up the plan tree, we've now put what we've encountered so far
             // into a stage.   We want to replace this plan now with an ArrowFlightReadExec