WIP commit for execution functionality

Author: robtandy

src/common/mod.rs

@@ -0,0 +1,2 @@
+pub mod result;
+pub mod util;

src/common/result.rs

@@ -0,0 +1 @@
+

src/common/util.rs

@@ -0,0 +1,36 @@
+use datafusion::error::Result;
+use datafusion::physical_plan::{displayable, ExecutionPlan, ExecutionPlanProperties};
+
+use std::fmt::Write;
+
+pub fn display_plan_with_partition_in_out(plan: &dyn ExecutionPlan) -> Result<String> {
+    let mut f = String::new();
+
+    fn visit(plan: &dyn ExecutionPlan, indent: usize, f: &mut String) -> Result<()> {
+        let output_partitions = plan.output_partitioning().partition_count();
+        let input_partitions = plan
+            .children()
+            .first()
+            .map(|child| child.output_partitioning().partition_count());
+
+        write!(
+            f,
+            "partitions [out:{:<3}{}]{} {}",
+            output_partitions,
+            input_partitions
+                .map(|p| format!("<-- in:{:<3}", p))
+                .unwrap_or("          ".to_string()),
+            " ".repeat(indent),
+            displayable(plan).one_line()
+        )?;
+
+        plan.children()
+            .iter()
+            .try_for_each(|input| visit(input.as_ref(), indent + 2, f))?;
+
+        Ok(())
+    }
+
+    visit(plan, 0, &mut f)?;
+    Ok(f)
+}

src/exec/distributed.rs

@@ -0,0 +1,98 @@
+use std::{any::Any, fmt::Formatter, pin::Pin, sync::Arc};
+
+use datafusion::{
+    execution::{RecordBatchStream, TaskContext},
+    physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan},
+};
+
+use crate::{remote::WorkerDiscovery, stage::ExecutionStage, task::WorkerAddr};
+
+/// This execution plan provides the capability to execute its input plan in
+/// a distributed manner across multiple nodes in a cluster.
+///
+/// When execute() is called, it will use the provided `WorkerDiscovery` to
+/// discover worker nodes and distribute the execution of the input plan
+/// across those nodes.
+#[derive(Debug)]
+pub struct DistributedExec {
+    /// This is the unmodified input plan that this execution plan will execute
+    /// We keep it for display purporses
+    input: Arc<dyn ExecutionPlan>,
+    /// This is the worker discovery mechanism that will be used to provide
+    /// worker addresses
+    discovery: Arc<dyn WorkerDiscovery>,
+    /// This contains stages as defined by the `crate::stage::planner::StageOptimizer`
+    stage: Arc<ExecutionStage>,
+}
+
+impl DistributedExec {
+    /// Creates a new `DistributedExec` with the given input plan.
+    pub fn new(
+        input: Arc<dyn ExecutionPlan>,
+        discovery: Arc<dyn WorkerDiscovery>,
+        stage: Arc<ExecutionStage>,
+    ) -> Self {
+        DistributedExec {
+            input,
+            discovery,
+            stage,
+        }
+    }
+}
+
+impl DisplayAs for DistributedExec {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+        match t {
+            DisplayFormatType::Default => write!(f, "DistributedExec: {}", self.input.name()),
+            DisplayFormatType::Verbose => write!(
+                f,
+                "DistributedExec: input = {}, properties = {:?}",
+                self.input.name(),
+                self.input.properties()
+            ),
+            DisplayFormatType::TreeRender => write!(f, ""),
+        }
+    }
+}
+
+impl ExecutionPlan for DistributedExec {
+    fn name(&self) -> &str {
+        "DistributedExec"
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn properties(&self) -> &datafusion::physical_plan::PlanProperties {
+        &self.input.properties()
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![&self.input]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
+        if children.len() != 1 {
+            return Err(datafusion::error::DataFusionError::Internal(
+                "DistributedExec requires exactly one child".to_string(),
+            ));
+        }
+        Ok(Arc::new(DistributedExec {
+            input: children[0].clone(),
+            discovery: self.discovery.clone(),
+            stage: self.stage.clone(),
+        }))
+    }
+
+    fn execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> datafusion::error::Result<Pin<Box<dyn RecordBatchStream + Send>>> {
+        todo!()
+    }
+}

src/exec/isolator.rs

@@ -0,0 +1,112 @@
+use std::{fmt::Formatter, sync::Arc};
+
+use datafusion::{
+    common::internal_datafusion_err,
+    error::Result,
+    execution::SendableRecordBatchStream,
+    physical_plan::{
+        DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan,
+        ExecutionPlanProperties, Partitioning, PlanProperties,
+    },
+};
+
+// TODO: REFACTOR NOT COMPLETE, just copied here and gutted uneccessary parts for now
+// IGNORE
+
+/// This is a simple execution plan that isolates a partition from the input
+/// plan It will advertise that it has a single partition and when
+/// asked to execute, it will execute a particular partition from the child
+/// input plan.
+///
+/// This allows us to execute Repartition Exec's on different processes
+/// by showing each one only a single child partition
+#[derive(Debug)]
+pub struct PartitionIsolatorExec {
+    pub input: Arc<dyn ExecutionPlan>,
+    properties: PlanProperties,
+    pub partition_count: usize,
+}
+
+impl PartitionIsolatorExec {
+    pub fn new(input: Arc<dyn ExecutionPlan>, partition_count: usize) -> Self {
+        // We advertise that we only have partition_count partitions
+        let properties = input
+            .properties()
+            .clone()
+            .with_partitioning(Partitioning::UnknownPartitioning(partition_count));
+
+        Self {
+            input,
+            properties,
+            partition_count,
+        }
+    }
+}
+
+impl DisplayAs for PartitionIsolatorExec {
+    fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+        write!(
+            f,
+            "PartitionIsolatorExec [providing upto {} partitions]",
+            self.partition_count
+        )
+    }
+}
+
+impl ExecutionPlan for PartitionIsolatorExec {
+    fn name(&self) -> &str {
+        "PartitionIsolatorExec"
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn properties(&self) -> &PlanProperties {
+        &self.properties
+    }
+
+    fn children(&self) -> Vec<&std::sync::Arc<dyn ExecutionPlan>> {
+        vec![&self.input]
+    }
+
+    fn with_new_children(
+        self: std::sync::Arc<Self>,
+        children: Vec<std::sync::Arc<dyn ExecutionPlan>>,
+    ) -> Result<std::sync::Arc<dyn ExecutionPlan>> {
+        // TODO: generalize this
+        assert_eq!(children.len(), 1);
+        Ok(Arc::new(Self::new(
+            children[0].clone(),
+            self.partition_count,
+        )))
+    }
+
+    fn execute(
+        &self,
+        partition: usize,
+        context: std::sync::Arc<datafusion::execution::TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        let config = context.session_config();
+        let partition_group = &[0, 1];
+
+        let partitions_in_input = self.input.output_partitioning().partition_count() as u64;
+
+        let output_stream = match partition_group.get(partition) {
+            Some(actual_partition_number) => {
+                if *actual_partition_number >= partitions_in_input {
+                    //trace!("{} returning empty stream", ctx_name);
+                    Ok(Box::pin(EmptyRecordBatchStream::new(self.input.schema()))
+                        as SendableRecordBatchStream)
+                } else {
+                    //trace!("{} returning actual stream", ctx_name);
+                    self.input
+                        .execute(*actual_partition_number as usize, context)
+                }
+            }
+            None => Ok(Box::pin(EmptyRecordBatchStream::new(self.input.schema()))
+                as SendableRecordBatchStream),
+        };
+        output_stream
+    }
+}

src/exec/mod.rs

@@ -0,0 +1,2 @@
+mod distributed;
+mod isolator;

src/lib.rs

@@ -1,13 +1,16 @@
 mod channel_manager;
-mod composed_extension_codec;
-mod flight_service;
+mod common;
+mod exec;
+//mod flight_service;
 mod plan;
-mod stage_delegation;
-#[cfg(test)]
-pub mod test_utils;
+mod remote;
+mod stage;
+pub mod task;
+mod test_utils;
 
 pub use channel_manager::{
     ArrowFlightChannel, BoxCloneSyncChannel, ChannelManager, ChannelResolver,
 };
-pub use flight_service::{ArrowFlightEndpoint, SessionBuilder};
+
 pub use plan::ArrowFlightReadExec;
+pub use stage::{display_stage, display_stage_tree, ExecutionStage, StagePlanner};