diff --git a/hamilton/plugins/polars_lazyframe_extensions.py b/hamilton/plugins/polars_lazyframe_extensions.py index b06f2db94..342171e6d 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 @@ -183,6 +183,117 @@ def name(cls) -> str: return "csv" +@dataclasses.dataclass +class PolarsSinkCSVWriter(DataSaver): + """Class to handle sinking a Polars LazyFrame to a CSV file using streaming. + + Calls LazyFrame.sink_csv() directly, avoiding collect() for better performance. + Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_csv.html + + Note: ``lazy=True`` is intentionally excluded because ``save_data()`` expects + the file to exist immediately after the call returns. + """ + + file: str | Path + # kwargs: + include_bom: bool | None = None + compression: str | None = None + compression_level: int | None = None + check_extension: bool | None = None + include_header: bool = True + separator: str = "," + line_terminator: str = "\n" + quote_char: str = '"' + batch_size: int = 1024 + datetime_format: str | None = None + date_format: str | None = None + time_format: str | None = None + float_scientific: bool | None = None + float_precision: int | None = None + decimal_comma: bool | None = None + null_value: str | None = None + quote_style: Any = None + maintain_order: bool = True + storage_options: dict[str, Any] | None = None + credential_provider: Any = None + retries: int | None = None + sync_on_close: Any = None + mkdir: bool | None = None + engine: Any = None + optimizations: Any = None + extra_kwargs: dict[str, Any] | None = None + + def _get_saving_kwargs(self) -> dict[str, Any]: + kwargs = {} + if self.include_bom is not None: + kwargs["include_bom"] = self.include_bom + if self.compression is not None: + kwargs["compression"] = self.compression + if self.compression_level is not None: + kwargs["compression_level"] = self.compression_level + if self.check_extension is not None: + kwargs["check_extension"] = self.check_extension + if self.include_header is not None: + kwargs["include_header"] = self.include_header + if self.separator is not None: + kwargs["separator"] = self.separator + if self.line_terminator is not None: + kwargs["line_terminator"] = self.line_terminator + if self.quote_char is not None: + kwargs["quote_char"] = self.quote_char + if self.batch_size is not None: + kwargs["batch_size"] = self.batch_size + if self.datetime_format is not None: + kwargs["datetime_format"] = self.datetime_format + if self.date_format is not None: + kwargs["date_format"] = self.date_format + if self.time_format is not None: + kwargs["time_format"] = self.time_format + if self.float_scientific is not None: + kwargs["float_scientific"] = self.float_scientific + if self.float_precision is not None: + kwargs["float_precision"] = self.float_precision + if self.decimal_comma is not None: + kwargs["decimal_comma"] = self.decimal_comma + if self.null_value is not None: + kwargs["null_value"] = self.null_value + if self.quote_style is not None: + kwargs["quote_style"] = self.quote_style + if self.maintain_order is not None: + kwargs["maintain_order"] = self.maintain_order + if self.storage_options is not None: + kwargs["storage_options"] = self.storage_options + if self.credential_provider is not None: + kwargs["credential_provider"] = self.credential_provider + if self.retries is not None: + kwargs["retries"] = self.retries + if self.sync_on_close is not None: + kwargs["sync_on_close"] = self.sync_on_close + if self.mkdir is not None: + kwargs["mkdir"] = self.mkdir + if self.engine is not None: + kwargs["engine"] = self.engine + if self.optimizations is not None: + kwargs["optimizations"] = self.optimizations + if self.extra_kwargs is not None: + if self.extra_kwargs.get("lazy", False): + raise ValueError("lazy=True is incompatible with synchronous data savers.") + kwargs.update(self.extra_kwargs) + return kwargs + + @classmethod + def applicable_types(cls) -> Collection[type]: + return [DATAFRAME_TYPE] + + def save_data(self, data: pl.LazyFrame) -> dict[str, Any]: + data.sink_csv(self.file, **self._get_saving_kwargs()) + return utils.get_file_metadata(self.file) + + @classmethod + def name(cls) -> str: + return "csv" + + @dataclasses.dataclass class PolarsScanParquetReader(DataLoader): """Class specifically to handle loading parquet files with polars @@ -239,6 +350,87 @@ def name(cls) -> str: return "parquet" +@dataclasses.dataclass +class PolarsSinkParquetWriter(DataSaver): + """Class to handle sinking a Polars LazyFrame to a Parquet file using streaming. + + Calls LazyFrame.sink_parquet() directly, avoiding collect() for better performance. + Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_parquet.html + + Note: ``lazy=True`` is intentionally excluded because ``save_data()`` expects + the file to exist immediately after the call returns. + """ + + file: str | Path + # kwargs: + compression: str = "zstd" + compression_level: int | None = None + statistics: bool | str | dict[str, bool] = True + row_group_size: int | None = None + data_page_size: int | None = None + maintain_order: bool = True + storage_options: dict[str, Any] | None = None + credential_provider: Any = None + retries: int | None = None + sync_on_close: Any = None + metadata: Any = None + arrow_schema: Any = None + mkdir: bool | None = None + engine: Any = None + optimizations: Any = None + extra_kwargs: dict[str, Any] | None = None + + def _get_saving_kwargs(self) -> dict[str, Any]: + kwargs = {} + if self.compression is not None: + kwargs["compression"] = self.compression + if self.compression_level is not None: + kwargs["compression_level"] = self.compression_level + if self.statistics is not None: + kwargs["statistics"] = self.statistics + if self.row_group_size is not None: + kwargs["row_group_size"] = self.row_group_size + if self.data_page_size is not None: + kwargs["data_page_size"] = self.data_page_size + if self.maintain_order is not None: + kwargs["maintain_order"] = self.maintain_order + if self.storage_options is not None: + kwargs["storage_options"] = self.storage_options + if self.credential_provider is not None: + kwargs["credential_provider"] = self.credential_provider + if self.retries is not None: + kwargs["retries"] = self.retries + if self.sync_on_close is not None: + kwargs["sync_on_close"] = self.sync_on_close + if self.metadata is not None: + kwargs["metadata"] = self.metadata + if self.arrow_schema is not None: + kwargs["arrow_schema"] = self.arrow_schema + if self.mkdir is not None: + kwargs["mkdir"] = self.mkdir + if self.engine is not None: + kwargs["engine"] = self.engine + if self.optimizations is not None: + kwargs["optimizations"] = self.optimizations + if self.extra_kwargs is not None: + if self.extra_kwargs.get("lazy", False): + raise ValueError("lazy=True is incompatible with synchronous data savers.") + kwargs.update(self.extra_kwargs) + return kwargs + + @classmethod + def applicable_types(cls) -> Collection[type]: + return [DATAFRAME_TYPE] + + def save_data(self, data: pl.LazyFrame) -> dict[str, Any]: + data.sink_parquet(self.file, **self._get_saving_kwargs()) + return utils.get_file_metadata(self.file) + + @classmethod + def name(cls) -> str: + return "parquet" + + @dataclasses.dataclass class PolarsScanFeatherReader(DataLoader): """ @@ -289,12 +481,154 @@ def name(cls) -> str: return "feather" +@dataclasses.dataclass +class PolarsSinkFeatherWriter(DataSaver): + """Class to handle sinking a Polars LazyFrame to an IPC/Feather file using streaming. + + Calls LazyFrame.sink_ipc() directly, avoiding collect() for better performance. + Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_ipc.html + + Note: ``lazy=True`` is intentionally excluded because ``save_data()`` expects + the file to exist immediately after the call returns. + """ + + file: str | Path + # kwargs: + compression: str | None = None + compat_level: Any = None + record_batch_size: int | None = None + maintain_order: bool = True + storage_options: dict[str, Any] | None = None + credential_provider: Any = None + retries: int | None = None + sync_on_close: Any = None + mkdir: bool | None = None + engine: Any = None + optimizations: Any = None + extra_kwargs: dict[str, Any] | None = None + + def _get_saving_kwargs(self) -> dict[str, Any]: + kwargs = {} + if self.compression is not None: + kwargs["compression"] = self.compression + if self.compat_level is not None: + kwargs["compat_level"] = self.compat_level + if self.record_batch_size is not None: + kwargs["record_batch_size"] = self.record_batch_size + if self.maintain_order is not None: + kwargs["maintain_order"] = self.maintain_order + if self.storage_options is not None: + kwargs["storage_options"] = self.storage_options + if self.credential_provider is not None: + kwargs["credential_provider"] = self.credential_provider + if self.retries is not None: + kwargs["retries"] = self.retries + if self.sync_on_close is not None: + kwargs["sync_on_close"] = self.sync_on_close + if self.mkdir is not None: + kwargs["mkdir"] = self.mkdir + if self.engine is not None: + kwargs["engine"] = self.engine + if self.optimizations is not None: + kwargs["optimizations"] = self.optimizations + if self.extra_kwargs is not None: + if self.extra_kwargs.get("lazy", False): + raise ValueError("lazy=True is incompatible with synchronous data savers.") + kwargs.update(self.extra_kwargs) + return kwargs + + @classmethod + def applicable_types(cls) -> Collection[type]: + return [DATAFRAME_TYPE] + + def save_data(self, data: pl.LazyFrame) -> dict[str, Any]: + data.sink_ipc(self.file, **self._get_saving_kwargs()) + return utils.get_file_metadata(self.file) + + @classmethod + def name(cls) -> str: + return "feather" + + +@dataclasses.dataclass +class PolarsSinkNDJSONWriter(DataSaver): + """Class to handle sinking a Polars LazyFrame to an NDJSON file using streaming. + + Calls LazyFrame.sink_ndjson() directly, avoiding collect() for better performance. + Should map to https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.sink_ndjson.html + + Note: ``lazy=True`` is intentionally excluded because ``save_data()`` expects + the file to exist immediately after the call returns. + """ + + file: str | Path + # kwargs: + compression: str | None = None + compression_level: int | None = None + check_extension: bool | None = None + maintain_order: bool = True + storage_options: dict[str, Any] | None = None + credential_provider: Any = None + retries: int | None = None + sync_on_close: Any = None + mkdir: bool | None = None + engine: Any = None + optimizations: Any = None + extra_kwargs: dict[str, Any] | None = None + + def _get_saving_kwargs(self) -> dict[str, Any]: + kwargs = {} + if self.compression is not None: + kwargs["compression"] = self.compression + if self.compression_level is not None: + kwargs["compression_level"] = self.compression_level + if self.check_extension is not None: + kwargs["check_extension"] = self.check_extension + if self.maintain_order is not None: + kwargs["maintain_order"] = self.maintain_order + if self.storage_options is not None: + kwargs["storage_options"] = self.storage_options + if self.credential_provider is not None: + kwargs["credential_provider"] = self.credential_provider + if self.retries is not None: + kwargs["retries"] = self.retries + if self.sync_on_close is not None: + kwargs["sync_on_close"] = self.sync_on_close + if self.mkdir is not None: + kwargs["mkdir"] = self.mkdir + if self.engine is not None: + kwargs["engine"] = self.engine + if self.optimizations is not None: + kwargs["optimizations"] = self.optimizations + if self.extra_kwargs is not None: + if self.extra_kwargs.get("lazy", False): + raise ValueError("lazy=True is incompatible with synchronous data savers.") + kwargs.update(self.extra_kwargs) + return kwargs + + @classmethod + def applicable_types(cls) -> Collection[type]: + return [DATAFRAME_TYPE] + + def save_data(self, data: pl.LazyFrame) -> dict[str, Any]: + data.sink_ndjson(self.file, **self._get_saving_kwargs()) + return utils.get_file_metadata(self.file) + + @classmethod + def name(cls) -> str: + return "ndjson" + + def register_data_loaders(): """Function to register the data loaders for this extension.""" for loader in [ PolarsScanCSVReader, + PolarsSinkCSVWriter, PolarsScanParquetReader, + PolarsSinkParquetWriter, PolarsScanFeatherReader, + PolarsSinkFeatherWriter, + PolarsSinkNDJSONWriter, ]: registry.register_adapter(loader) diff --git a/hamilton/plugins/polars_post_1_0_0_extensions.py b/hamilton/plugins/polars_post_1_0_0_extensions.py index dab1a3973..2f4d7131e 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 = {} @@ -318,7 +318,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 = {} @@ -339,7 +339,6 @@ def _get_saving_kwargs(self): 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) @@ -414,7 +413,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 = {} @@ -600,12 +599,11 @@ class PolarsNDJSONWriter(DataSaver): @classmethod def applicable_types(cls) -> Collection[type]: - return [DATAFRAME_TYPE, pl.LazyFrame] + return [DATAFRAME_TYPE] 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 4dd92dc43..bc01f13c6 100644 --- a/hamilton/plugins/polars_pre_1_0_0_extension.py +++ b/hamilton/plugins/polars_pre_1_0_0_extension.py @@ -44,9 +44,12 @@ # 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: @@ -206,7 +209,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 = {} @@ -327,7 +330,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 = {} @@ -348,7 +351,6 @@ def _get_saving_kwargs(self): 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) @@ -423,7 +425,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 = {} 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 a35929654..c1ac7743d 100644 --- a/tests/plugins/test_polars_lazyframe_extensions.py +++ b/tests/plugins/test_polars_lazyframe_extensions.py @@ -23,10 +23,17 @@ 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, PolarsScanParquetReader, + PolarsSinkCSVWriter, + PolarsSinkFeatherWriter, + PolarsSinkNDJSONWriter, + PolarsSinkParquetWriter, ) from hamilton.plugins.polars_post_1_0_0_extensions import ( PolarsAvroReader, @@ -43,6 +50,15 @@ 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 @@ -53,15 +69,15 @@ 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) 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 PolarsSinkCSVWriter.applicable_types() == [pl.LazyFrame] assert PolarsScanCSVReader.applicable_types() == [pl.LazyFrame] assert kwargs1["separator"] == "," assert kwargs2["has_header"] is True @@ -71,15 +87,15 @@ 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) 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 PolarsSinkParquetWriter.applicable_types() == [pl.LazyFrame] assert PolarsScanParquetReader.applicable_types() == [pl.LazyFrame] assert kwargs1["compression"] == "zstd" assert kwargs2["n_rows"] == 2 @@ -101,7 +117,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) @@ -122,7 +138,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] @@ -138,7 +154,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] @@ -149,14 +165,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) + 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 @@ -179,7 +195,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] @@ -211,3 +227,241 @@ 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(file=file) + kwargs = sink._get_saving_kwargs() + 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) + + +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(file=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().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(file=file) + kwargs = sink._get_saving_kwargs() + 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(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"]) + + assert kwargs["separator"] == ";" + assert kwargs["include_header"] is False + 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 = PolarsSinkFeatherWriter(file=file) + kwargs = sink._get_saving_kwargs() + metadata = sink.save_data(df) + df2 = pl.read_ipc(file) + + assert PolarsSinkFeatherWriter.applicable_types() == [pl.LazyFrame] + assert file.exists() + assert "compression" not in kwargs + 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(file=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().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(file=file) + kwargs = sink._get_saving_kwargs() + 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) + + +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(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().sort(["a", "b"]), df2.sort(["a", "b"])) + + +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" + + +@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 + + +@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.""" + + 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())