feat(datasets): support Polars lazy evaluation

[MatthiasRoels]

Description

For optimal data processing, it is recommended to use Polar's Lazy API (instead of the Eager one) so it is only natural that kedro-datasets supports it too.

Closes #224

Development notes

Checklist

• Opened this PR as a 'Draft Pull Request' if it is work-in-progress
• Updated the documentation to reflect the code changes
• Added a description of this change in the relevant RELEASE.md file
• Added tests to cover my changes

[noklam]

You are killing it with this rate of PR 🔥

[MatthiasRoels]

And just as I created this PR, Polars releases a new version with native support to read parquet files from AWS/GCP/Azure... From what I can see, it should also work for the lazy API.

So, this brings us to the next question: how do we proceed with this PR?

1. Should we just finish it as-is (allowing versions < 0.19.4) and use arrow to read data AWS/GCP/Azure? I have been using this since beginning of summer and never ran into issue, so I know it works properly!
2. Bump polars version to 0.19.4 to allow native reading of parquet files from Cloud? Pro: easier to implement. Con: it only works for parquet for now and it looks like it is experimental...

[astrojuanlu]

Maybe wait a couple of releases to see if it goes from experimental to mature? If we are going to throw away this code in a couple of months, I wouldn't say it makes a lot of sense to continue working on it.

They also said that the first release candidate of polars 1.0 is coming before the end of the year pola-rs/polars#6616 (comment) so maybe things will estabilize soon :)

[MatthiasRoels]

I gave it a little thought and if anyone is willing to review, I would like to finish it, despite Polars’ native support for reading/writing to object store. The reason is simple; if this functionality stabilises on Polars’ side, we only have change (read, simplify) our load/save methods!

The way I see it moving forward is that we need 2 datasets: one for eager loading and one for lazy loading. So I would keep the GenericDataset (potentially rename it to e.g. EagerDataset) and rename mine to LazyDataset. The reason I want to keep the two is that there are some file extensions (e.g. Excel) that don’t have a Lazy loading option.

Curious to hear other opinions!

[astrojuanlu]

There's two options: LazyPolarsGenericDataset (ugh, a bit long) or

if lazy:
  pl.scan_csv(xxx)
else:
  pl.read_csv(xxx)

as @noklam suggested in #224 (comment)

Since the return type is going to be completely different, I'd rather have a different dataset indeed.

[MatthiasRoels]

Is it ok if I:

1. Rename my dataset to LazyDataset (other suggestions welcome)
2. Rename GenericDataset to EagerDataset (other suggestions welcome)

To make it clear that one uses the eager API while the other one uses the Lazy API

Shouldn't we then also consider removing (or mark as deprecated) the CSVDataset? And as a follow-up: should we do the same with the pandas datasets? Basically, the pandas.GenericDataset replaces the CSV, Excel, Feather, JSON, Parquet and XML ones, which should make the code base leaner. This is beyond the scope of this PR obviously...

[astrojuanlu]

Shouldn't we then also consider removing (or mark as deprecated) the CSVDataset? And as a follow-up: should we do the same with the pandas datasets? Basically, the pandas.GenericDataset replaces the CSV, Excel, Feather, JSON, Parquet and XML ones, which should make the code base leaner. This is beyond the scope of this PR obviously...

This is a discussion worth having but it has some implications in the docs, starters etc, could you open a separate issue about it?

[merelcht]

Thanks a lot for this contribution @MatthiasRoels! I've left some comments, the main being around the naming syntax (classes ending in DataSet are old, any new ones should end in Dataset).

In terms of naming the datasets, how about:

• LazyPolarsDataset
• EagerPolarsDataset

[MatthiasRoels]

Thanks a lot for this contribution @MatthiasRoels! I've left some comments, the main being around the naming syntax (classes ending in DataSet are old, any new ones should end in Dataset).

No worries, I am always happy to contribute!

In terms of naming the datasets, how about:

• LazyPolarsDataset
• EagerPolarsDataset

@merelcht That a nice suggestion! But that would introduce a breaking change (renaming the GenericDataset to EagerPolarsDataset). Should I already rename it while keeping the old names with deprecation warnings (like we did for DataSet vs Dataset)?

[merelcht]

@merelcht That a nice suggestion! But that would introduce a breaking change (renaming the GenericDataset to EagerPolarsDataset). Should I already rename it while keeping the old names with deprecation warnings (like we did for DataSet vs Dataset)?

Yes that sounds good! And then we can remove the alias in the next breaking kedro-datasets release 2.0.0 🙂

[merelcht]

I haven't tried running this, but the code looks good! Can you also add this change to the release note + mention that polars.GenericDataSet will be deprecated and be replaced by polars.EagerPolarsDataset?

[merelcht]

Should EagerPolarsDataset be added here as well? If someone uses that directly I'm guessing the requirements would now not be picked up?

[MatthiasRoels]

Read the Docs build failed because of a dependency conflict with Dask in test vs docs. Weird that this does not occur in other PR's (or is it?)...

[merelcht]

Read the Docs build failed because of a dependency conflict with Dask in test vs docs. Weird that this does not occur in other PR's (or is it?)...

We are aware of this issue, it has nothing to do with your PR!

[MatthiasRoels]

Read the Docs build failed because of a dependency conflict with Dask in test vs docs. Weird that this does not occur in other PR's (or is it?)...

We are aware of this issue, it has nothing to do with your PR!

Is there an open issue for this already?

[astrojuanlu]

It got fixed already in #396 ✨

[astrojuanlu]

Tried with a local CSV and it works, but the same CSV on a Minio bucket failed:

In [1]: import polars as pl

In [2]: from kedro_datasets.polars import LazyPolarsDataset

In [3]: ds = LazyPolarsDataset(
   ...:   filepath="s3://temp-openrepair/OpenRepairData_v0.3_aggregate_202210.csv",
   ...:   file_format="csv",
   ...:   load_args=dict(dtypes=dict(product_age=pl.Float64, group_identifier=pl.Utf8), try_parse_dates=True),
   ...: )

In [4]: df_l = ds.load()

In [5]: df_l
Out[5]: <LazyFrame [14 cols, {"id": Utf8 … "problem": Utf8}] at 0x107F1ABB0>

In [6]: df_l.collect().head()
---------------------------------------------------------------------------
ComputeError                              Traceback (most recent call last)
Cell In[6], line 1
----> 1 df_l.collect().head()

File ~/.micromamba/envs/kedro38-dev2/lib/python3.8/site-packages/polars/utils/deprecation.py:96, in deprecate_renamed_parameter.<locals>.decorate.<locals>.wrapper(*args, **kwargs)
     91 @wraps(function)
     92 def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
     93     _rename_keyword_argument(
     94         old_name, new_name, kwargs, function.__name__, version
     95     )
---> 96     return function(*args, **kwargs)

File ~/.micromamba/envs/kedro38-dev2/lib/python3.8/site-packages/polars/lazyframe/frame.py:1787, in LazyFrame.collect(self, type_coercion, predicate_pushdown, projection_pushdown, simplify_expression, slice_pushdown, comm_subplan_elim, comm_subexpr_elim, no_optimization, streaming, _eager)
   1774     comm_subplan_elim = False
   1776 ldf = self._ldf.optimization_toggle(
   1777     type_coercion,
   1778     predicate_pushdown,
   (...)
   1785     _eager,
   1786 )
-> 1787 return wrap_df(ldf.collect())

ComputeError: ArrowInvalid: In CSV column #11: Row #9889: CSV conversion error to int64: invalid value 'Fixit Clinic'

Notice that it's trying to use int64 for the group_identifier column, despite having specified pl.Utf8. Maybe load_args is not being properly passed for fsspec files?

The EagerPolarsDataset doesn't have this problem:

In [16]: ds = EagerPolarsDataset(
    ...:   filepath="s3://temp-openrepair/OpenRepairData_v0.3_aggregate_202210.csv",
    ...:   file_format="csv",
    ...:   load_args=dict(dtypes=dict(product_age=pl.Float64, group_identifier=pl.Utf8), try_parse_dates=True),
    ...: )

In [17]: ds.load().head()
Out[17]: 
shape: (5, 14)
┌─────────────────┬───────────────┬─────────┬─────────────────┬───┬─────────────────┬─────────────────┬────────────┬─────────────────┐
│ id              ┆ data_provider ┆ country ┆ partner_product ┆ … ┆ repair_barrier_ ┆ group_identifie ┆ event_date ┆ problem         │
│ ---             ┆ ---           ┆ ---     ┆ _category       ┆   ┆ if_end_of_life  ┆ r               ┆ ---        ┆ ---             │
│ str             ┆ str           ┆ str     ┆ ---             ┆   ┆ ---             ┆ ---             ┆ date       ┆ str             │
│                 ┆               ┆         ┆ str             ┆   ┆ str             ┆ str             ┆            ┆                 │
╞═════════════════╪═══════════════╪═════════╪═════════════════╪═══╪═════════════════╪═════════════════╪════════════╪═════════════════╡
│ anstiftung_2749 ┆ anstiftung    ┆ DEU     ┆ Elektro divers  ┆ … ┆ null            ┆ 5073            ┆ 2012-06-20 ┆ Funktionierte   │
│                 ┆               ┆         ┆ ~ Nähmaschine   ┆   ┆                 ┆                 ┆            ┆ nicht mehr.     │
│                 ┆               ┆         ┆                 ┆   ┆                 ┆                 ┆            ┆ Fehler…         │
│ anstiftung_2750 ┆ anstiftung    ┆ DEU     ┆ Computer ~      ┆ … ┆ null            ┆ 5073            ┆ 2012-06-20 ┆ Wurde schnell   │
│                 ┆               ┆         ┆ Laptop          ┆   ┆                 ┆                 ┆            ┆ heiß. Der       │
│                 ┆               ┆         ┆                 ┆   ┆                 ┆                 ┆            ┆ Lüfter w…       │
│ anstiftung_2746 ┆ anstiftung    ┆ DEU     ┆ Computer ~      ┆ … ┆ null            ┆ 5073            ┆ 2012-06-20 ┆ Funktionierte   │
│                 ┆               ┆         ┆ Drucker         ┆   ┆                 ┆                 ┆            ┆ nicht mehr.     │
│                 ┆               ┆         ┆                 ┆   ┆                 ┆                 ┆            ┆ Fehler…         │
│ anstiftung_2747 ┆ anstiftung    ┆ DEU     ┆ Unterhaltungsel ┆ … ┆ null            ┆ 5073            ┆ 2012-06-20 ┆ Funktionierte   │
│                 ┆               ┆         ┆ ektronik ~      ┆   ┆                 ┆                 ┆            ┆ nicht mehr.     │
│                 ┆               ┆         ┆ Kopfhö…         ┆   ┆                 ┆                 ┆            ┆ Fehler…         │
│ anstiftung_2742 ┆ anstiftung    ┆ DEU     ┆ Haushaltsgeräte ┆ … ┆ null            ┆ 5073            ┆ 2012-09-19 ┆ Die Beine der   │
│                 ┆               ┆         ┆ ~ Spielzeug     ┆   ┆                 ┆                 ┆            ┆ Puppe waren ab. │
│                 ┆               ┆         ┆                 ┆   ┆                 ┆                 ┆            ┆ Si…             │
└─────────────────┴───────────────┴─────────┴─────────────────┴───┴─────────────────┴─────────────────┴────────────┴─────────────────┘

[astrojuanlu]

This is using https://github.com/astrojuanlu/workshop-jupyter-kedro/tree/workshop-steps/data by the way

[MatthiasRoels]

@astrojuanlu: You were correct that the load args were not properly passed when loading files from object stores. However, this will not fix the issue you had because for object stores, we leverage Arrow to load the dataset. This means that you actually have to pass an Arrow schema (instead of polars dtypes). Unfortunately, you have to pass in the full schema, not just the types you want to manually specify... This is downside from the fact that older Polars version could not read directly from Object stores (which, in the newer versions, is in beta).

So you should be able to do the following now:

import polars as pl
import pyarrow as pa 

from kedro_datasets.polars import LazyPolarsDataset

pa_schema = pa.schema(
    [
        ("id", pa.string()),
        ("data_provider", pa.string()),
        ("country", pa.string()),
        ("partner_product_category", pa.string()),
        ("product_category", pa.string()),
        ("product_category_id", pa.int64()),
        ("brand", pa.string()),
        ("year_of_manufacture", pa.int64()),
        ("product_age", pa.float64()),
        ("repair_status", pa.string()),
        ("repair_barrier_if_end_of_life", pa.string()),
        ("group_identifier", pa.string()),
        ("event_date", pa.date32()),
        ("problem", pa.string()),
    ]
)

ds = LazyPolarsDataset(
    filepath="s3://temp-openrepair/OpenRepairData_v0.3_aggregate_202210.csv",
    file_format="csv",
    load_args={"schema": pa_schema},
)

[astrojuanlu]

I see, thanks! it's a bit annoying to specify the full schema for remote filepaths, do you think this is enough reason to embrace the new method, even if it's beta?

[MatthiasRoels]

@astrojuanlu: I think there are pro's and cons of both approaches

Current approach
Pros:

• It is based on PyArrow which is well-established
• I have been using it for 3 months with 100+ daily kedro runs on different datasets daily (well tested?)

Cons:

• you cannot use convenient Polars methods and you have to in the Arrow docs if you want to do something special like overwriting the schema
• load_args look different when you switch from a file stored to locally to one stored in an object store
• If we change later on, this change could be considered breaking (everyone will potentially need to change load_args).

Approach when we force a newer version of Polars
Pros:

• Uniform syntax
• No further (breaking?) changes required in the future

Cons:

• We have to force users to use a more recent version of Polars (not sure if that's necessarily a bad thing?)
• reading from cloud storage using polars is experimental so we might encounter unexpected behaviour (I have also never tested it so I have no idea what to expect).

Given the pros/cons of each, we need to decide how we proceed. I think it comes down to stability vs convenience, no?

[astrojuanlu]

Good assessment, thanks for the writeup. I keep coming back and forth about this, because I'm pretty confident the new approach will stabilise, but to be honest there's no 100 % guarantee.

Let's proceed with what you have created here, and we can revisit when Polars 1.0 is out 👍🏽

[astrojuanlu]

🚀

[merelcht]

Tiny comment, but otherwise LGTM! Thank you so much for your contribution @MatthiasRoels it's truly awesome ⭐ 😄

[astrojuanlu]

Congrats on your first contribution @MatthiasRoels ! 🎉