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
@@ -287,6 +287,171 @@ def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]:
     @classmethod
     def name(cls) -> str:
         return "feather"
+
+
+####
+
+@dataclasses.dataclass
+class PolarsSinkParquetWriter(DataSaver):
+    """Class to handle sinking a Polars LazyFrame to a Parquet file.
+    Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_parquet.html
+    """
+    path: str | Path
+    # kwargs:
+    compression: str = "zstd"
+    compression_level: int | None = None
+    statistics: bool = True
+    row_group_size: int | None = None
+    data_page_size: int | None = None
+    maintain_order: bool = True
+
+    def _get_saving_kwargs(self):
+        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
+        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.path, **self._get_saving_kwargs())
+        return utils.get_file_metadata(self.path)
+
+    @classmethod
+    def name(cls) -> str:
+        return "parquet"
+
+
+@dataclasses.dataclass
+class PolarsSinkCSVWriter(DataSaver):
+    """Class to handle sinking a Polars LazyFrame to a CSV file.
+    Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_csv.html
+    """
+    path: str | Path
+    # kwargs:
+    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
+    null_value: str = ""
+    quote_style: str = "necessary"
+    maintain_order: bool = True
+
+    def _get_saving_kwargs(self):
+        kwargs = {}
+        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.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
+        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.path, **self._get_saving_kwargs())
+        return utils.get_file_metadata(self.path)
+
+    @classmethod
+    def name(cls) -> str:
+        return "csv"
+
+
+@dataclasses.dataclass
+class PolarsSinkFeatherWriter(DataSaver):
+    """Class to handle sinking a Polars LazyFrame to an IPC/Feather file.
+    Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_ipc.html
+    """
+    path: str | Path
+    # kwargs:
+    compression: str = "uncompressed"
+    maintain_order: bool = True
+
+    def _get_saving_kwargs(self):
+        kwargs = {}
+        if self.compression is not None:
+            kwargs["compression"] = self.compression
+        if self.maintain_order is not None:
+            kwargs["maintain_order"] = self.maintain_order
+        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.path, **self._get_saving_kwargs())
+        return utils.get_file_metadata(self.path)
+
+    @classmethod
+    def name(cls) -> str:
+        return "ipc"
+
+
+@dataclasses.dataclass
+class PolarsSinkNDJSONWriter(DataSaver):
+    """Class to handle sinking a Polars LazyFrame to an NDJSON file.
+    Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_ndjson.html
+    """
+    path: str | Path
+    # kwargs:
+    maintain_order: bool = True
+
+    def _get_saving_kwargs(self):
+        kwargs = {}
+        if self.maintain_order is not None:
+            kwargs["maintain_order"] = self.maintain_order
+        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.path, **self._get_saving_kwargs())
+        return utils.get_file_metadata(self.path)
+
+    @classmethod
+    def name(cls) -> str:
+        return "ndjson"
 
 
 def register_data_loaders():
@@ -295,8 +460,108 @@ def register_data_loaders():
         PolarsScanCSVReader,
         PolarsScanParquetReader,
         PolarsScanFeatherReader,
+        PolarsSinkParquetWriter,
+        PolarsSinkCSVWriter,
+        PolarsSinkFeatherWriter,
+        PolarsSinkNDJSONWriter,
     ]:
         registry.register_adapter(loader)
 
 
 register_data_loaders()
+
+# @dataclasses.dataclass
+# class PolarsLazyFrameSinkParquet(DataSaver):
+#     """Class to handle sinking a Polars LazyFrame to a Parquet file.
+#     Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_parquet.html
+#     """
+#     path: str | Path
+
+#     @classmethod
+#     def applicable_types(cls) -> Collection[type]:
+#         return [DATAFRAME_TYPE]
+
+#     def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+#         data.sink_parquet(self.path)
+#         return utils.get_file_metadata(self.path)
+
+#     @classmethod
+#     def name(cls) -> str:
+#         return "parquet"
+
+
+# @dataclasses.dataclass
+# class PolarsLazyFrameSinkCSV(DataSaver):
+#     """Class to handle sinking a Polars LazyFrame to a CSV file.
+#     Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_csv.html
+#     """
+#     path: str | Path
+
+#     @classmethod
+#     def applicable_types(cls) -> Collection[type]:
+#         return [DATAFRAME_TYPE]
+
+#     def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+#         data.sink_csv(self.path)
+#         return utils.get_file_metadata(self.path)
+
+#     @classmethod
+#     def name(cls) -> str:
+#         return "csv"
+
+
+# @dataclasses.dataclass
+# class PolarsLazyFrameSinkIPC(DataSaver):
+#     """Class to handle sinking a Polars LazyFrame to an IPC/Feather file.
+#     Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_ipc.html
+#     """
+#     path: str | Path
+
+#     @classmethod
+#     def applicable_types(cls) -> Collection[type]:
+#         return [DATAFRAME_TYPE]
+
+#     def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+#         data.sink_ipc(self.path)
+#         return utils.get_file_metadata(self.path)
+
+#     @classmethod
+#     def name(cls) -> str:
+#         return "ipc"
+
+
+# @dataclasses.dataclass
+# class PolarsLazyFrameSinkNDJSON(DataSaver):
+#     """Class to handle sinking a Polars LazyFrame to an NDJSON file.
+#     Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_ndjson.html
+#     """
+#     path: str | Path
+
+#     @classmethod
+#     def applicable_types(cls) -> Collection[type]:
+#         return [DATAFRAME_TYPE]
+
+#     def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+#         data.sink_ndjson(self.path)
+#         return utils.get_file_metadata(self.path)
+
+#     @classmethod
+#     def name(cls) -> str:
+#         return "ndjson"
+
+
+# def register_data_loaders():
+#     """Function to register the data loaders for this extension."""
+#     for loader in [
+#         PolarsScanCSVReader,
+#         PolarsScanParquetReader,
+#         PolarsScanFeatherReader,
+#         PolarsLazyFrameSinkParquet,
+#         PolarsLazyFrameSinkCSV,
+#         PolarsLazyFrameSinkIPC,
+#         PolarsLazyFrameSinkNDJSON,
+#     ]:
+#         registry.register_adapter(loader)
+
+
+# register_data_loaders()

tests/plugins/test_polars_lazyframe_extensions.py

@@ -23,10 +23,15 @@
 from polars.testing import assert_frame_equal
 from sqlalchemy import create_engine
 
+
 from hamilton.plugins.polars_lazyframe_extensions import (
     PolarsScanCSVReader,
     PolarsScanFeatherReader,
     PolarsScanParquetReader,
+    PolarsSinkParquetWriter,
+    PolarsSinkCSVWriter,
+    PolarsSinkFeatherWriter,
+    PolarsSinkNDJSONWriter,
 )
 from hamilton.plugins.polars_post_1_0_0_extensions import (
     PolarsAvroReader,
@@ -211,3 +216,93 @@ def test_polars_spreadsheet(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None:
     assert write_kwargs["include_header"] is True
     assert "raise_if_empty" in read_kwargs
     assert read_kwargs["raise_if_empty"] is True
+
+#####
+
+def test_polars_lazyframe_sink_parquet(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None:
+    file = tmp_path / "test.parquet"
+    sink = PolarsSinkParquetWriter(path=file)
+    kwargs = sink._get_saving_kwargs()
+    sink.save_data(df)
+    df2 = pl.read_parquet(file)
+    assert PolarsSinkParquetWriter.applicable_types() == [pl.LazyFrame]
+    assert file.exists()
+    assert kwargs["compression"] == "zstd"
+    assert_frame_equal(df.collect(), df2)
+
+
+def test_polars_lazyframe_sink_csv(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None:
+    file = tmp_path / "test.csv"
+    sink = PolarsSinkCSVWriter(path=file)
+    kwargs = sink._get_saving_kwargs()
+    sink.save_data(df)
+    df2 = pl.read_csv(file)
+    assert PolarsSinkCSVWriter.applicable_types() == [pl.LazyFrame]
+    assert file.exists()
+    assert kwargs["separator"] == ","
+    assert_frame_equal(df.collect(), df2)
+
+
+def test_polars_lazyframe_sink_ipc(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None:
+    file = tmp_path / "test.ipc"
+    sink = PolarsSinkFeatherWriter(path=file)
+    kwargs = sink._get_saving_kwargs()
+    sink.save_data(df)
+    df2 = pl.read_ipc(file)
+    assert PolarsSinkFeatherWriter.applicable_types() == [pl.LazyFrame]
+    assert file.exists()
+    assert kwargs["compression"] == "uncompressed"
+    assert_frame_equal(df.collect(), df2)
+
+
+def test_polars_lazyframe_sink_ndjson(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None:
+    file = tmp_path / "test.ndjson"
+    sink = PolarsSinkNDJSONWriter(path=file)
+    kwargs = sink._get_saving_kwargs()
+    sink.save_data(df)
+    df2 = pl.read_ndjson(file)
+    assert PolarsSinkNDJSONWriter.applicable_types() == [pl.LazyFrame]
+    assert file.exists()
+    assert kwargs["maintain_order"] is True
+    assert_frame_equal(df.collect(), df2)
+
+
+# def test_polars_lazyframe_sink_parquet(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None:
+#     file = tmp_path / "test.parquet"
+#     sink = PolarsLazyFrameSinkParquet(path=file)
+#     metadata = sink.save_data(df)
+#     df2 = pl.read_parquet(file)
+#     assert PolarsLazyFrameSinkParquet.applicable_types() == [pl.LazyFrame]
+#     assert file.exists()
+#     assert_frame_equal(df.collect(), df2)
+
+
+# def test_polars_lazyframe_sink_csv(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None:
+#     file = tmp_path / "test.csv"
+#     sink = PolarsLazyFrameSinkCSV(path=file)
+#     metadata = sink.save_data(df)
+#     df2 = pl.read_csv(file)
+#     assert PolarsLazyFrameSinkCSV.applicable_types() == [pl.LazyFrame]
+#     assert file.exists()
+#     assert_frame_equal(df.collect(), df2)
+
+
+# def test_polars_lazyframe_sink_ipc(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None:
+#     file = tmp_path / "test.ipc"
+#     sink = PolarsLazyFrameSinkIPC(path=file)
+#     metadata = sink.save_data(df)
+#     df2 = pl.read_ipc(file)
+#     assert PolarsLazyFrameSinkIPC.applicable_types() == [pl.LazyFrame]
+#     assert file.exists()
+#     assert_frame_equal(df.collect(), df2)
+
+
+# def test_polars_lazyframe_sink_ndjson(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None:
+#     file = tmp_path / "test.ndjson"
+#     sink = PolarsLazyFrameSinkNDJSON(path=file)
+#     metadata = sink.save_data(df)
+#     df2 = pl.read_ndjson(file)
+#     assert PolarsLazyFrameSinkNDJSON.applicable_types() == [pl.LazyFrame]
+#     assert file.exists()
+#     assert_frame_equal(df.collect(), df2)
+