From aa1a54e495b34894de7e26d1811c1a0622744636 Mon Sep 17 00:00:00 2001 From: sonalishintre <42985737+sonalishintre@users.noreply.github.com> Date: Mon, 22 Jun 2026 08:51:16 +0530 Subject: [PATCH 1/7] feat: add LazyFrame sink classes for parquet, csv, ipc, ndjson --- .../plugins/polars_lazyframe_extensions.py | 85 ++++++++++++++++++- 1 file changed, 84 insertions(+), 1 deletion(-) diff --git a/hamilton/plugins/polars_lazyframe_extensions.py b/hamilton/plugins/polars_lazyframe_extensions.py index b06f2db94..a300195a7 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 @@ -288,6 +288,85 @@ def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]: def name(cls) -> str: return "feather" +@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.""" @@ -295,6 +374,10 @@ def register_data_loaders(): PolarsScanCSVReader, PolarsScanParquetReader, PolarsScanFeatherReader, + PolarsLazyFrameSinkParquet, + PolarsLazyFrameSinkCSV, + PolarsLazyFrameSinkIPC, + PolarsLazyFrameSinkNDJSON, ]: registry.register_adapter(loader) From 0891300cff25afec54ee5ca1fe091c0200730570 Mon Sep 17 00:00:00 2001 From: sonalishintre <42985737+sonalishintre@users.noreply.github.com> Date: Mon, 29 Jun 2026 05:55:07 +0530 Subject: [PATCH 2/7] test: add tests for LazyFrame sink classes --- .../test_polars_lazyframe_extensions.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/plugins/test_polars_lazyframe_extensions.py b/tests/plugins/test_polars_lazyframe_extensions.py index a35929654..dedb52bbe 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, + PolarsLazyFrameSinkParquet, + PolarsLazyFrameSinkCSV, + PolarsLazyFrameSinkIPC, + PolarsLazyFrameSinkNDJSON, ) from hamilton.plugins.polars_post_1_0_0_extensions import ( PolarsAvroReader, @@ -211,3 +216,45 @@ 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 = 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) + From 2bb229d65b8c02d14417553d520237b922e5a9e6 Mon Sep 17 00:00:00 2001 From: sonalishintre <42985737+sonalishintre@users.noreply.github.com> Date: Tue, 21 Jul 2026 04:56:56 +0530 Subject: [PATCH 3/7] refactor: rename sink classes and add kwargs support per reviewer feedback --- .../plugins/polars_lazyframe_extensions.py | 206 +++++++++++++++++- .../test_polars_lazyframe_extensions.py | 80 +++++-- 2 files changed, 258 insertions(+), 28 deletions(-) diff --git a/hamilton/plugins/polars_lazyframe_extensions.py b/hamilton/plugins/polars_lazyframe_extensions.py index a300195a7..73db28ba0 100644 --- a/hamilton/plugins/polars_lazyframe_extensions.py +++ b/hamilton/plugins/polars_lazyframe_extensions.py @@ -287,20 +287,46 @@ def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]: @classmethod def name(cls) -> str: return "feather" + + +#### @dataclasses.dataclass -class PolarsLazyFrameSinkParquet(DataSaver): +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) + data.sink_parquet(self.path, **self._get_saving_kwargs()) return utils.get_file_metadata(self.path) @classmethod @@ -309,18 +335,59 @@ def name(cls) -> str: @dataclasses.dataclass -class PolarsLazyFrameSinkCSV(DataSaver): +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) + data.sink_csv(self.path, **self._get_saving_kwargs()) return utils.get_file_metadata(self.path) @classmethod @@ -329,18 +396,29 @@ def name(cls) -> str: @dataclasses.dataclass -class PolarsLazyFrameSinkIPC(DataSaver): +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) + data.sink_ipc(self.path, **self._get_saving_kwargs()) return utils.get_file_metadata(self.path) @classmethod @@ -349,18 +427,26 @@ def name(cls) -> str: @dataclasses.dataclass -class PolarsLazyFrameSinkNDJSON(DataSaver): +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) + data.sink_ndjson(self.path, **self._get_saving_kwargs()) return utils.get_file_metadata(self.path) @classmethod @@ -374,12 +460,108 @@ def register_data_loaders(): PolarsScanCSVReader, PolarsScanParquetReader, PolarsScanFeatherReader, - PolarsLazyFrameSinkParquet, - PolarsLazyFrameSinkCSV, - PolarsLazyFrameSinkIPC, - PolarsLazyFrameSinkNDJSON, + 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 dedb52bbe..dbba2429e 100644 --- a/tests/plugins/test_polars_lazyframe_extensions.py +++ b/tests/plugins/test_polars_lazyframe_extensions.py @@ -28,10 +28,10 @@ PolarsScanCSVReader, PolarsScanFeatherReader, PolarsScanParquetReader, - PolarsLazyFrameSinkParquet, - PolarsLazyFrameSinkCSV, - PolarsLazyFrameSinkIPC, - PolarsLazyFrameSinkNDJSON, + PolarsSinkParquetWriter, + PolarsSinkCSVWriter, + PolarsSinkFeatherWriter, + PolarsSinkNDJSONWriter, ) from hamilton.plugins.polars_post_1_0_0_extensions import ( PolarsAvroReader, @@ -221,40 +221,88 @@ def test_polars_spreadsheet(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: 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) + sink = PolarsSinkParquetWriter(path=file) + kwargs = sink._get_saving_kwargs() + sink.save_data(df) df2 = pl.read_parquet(file) - assert PolarsLazyFrameSinkParquet.applicable_types() == [pl.LazyFrame] + 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 = PolarsLazyFrameSinkCSV(path=file) - metadata = sink.save_data(df) + sink = PolarsSinkCSVWriter(path=file) + kwargs = sink._get_saving_kwargs() + sink.save_data(df) df2 = pl.read_csv(file) - assert PolarsLazyFrameSinkCSV.applicable_types() == [pl.LazyFrame] + 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 = PolarsLazyFrameSinkIPC(path=file) - metadata = sink.save_data(df) + sink = PolarsSinkFeatherWriter(path=file) + kwargs = sink._get_saving_kwargs() + sink.save_data(df) df2 = pl.read_ipc(file) - assert PolarsLazyFrameSinkIPC.applicable_types() == [pl.LazyFrame] + 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 = PolarsLazyFrameSinkNDJSON(path=file) - metadata = sink.save_data(df) + sink = PolarsSinkNDJSONWriter(path=file) + kwargs = sink._get_saving_kwargs() + sink.save_data(df) df2 = pl.read_ndjson(file) - assert PolarsLazyFrameSinkNDJSON.applicable_types() == [pl.LazyFrame] + 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) + From a75dc1197c3a18a51ecced44deda568c3df93659 Mon Sep 17 00:00:00 2001 From: sonalishintre <42985737+sonalishintre@users.noreply.github.com> Date: Thu, 13 Aug 2026 05:34:17 +0530 Subject: [PATCH 4/7] fix: address reviewer feedback - feather name, kwargs, ordering, formatting --- .../plugins/polars_lazyframe_extensions.py | 333 +++++++----------- .../test_polars_lazyframe_extensions.py | 135 ++++--- 2 files changed, 219 insertions(+), 249 deletions(-) diff --git a/hamilton/plugins/polars_lazyframe_extensions.py b/hamilton/plugins/polars_lazyframe_extensions.py index 73db28ba0..a0b476aaf 100644 --- a/hamilton/plugins/polars_lazyframe_extensions.py +++ b/hamilton/plugins/polars_lazyframe_extensions.py @@ -183,6 +183,76 @@ 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. + """ + + path: str | Path + # kwargs: + 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 + null_value: str = "" + quote_style: str = "necessary" + maintain_order: bool = True + + def _get_saving_kwargs(self): + kwargs = {} + 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.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 PolarsScanParquetReader(DataLoader): """Class specifically to handle loading parquet files with polars @@ -239,63 +309,17 @@ def name(cls) -> str: return "parquet" -@dataclasses.dataclass -class PolarsScanFeatherReader(DataLoader): - """ - Class specifically to handle loading Feather/Arrow IPC files with Polars. - Should map to https://pola-rs.github.io/polars/py-polars/html/reference/api/polars.read_ipc.html - """ - - source: str | BinaryIO | BytesIO | Path | bytes - # kwargs: - columns: list[str] | list[int] | None = None - n_rows: int | None = None - use_pyarrow: bool = False - memory_map: bool = True - storage_options: dict[str, Any] | None = None - row_count_name: str | None = None - row_count_offset: int = 0 - rechunk: bool = True - - @classmethod - def applicable_types(cls) -> Collection[type]: - return [DATAFRAME_TYPE] - - def _get_loading_kwargs(self): - kwargs = {} - if self.columns is not None: - kwargs["columns"] = self.columns - if self.n_rows is not None: - kwargs["n_rows"] = self.n_rows - if self.memory_map is not None: - kwargs["memory_map"] = self.memory_map - if self.storage_options is not None: - kwargs["storage_options"] = self.storage_options - if self.row_count_name is not None: - kwargs["row_count_name"] = self.row_count_name - if self.row_count_offset is not None: - kwargs["row_count_offset"] = self.row_count_offset - if self.rechunk is not None: - kwargs["rechunk"] = self.rechunk - return kwargs - - def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]: - df = pl.scan_ipc(self.source, **self._get_loading_kwargs()) - metadata = utils.get_file_metadata(self.source) - return df, metadata - - @classmethod - def name(cls) -> str: - return "feather" - - -#### - @dataclasses.dataclass class PolarsSinkParquetWriter(DataSaver): - """Class to handle sinking a Polars LazyFrame to a Parquet file. + """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. """ + path: str | Path # kwargs: compression: str = "zstd" @@ -335,71 +359,66 @@ def name(cls) -> str: @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 +class PolarsScanFeatherReader(DataLoader): + """ + Class specifically to handle loading Feather/Arrow IPC files with Polars. + Should map to https://pola-rs.github.io/polars/py-polars/html/reference/api/polars.read_ipc.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 + source: str | BinaryIO | BytesIO | Path | bytes + # kwargs: + columns: list[str] | list[int] | None = None + n_rows: int | None = None + use_pyarrow: bool = False + memory_map: bool = True + storage_options: dict[str, Any] | None = None + row_count_name: str | None = None + row_count_offset: int = 0 + rechunk: bool = True @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) + def _get_loading_kwargs(self): + kwargs = {} + if self.columns is not None: + kwargs["columns"] = self.columns + if self.n_rows is not None: + kwargs["n_rows"] = self.n_rows + if self.memory_map is not None: + kwargs["memory_map"] = self.memory_map + if self.storage_options is not None: + kwargs["storage_options"] = self.storage_options + if self.row_count_name is not None: + kwargs["row_count_name"] = self.row_count_name + if self.row_count_offset is not None: + kwargs["row_count_offset"] = self.row_count_offset + if self.rechunk is not None: + kwargs["rechunk"] = self.rechunk + return kwargs + + def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]: + df = pl.scan_ipc(self.source, **self._get_loading_kwargs()) + metadata = utils.get_file_metadata(self.source) + return df, metadata @classmethod def name(cls) -> str: - return "csv" + return "feather" @dataclasses.dataclass class PolarsSinkFeatherWriter(DataSaver): - """Class to handle sinking a Polars LazyFrame to an IPC/Feather file. + """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. """ + path: str | Path # kwargs: compression: str = "uncompressed" @@ -423,14 +442,20 @@ def save_data(self, data: pl.LazyFrame) -> dict[str, Any]: @classmethod def name(cls) -> str: - return "ipc" + return "feather" @dataclasses.dataclass class PolarsSinkNDJSONWriter(DataSaver): - """Class to handle sinking a Polars LazyFrame to an NDJSON file. + """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. """ + path: str | Path # kwargs: maintain_order: bool = True @@ -458,10 +483,10 @@ def register_data_loaders(): """Function to register the data loaders for this extension.""" for loader in [ PolarsScanCSVReader, + PolarsSinkCSVWriter, PolarsScanParquetReader, - PolarsScanFeatherReader, PolarsSinkParquetWriter, - PolarsSinkCSVWriter, + PolarsScanFeatherReader, PolarsSinkFeatherWriter, PolarsSinkNDJSONWriter, ]: @@ -469,99 +494,3 @@ def register_data_loaders(): 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 dbba2429e..b909584ed 100644 --- a/tests/plugins/test_polars_lazyframe_extensions.py +++ b/tests/plugins/test_polars_lazyframe_extensions.py @@ -23,15 +23,14 @@ 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, + PolarsSinkParquetWriter, ) from hamilton.plugins.polars_post_1_0_0_extensions import ( PolarsAvroReader, @@ -64,7 +63,7 @@ def test_lazy_polars_lazyframe_csv(df: pl.LazyFrame, tmp_path: pathlib.Path) -> reader = PolarsScanCSVReader(file=file) kwargs2 = reader._get_loading_kwargs() - df2, metadata = reader.load_data(pl.LazyFrame) + df2, _metadata = reader.load_data(pl.LazyFrame) assert PolarsCSVWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] assert PolarsScanCSVReader.applicable_types() == [pl.LazyFrame] @@ -82,7 +81,7 @@ def test_lazy_polars_parquet(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: reader = PolarsScanParquetReader(file=file, n_rows=2) kwargs2 = reader._get_loading_kwargs() - df2, metadata = reader.load_data(pl.LazyFrame) + df2, _metadata = reader.load_data(pl.LazyFrame) assert PolarsParquetWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] assert PolarsScanParquetReader.applicable_types() == [pl.LazyFrame] @@ -127,7 +126,7 @@ def test_lazy_polars_avro(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: reader = PolarsAvroReader(file=file, n_rows=2) kwargs2 = reader._get_loading_kwargs() - df2, metadata = reader.load_data(pl.DataFrame) + df2, _metadata = reader.load_data(pl.DataFrame) assert PolarsAvroWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] assert PolarsAvroReader.applicable_types() == [pl.DataFrame] @@ -143,7 +142,7 @@ def test_polars_json(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: reader = PolarsJSONReader(source=file) kwargs2 = reader._get_loading_kwargs() - df2, metadata = reader.load_data(pl.DataFrame) + df2, _metadata = reader.load_data(pl.DataFrame) assert PolarsJSONWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] assert PolarsJSONReader.applicable_types() == [pl.DataFrame] @@ -159,7 +158,7 @@ def test_polars_ndjson(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: reader = PolarsNDJSONReader(source=file) kwargs2 = reader._get_loading_kwargs() - df2, metadata = reader.load_data(pl.DataFrame) + df2, _metadata = reader.load_data(pl.DataFrame) assert PolarsNDJSONWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] assert PolarsNDJSONReader.applicable_types() == [pl.DataFrame] @@ -184,7 +183,7 @@ def test_polars_database(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: reader = PolarsDatabaseReader(query=f"SELECT * FROM {table_name}", connection=connector) kwargs2 = reader._get_loading_kwargs() - df2, metadata = reader.load_data(pl.DataFrame) + df2, _metadata = reader.load_data(pl.DataFrame) assert PolarsDatabaseWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] assert PolarsDatabaseReader.applicable_types() == [pl.DataFrame] @@ -217,7 +216,6 @@ def test_polars_spreadsheet(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: 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" @@ -225,21 +223,54 @@ def test_polars_lazyframe_sink_parquet(df: pl.LazyFrame, tmp_path: pathlib.Path) 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_parquet_custom_kwargs( + df: pl.LazyFrame, tmp_path: pathlib.Path +) -> None: + """Test that non-default kwargs are passed through correctly.""" + file = tmp_path / "test.parquet" + sink = PolarsSinkParquetWriter(path=file, compression="snappy", maintain_order=False) + kwargs = sink._get_saving_kwargs() + sink.save_data(df) + df2 = pl.read_parquet(file) + + assert kwargs["compression"] == "snappy" + assert kwargs["maintain_order"] is False + 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 = 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 kwargs["include_header"] is True + assert_frame_equal(df.collect(), df2) + + +def test_polars_lazyframe_sink_csv_custom_kwargs(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: + """Test that non-default kwargs are passed through correctly.""" + file = tmp_path / "test.csv" + sink = PolarsSinkCSVWriter(path=file, separator=";", include_header=False) + kwargs = sink._get_saving_kwargs() + sink.save_data(df) + df2 = pl.read_csv(file, separator=";", has_header=False, new_columns=["a", "b"]) + + assert kwargs["separator"] == ";" + assert kwargs["include_header"] is False + assert file.exists() assert_frame_equal(df.collect(), df2) @@ -249,60 +280,70 @@ def test_polars_lazyframe_sink_ipc(df: pl.LazyFrame, tmp_path: pathlib.Path) -> 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_ipc_custom_kwargs(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: + """Test that non-default kwargs are passed through correctly.""" + file = tmp_path / "test.ipc" + sink = PolarsSinkFeatherWriter(path=file, compression="lz4", maintain_order=False) + kwargs = sink._get_saving_kwargs() + sink.save_data(df) + df2 = pl.read_ipc(file) + + assert kwargs["compression"] == "lz4" + assert kwargs["maintain_order"] is False + 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 = 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) +def test_polars_lazyframe_sink_ndjson_custom_kwargs( + df: pl.LazyFrame, tmp_path: pathlib.Path +) -> None: + """Test that non-default kwargs are passed through correctly.""" + file = tmp_path / "test.ndjson" + sink = PolarsSinkNDJSONWriter(path=file, maintain_order=False) + kwargs = sink._get_saving_kwargs() + sink.save_data(df) + df2 = pl.read_ndjson(file) + + assert kwargs["maintain_order"] is False + assert file.exists() + assert_frame_equal(df.collect(), df2) + + +def test_polars_lazyframe_sink_feather_adapter_name() -> None: + """Test that PolarsSinkFeatherWriter is registered under 'feather', not 'ipc'.""" + assert PolarsSinkFeatherWriter.name() == "feather" + + +def test_polars_lazyframe_sink_csv_adapter_name() -> None: + """Test that PolarsSinkCSVWriter is registered under 'csv'.""" + assert PolarsSinkCSVWriter.name() == "csv" + + +def test_polars_lazyframe_sink_parquet_adapter_name() -> None: + """Test that PolarsSinkParquetWriter is registered under 'parquet'.""" + assert PolarsSinkParquetWriter.name() == "parquet" + +def test_polars_lazyframe_sink_ndjson_adapter_name() -> None: + """Test that PolarsSinkNDJSONWriter is registered under 'ndjson'.""" + assert PolarsSinkNDJSONWriter.name() == "ndjson" From 23601dd4ba87a83a3c0dee7424349ce98eeadd85 Mon Sep 17 00:00:00 2001 From: jernejfrank Date: Thu, 13 Aug 2026 10:03:16 +0100 Subject: [PATCH 5/7] fix(polars): close LazyFrame sink gaps Make streaming sinks the sole adapters for supported LazyFrame formats while preserving eager writers for DataFrames. Align sink options with Polars and cover registry and materializer resolution. Refs apache/hamilton#1653 --- .../plugins/polars_lazyframe_extensions.py | 178 +++++++++++++++--- .../plugins/polars_post_1_0_0_extensions.py | 26 +-- .../plugins/polars_pre_1_0_0_extension.py | 19 +- tests/plugins/test_polars_extensions.py | 8 +- .../test_polars_lazyframe_extensions.py | 129 ++++++++++--- 5 files changed, 277 insertions(+), 83 deletions(-) diff --git a/hamilton/plugins/polars_lazyframe_extensions.py b/hamilton/plugins/polars_lazyframe_extensions.py index a0b476aaf..1c57188b7 100644 --- a/hamilton/plugins/polars_lazyframe_extensions.py +++ b/hamilton/plugins/polars_lazyframe_extensions.py @@ -194,8 +194,12 @@ class PolarsSinkCSVWriter(DataSaver): the file to exist immediately after the call returns. """ - path: str | Path + 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" @@ -206,12 +210,29 @@ class PolarsSinkCSVWriter(DataSaver): time_format: str | None = None float_scientific: bool | None = None float_precision: int | None = None - null_value: str = "" - quote_style: str = "necessary" + decimal_comma: bool | None = None + null_value: str | None = None + quote_style: Any = None maintain_order: bool = True - - def _get_saving_kwargs(self): + 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: @@ -232,12 +253,30 @@ def _get_saving_kwargs(self): 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: + kwargs.update(self.extra_kwargs) return kwargs @classmethod @@ -245,8 +284,8 @@ 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) + data.sink_csv(self.file, **self._get_saving_kwargs()) + return utils.get_file_metadata(self.file) @classmethod def name(cls) -> str: @@ -320,16 +359,26 @@ class PolarsSinkParquetWriter(DataSaver): the file to exist immediately after the call returns. """ - path: str | Path + file: str | Path # kwargs: compression: str = "zstd" compression_level: int | None = None - statistics: bool = True + statistics: bool | str | dict[str, bool] = True row_group_size: int | None = None data_page_size: int | None = None maintain_order: bool = True - - def _get_saving_kwargs(self): + 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 @@ -343,6 +392,26 @@ def _get_saving_kwargs(self): 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: + kwargs.update(self.extra_kwargs) return kwargs @classmethod @@ -350,8 +419,8 @@ 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) + data.sink_parquet(self.file, **self._get_saving_kwargs()) + return utils.get_file_metadata(self.file) @classmethod def name(cls) -> str: @@ -419,17 +488,47 @@ class PolarsSinkFeatherWriter(DataSaver): the file to exist immediately after the call returns. """ - path: str | Path + file: str | Path # kwargs: compression: str = "uncompressed" + compat_level: Any = None + record_batch_size: int | None = None maintain_order: bool = True - - def _get_saving_kwargs(self): + 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: + kwargs.update(self.extra_kwargs) return kwargs @classmethod @@ -437,8 +536,8 @@ 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) + data.sink_ipc(self.file, **self._get_saving_kwargs()) + return utils.get_file_metadata(self.file) @classmethod def name(cls) -> str: @@ -456,14 +555,47 @@ class PolarsSinkNDJSONWriter(DataSaver): the file to exist immediately after the call returns. """ - path: str | Path + file: str | Path # kwargs: + compression: str | None = None + compression_level: int | None = None + check_extension: bool | None = None maintain_order: bool = True - - def _get_saving_kwargs(self): + 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: + kwargs.update(self.extra_kwargs) return kwargs @classmethod @@ -471,8 +603,8 @@ 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) + data.sink_ndjson(self.file, **self._get_saving_kwargs()) + return utils.get_file_metadata(self.file) @classmethod def name(cls) -> str: diff --git a/hamilton/plugins/polars_post_1_0_0_extensions.py b/hamilton/plugins/polars_post_1_0_0_extensions.py index dab1a3973..44b379646 100644 --- a/hamilton/plugins/polars_post_1_0_0_extensions.py +++ b/hamilton/plugins/polars_post_1_0_0_extensions.py @@ -197,7 +197,7 @@ class PolarsCSVWriter(DataSaver): @classmethod def applicable_types(cls) -> Collection[type]: - return [DATAFRAME_TYPE, pl.LazyFrame] + return [DATAFRAME_TYPE] def _get_saving_kwargs(self): kwargs = {} @@ -227,9 +227,7 @@ def _get_saving_kwargs(self): kwargs["quote_style"] = self.quote_style return kwargs - def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: - if isinstance(data, pl.LazyFrame): - data = data.collect() + def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: data.write_csv(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) @@ -318,7 +316,7 @@ class PolarsParquetWriter(DataSaver): @classmethod def applicable_types(cls) -> Collection[type]: - return [DATAFRAME_TYPE, pl.LazyFrame] + return [DATAFRAME_TYPE] def _get_saving_kwargs(self): kwargs = {} @@ -336,10 +334,7 @@ def _get_saving_kwargs(self): kwargs["pyarrow_options"] = self.pyarrow_options return kwargs - def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: - if isinstance(data, pl.LazyFrame): - data = data.collect() - + def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: data.write_parquet(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) @@ -414,7 +409,7 @@ class PolarsFeatherWriter(DataSaver): @classmethod def applicable_types(cls) -> Collection[type]: - return [DATAFRAME_TYPE, pl.LazyFrame] + return [DATAFRAME_TYPE] def _get_saving_kwargs(self): kwargs = {} @@ -422,9 +417,7 @@ def _get_saving_kwargs(self): kwargs["compression"] = self.compression return kwargs - def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: - if isinstance(data, pl.LazyFrame): - data = data.collect() + def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: data.write_ipc(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) @@ -600,12 +593,9 @@ class PolarsNDJSONWriter(DataSaver): @classmethod def applicable_types(cls) -> Collection[type]: - return [DATAFRAME_TYPE, pl.LazyFrame] - - def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: - if isinstance(data, pl.LazyFrame): - data = data.collect() + return [DATAFRAME_TYPE] + def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: data.write_ndjson(self.file) return utils.get_file_and_dataframe_metadata(self.file, data) diff --git a/hamilton/plugins/polars_pre_1_0_0_extension.py b/hamilton/plugins/polars_pre_1_0_0_extension.py index 4dd92dc43..bd3b54d2b 100644 --- a/hamilton/plugins/polars_pre_1_0_0_extension.py +++ b/hamilton/plugins/polars_pre_1_0_0_extension.py @@ -206,7 +206,7 @@ class PolarsCSVWriter(DataSaver): @classmethod def applicable_types(cls) -> Collection[type]: - return [DATAFRAME_TYPE, pl.LazyFrame] + return [DATAFRAME_TYPE] def _get_saving_kwargs(self): kwargs = {} @@ -236,9 +236,7 @@ def _get_saving_kwargs(self): kwargs["quote_style"] = self.quote_style return kwargs - def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: - if isinstance(data, pl.LazyFrame): - data = data.collect() + def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: data.write_csv(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) @@ -327,7 +325,7 @@ class PolarsParquetWriter(DataSaver): @classmethod def applicable_types(cls) -> Collection[type]: - return [DATAFRAME_TYPE, pl.LazyFrame] + return [DATAFRAME_TYPE] def _get_saving_kwargs(self): kwargs = {} @@ -345,10 +343,7 @@ def _get_saving_kwargs(self): kwargs["pyarrow_options"] = self.pyarrow_options return kwargs - def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: - if isinstance(data, pl.LazyFrame): - data = data.collect() - + def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: data.write_parquet(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) @@ -423,7 +418,7 @@ class PolarsFeatherWriter(DataSaver): @classmethod def applicable_types(cls) -> Collection[type]: - return [DATAFRAME_TYPE, pl.LazyFrame] + return [DATAFRAME_TYPE] def _get_saving_kwargs(self): kwargs = {} @@ -431,9 +426,7 @@ def _get_saving_kwargs(self): kwargs["compression"] = self.compression return kwargs - def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: - if isinstance(data, pl.LazyFrame): - data = data.collect() + def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: data.write_ipc(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) diff --git a/tests/plugins/test_polars_extensions.py b/tests/plugins/test_polars_extensions.py index 2709dd85c..1d4b9882e 100644 --- a/tests/plugins/test_polars_extensions.py +++ b/tests/plugins/test_polars_extensions.py @@ -67,7 +67,7 @@ def test_polars_csv(df: pl.DataFrame, tmp_path: pathlib.Path) -> None: kwargs2 = reader._get_loading_kwargs() df2, metadata = reader.load_data(pl.DataFrame) - assert PolarsCSVWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] + assert PolarsCSVWriter.applicable_types() == [pl.DataFrame] assert PolarsCSVReader.applicable_types() == [pl.DataFrame] assert kwargs1["separator"] == "," assert kwargs2["has_header"] is True @@ -85,7 +85,7 @@ def test_polars_parquet(df: pl.DataFrame, tmp_path: pathlib.Path) -> None: kwargs2 = reader._get_loading_kwargs() df2, metadata = reader.load_data(pl.DataFrame) - assert PolarsParquetWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] + assert PolarsParquetWriter.applicable_types() == [pl.DataFrame] assert PolarsParquetReader.applicable_types() == [pl.DataFrame] assert kwargs1["compression"] == "zstd" assert kwargs2["n_rows"] == 2 @@ -107,7 +107,7 @@ def test_polars_feather(tmp_path: pathlib.Path) -> None: assert "n_rows" not in read_kwargs assert df.shape == (4, 3) - assert PolarsFeatherWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] + assert PolarsFeatherWriter.applicable_types() == [pl.DataFrame] assert "compression" in write_kwargs assert file_path.exists() assert metadata["file_metadata"]["path"] == str(file_path) @@ -140,7 +140,7 @@ def test_polars_ndjson(df: pl.DataFrame, tmp_path: pathlib.Path) -> None: kwargs2 = reader._get_loading_kwargs() df2, metadata = reader.load_data(pl.DataFrame) - assert PolarsNDJSONWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] + assert PolarsNDJSONWriter.applicable_types() == [pl.DataFrame] assert PolarsNDJSONReader.applicable_types() == [pl.DataFrame] assert df2.shape == (2, 2) assert "schema" not in kwargs2 diff --git a/tests/plugins/test_polars_lazyframe_extensions.py b/tests/plugins/test_polars_lazyframe_extensions.py index b909584ed..c4b6a7411 100644 --- a/tests/plugins/test_polars_lazyframe_extensions.py +++ b/tests/plugins/test_polars_lazyframe_extensions.py @@ -23,6 +23,9 @@ from polars.testing import assert_frame_equal from sqlalchemy import create_engine +from hamilton import ad_hoc_utils, driver, registry +from hamilton.function_modifiers.adapters import resolve_adapter_class +from hamilton.io.materialization import to from hamilton.plugins.polars_lazyframe_extensions import ( PolarsScanCSVReader, PolarsScanFeatherReader, @@ -35,18 +38,24 @@ from hamilton.plugins.polars_post_1_0_0_extensions import ( PolarsAvroReader, PolarsAvroWriter, - PolarsCSVWriter, PolarsDatabaseReader, PolarsDatabaseWriter, PolarsFeatherWriter, PolarsJSONReader, PolarsJSONWriter, PolarsNDJSONReader, - PolarsNDJSONWriter, - PolarsParquetWriter, PolarsSpreadsheetReader, PolarsSpreadsheetWriter, ) +from hamilton.plugins.polars_pre_1_0_0_extension import ( + PolarsCSVWriter as Pre1PolarsCSVWriter, +) +from hamilton.plugins.polars_pre_1_0_0_extension import ( + PolarsFeatherWriter as Pre1PolarsFeatherWriter, +) +from hamilton.plugins.polars_pre_1_0_0_extension import ( + PolarsParquetWriter as Pre1PolarsParquetWriter, +) @pytest.fixture @@ -57,7 +66,7 @@ def df(): def test_lazy_polars_lazyframe_csv(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: file = tmp_path / "test.csv" - writer = PolarsCSVWriter(file=file) + writer = PolarsSinkCSVWriter(file=file) kwargs1 = writer._get_saving_kwargs() writer.save_data(df) @@ -65,7 +74,7 @@ def test_lazy_polars_lazyframe_csv(df: pl.LazyFrame, tmp_path: pathlib.Path) -> kwargs2 = reader._get_loading_kwargs() df2, _metadata = reader.load_data(pl.LazyFrame) - assert PolarsCSVWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] + assert PolarsSinkCSVWriter.applicable_types() == [pl.LazyFrame] assert PolarsScanCSVReader.applicable_types() == [pl.LazyFrame] assert kwargs1["separator"] == "," assert kwargs2["has_header"] is True @@ -75,7 +84,7 @@ def test_lazy_polars_lazyframe_csv(df: pl.LazyFrame, tmp_path: pathlib.Path) -> def test_lazy_polars_parquet(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: file = tmp_path / "test.parquet" - writer = PolarsParquetWriter(file=file) + writer = PolarsSinkParquetWriter(file=file) kwargs1 = writer._get_saving_kwargs() writer.save_data(df) @@ -83,7 +92,7 @@ def test_lazy_polars_parquet(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: kwargs2 = reader._get_loading_kwargs() df2, _metadata = reader.load_data(pl.LazyFrame) - assert PolarsParquetWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] + assert PolarsSinkParquetWriter.applicable_types() == [pl.LazyFrame] assert PolarsScanParquetReader.applicable_types() == [pl.LazyFrame] assert kwargs1["compression"] == "zstd" assert kwargs2["n_rows"] == 2 @@ -105,7 +114,7 @@ def test_lazy_polars_feather(tmp_path: pathlib.Path) -> None: assert "n_rows" not in read_kwargs assert df.collect().shape == (4, 3) - assert PolarsFeatherWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] + assert PolarsFeatherWriter.applicable_types() == [pl.DataFrame] assert "compression" in write_kwargs assert file_path.exists() assert metadata["file_metadata"]["path"] == str(file_path) @@ -153,14 +162,14 @@ def test_polars_json(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: def test_polars_ndjson(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: file = tmp_path / "test.ndjson" - writer = PolarsNDJSONWriter(file=file) + writer = PolarsSinkNDJSONWriter(file=file) writer.save_data(df) reader = PolarsNDJSONReader(source=file) kwargs2 = reader._get_loading_kwargs() df2, _metadata = reader.load_data(pl.DataFrame) - assert PolarsNDJSONWriter.applicable_types() == [pl.DataFrame, pl.LazyFrame] + assert PolarsSinkNDJSONWriter.applicable_types() == [pl.LazyFrame] assert PolarsNDJSONReader.applicable_types() == [pl.DataFrame] assert df2.shape == (2, 2) assert "schema" not in kwargs2 @@ -219,14 +228,15 @@ def test_polars_spreadsheet(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: def test_polars_lazyframe_sink_parquet(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: file = tmp_path / "test.parquet" - sink = PolarsSinkParquetWriter(path=file) + sink = PolarsSinkParquetWriter(file=file) kwargs = sink._get_saving_kwargs() - sink.save_data(df) + metadata = sink.save_data(df) df2 = pl.read_parquet(file) assert PolarsSinkParquetWriter.applicable_types() == [pl.LazyFrame] assert file.exists() assert kwargs["compression"] == "zstd" + assert metadata["file_metadata"]["path"] == str(file) assert_frame_equal(df.collect(), df2) @@ -235,7 +245,7 @@ def test_polars_lazyframe_sink_parquet_custom_kwargs( ) -> None: """Test that non-default kwargs are passed through correctly.""" file = tmp_path / "test.parquet" - sink = PolarsSinkParquetWriter(path=file, compression="snappy", maintain_order=False) + sink = PolarsSinkParquetWriter(file=file, compression="snappy", maintain_order=False) kwargs = sink._get_saving_kwargs() sink.save_data(df) df2 = pl.read_parquet(file) @@ -243,27 +253,28 @@ def test_polars_lazyframe_sink_parquet_custom_kwargs( assert kwargs["compression"] == "snappy" assert kwargs["maintain_order"] is False assert file.exists() - assert_frame_equal(df.collect(), df2) + assert_frame_equal(df.collect().sort(["a", "b"]), df2.sort(["a", "b"])) def test_polars_lazyframe_sink_csv(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: file = tmp_path / "test.csv" - sink = PolarsSinkCSVWriter(path=file) + sink = PolarsSinkCSVWriter(file=file) kwargs = sink._get_saving_kwargs() - sink.save_data(df) + metadata = sink.save_data(df) df2 = pl.read_csv(file) assert PolarsSinkCSVWriter.applicable_types() == [pl.LazyFrame] assert file.exists() assert kwargs["separator"] == "," assert kwargs["include_header"] is True + assert metadata["file_metadata"]["path"] == str(file) assert_frame_equal(df.collect(), df2) def test_polars_lazyframe_sink_csv_custom_kwargs(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: """Test that non-default kwargs are passed through correctly.""" file = tmp_path / "test.csv" - sink = PolarsSinkCSVWriter(path=file, separator=";", include_header=False) + sink = PolarsSinkCSVWriter(file=file, separator=";", include_header=False) kwargs = sink._get_saving_kwargs() sink.save_data(df) df2 = pl.read_csv(file, separator=";", has_header=False, new_columns=["a", "b"]) @@ -276,21 +287,22 @@ def test_polars_lazyframe_sink_csv_custom_kwargs(df: pl.LazyFrame, tmp_path: pat def test_polars_lazyframe_sink_ipc(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: file = tmp_path / "test.ipc" - sink = PolarsSinkFeatherWriter(path=file) + sink = PolarsSinkFeatherWriter(file=file) kwargs = sink._get_saving_kwargs() - sink.save_data(df) + metadata = sink.save_data(df) df2 = pl.read_ipc(file) assert PolarsSinkFeatherWriter.applicable_types() == [pl.LazyFrame] assert file.exists() assert kwargs["compression"] == "uncompressed" + assert metadata["file_metadata"]["path"] == str(file) assert_frame_equal(df.collect(), df2) def test_polars_lazyframe_sink_ipc_custom_kwargs(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: """Test that non-default kwargs are passed through correctly.""" file = tmp_path / "test.ipc" - sink = PolarsSinkFeatherWriter(path=file, compression="lz4", maintain_order=False) + sink = PolarsSinkFeatherWriter(file=file, compression="lz4", maintain_order=False) kwargs = sink._get_saving_kwargs() sink.save_data(df) df2 = pl.read_ipc(file) @@ -298,19 +310,20 @@ def test_polars_lazyframe_sink_ipc_custom_kwargs(df: pl.LazyFrame, tmp_path: pat assert kwargs["compression"] == "lz4" assert kwargs["maintain_order"] is False assert file.exists() - assert_frame_equal(df.collect(), df2) + assert_frame_equal(df.collect().sort(["a", "b"]), df2.sort(["a", "b"])) def test_polars_lazyframe_sink_ndjson(df: pl.LazyFrame, tmp_path: pathlib.Path) -> None: file = tmp_path / "test.ndjson" - sink = PolarsSinkNDJSONWriter(path=file) + sink = PolarsSinkNDJSONWriter(file=file) kwargs = sink._get_saving_kwargs() - sink.save_data(df) + metadata = 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 metadata["file_metadata"]["path"] == str(file) assert_frame_equal(df.collect(), df2) @@ -319,14 +332,14 @@ def test_polars_lazyframe_sink_ndjson_custom_kwargs( ) -> None: """Test that non-default kwargs are passed through correctly.""" file = tmp_path / "test.ndjson" - sink = PolarsSinkNDJSONWriter(path=file, maintain_order=False) + sink = PolarsSinkNDJSONWriter(file=file, maintain_order=False) kwargs = sink._get_saving_kwargs() sink.save_data(df) df2 = pl.read_ndjson(file) assert kwargs["maintain_order"] is False assert file.exists() - assert_frame_equal(df.collect(), df2) + assert_frame_equal(df.collect().sort(["a", "b"]), df2.sort(["a", "b"])) def test_polars_lazyframe_sink_feather_adapter_name() -> None: @@ -347,3 +360,69 @@ def test_polars_lazyframe_sink_parquet_adapter_name() -> None: def test_polars_lazyframe_sink_ndjson_adapter_name() -> None: """Test that PolarsSinkNDJSONWriter is registered under 'ndjson'.""" assert PolarsSinkNDJSONWriter.name() == "ndjson" + + +@pytest.mark.parametrize( + ("adapter_name", "expected_saver"), + [ + ("csv", PolarsSinkCSVWriter), + ("parquet", PolarsSinkParquetWriter), + ("feather", PolarsSinkFeatherWriter), + ("ndjson", PolarsSinkNDJSONWriter), + ], +) +def test_lazyframe_sink_registry_resolution(adapter_name, expected_saver) -> None: + """Each format resolves one LazyFrame saver, independent of registration order.""" + registered_savers = registry.SAVER_REGISTRY[adapter_name] + applicable_savers = [saver for saver in registered_savers if saver.applies_to(pl.LazyFrame)] + + assert applicable_savers == [expected_saver] + assert resolve_adapter_class(pl.LazyFrame, registered_savers) is expected_saver + assert resolve_adapter_class(pl.LazyFrame, list(reversed(registered_savers))) is expected_saver + + +@pytest.mark.parametrize( + ("eager_saver", "streaming_saver"), + [ + (Pre1PolarsCSVWriter, PolarsSinkCSVWriter), + (Pre1PolarsParquetWriter, PolarsSinkParquetWriter), + (Pre1PolarsFeatherWriter, PolarsSinkFeatherWriter), + ], +) +def test_pre_1_0_lazyframe_sink_resolution_is_unambiguous(eager_saver, streaming_saver) -> None: + assert eager_saver.applicable_types() == [pl.DataFrame] + assert resolve_adapter_class(pl.LazyFrame, [eager_saver, streaming_saver]) is streaming_saver + assert resolve_adapter_class(pl.LazyFrame, [streaming_saver, eager_saver]) is streaming_saver + + +def test_lazyframe_sink_csv_complete_kwargs(tmp_path: pathlib.Path) -> None: + sink = PolarsSinkCSVWriter( + file=tmp_path / "output.csv", + include_bom=True, + decimal_comma=True, + storage_options={"key": "value"}, + extra_kwargs={"retries": 7}, + ) + + assert sink._get_saving_kwargs()["include_bom"] is True + assert sink._get_saving_kwargs()["decimal_comma"] is True + assert sink._get_saving_kwargs()["storage_options"] == {"key": "value"} + assert sink._get_saving_kwargs()["retries"] == 7 + + +def test_lazyframe_sink_materializer_preserves_file_argument(tmp_path: pathlib.Path) -> None: + """The streaming saver must remain compatible with the existing Polars ``file`` API.""" + + def lazy_data() -> pl.LazyFrame: + return pl.LazyFrame({"a": [1, 2], "b": [3, 4]}) + + output_file = tmp_path / "output.csv" + module = ad_hoc_utils.create_temporary_module(lazy_data) + dr = driver.Driver({}, module) + + materialization_result, _ = dr.materialize( + to.csv(id="save_lazy_data", dependencies=["lazy_data"], file=output_file) + ) + + assert materialization_result["save_lazy_data"]["file_metadata"]["path"] == str(output_file) + assert_frame_equal(pl.read_csv(output_file), lazy_data().collect()) From 1a4577f4f1e5619ca2fa6e75e8871f773957a620 Mon Sep 17 00:00:00 2001 From: sonalishintre <42985737+sonalishintre@users.noreply.github.com> Date: Fri, 14 Aug 2026 05:00:37 +0530 Subject: [PATCH 6/7] fix: fix SchemaDefinition import error in polars_pre_1_0_0_extension --- hamilton/plugins/polars_pre_1_0_0_extension.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/hamilton/plugins/polars_pre_1_0_0_extension.py b/hamilton/plugins/polars_pre_1_0_0_extension.py index bd3b54d2b..5c4020826 100644 --- a/hamilton/plugins/polars_pre_1_0_0_extension.py +++ b/hamilton/plugins/polars_pre_1_0_0_extension.py @@ -44,9 +44,11 @@ # for polars 0.18.0 we need to check what to do. if has_alias and hasattr(pl.type_aliases, "CsvEncoding"): - from polars.type_aliases import CsvEncoding, SchemaDefinition + from polars.type_aliases import CsvEncoding + SchemaDefinition = type else: CsvEncoding = type + SchemaDefinition = type if has_alias and hasattr(pl.type_aliases, "CsvQuoteStyle"): from polars.type_aliases import CsvQuoteStyle else: From 876d041cd87e75800ef411834b68ef8e4a7fb6ff Mon Sep 17 00:00:00 2001 From: jernejfrank Date: Mon, 17 Aug 2026 10:58:06 +0100 Subject: [PATCH 7/7] Minor fixes based on CI tests --- .../plugins/polars_lazyframe_extensions.py | 10 ++++- .../plugins/polars_post_1_0_0_extensions.py | 16 ++++++-- .../plugins/polars_pre_1_0_0_extension.py | 13 ++++-- .../test_polars_lazyframe_extensions.py | 41 ++++++++++++++++++- 4 files changed, 71 insertions(+), 9 deletions(-) diff --git a/hamilton/plugins/polars_lazyframe_extensions.py b/hamilton/plugins/polars_lazyframe_extensions.py index 1c57188b7..342171e6d 100644 --- a/hamilton/plugins/polars_lazyframe_extensions.py +++ b/hamilton/plugins/polars_lazyframe_extensions.py @@ -276,6 +276,8 @@ def _get_saving_kwargs(self) -> dict[str, Any]: 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 @@ -411,6 +413,8 @@ def _get_saving_kwargs(self) -> dict[str, Any]: 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 @@ -490,7 +494,7 @@ class PolarsSinkFeatherWriter(DataSaver): file: str | Path # kwargs: - compression: str = "uncompressed" + compression: str | None = None compat_level: Any = None record_batch_size: int | None = None maintain_order: bool = True @@ -528,6 +532,8 @@ def _get_saving_kwargs(self) -> dict[str, Any]: 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 @@ -595,6 +601,8 @@ def _get_saving_kwargs(self) -> dict[str, Any]: 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 diff --git a/hamilton/plugins/polars_post_1_0_0_extensions.py b/hamilton/plugins/polars_post_1_0_0_extensions.py index 44b379646..2f4d7131e 100644 --- a/hamilton/plugins/polars_post_1_0_0_extensions.py +++ b/hamilton/plugins/polars_post_1_0_0_extensions.py @@ -227,7 +227,9 @@ def _get_saving_kwargs(self): kwargs["quote_style"] = self.quote_style return kwargs - def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: + def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: + if isinstance(data, pl.LazyFrame): + data = data.collect() data.write_csv(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) @@ -334,7 +336,9 @@ def _get_saving_kwargs(self): kwargs["pyarrow_options"] = self.pyarrow_options return kwargs - def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: + def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: + if isinstance(data, pl.LazyFrame): + data = data.collect() data.write_parquet(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) @@ -417,7 +421,9 @@ def _get_saving_kwargs(self): kwargs["compression"] = self.compression return kwargs - def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: + def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: + if isinstance(data, pl.LazyFrame): + data = data.collect() data.write_ipc(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) @@ -595,7 +601,9 @@ class PolarsNDJSONWriter(DataSaver): def applicable_types(cls) -> Collection[type]: return [DATAFRAME_TYPE] - def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: + def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: + if isinstance(data, pl.LazyFrame): + data = data.collect() data.write_ndjson(self.file) return utils.get_file_and_dataframe_metadata(self.file, data) diff --git a/hamilton/plugins/polars_pre_1_0_0_extension.py b/hamilton/plugins/polars_pre_1_0_0_extension.py index 5c4020826..bc01f13c6 100644 --- a/hamilton/plugins/polars_pre_1_0_0_extension.py +++ b/hamilton/plugins/polars_pre_1_0_0_extension.py @@ -45,6 +45,7 @@ # for polars 0.18.0 we need to check what to do. if has_alias and hasattr(pl.type_aliases, "CsvEncoding"): from polars.type_aliases import CsvEncoding + SchemaDefinition = type else: CsvEncoding = type @@ -238,7 +239,9 @@ def _get_saving_kwargs(self): kwargs["quote_style"] = self.quote_style return kwargs - def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: + def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: + if isinstance(data, pl.LazyFrame): + data = data.collect() data.write_csv(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) @@ -345,7 +348,9 @@ def _get_saving_kwargs(self): kwargs["pyarrow_options"] = self.pyarrow_options return kwargs - def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: + def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: + if isinstance(data, pl.LazyFrame): + data = data.collect() data.write_parquet(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) @@ -428,7 +433,9 @@ def _get_saving_kwargs(self): kwargs["compression"] = self.compression return kwargs - def save_data(self, data: DATAFRAME_TYPE) -> dict[str, Any]: + def save_data(self, data: DATAFRAME_TYPE | pl.LazyFrame) -> dict[str, Any]: + if isinstance(data, pl.LazyFrame): + data = data.collect() data.write_ipc(self.file, **self._get_saving_kwargs()) return utils.get_file_and_dataframe_metadata(self.file, data) diff --git a/tests/plugins/test_polars_lazyframe_extensions.py b/tests/plugins/test_polars_lazyframe_extensions.py index c4b6a7411..c1ac7743d 100644 --- a/tests/plugins/test_polars_lazyframe_extensions.py +++ b/tests/plugins/test_polars_lazyframe_extensions.py @@ -38,12 +38,15 @@ from hamilton.plugins.polars_post_1_0_0_extensions import ( PolarsAvroReader, PolarsAvroWriter, + PolarsCSVWriter, PolarsDatabaseReader, PolarsDatabaseWriter, PolarsFeatherWriter, PolarsJSONReader, PolarsJSONWriter, PolarsNDJSONReader, + PolarsNDJSONWriter, + PolarsParquetWriter, PolarsSpreadsheetReader, PolarsSpreadsheetWriter, ) @@ -294,7 +297,7 @@ def test_polars_lazyframe_sink_ipc(df: pl.LazyFrame, tmp_path: pathlib.Path) -> assert PolarsSinkFeatherWriter.applicable_types() == [pl.LazyFrame] assert file.exists() - assert kwargs["compression"] == "uncompressed" + assert "compression" not in kwargs assert metadata["file_metadata"]["path"] == str(file) assert_frame_equal(df.collect(), df2) @@ -410,6 +413,42 @@ def test_lazyframe_sink_csv_complete_kwargs(tmp_path: pathlib.Path) -> None: assert sink._get_saving_kwargs()["retries"] == 7 +@pytest.mark.parametrize( + "sink_class", + [ + PolarsSinkCSVWriter, + PolarsSinkParquetWriter, + PolarsSinkFeatherWriter, + PolarsSinkNDJSONWriter, + ], +) +def test_lazyframe_sinks_reject_lazy_extra_kwarg(sink_class, tmp_path: pathlib.Path) -> None: + sink = sink_class(file=tmp_path / "output", extra_kwargs={"lazy": True}) + + with pytest.raises(ValueError, match="lazy=True"): + sink._get_saving_kwargs() + + +@pytest.mark.parametrize( + ("writer_class", "extension"), + [ + (PolarsCSVWriter, "csv"), + (PolarsParquetWriter, "parquet"), + (PolarsFeatherWriter, "ipc"), + (PolarsNDJSONWriter, "ndjson"), + ], +) +def test_eager_writers_preserve_direct_lazyframe_support( + writer_class, extension: str, df: pl.LazyFrame, tmp_path: pathlib.Path +) -> None: + output_file = tmp_path / f"output.{extension}" + + metadata = writer_class(file=output_file).save_data(df) + + assert output_file.exists() + assert metadata["dataframe_metadata"]["rows"] == 2 + + def test_lazyframe_sink_materializer_preserves_file_argument(tmp_path: pathlib.Path) -> None: """The streaming saver must remain compatible with the existing Polars ``file`` API."""