Fix issue 791: add LazyFrame sink classes for parquet, csv, ipc, ndjson

hamilton/plugins/polars_lazyframe_extensions.py

@@ -47,7 +47,7 @@
 
 from hamilton import registry
 from hamilton.io import utils
-from hamilton.io.data_adapters import DataLoader
+from hamilton.io.data_adapters import DataLoader, DataSaver
 
 DATAFRAME_TYPE = pl.LazyFrame
 COLUMN_TYPE = pl.Expr
@@ -183,6 +183,117 @@ def name(cls) -> str:
         return "csv"
 
 
+@dataclasses.dataclass
+class PolarsSinkCSVWriter(DataSaver):
+    """Class to handle sinking a Polars LazyFrame to a CSV file using streaming.
+
+    Calls LazyFrame.sink_csv() directly, avoiding collect() for better performance.
+    Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_csv.html
+
+    Note: ``lazy=True`` is intentionally excluded because ``save_data()`` expects
+    the file to exist immediately after the call returns.
+    """
+
+    file: str | Path
+    # kwargs:
+    include_bom: bool | None = None
+    compression: str | None = None
+    compression_level: int | None = None
+    check_extension: bool | None = None
+    include_header: bool = True
+    separator: str = ","
+    line_terminator: str = "\n"
+    quote_char: str = '"'
+    batch_size: int = 1024
+    datetime_format: str | None = None
+    date_format: str | None = None
+    time_format: str | None = None
+    float_scientific: bool | None = None
+    float_precision: int | None = None
+    decimal_comma: bool | None = None
+    null_value: str | None = None
+    quote_style: Any = None
+    maintain_order: bool = True
+    storage_options: dict[str, Any] | None = None
+    credential_provider: Any = None
+    retries: int | None = None
+    sync_on_close: Any = None
+    mkdir: bool | None = None
+    engine: Any = None
+    optimizations: Any = None
+    extra_kwargs: dict[str, Any] | None = None
+
+    def _get_saving_kwargs(self) -> dict[str, Any]:
+        kwargs = {}
+        if self.include_bom is not None:
+            kwargs["include_bom"] = self.include_bom
+        if self.compression is not None:
+            kwargs["compression"] = self.compression
+        if self.compression_level is not None:
+            kwargs["compression_level"] = self.compression_level
+        if self.check_extension is not None:
+            kwargs["check_extension"] = self.check_extension
+        if self.include_header is not None:
+            kwargs["include_header"] = self.include_header
+        if self.separator is not None:
+            kwargs["separator"] = self.separator
+        if self.line_terminator is not None:
+            kwargs["line_terminator"] = self.line_terminator
+        if self.quote_char is not None:
+            kwargs["quote_char"] = self.quote_char
+        if self.batch_size is not None:
+            kwargs["batch_size"] = self.batch_size
+        if self.datetime_format is not None:
+            kwargs["datetime_format"] = self.datetime_format
+        if self.date_format is not None:
+            kwargs["date_format"] = self.date_format
+        if self.time_format is not None:
+            kwargs["time_format"] = self.time_format
+        if self.float_scientific is not None:
+            kwargs["float_scientific"] = self.float_scientific
+        if self.float_precision is not None:
+            kwargs["float_precision"] = self.float_precision
+        if self.decimal_comma is not None:
+            kwargs["decimal_comma"] = self.decimal_comma
+        if self.null_value is not None:
+            kwargs["null_value"] = self.null_value
+        if self.quote_style is not None:
+            kwargs["quote_style"] = self.quote_style
+        if self.maintain_order is not None:
+            kwargs["maintain_order"] = self.maintain_order
+        if self.storage_options is not None:
+            kwargs["storage_options"] = self.storage_options
+        if self.credential_provider is not None:
+            kwargs["credential_provider"] = self.credential_provider
+        if self.retries is not None:
+            kwargs["retries"] = self.retries
+        if self.sync_on_close is not None:
+            kwargs["sync_on_close"] = self.sync_on_close
+        if self.mkdir is not None:
+            kwargs["mkdir"] = self.mkdir
+        if self.engine is not None:
+            kwargs["engine"] = self.engine
+        if self.optimizations is not None:
+            kwargs["optimizations"] = self.optimizations
+        if self.extra_kwargs is not None:
+            if self.extra_kwargs.get("lazy", False):
+                raise ValueError("lazy=True is incompatible with synchronous data savers.")
+            kwargs.update(self.extra_kwargs)
+        return kwargs
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [DATAFRAME_TYPE]
+
+    def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+        data.sink_csv(self.file, **self._get_saving_kwargs())
+        return utils.get_file_metadata(self.file)
+
+    @classmethod
+    def name(cls) -> str:
+        return "csv"
+
+
 @dataclasses.dataclass
 class PolarsScanParquetReader(DataLoader):
     """Class specifically to handle loading parquet files with polars
@@ -239,6 +350,87 @@ def name(cls) -> str:
         return "parquet"
 
 
+@dataclasses.dataclass
+class PolarsSinkParquetWriter(DataSaver):
+    """Class to handle sinking a Polars LazyFrame to a Parquet file using streaming.
+
+    Calls LazyFrame.sink_parquet() directly, avoiding collect() for better performance.
+    Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_parquet.html
+
+    Note: ``lazy=True`` is intentionally excluded because ``save_data()`` expects
+    the file to exist immediately after the call returns.
+    """
+
+    file: str | Path
+    # kwargs:
+    compression: str = "zstd"
+    compression_level: int | None = None
+    statistics: bool | str | dict[str, bool] = True
+    row_group_size: int | None = None
+    data_page_size: int | None = None
+    maintain_order: bool = True
+    storage_options: dict[str, Any] | None = None
+    credential_provider: Any = None
+    retries: int | None = None
+    sync_on_close: Any = None
+    metadata: Any = None
+    arrow_schema: Any = None
+    mkdir: bool | None = None
+    engine: Any = None
+    optimizations: Any = None
+    extra_kwargs: dict[str, Any] | None = None
+
+    def _get_saving_kwargs(self) -> dict[str, Any]:
+        kwargs = {}
+        if self.compression is not None:
+            kwargs["compression"] = self.compression
+        if self.compression_level is not None:
+            kwargs["compression_level"] = self.compression_level
+        if self.statistics is not None:
+            kwargs["statistics"] = self.statistics
+        if self.row_group_size is not None:
+            kwargs["row_group_size"] = self.row_group_size
+        if self.data_page_size is not None:
+            kwargs["data_page_size"] = self.data_page_size
+        if self.maintain_order is not None:
+            kwargs["maintain_order"] = self.maintain_order
+        if self.storage_options is not None:
+            kwargs["storage_options"] = self.storage_options
+        if self.credential_provider is not None:
+            kwargs["credential_provider"] = self.credential_provider
+        if self.retries is not None:
+            kwargs["retries"] = self.retries
+        if self.sync_on_close is not None:
+            kwargs["sync_on_close"] = self.sync_on_close
+        if self.metadata is not None:
+            kwargs["metadata"] = self.metadata
+        if self.arrow_schema is not None:
+            kwargs["arrow_schema"] = self.arrow_schema
+        if self.mkdir is not None:
+            kwargs["mkdir"] = self.mkdir
+        if self.engine is not None:
+            kwargs["engine"] = self.engine
+        if self.optimizations is not None:
+            kwargs["optimizations"] = self.optimizations
+        if self.extra_kwargs is not None:
+            if self.extra_kwargs.get("lazy", False):
+                raise ValueError("lazy=True is incompatible with synchronous data savers.")
+            kwargs.update(self.extra_kwargs)
+        return kwargs
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [DATAFRAME_TYPE]
+
+    def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+        data.sink_parquet(self.file, **self._get_saving_kwargs())
+        return utils.get_file_metadata(self.file)
+
+    @classmethod
+    def name(cls) -> str:
+        return "parquet"
+
+
 @dataclasses.dataclass
 class PolarsScanFeatherReader(DataLoader):
     """
@@ -289,12 +481,154 @@ def name(cls) -> str:
         return "feather"
 
 
+@dataclasses.dataclass
+class PolarsSinkFeatherWriter(DataSaver):
+    """Class to handle sinking a Polars LazyFrame to an IPC/Feather file using streaming.
+
+    Calls LazyFrame.sink_ipc() directly, avoiding collect() for better performance.
+    Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_ipc.html
+
+    Note: ``lazy=True`` is intentionally excluded because ``save_data()`` expects
+    the file to exist immediately after the call returns.
+    """
+
+    file: str | Path
+    # kwargs:
+    compression: str | None = None
+    compat_level: Any = None
+    record_batch_size: int | None = None
+    maintain_order: bool = True
+    storage_options: dict[str, Any] | None = None
+    credential_provider: Any = None
+    retries: int | None = None
+    sync_on_close: Any = None
+    mkdir: bool | None = None
+    engine: Any = None
+    optimizations: Any = None
+    extra_kwargs: dict[str, Any] | None = None
+
+    def _get_saving_kwargs(self) -> dict[str, Any]:
+        kwargs = {}
+        if self.compression is not None:
+            kwargs["compression"] = self.compression
+        if self.compat_level is not None:
+            kwargs["compat_level"] = self.compat_level
+        if self.record_batch_size is not None:
+            kwargs["record_batch_size"] = self.record_batch_size
+        if self.maintain_order is not None:
+            kwargs["maintain_order"] = self.maintain_order
+        if self.storage_options is not None:
+            kwargs["storage_options"] = self.storage_options
+        if self.credential_provider is not None:
+            kwargs["credential_provider"] = self.credential_provider
+        if self.retries is not None:
+            kwargs["retries"] = self.retries
+        if self.sync_on_close is not None:
+            kwargs["sync_on_close"] = self.sync_on_close
+        if self.mkdir is not None:
+            kwargs["mkdir"] = self.mkdir
+        if self.engine is not None:
+            kwargs["engine"] = self.engine
+        if self.optimizations is not None:
+            kwargs["optimizations"] = self.optimizations
+        if self.extra_kwargs is not None:
+            if self.extra_kwargs.get("lazy", False):
+                raise ValueError("lazy=True is incompatible with synchronous data savers.")
+            kwargs.update(self.extra_kwargs)
+        return kwargs
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [DATAFRAME_TYPE]
+
+    def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+        data.sink_ipc(self.file, **self._get_saving_kwargs())
+        return utils.get_file_metadata(self.file)
+
+    @classmethod
+    def name(cls) -> str:
+        return "feather"
+
+
+@dataclasses.dataclass
+class PolarsSinkNDJSONWriter(DataSaver):
+    """Class to handle sinking a Polars LazyFrame to an NDJSON file using streaming.
+
+    Calls LazyFrame.sink_ndjson() directly, avoiding collect() for better performance.
+    Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_ndjson.html
+
+    Note: ``lazy=True`` is intentionally excluded because ``save_data()`` expects
+    the file to exist immediately after the call returns.
+    """
+
+    file: str | Path
+    # kwargs:
+    compression: str | None = None
+    compression_level: int | None = None
+    check_extension: bool | None = None
+    maintain_order: bool = True
+    storage_options: dict[str, Any] | None = None
+    credential_provider: Any = None
+    retries: int | None = None
+    sync_on_close: Any = None
+    mkdir: bool | None = None
+    engine: Any = None
+    optimizations: Any = None
+    extra_kwargs: dict[str, Any] | None = None
+
+    def _get_saving_kwargs(self) -> dict[str, Any]:
+        kwargs = {}
+        if self.compression is not None:
+            kwargs["compression"] = self.compression
+        if self.compression_level is not None:
+            kwargs["compression_level"] = self.compression_level
+        if self.check_extension is not None:
+            kwargs["check_extension"] = self.check_extension
+        if self.maintain_order is not None:
+            kwargs["maintain_order"] = self.maintain_order
+        if self.storage_options is not None:
+            kwargs["storage_options"] = self.storage_options
+        if self.credential_provider is not None:
+            kwargs["credential_provider"] = self.credential_provider
+        if self.retries is not None:
+            kwargs["retries"] = self.retries
+        if self.sync_on_close is not None:
+            kwargs["sync_on_close"] = self.sync_on_close
+        if self.mkdir is not None:
+            kwargs["mkdir"] = self.mkdir
+        if self.engine is not None:
+            kwargs["engine"] = self.engine
+        if self.optimizations is not None:
+            kwargs["optimizations"] = self.optimizations
+        if self.extra_kwargs is not None:
+            if self.extra_kwargs.get("lazy", False):
+                raise ValueError("lazy=True is incompatible with synchronous data savers.")
+            kwargs.update(self.extra_kwargs)
+        return kwargs
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [DATAFRAME_TYPE]
+
+    def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+        data.sink_ndjson(self.file, **self._get_saving_kwargs())
+        return utils.get_file_metadata(self.file)
+
+    @classmethod
+    def name(cls) -> str:
+        return "ndjson"
+
+
 def register_data_loaders():
     """Function to register the data loaders for this extension."""
     for loader in [
         PolarsScanCSVReader,
+        PolarsSinkCSVWriter,
         PolarsScanParquetReader,
+        PolarsSinkParquetWriter,
         PolarsScanFeatherReader,
+        PolarsSinkFeatherWriter,
+        PolarsSinkNDJSONWriter,
     ]:
         registry.register_adapter(loader)