diff --git a/hamilton/plugins/polars_lazyframe_extensions.py b/hamilton/plugins/polars_lazyframe_extensions.py index b06f2db94..73db28ba0 100644 --- a/hamilton/plugins/polars_lazyframe_extensions.py +++ b/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() diff --git a/tests/plugins/test_polars_lazyframe_extensions.py b/tests/plugins/test_polars_lazyframe_extensions.py index a35929654..dbba2429e 100644 --- a/tests/plugins/test_polars_lazyframe_extensions.py +++ b/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) +