Polars Lazyframe Support

[buggtb]

This PR is to aid for support of Polars LazyFrames in Hamilton.

Changes

I've currently stubbed out the CSV Reader / Writer to work both on Eager and Lazy mode in Polars.

How I tested this

There is an accompanying test that mimics the Eager test but using Lazyframes instead

Notes

I've also had to update the get_dataframe_metadata in utils.py to allow it to work with Lazyframes that don't have a row count. I abstracted all the lookups so that if others passed/failed in the future for support of other read/writers they would return what they can.

Checklist

• PR has an informative and human-readable title (this will be pulled into the release notes)
• Changes are limited to a single goal (no scope creep)
• Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
• Any change in functionality is tested
• New functions are documented (with a description, list of inputs, and expected output)
• Placeholder code is flagged / future TODOs are captured in comments
• Project documentation has been updated if adding/changing functionality.

[skrawcz]

Thanks @buggtb for driving this!

I think we can simplify a few things actually.

1. polars_extensions.py: All existing polars DataSavers can have the following applicable_types():

    @classmethod
    def applicable_types(cls) -> Collection[Type]:
        return [DATAFRAME_TYPE, pl.LazyFrame]

Then in the save_data() method, add union type annotation with pl.LazyFrame, and then in the body:

if isinstance(data, pl.LazyFrame):
    data = data.collect()
# continue with rest of function

2. h_polars.py should be modified to:
   (a) include a PolarsLazyDataFrameResult -- this won't call collect, it will just return a LazyFrame. (though not sure if it makes sense)
   (b) the existing PolarsDataFrameResult should have isinstance checks for lazyframe, and if found, should collect and move on.
   This means we can delete h_polars_lazyframe.py.

3. polars_lazyframe_extensions.py: I would use spark_extensions.py as the template. No need to copy all of what polars_extesions.py has. Otherwise it should house the scan_* and sink_* DataLoaders and DataSavers, e.g. scan_csv. Let me know if you need more guidance here.

4. Remove polars_shared.py.
   I think this is a little over engineering. I don't see the value in the abstraction right now -- I think a little duplication is more straightforward at this time.

5. io/utils.py. I think the changes here are okay. @elijahbenizzy and I can think about whether there's a nicer way to do this.

[buggtb]

Thanks for the review @skrawcz

A couple of follow up questions:

https://github.com/DAGWorks-Inc/hamilton/pull/775/files#diff-fbd89fcde3b2c949a6f81da8ecfd52a163490f1645f1eca8f9451047cf7785faR179

the save function is fair enough, the load is more problematic in terms of design, if I want to load into a lazy frame instead of a dataframe, there's nothing in there currently to allow that. I could extend that function to default to dataframe and let a user request a lazyframe I guess. I'm unsure how you'd want that pattern to look perhaps something like this, or something else:

    def load_data(self, type_: Type, frametype: Union[DATAFRAME_TYPE, pl.LazyFrame] = pl.DataFrame) -> Tuple[DATAFRAME_TYPE, Dict[str, Any]]:
        if isinstance(type_, pl.LazyFrame):
            df = pl.scan_csv(self.file, **self._get_loading_kwargs())
        else:
            df = pl.read_csv(self.file, **self._get_loading_kwargs())

        metadata = utils.get_file_and_dataframe_metadata(self.file, df)
        return df, metadata

also there are some minor differences in kwargs in the lazyframe when compared to the dataframe and so I'd have to add some conditional logic to the get_kwargs() functions to return the valid sets there.

Not entirely sure how the spark extensions example applies here as I suspect we'll need loaders and savers for each type, unless we can overload an existing class but that doesn't appear to be what the spark_extensions example does, but I'm probably missing something.

[skrawcz]

Thanks for the review @skrawcz

A couple of follow up questions:

https://github.com/DAGWorks-Inc/hamilton/pull/775/files#diff-fbd89fcde3b2c949a6f81da8ecfd52a163490f1645f1eca8f9451047cf7785faR179

the save function is fair enough, the load is more problematic in terms of design, if I want to load into a lazy frame instead of a dataframe, there's nothing in there currently to allow that. I could extend that function to default to dataframe and let a user request a lazyframe I guess. I'm unsure how you'd want that pattern to look perhaps something like this, or something else:

    def load_data(self, type_: Type, frametype: Union[DATAFRAME_TYPE, pl.LazyFrame] = pl.DataFrame) -> Tuple[DATAFRAME_TYPE, Dict[str, Any]]:
        if isinstance(type_, pl.LazyFrame):
            df = pl.scan_csv(self.file, **self._get_loading_kwargs())
        else:
            df = pl.read_csv(self.file, **self._get_loading_kwargs())

        metadata = utils.get_file_and_dataframe_metadata(self.file, df)
        return df, metadata

also there are some minor differences in kwargs in the lazyframe when compared to the dataframe and so I'd have to add some conditional logic to the get_kwargs() functions to return the valid sets there.

Yep, so I'm thinking don't couple them at all precisely for the reasons you mentioned. So polars_lazyframe_extensions.py would contain the "load" ones specific to creating lazyframes that pertain to scan_csv, scan_parquet etc. Does that clarify it more?

Not entirely sure how the spark extensions example applies here as I suspect we'll need loaders and savers for each type, unless we can overload an existing class but that doesn't appear to be what the spark_extensions example does, but I'm probably missing something.

I just mean to use it in terms of structure. Since pyspark doesn't make use of the column stuff. Lazyframes wouldn't make use of it either.

[buggtb]

Okay so I removed Polars shared, I then removed most from polars_lazyframe_exensions but kept the loader. Can you check to make sure it's as you imaged, before I go fill out the other save/loaders etc. I see how the save_to and load_from annotations work it makes sense now, I wasn't sure how you could blend loaders and savers across classes, all good.

In h_polars.py I took the same logic I had in the lazyframe version and I do run collect and return a dataframe. When I was running it, to me at least it made sense to collect at the end and return a dataframe as the pipeline is done, but now I think about it there are times, perhaps if you're running chained pipelines, subdags or whatever that you may not want to, I don't know, the jury is out there!

[skrawcz]

Yep looking good. Added code for what I would change.

[skrawcz]

In h_polars.py I took the same logic I had in the lazyframe version and I do run collect and return a dataframe. When I was running it, to me at least it made sense to collect at the end and return a dataframe as the pipeline is done, but now I think about it there are times, perhaps if you're running chained pipelines, subdags or whatever that you may not want to, I don't know, the jury is out there!

Yep! So that's a vote for a second one specific to LazyFrames IMO.

[skrawcz]

Looking great. A few minor things.

We should also have a lazyframe test for each of the savers. If code duplication is the simplest here, I think that's fine; not sure there's a DRY way that would be worth the time/effort since this is a one off.

[skrawcz]

Thanks @buggtb for the work! Looking good I think we just need an example to help people show how to use the new functionality you just added!

So two things:

1. Can you add an example exercising the scan variants? E.g. add something to examples/polars/lazyframe? I think we should add this to complete this PR.
2. There are two minor things (see suggested commits), I can do the commit before merging, or you can. Let me know.

Otherwise question, should I create an issue for the sink_* variants for lazyframe writing to csv, etc. to track that? No need to add it in this PR, but I can see that be a natural progression if we have scan_* support.

Thanks!

[buggtb]

Added a test which loads a CSV using the annotation, does a simple sum and dumps out the result as a LazyFrame, added a comment in my_script.py explaining how you could use either LazyFrame or DataFrame as a resultset depending on what you needed.

Added #792 and #791 on the issues lists for some missing sources and sinks

[skrawcz]

Looks great let's 🚢 it.

But, argh, I see there's a rebase that's needed. Mind doing that? Then we'll be able to squash merge.

[buggtb]

Pulled the latest main branch into my fork, should be good, I think.

[skrawcz]

hmm -- 🤔

[buggtb]

I'm not sure why I don't see the same

[skrawcz]

Done. Thanks @buggtb 🍾 . This will go out next week.