[FEAT] Streaming Local Writes for Native Executor

Cargo.toml

@@ -20,6 +20,7 @@ daft-local-execution = {path = "src/daft-local-execution", default-features = fa
 daft-micropartition = {path = "src/daft-micropartition", default-features = false}
 daft-minhash = {path = "src/daft-minhash", default-features = false}
 daft-parquet = {path = "src/daft-parquet", default-features = false}
+daft-physical-plan = {path = "src/daft-physical-plan", default-features = false}
 daft-plan = {path = "src/daft-plan", default-features = false}
 daft-scan = {path = "src/daft-scan", default-features = false}
 daft-scheduler = {path = "src/daft-scheduler", default-features = false}
@@ -47,6 +48,7 @@ python = [
   "daft-json/python",
   "daft-micropartition/python",
   "daft-parquet/python",
+  "daft-physical-plan/python",
   "daft-plan/python",
   "daft-scan/python",
   "daft-scheduler/python",

daft/io/writer.py

@@ -0,0 +1,79 @@
+import uuid
+from typing import Optional, Union
+
+from daft.daft import IOConfig
+from daft.dependencies import pa, pacsv, pq
+from daft.filesystem import (
+    _resolve_paths_and_filesystem,
+    canonicalize_protocol,
+    get_protocol_from_path,
+)
+from daft.table.micropartition import MicroPartition
+
+
+class FileWriterBase:
+    def __init__(
+        self,
+        root_dir: str,
+        file_idx: int,
+        file_format: str,
+        compression: Optional[str] = None,
+        io_config: Optional[IOConfig] = None,
+    ):
+        [self.resolved_path], self.fs = _resolve_paths_and_filesystem(root_dir, io_config=io_config)
+        protocol = get_protocol_from_path(root_dir)
+        canonicalized_protocol = canonicalize_protocol(protocol)
+        is_local_fs = canonicalized_protocol == "file"
+        if is_local_fs:
+            self.fs.create_dir(root_dir)
+
+        self.file_name = f"{uuid.uuid4()}-{file_idx}.{file_format}"
+        self.full_path = f"{self.resolved_path}/{self.file_name}"
+        self.compression = compression if compression is not None else "none"
+        self.current_writer: Optional[Union[pq.ParquetWriter, pacsv.CSVWriter]] = None
+
+    def _create_writer(self, schema: pa.Schema):
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    def write(self, table: MicroPartition):
+        if self.current_writer is None:
+            self.current_writer = self._create_writer(table.schema().to_pyarrow_schema())
+        self.current_writer.write_table(table.to_arrow())
+
+    def close(self) -> Optional[str]:
+        if self.current_writer is None:
+            return None
+        self.current_writer.close()
+        return self.full_path
+
+
+class ParquetFileWriter(FileWriterBase):
+    def __init__(
+        self,
+        root_dir: str,
+        file_idx: int,
+        compression: str = "none",
+        io_config: Optional[IOConfig] = None,
+    ):
+        super().__init__(root_dir, file_idx, "parquet", compression, io_config)
+
+    def _create_writer(self, schema: pa.Schema) -> pq.ParquetWriter:
+        return pq.ParquetWriter(
+            self.full_path,
+            schema,
+            compression=self.compression,
+            use_compliant_nested_type=False,
+            filesystem=self.fs,
+        )
+
+
+class CSVFileWriter(FileWriterBase):
+    def __init__(self, root_dir: str, file_idx: int, io_config: Optional[IOConfig] = None):
+        super().__init__(root_dir, file_idx, "csv", None, io_config)
+
+    def _create_writer(self, schema: pa.Schema) -> pacsv.CSVWriter:
+        file_path = f"{self.resolved_path}/{self.file_name}"
+        return pacsv.CSVWriter(
+            file_path,
+            schema,
+        )

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

@@ -0,0 +1,89 @@
+use std::{cmp::Ordering::*, collections::VecDeque, sync::Arc};
+
+use common_error::DaftResult;
+use daft_micropartition::MicroPartition;
+
+pub struct RowBasedBuffer {
+    pub buffer: VecDeque<Arc<MicroPartition>>,
+    pub curr_len: usize,
+    pub threshold: usize,
+}
+
+impl RowBasedBuffer {
+    pub fn new(threshold: usize) -> Self {
+        assert!(threshold > 0);
+        Self {
+            buffer: VecDeque::new(),
+            curr_len: 0,
+            threshold,
+        }
+    }
+
+    pub fn push(&mut self, part: Arc<MicroPartition>) {
+        self.curr_len += part.len();
+        self.buffer.push_back(part);
+    }
+
+    pub fn pop_enough(&mut self) -> DaftResult<Option<Vec<Arc<MicroPartition>>>> {
+        match self.curr_len.cmp(&self.threshold) {
+            Less => Ok(None),
+            Equal => {
+                if self.buffer.len() == 1 {
+                    let part = self.buffer.pop_front().unwrap();
+                    self.curr_len = 0;
+                    Ok(Some(vec![part]))
+                } else {
+                    let chunk = MicroPartition::concat(
+                        &std::mem::take(&mut self.buffer)
+                            .iter()
+                            .map(|x| x.as_ref())
+                            .collect::<Vec<_>>(),
+                    )?;
+                    self.curr_len = 0;
+                    Ok(Some(vec![Arc::new(chunk)]))
+                }
+            }
+            Greater => {
+                let num_ready_chunks = self.curr_len / self.threshold;
+                let concated = MicroPartition::concat(
+                    &std::mem::take(&mut self.buffer)
+                        .iter()
+                        .map(|x| x.as_ref())
+                        .collect::<Vec<_>>(),
+                )?;
+                let mut start = 0;
+                let mut parts_to_return = Vec::with_capacity(num_ready_chunks);
+                for _ in 0..num_ready_chunks {
+                    let end = start + self.threshold;
+                    let part = Arc::new(concated.slice(start, end)?);
+                    parts_to_return.push(part);
+                    start = end;
+                }
+                if start < concated.len() {
+                    let part = Arc::new(concated.slice(start, concated.len())?);
+                    self.curr_len = part.len();
+                    self.buffer.push_back(part);
+                } else {
+                    self.curr_len = 0;
+                }
+                Ok(Some(parts_to_return))
+            }
+        }
+    }
+
+    pub fn pop_all(&mut self) -> DaftResult<Option<Arc<MicroPartition>>> {
+        assert!(self.curr_len < self.threshold);
+        if self.buffer.is_empty() {
+            Ok(None)
+        } else {
+            let concated = MicroPartition::concat(
+                &std::mem::take(&mut self.buffer)
+                    .iter()
+                    .map(|x| x.as_ref())
+                    .collect::<Vec<_>>(),
+            )?;
+            self.curr_len = 0;
+            Ok(Some(Arc::new(concated)))
+        }
+    }
+}

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

@@ -5,8 +5,8 @@
 use daft_micropartition::MicroPartition;
 use tracing::{info_span, instrument};
 
-use super::buffer::OperatorBuffer;
 use crate::{
+    buffer::RowBasedBuffer,
     channel::{create_channel, PipelineChannel, Receiver, Sender},
     pipeline::{PipelineNode, PipelineResultType},
     runtime_stats::{CountingReceiver, CountingSender, RuntimeStatsContext},
@@ -135,26 +135,23 @@
         };
 
         for (idx, mut receiver) in receivers.into_iter().enumerate() {
-            let mut buffer = OperatorBuffer::new(morsel_size);
+            let mut buffer = RowBasedBuffer::new(morsel_size);
             while let Some(morsel) = receiver.recv().await {
                 if morsel.should_broadcast() {
                     for worker_sender in worker_senders.iter() {
                         let _ = worker_sender.send((idx, morsel.clone())).await;
                     }
                 } else {
                     buffer.push(morsel.as_data().clone());
-                    if let Some(ready) = buffer.try_clear() {
-                        let _ = send_to_next_worker(idx, ready?.into()).await;
+                    if let Some(ready) = buffer.pop_enough()? {
+                        for part in ready {
+                            let _ = send_to_next_worker(idx, part.into()).await;
+                        }
                     }
                 }
             }
-            // Buffer may still have some morsels left above the threshold
-            while let Some(ready) = buffer.try_clear() {
-                let _ = send_to_next_worker(idx, ready?.into()).await;
-            }
-            // Clear all remaining morsels
-            if let Some(last_morsel) = buffer.clear_all() {
-                let _ = send_to_next_worker(idx, last_morsel?.into()).await;
+            if let Some(ready) = buffer.pop_all()? {
+                let _ = send_to_next_worker(idx, ready.into()).await;
             }
         }
         Ok(())

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

@@ -1,6 +1,5 @@
 pub mod aggregate;
 pub mod anti_semi_hash_join_probe;
-pub mod buffer;
 pub mod filter;
 pub mod hash_join_probe;
 pub mod intermediate_op;

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

@@ -1,4 +1,5 @@
 #![feature(let_chains)]
+mod buffer;
 mod channel;
 mod intermediate_ops;
 mod pipeline;
@@ -14,15 +15,20 @@ lazy_static! {
     pub static ref NUM_CPUS: usize = std::thread::available_parallelism().unwrap().get();
 }
 
+pub(crate) type TaskSet<T> = tokio::task::JoinSet<T>;
+pub(crate) fn create_task_set<T>() -> TaskSet<T> {
+    TaskSet::new()
+}
+
 pub struct ExecutionRuntimeHandle {
-    worker_set: tokio::task::JoinSet<crate::Result<()>>,
+    worker_set: TaskSet<crate::Result<()>>,
     default_morsel_size: usize,
 }
 
 impl ExecutionRuntimeHandle {
     pub fn new(default_morsel_size: usize) -> Self {
         Self {
-            worker_set: tokio::task::JoinSet::new(),
+            worker_set: create_task_set(),
             default_morsel_size,
         }
     }