[FEAT] Streaming physical writes for native executor

Cargo.toml

@@ -1,4 +1,5 @@
 [dependencies]
+common-buffer = {path = "src/common/buffer", default-features = false}
 common-daft-config = {path = "src/common/daft-config", default-features = false}
 common-display = {path = "src/common/display", default-features = false}
 common-file-formats = {path = "src/common/file-formats", default-features = false}
@@ -21,12 +22,14 @@ 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}
 daft-sql = {path = "src/daft-sql", default-features = false}
 daft-stats = {path = "src/daft-stats", default-features = false}
 daft-table = {path = "src/daft-table", default-features = false}
+daft-writers = {path = "src/daft-writers", default-features = false}
 lazy_static = {workspace = true}
 log = {workspace = true}
 lzma-sys = {version = "*", features = ["static"]}
@@ -48,6 +51,7 @@ python = [
   "daft-json/python",
   "daft-micropartition/python",
   "daft-parquet/python",
+  "daft-physical-plan/python",
   "daft-plan/python",
   "daft-scan/python",
   "daft-scheduler/python",
@@ -56,6 +60,7 @@ python = [
   "daft-table/python",
   "daft-functions/python",
   "daft-functions-json/python",
+  "daft-writers/python",
   "common-daft-config/python",
   "common-system-info/python",
   "common-display/python",
@@ -138,6 +143,7 @@ members = [
   "src/daft-functions",
   "src/daft-functions-json",
   "src/daft-sql",
+  "src/daft-writers",
   "src/hyperloglog"
 ]
 

daft/io/writer.py

@@ -0,0 +1,152 @@
+import uuid
+from abc import ABC, abstractmethod
+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.series import Series
+from daft.table.micropartition import MicroPartition
+from daft.table.partitioning import (
+    partition_strings_to_path,
+    partition_values_to_str_mapping,
+)
+from daft.table.table import Table
+
+
+class FileWriterBase(ABC):
+    def __init__(
+        self,
+        root_dir: str,
+        file_idx: int,
+        file_format: str,
+        partition_values: Optional[Table] = None,
+        compression: Optional[str] = None,
+        io_config: Optional[IOConfig] = None,
+        default_partition_fallback: str = "__HIVE_DEFAULT_PARTITION__",
+    ):
+        [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"
+
+        self.file_name = f"{uuid.uuid4()}-{file_idx}.{file_format}"
+        self.partition_values = partition_values
+        if self.partition_values is not None:
+            partition_strings = {
+                key: values.to_pylist()[0]
+                for key, values in partition_values_to_str_mapping(self.partition_values).items()
+            }
+            self.dir_path = partition_strings_to_path(self.resolved_path, partition_strings, default_partition_fallback)
+        else:
+            self.dir_path = f"{self.resolved_path}"
+
+        self.full_path = f"{self.dir_path}/{self.file_name}"
+        if is_local_fs:
+            self.fs.create_dir(self.dir_path, recursive=True)
+
+        self.compression = compression if compression is not None else "none"
+        self.current_writer: Optional[Union[pq.ParquetWriter, pacsv.CSVWriter]] = None
+
+    @abstractmethod
+    def write(self, table: MicroPartition) -> None:
+        """Write data to the file using the appropriate writer.
+
+        Args:
+            table: MicroPartition containing the data to be written.
+        """
+        pass
+
+    @abstractmethod
+    def close(self) -> Table:
+        """Close the writer and return metadata about the written file.
+
+        Returns:
+            Table containing metadata about the written file, including path and partition values.
+        """
+        pass
+
+
+class ParquetFileWriter(FileWriterBase):
+    def __init__(
+        self,
+        root_dir: str,
+        file_idx: int,
+        partition_values: Optional[Table] = None,
+        compression: str = "none",
+        io_config: Optional[IOConfig] = None,
+    ):
+        super().__init__(
+            root_dir=root_dir,
+            file_idx=file_idx,
+            file_format="parquet",
+            partition_values=partition_values,
+            compression=compression,
+            io_config=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,
+        )
+
+    def write(self, table: MicroPartition) -> None:
+        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) -> Table:
+        if self.current_writer is not None:
+            self.current_writer.close()
+
+        metadata = {"path": Series.from_pylist([self.full_path])}
+        if self.partition_values is not None:
+            for col_name in self.partition_values.column_names():
+                metadata[col_name] = self.partition_values.get_column(col_name)
+        return Table.from_pydict(metadata)
+
+
+class CSVFileWriter(FileWriterBase):
+    def __init__(
+        self,
+        root_dir: str,
+        file_idx: int,
+        partition_values: Optional[Table] = None,
+        io_config: Optional[IOConfig] = None,
+    ):
+        super().__init__(
+            root_dir=root_dir,
+            file_idx=file_idx,
+            file_format="csv",
+            partition_values=partition_values,
+            io_config=io_config,
+        )
+
+    def _create_writer(self, schema: pa.Schema) -> pacsv.CSVWriter:
+        return pacsv.CSVWriter(
+            self.full_path,
+            schema,
+        )
+
+    def write(self, table: MicroPartition) -> None:
+        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) -> Table:
+        if self.current_writer is not None:
+            self.current_writer.close()
+
+        metadata = {"path": Series.from_pylist([self.full_path])}
+        if self.partition_values is not None:
+            for col_name in self.partition_values.column_names():
+                metadata[col_name] = self.partition_values.get_column(col_name)
+        return Table.from_pydict(metadata)

daft/table/partitioning.py

@@ -1,20 +1,42 @@
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from daft import Series
 from daft.expressions import ExpressionsProjection
+from daft.table.table import Table
 
 from .micropartition import MicroPartition
 
 
 def partition_strings_to_path(
-    root_path: str, parts: Dict[str, str], partition_null_fallback: str = "__HIVE_DEFAULT_PARTITION__"
+    root_path: str,
+    parts: Dict[str, str],
+    partition_null_fallback: str = "__HIVE_DEFAULT_PARTITION__",
 ) -> str:
     keys = parts.keys()
     values = [partition_null_fallback if value is None else value for value in parts.values()]
     postfix = "/".join(f"{k}={v}" for k, v in zip(keys, values))
     return f"{root_path}/{postfix}"
 
 
+def partition_values_to_str_mapping(
+    partition_values: Union[MicroPartition, Table],
+) -> Dict[str, Series]:
+    null_part = Series.from_pylist(
+        [None]
+    )  # This is to ensure that the null values are replaced with the default_partition_fallback value
+    pkey_names = partition_values.column_names()
+
+    partition_strings = {}
+
+    for c in pkey_names:
+        column = partition_values.get_column(c)
+        string_names = column._to_str_values()
+        null_filled = column.is_null().if_else(null_part, string_names)
+        partition_strings[c] = null_filled
+
+    return partition_strings
+
+
 class PartitionedTable:
     def __init__(self, table: MicroPartition, partition_keys: Optional[ExpressionsProjection]):
         self.table = table
@@ -56,20 +78,10 @@ def partition_values_str(self) -> Optional[MicroPartition]:
 
         If the table is not partitioned, returns None.
         """
-        null_part = Series.from_pylist([None])
         partition_values = self.partition_values()
 
         if partition_values is None:
             return None
         else:
-            pkey_names = partition_values.column_names()
-
-            partition_strings = {}
-
-            for c in pkey_names:
-                column = partition_values.get_column(c)
-                string_names = column._to_str_values()
-                null_filled = column.is_null().if_else(null_part, string_names)
-                partition_strings[c] = null_filled
-
+            partition_strings = partition_values_to_str_mapping(partition_values)
             return MicroPartition.from_pydict(partition_strings)

src/common/buffer/Cargo.toml

@@ -0,0 +1,10 @@
+[dependencies]
+common-error = {path = "../error", default-features = false}
+
+[lints]
+workspace = true
+
+[package]
+edition = {workspace = true}
+name = "common-buffer"
+version = {workspace = true}