[PERF] Swordfish Dynamic Pipelines (#2599)

src/daft-core/src/array/growable/mod.rs

@@ -43,7 +43,7 @@ use crate::datatypes::PythonArray;
 /// * `capacity` - Helps pre-allocate memory to the [`Growable`] by providing a capacity up-front. Note that variable-length types
 ///     such as [`ListArray`] only understands this as the "top-level" capacity, but the capacity of nested children arrays cannot be specified
 ///     through this [`make_growable`] API. Instead, you may wish to instantiate and use the [`nested_growable::ListGrowable`] directly if
-///     this is important to your use0-case.
+///     this is important to your use-case.
 pub fn make_growable<'a>(
     name: &str,
     dtype: &DataType,

src/daft-local-execution/Cargo.toml

@@ -1,5 +1,6 @@
 [dependencies]
 async-stream = {workspace = true}
+async-trait = {workspace = true}
 common-error = {path = "../common/error", default-features = false}
 common-tracing = {path = "../common/tracing", default-features = false}
 daft-core = {path = "../daft-core", default-features = false}
@@ -14,7 +15,6 @@ daft-plan = {path = "../daft-plan", default-features = false}
 daft-scan = {path = "../daft-scan", default-features = false}
 daft-stats = {path = "../daft-stats", default-features = false}
 daft-table = {path = "../daft-table", default-features = false}
-dyn-clone = {workspace = true}
 futures = {workspace = true}
 lazy_static = {workspace = true}
 pyo3 = {workspace = true, optional = true}

src/daft-local-execution/src/intermediate_ops/intermediate_op.rs

@@ -1,23 +1,35 @@
-use std::sync::Arc;
+use std::{env, sync::Arc};
 
 use common_error::DaftResult;
 use daft_micropartition::MicroPartition;
 use tracing::{info_span, instrument};
 
+use async_trait::async_trait;
+
 use crate::{
     channel::{
         create_channel, create_single_channel, spawn_compute_task, MultiReceiver, MultiSender,
         SingleReceiver, SingleSender,
     },
+    pipeline::PipelineNode,
     DEFAULT_MORSEL_SIZE, NUM_CPUS,
 };
 
-pub trait IntermediateOperator: dyn_clone::DynClone + Send + Sync {
+pub trait IntermediateOperator: Send + Sync {
     fn execute(&self, input: &Arc<MicroPartition>) -> DaftResult<Arc<MicroPartition>>;
+    #[allow(dead_code)]
+
     fn name(&self) -> &'static str;
 }
 
-dyn_clone::clone_trait_object!(IntermediateOperator);
+/// The number of rows that will trigger an intermediate operator to output its data.
+#[allow(dead_code)]
+fn get_output_threshold() -> usize {
+    env::var("OUTPUT_THRESHOLD")
+        .unwrap_or_else(|_| "1000".to_string())
+        .parse()
+        .expect("OUTPUT_THRESHOLD must be a number")
+}
 
 /// State of an operator task, used to buffer data and output it when a threshold is reached.
 pub struct OperatorTaskState {
@@ -55,6 +67,10 @@ impl OperatorTaskState {
         if self.buffer.is_empty() {
             return None;
         }
+        assert!(
+            !self.buffer.is_empty(),
+            "We can not run concat with no data"
+        );
         let concated =
             MicroPartition::concat(&self.buffer.iter().map(|x| x.as_ref()).collect::<Vec<_>>())
                 .map(Arc::new);
@@ -70,12 +86,12 @@ impl OperatorTaskState {
 pub struct IntermediateOpActor {
     sender: MultiSender,
     receiver: MultiReceiver,
-    op: Box<dyn IntermediateOperator>,
+    op: Arc<dyn IntermediateOperator>,
 }
 
 impl IntermediateOpActor {
     pub fn new(
-        op: Box<dyn IntermediateOperator>,
+        op: Arc<dyn IntermediateOperator>,
         receiver: MultiReceiver,
         sender: MultiSender,
     ) -> Self {
@@ -91,7 +107,7 @@ impl IntermediateOpActor {
     async fn run_single(
         mut receiver: SingleReceiver,
         sender: SingleSender,
-        op: Box<dyn IntermediateOperator>,
+        op: Arc<dyn IntermediateOperator>,
     ) -> DaftResult<()> {
         let mut state = OperatorTaskState::new();
         let span = info_span!("IntermediateOp::execute");
@@ -139,11 +155,51 @@ impl IntermediateOpActor {
     }
 }
 
-pub fn run_intermediate_op(op: Box<dyn IntermediateOperator>, send_to: MultiSender) -> MultiSender {
-    let (sender, receiver) = create_channel(*NUM_CPUS, send_to.in_order());
-    let mut actor = IntermediateOpActor::new(op, receiver, send_to);
-    tokio::spawn(async move {
-        let _ = actor.run_parallel().await;
-    });
-    sender
+pub(crate) struct IntermediateNode {
+    intermediate_op: Arc<dyn IntermediateOperator>,
+    children: Vec<Box<dyn PipelineNode>>,
+}
+
+impl IntermediateNode {
+    pub(crate) fn new(
+        intermediate_op: Arc<dyn IntermediateOperator>,
+        children: Vec<Box<dyn PipelineNode>>,
+    ) -> Self {
+        IntermediateNode {
+            intermediate_op,
+            children,
+        }
+    }
+    pub(crate) fn boxed(self) -> Box<dyn PipelineNode> {
+        Box::new(self)
+    }
+}
+
+#[async_trait]
+impl PipelineNode for IntermediateNode {
+    fn children(&self) -> Vec<&dyn PipelineNode> {
+        self.children.iter().map(|v| v.as_ref()).collect()
+    }
+
+    async fn start(&mut self, destination: MultiSender) -> DaftResult<()> {
+        assert_eq!(
+            self.children.len(),
+            1,
+            "we only support 1 child for Intermediate Node for now"
+        );
+
+        let (sender, receiver) = create_channel(*NUM_CPUS, destination.in_order());
+
+        let child = self
+            .children
+            .get_mut(0)
+            .expect("we should only have 1 child");
+        child.start(sender).await?;
+
+        let mut actor =
+            IntermediateOpActor::new(self.intermediate_op.clone(), receiver, destination);
+        // this should ideally be in the actor
+        spawn_compute_task(async move { actor.run_parallel().await });
+        Ok(())
+    }
 }

src/daft-local-execution/src/lib.rs

@@ -1,3 +1,4 @@
+#![feature(let_chains)]
 mod channel;
 mod intermediate_ops;
 mod pipeline;