Fix issue 791: add LazyFrame sink classes for parquet, csv, ipc, ndjson by sonalishintre · Pull Request #1653 · apache/hamilton

sonalishintre · 2026-06-29T00:31:31Z

Summary

Closes #791

Adds data sink support for Polars LazyFrames, allowing users to write
LazyFrames directly to disk without calling lf.collect() first, which
is more performant.

Changes

Added PolarsLazyFrameSinkParquet class
Added PolarsLazyFrameSinkCSV class
Added PolarsLazyFrameSinkIPC class
Added PolarsLazyFrameSinkNDJSON class
Registered all new sinks in register_data_loaders()

How I tested this

Added 4 new tests in tests/plugins/test_polars_lazyframe_extensions.py:

test_polars_lazyframe_sink_parquet ✅
test_polars_lazyframe_sink_csv ✅
test_polars_lazyframe_sink_ipc ✅
test_polars_lazyframe_sink_ndjson ✅

Each test creates a LazyFrame, sinks it to a temp file, reads it back,
and verifies the data matches.

Notes

sink_ndjson has no corresponding loader yet as noted in the issue

Checklist

PR has an informative and human-readable title
Changes are limited to a single goal (no scope creep)
Any change in functionality is tested
New functions are documented (with a description and docstring)

sonalishintre · 2026-06-29T00:44:01Z

Hi @skrawcz,
this is my first contribution to this project.
I've implemented the LazyFrame sink classes requested in #791
(sink_parquet, sink_csv, sink_ipc, sink_ndjson) with tests for each.
Would appreciate a review when you have time!

jernejfrank

Hi @sonalishintre and welcome. Thanks for taking the time to contribute. In general looks good just a couple of remarks to keep it consistent with the rest of the existing implementations (see below).

If you look at the https://github.com/apache/hamilton/blob/main/hamilton/plugins/polars_post_1_0_0_extensions.py you will see the same implementations for the eager frame and can see the class naming conventions and order in which the classes appear. Could you try to mirror that?

Otherwise looking good and great job!

jernejfrank · 2026-07-02T10:36:12Z

+        return [DATAFRAME_TYPE]
+
+    def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+        data.sink_parquet(self.path)


Could you add the ability to add the lazyframe.sink_parquet kwargs? https://docs.pola.rs/api/python/stable/reference/api/polars.LazyFrame.sink_parquet.html

jernejfrank · 2026-07-02T10:37:27Z

        return "feather"

+@dataclasses.dataclass
+class PolarsLazyFrameSinkParquet(DataSaver):


Can we name this PolarsSinkParquetWriter?

jernejfrank · 2026-07-02T10:38:28Z

+        return [DATAFRAME_TYPE]
+
+    def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+        data.sink_ipc(self.path)


same here adding kwargs: https://docs.pola.rs/api/python/dev/reference/api/polars.LazyFrame.sink_ipc.html

jernejfrank · 2026-07-02T10:39:47Z

+
+
+@dataclasses.dataclass
+class PolarsLazyFrameSinkIPC(DataSaver):


Same rename style: PolarsSinkFeatherWriter

jernejfrank · 2026-07-02T10:41:01Z

+
+
+@dataclasses.dataclass
+class PolarsLazyFrameSinkNDJSON(DataSaver):


Same rename: PolarsSinkNDJSONWriter

jernejfrank · 2026-07-02T10:41:10Z

+        return [DATAFRAME_TYPE]
+
+    def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+        data.sink_ndjson(self.path)


Adding kwargs: https://docs.pola.rs/api/python/dev/reference/api/polars.LazyFrame.sink_ndjson.html

jernejfrank · 2026-07-02T10:42:25Z

+
+
+@dataclasses.dataclass
+class PolarsLazyFrameSinkCSV(DataSaver):


Same rename: PolarsSinkCSVWriter

jernejfrank · 2026-07-02T10:42:47Z

+        return [DATAFRAME_TYPE]
+
+    def save_data(self, data: pl.LazyFrame) -> dict[str, Any]:
+        data.sink_csv(self.path)


And adding kwargs: https://docs.pola.rs/api/python/stable/reference/api/polars.LazyFrame.sink_csv.html

…dback

sonalishintre · 2026-07-20T23:29:46Z

Hi @jernejfrank — I've made the requested changes:

Renamed classes to match existing naming convention
(PolarsSinkParquetWriter, PolarsSinkCSVWriter,
PolarsSinkFeatherWriter, PolarsSinkNDJSONWriter)
Added kwargs support with _get_saving_kwargs() for all 4 sink classes
All 4 tests passing

Please let me know if any further changes are needed!
Thank You

jernejfrank · 2026-07-23T20:11:28Z

Thanks for the update. I took another pass through the latest revision and found a few things that should be addressed before merging:

Feather saver is registered under the wrong adapter name

PolarsSinkFeatherWriter.name() currently returns "ipc", whereas PolarsScanFeatherReader and the existing PolarsFeatherWriter both use "feather". As a result, to.feather(...) will not select this new streaming saver and can continue selecting the eager writer, which collects the LazyFrame. Could this return "feather", with an integration test confirming that a LazyFrame passed to to.feather selects PolarsSinkFeatherWriter?
LazyFrame saver selection is ambiguous and registration-order dependent

The existing eager CSV, Parquet, Feather, and NDJSON writers already advertise support for pl.LazyFrame. The new sink writers advertise the same type under the same adapter names, so Hamilton reports more than one applicable adapter and chooses based on registration order. A different manual import order could therefore select the eager writer and call collect() again. Could we make LazyFrame selection unambiguous and test adapter resolution through the registry/materializer API?
The kwargs support is incomplete

The dataclass fields plus _get_saving_kwargs() pattern is consistent with the other implementations, but not all intended sink options are represented. Most clearly, LazyFrame.sink_csv() supports include_header=True, and the existing PolarsCSVWriter exposes it too, but PolarsSinkCSVWriter does not. IPC and NDJSON also expose only a small subset of their documented options. Could we cover the compatible options we intend to support, or document why specific ones are intentionally omitted? (lazy=True may need to remain excluded or receive special handling because it conflicts with save_data() expecting the file to exist immediately.)
The tests do not verify custom kwargs or adapter integration

The tests currently inspect default values such as separator == ",", but do not instantiate savers with non-default kwargs and verify their effect. They also instantiate the classes directly, which misses the "ipc"/"feather" registration issue and duplicate saver selection. Could we add at least one non-default kwarg test per relevant saver and a registry/materializer resolution test?
The superseded implementation and tests are commented out rather than removed

The old four saver classes and duplicate registration function remain commented out near the end of polars_lazyframe_extensions.py, and the old four tests remain commented out near the end of the test file. Could these dead blocks be deleted before merging?
The surrounding-code ordering was not fully mirrored

The established Polars module uses CSV, Parquet, Feather, etc., generally placing each reader and writer together. This revision puts Parquet before CSV and groups all new writers after the readers. Could we mirror the surrounding module's ordering as requested in the previous review?
Formatting checks currently fail

Ruff reports trailing whitespace in polars_lazyframe_extensions.py, an unsorted import block in the test file, and both changed files require formatting. git diff --check also reports an extra blank line at EOF. There are also temporary-looking #### separators. Please run the configured formatting/lint hooks before the next review.

For reference, the functional test file passes locally (11 passed, 1 skipped), but the current tests do not cover the registry and kwargs issues above.

jernejfrank · 2026-07-23T20:13:33Z

@sonalishintre let me know if any of the above are unclear or you need help with any of the points. Happy to sketch it out / point to the right places so you can get this over the line

sonalishintre added 2 commits June 22, 2026 08:52

feat: add LazyFrame sink classes for parquet, csv, ipc, ndjson

d8e3e28

test: add tests for LazyFrame sink classes

dbb400c

jernejfrank reviewed Jul 2, 2026

View reviewed changes

refactor: rename sink classes and add kwargs support per reviewer fee…

4ed01ea

…dback



		@dataclasses.dataclass
		class PolarsLazyFrameSinkIPC(DataSaver):



		@dataclasses.dataclass
		class PolarsLazyFrameSinkNDJSON(DataSaver):



		@dataclasses.dataclass
		class PolarsLazyFrameSinkCSV(DataSaver):

Uh oh!

Conversation

sonalishintre commented Jun 29, 2026

Summary

Changes

How I tested this

Notes

Checklist

Uh oh!

sonalishintre commented Jun 29, 2026

Uh oh!

jernejfrank left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sonalishintre commented Jul 20, 2026

Uh oh!

jernejfrank commented Jul 23, 2026

Uh oh!

jernejfrank commented Jul 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants