feat: Add generic interface for custom data sinks (#4244)

## Changes Made

1. Add a generic `DataSink` interface. Users can use this to write
custom write sinks that have optional `.start()`, `.write()`,
`.finalize()` methods.
2. Add `DataFrame.write()` that takes in a custom data sink and writes
the dataframe using the sink.
3. Add `LanceDataSink` as an example of implementing `DataSink` for
LanceDB.
4. Modify `DataFrame.write_lance()` to use this new `LanceDataSink`.

## Checklist

- [x] Documented in API Docs (if applicable)
- [x] Documented in User Guide (if applicable)
- [x] If adding a new documentation page, doc is added to
`docs/mkdocs.yml` navigation
- [x] Documentation builds and is formatted properly (tag @/ccmao1130
for docs review)

Author: desmondcheongzx

daft/daft/__init__.pyi

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Iterator, Litera
 from daft.catalog import Catalog, Table
 from daft.dataframe.display import MermaidOptions
 from daft.execution import physical_plan
+from daft.io import DataSink
 from daft.io.scan import ScanOperator
 from daft.plan_scheduler.physical_plan_scheduler import PartitionT
 from daft.runners.partitioning import PartitionCacheEntry
@@ -1821,6 +1822,7 @@ class LogicalPlanBuilder:
         io_config: IOConfig | None = None,
         kwargs: dict[str, Any] | None = None,
     ) -> LogicalPlanBuilder: ...
+    def datasink_write(self, name: str, sink: DataSink) -> LogicalPlanBuilder: ...
     def schema(self) -> PySchema: ...
     def describe(self) -> LogicalPlanBuilder: ...
     def summarize(self) -> LogicalPlanBuilder: ...

daft/dataframe/dataframe.py

@@ -54,7 +54,7 @@
     import ray
     import torch
 
-    from daft.io import DataCatalogTable
+    from daft.io import DataCatalogTable, DataSink
     from daft.unity_catalog import UnityCatalogTable
 
 if sys.version_info < (3, 10):
@@ -65,8 +65,8 @@
 from daft.schema import Schema
 
 UDFReturnType = TypeVar("UDFReturnType", covariant=True)
-
 T = TypeVar("T")
+R = TypeVar("R")
 P = ParamSpec("P")
 
 
@@ -1178,6 +1178,35 @@ def _create_metadata_param(metadata: Optional[Dict[str, str]]):
 
         return with_operations
 
+    @DataframePublicAPI
+    def write_sink(self, sink: "DataSink[T]") -> "DataFrame":
+        """Writes the DataFrame to the given DataSink.
+
+        Args:
+            sink: The DataSink to write to.
+
+        Returns:
+            DataFrame: A dataframe from the micropartition returned by the DataSink's `.finalize()` method.
+        """
+        sink.start()
+
+        builder = self._builder.write_datasink(sink.name(), sink)
+        write_df = DataFrame(builder)
+        write_df.collect()
+
+        results = write_df.to_pydict()
+        assert "write_results" in results
+        micropartition = sink.finalize(results["write_results"])
+        if micropartition.schema() != sink.schema():
+            raise ValueError(
+                f"Schema mismatch between the data sink's schema and the result's schema:\nSink schema:\n{sink.schema()}\nResult schema:\n{micropartition.schema()}"
+            )
+        # TODO(desmond): Connect the old and new logical plan builders so that a .explain() shows the
+        # plan from the source all the way to the sink to the sink's results. In theory we can do this
+        # for all other sinks too.
+        write_plan_builder = to_logical_plan_builder(micropartition)
+        return DataFrame(write_plan_builder)
+
     @DataframePublicAPI
     def write_lance(
         self,
@@ -1239,75 +1268,10 @@ def write_lance(
             <BLANKLINE>
             (Showing first 1 of 1 rows)
         """
-        from daft import from_pydict
-        from daft.io.object_store_options import io_config_to_storage_options
-
-        try:
-            import lance
-            import pyarrow as pa
+        from daft.dataframe.lance_data_sink import LanceDataSink
 
-        except ImportError:
-            raise ImportError("lance is not installed. Please install lance using `pip install daft[lance]`")
-
-        io_config = get_context().daft_planning_config.default_io_config if io_config is None else io_config
-
-        if isinstance(uri, (str, pathlib.Path)):
-            if isinstance(uri, str):
-                table_uri = uri
-            elif isinstance(uri, pathlib.Path):
-                table_uri = str(uri)
-            else:
-                table_uri = uri
-        pyarrow_schema = pa.schema((f.name, f.dtype.to_arrow_dtype()) for f in self.schema())
-
-        storage_options = io_config_to_storage_options(io_config, table_uri)
-
-        try:
-            table = lance.dataset(table_uri, storage_options=storage_options)
-
-        except ValueError:
-            table = None
-
-        version = 0
-        if table:
-            table_schema = table.schema
-            version = table.latest_version
-            if pyarrow_schema != table_schema and not (mode == "overwrite"):
-                raise ValueError(
-                    "Schema of data does not match table schema\n"
-                    f"Data schema:\n{pyarrow_schema}\nTable Schema:\n{table_schema}"
-                )
-
-        builder = self._builder.write_lance(
-            table_uri,
-            mode,
-            io_config=io_config,
-            kwargs=kwargs,
-        )
-        write_df = DataFrame(builder)
-        write_df.collect()
-
-        write_result = write_df.to_pydict()
-        assert "fragments" in write_result
-        fragments = write_result["fragments"]
-
-        if mode == "create" or mode == "overwrite":
-            operation = lance.LanceOperation.Overwrite(pyarrow_schema, fragments)
-        elif mode == "append":
-            operation = lance.LanceOperation.Append(fragments)
-
-        dataset = lance.LanceDataset.commit(table_uri, operation, read_version=version, storage_options=storage_options)
-        stats = dataset.stats.dataset_stats()
-
-        tbl = from_pydict(
-            {
-                "num_fragments": pa.array([stats["num_fragments"]], type=pa.int64()),
-                "num_deleted_rows": pa.array([stats["num_deleted_rows"]], type=pa.int64()),
-                "num_small_files": pa.array([stats["num_small_files"]], type=pa.int64()),
-                "version": pa.array([dataset.version], type=pa.int64()),
-            }
-        )
-        return tbl
+        sink = LanceDataSink(uri, self.schema(), mode, io_config, **kwargs)
+        return self.write_sink(sink)
 
     ###
     # DataFrame operations

daft/dataframe/lance_data_sink.py

@@ -0,0 +1,126 @@
+import pathlib
+from itertools import chain
+from typing import Iterator, List, Literal, Optional, Union
+
+import lance
+
+from daft.context import get_context
+from daft.daft import IOConfig
+from daft.datatype import DataType
+from daft.io import DataSink, WriteOutput
+from daft.recordbatch import MicroPartition
+from daft.schema import Schema
+
+
+class LanceDataSink(DataSink[list[lance.FragmentMetadata]]):
+    """WriteSink for writing data to a Lance dataset."""
+
+    def _import_lance(self):
+        try:
+            import lance
+
+            return lance
+        except ImportError:
+            raise ImportError("lance is not installed. Please install lance using `pip install daft[lance]`")
+
+    def __init__(
+        self,
+        uri: Union[str, pathlib.Path],
+        schema: Schema,
+        mode: Literal["create", "append", "overwrite"],
+        io_config: Optional[IOConfig] = None,
+        **kwargs,
+    ):
+        from daft.dependencies import pa
+        from daft.io.object_store_options import io_config_to_storage_options
+
+        lance = self._import_lance()
+
+        if not isinstance(uri, (str, pathlib.Path)):
+            raise TypeError(f"Expected URI to be str or pathlib.Path, got {type(uri)}")
+        self._table_uri = str(uri)
+        self._mode = mode
+        self._io_config = get_context().daft_planning_config.default_io_config if io_config is None else io_config
+        self._args = kwargs
+
+        self._storage_options = io_config_to_storage_options(self._io_config, self._table_uri)
+
+        self._pyarrow_schema = pa.schema((f.name, f.dtype.to_arrow_dtype()) for f in schema)
+
+        try:
+            table = lance.dataset(self._table_uri, storage_options=self._storage_options)
+
+        except ValueError:
+            table = None
+
+        self._version = 0
+        if table:
+            table_schema = table.schema
+            self._version = table.latest_version
+            if self._pyarrow_schema != table_schema and not (self._mode == "overwrite"):
+                raise ValueError(
+                    "Schema of data does not match table schema\n"
+                    f"Data schema:\n{self._pyarrow_schema}\nTable Schema:\n{table_schema}"
+                )
+
+        self._schema = Schema._from_field_name_and_types(
+            [
+                ("num_fragments", DataType.int64()),
+                ("num_deleted_rows", DataType.int64()),
+                ("num_small_files", DataType.int64()),
+                ("version", DataType.int64()),
+            ]
+        )
+
+    def schema(self) -> Schema:
+        return self._schema
+
+    def write(self, micropartitions: Iterator[MicroPartition]) -> Iterator[WriteOutput[list[lance.FragmentMetadata]]]:
+        """Writes fragments from the given micropartitions."""
+        lance = self._import_lance()
+
+        for micropartition in micropartitions:
+            arrow_table = micropartition.to_arrow()
+            bytes_written = arrow_table.nbytes
+            rows_written = arrow_table.num_rows
+
+            fragments = lance.fragment.write_fragments(
+                arrow_table,
+                dataset_uri=self._table_uri,
+                mode=self._mode,
+                storage_options=self._storage_options,
+                **self._args,
+            )
+            yield WriteOutput(
+                output=fragments,
+                bytes_written=bytes_written,
+                rows_written=rows_written,
+            )
+
+    def finalize(self, write_outputs: List[WriteOutput[list[lance.FragmentMetadata]]]) -> MicroPartition:
+        """Commits the fragments to the Lance dataset. Returns a DataFrame with the stats of the dataset."""
+        from daft.dependencies import pa
+
+        lance = self._import_lance()
+
+        fragments = list(chain.from_iterable(write_output.output for write_output in write_outputs))
+
+        if self._mode == "create" or self._mode == "overwrite":
+            operation = lance.LanceOperation.Overwrite(self._pyarrow_schema, fragments)
+        elif self._mode == "append":
+            operation = lance.LanceOperation.Append(fragments)
+
+        dataset = lance.LanceDataset.commit(
+            self._table_uri, operation, read_version=self._version, storage_options=self._storage_options
+        )
+        stats = dataset.stats.dataset_stats()
+
+        tbl = MicroPartition.from_pydict(
+            {
+                "num_fragments": pa.array([stats["num_fragments"]], type=pa.int64()),
+                "num_deleted_rows": pa.array([stats["num_deleted_rows"]], type=pa.int64()),
+                "num_small_files": pa.array([stats["num_small_files"]], type=pa.int64()),
+                "version": pa.array([dataset.version], type=pa.int64()),
+            }
+        )
+        return tbl

daft/execution/execution_step.py

@@ -5,17 +5,18 @@
 from typing import TYPE_CHECKING, Generic, Protocol
 
 from daft.context import get_context
-from daft.daft import JoinSide, ResourceRequest
+from daft.daft import JoinSide, PyRecordBatch, ResourceRequest
 from daft.expressions import Expression, ExpressionsProjection, col
 from daft.filesystem import overwrite_files
-from daft.recordbatch import MicroPartition, recordbatch_io
+from daft.recordbatch import MicroPartition, RecordBatch, recordbatch_io
 from daft.runners.partitioning import (
     Boundaries,
     MaterializedResult,
     PartialPartitionMetadata,
     PartitionMetadata,
     PartitionT,
 )
+from daft.series import Series
 
 if TYPE_CHECKING:
     import pathlib
@@ -24,6 +25,7 @@
     from pyiceberg.table import TableProperties as IcebergTableProperties
 
     from daft.daft import FileFormat, IOConfig, JoinType, ScanTask
+    from daft.io import DataSink
     from daft.logical.map_partition_ops import MapPartitionOp
     from daft.logical.schema import Schema
 
@@ -578,6 +580,30 @@ def _handle_file_write(self, input: MicroPartition) -> MicroPartition:
         )
 
 
+@dataclass(frozen=True)
+class DataSinkWrite(SingleOutputInstruction):
+    sink: DataSink
+
+    def run(self, inputs: list[MicroPartition]) -> list[MicroPartition]:
+        result_field_name = "write_results"
+        results = list(self.sink.write(iter(inputs)))
+        results_series = Series.from_pylist(results, result_field_name, pyobj="force")
+        series_dict = {result_field_name: results_series._series}
+        rb = RecordBatch._from_pyrecordbatch(PyRecordBatch.from_pylist_series(series_dict))
+        mp = MicroPartition._from_record_batches([rb])
+        return [mp]
+
+    def run_partial_metadata(self, input_metadatas: list[PartialPartitionMetadata]) -> list[PartialPartitionMetadata]:
+        # TODO(desmond): We can potentially do something more useful here. For now, copy the implementation for the other writers.
+        assert len(input_metadatas) == 1
+        return [
+            PartialPartitionMetadata(
+                num_rows=None,  # we can write more than 1 file per partition
+                size_bytes=None,
+            )
+        ]
+
+
 @dataclass(frozen=True)
 class Filter(SingleOutputInstruction):
     predicate: ExpressionsProjection

daft/execution/physical_plan.py

@@ -60,6 +60,7 @@
     from pyiceberg.table import TableProperties as IcebergTableProperties
 
     from daft.daft import FileFormat, IOConfig, JoinType
+    from daft.io import DataSink
     from daft.logical.schema import Schema
 
 
@@ -224,6 +225,17 @@ def lance_write(
     )
 
 
+def data_sink_write(
+    child_plan: InProgressPhysicalPlan[PartitionT],
+    sink: DataSink,
+) -> InProgressPhysicalPlan[PartitionT]:
+    """Write the results of `child_plan` into a custom write sink described by `sink`."""
+    yield from (
+        step.add_instruction(execution_step.DataSinkWrite(sink)) if isinstance(step, PartitionTaskBuilder) else step
+        for step in child_plan
+    )
+
+
 def pipeline_instruction(
     child_plan: InProgressPhysicalPlan[PartitionT],
     pipeable_instruction: Instruction,

daft/execution/rust_physical_plan_shim.py

@@ -23,6 +23,7 @@
     from pyiceberg.schema import Schema as IcebergSchema
     from pyiceberg.table import TableProperties as IcebergTableProperties
 
+    from daft.io import DataSink
     from daft.recordbatch import MicroPartition
 
 
@@ -415,3 +416,10 @@ def write_lance(
     kwargs: dict | None,
 ) -> physical_plan.InProgressPhysicalPlan[PartitionT]:
     return physical_plan.lance_write(input, path, mode, io_config, kwargs)
+
+
+def write_data_sink(
+    input: physical_plan.InProgressPhysicalPlan[PartitionT],
+    sink: DataSink,
+) -> physical_plan.InProgressPhysicalPlan[PartitionT]:
+    return physical_plan.data_sink_write(input, sink)