feat: Allow specifying the engine for validate and filter (#364)

Author: borchero

dataframely/_polars.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import datetime as dt
-from typing import TypeVar
+from typing import Any, TypeVar
 
 import polars as pl
 from polars.datatypes import DataTypeClass
@@ -41,18 +41,18 @@ def timedelta_matches_resolution(d: dt.timedelta, resolution: str) -> bool:
     return datetime_matches_resolution(EPOCH_DATETIME + d, resolution)
 
 
-def collect_if(lf: pl.LazyFrame, condition: bool) -> pl.LazyFrame:
+def collect_if(lf: pl.LazyFrame, condition: bool, **kwargs: Any) -> pl.LazyFrame:
     """Collect a lazy frame based on `condition`."""
     if condition:
-        return lf.collect().lazy()
+        return lf.collect(**kwargs).lazy()
     return lf
 
 
 def collect_all_if(
-    lfs: dict[str, pl.LazyFrame], condition: bool
+    lfs: dict[str, pl.LazyFrame], condition: bool, **kwargs: Any
 ) -> dict[str, pl.LazyFrame]:
     """Collect the lazy frames in the dictionary based on `condition`."""
     if condition:
-        dfs = pl.collect_all(lfs.values())
+        dfs = pl.collect_all(lfs.values(), **kwargs)
         return {k: v.lazy() for k, v in zip(lfs.keys(), dfs)}
     return lfs

dataframely/collection/collection.py

@@ -380,7 +380,13 @@ def _preprocess_sample(
 
     @classmethod
     def validate(
-        cls, data: Mapping[str, FrameType], /, *, cast: bool = False, eager: bool = True
+        cls,
+        data: Mapping[str, FrameType],
+        /,
+        *,
+        cast: bool = False,
+        eager: bool = True,
+        **kwargs: Any,
     ) -> Self:
         """Validate that a set of data frames satisfy the collection's invariants.
 
@@ -400,6 +406,8 @@ def validate(
                 :meth:`~polars.LazyFrame.collect` on the individual member or
                 :meth:`collect_all` on the collection. Note that, in the latter case,
                 information from error messages is limited.
+            kwargs: Keyword arguments passed directly to :meth:`polars.collect_all` and
+                :meth:`polars.LazyFrame.collect` when `eager=True`.
 
         Raises:
             ValueError: If an insufficient set of input data frames is provided, i.e. if
@@ -421,7 +429,7 @@ def validate(
         if eager:
             # If we perform the validation eagerly, we call filter and check the failure
             # information to properly construct a useful error message.
-            filtered, failures = cls.filter(data, cast=cast, eager=True)
+            filtered, failures = cls.filter(data, cast=cast, eager=True, **kwargs)
             if any(len(failure) > 0 for failure in failures.values()):
                 errors: dict[str, str] = {}
                 for member, failure in failures.items():
@@ -489,7 +497,9 @@ def validate(
             return cls._init(members)
 
     @classmethod
-    def is_valid(cls, data: Mapping[str, FrameType], /, *, cast: bool = False) -> bool:
+    def is_valid(
+        cls, data: Mapping[str, FrameType], /, *, cast: bool = False, **kwargs: Any
+    ) -> bool:
         """Utility method to check whether :meth:`validate` raises an exception.
 
         Args:
@@ -498,6 +508,8 @@ def is_valid(cls, data: Mapping[str, FrameType], /, *, cast: bool = False) -> bo
                 the member as key.
             cast: Whether columns with a wrong data type in the member data frame are
                 cast to their schemas' defined data types if possible.
+            kwargs: Keyword arguments passed directly to :meth:`polars.collect_all` and
+                :meth:`polars.LazyFrame.collect`.
 
         Returns:
             Whether the provided members satisfy the invariants of the collection.
@@ -512,7 +524,7 @@ def is_valid(cls, data: Mapping[str, FrameType], /, *, cast: bool = False) -> bo
         members: dict[str, pl.LazyFrame] = {}
         for member, schema in cls.member_schemas().items():
             if member in data:
-                if not schema.is_valid(data[member], cast=cast):
+                if not schema.is_valid(data[member], cast=cast, **kwargs):
                     return False
                 members[member] = data[member].lazy()
 
@@ -523,9 +535,12 @@ def is_valid(cls, data: Mapping[str, FrameType], /, *, cast: bool = False) -> bo
             keep = [filter.logic(result_cls).select(primary_key) for filter in filters]
             joined = _join_all(*keep, on=primary_key, how="inner")
             removed_rows = pl.collect_all(
-                data[member].lazy().join(joined, on=primary_key, how="anti")
-                for member in cls.members()
-                if member in data
+                (
+                    data[member].lazy().join(joined, on=primary_key, how="anti")
+                    for member in cls.members()
+                    if member in data
+                ),
+                **kwargs,
             )
             return all(df.is_empty() for df in removed_rows)
 
@@ -535,7 +550,13 @@ def is_valid(cls, data: Mapping[str, FrameType], /, *, cast: bool = False) -> bo
 
     @classmethod
     def filter(
-        cls, data: Mapping[str, FrameType], /, *, cast: bool = False, eager: bool = True
+        cls,
+        data: Mapping[str, FrameType],
+        /,
+        *,
+        cast: bool = False,
+        eager: bool = True,
+        **kwargs: Any,
     ) -> CollectionFilterResult[Self]:
         """Filter the members data frame by their schemas and the collection's filters.
 
@@ -551,6 +572,8 @@ def filter(
             eager: Whether the filter operation should be performed eagerly.
                 Note that until https://github.com/pola-rs/polars/pull/24129 is
                 released, eagerly filtering can provide significant speedups.
+            kwargs: Keyword arguments passed directly to :meth:`polars.collect_all` and
+                :meth:`polars.LazyFrame.collect` when `eager=True`.
 
         Returns:
             A named tuple with fields `result` and `failure`. The `result` field
@@ -593,7 +616,7 @@ class HospitalInvoiceData(dy.Collection):
                 continue
 
             member_result, failures[member_name] = member.schema.filter(
-                data[member_name].lazy(), cast=cast, eager=eager
+                data[member_name].lazy(), cast=cast, eager=eager, **kwargs
             )
             results[member_name] = member_result.lazy()
 
@@ -609,7 +632,7 @@ class HospitalInvoiceData(dy.Collection):
                 name: filter.logic(result_cls).select(primary_key)
                 for name, filter in filters.items()
             }
-            keep = collect_all_if(keep, eager)
+            keep = collect_all_if(keep, eager, **kwargs)
 
             drop: dict[str, pl.LazyFrame] = {
                 f"{failure_propagating_member}|failure_propagation": (
@@ -619,7 +642,7 @@ class HospitalInvoiceData(dy.Collection):
                 )
                 for failure_propagating_member in failure_propagating_members
             }
-            drop = collect_all_if(drop, eager)
+            drop = collect_all_if(drop, eager, **kwargs)
 
             # Now we can iterate over the results and left-join onto each individual
             # filter to obtain independent boolean indicators of whether to keep the row.
@@ -647,7 +670,7 @@ class HospitalInvoiceData(dy.Collection):
 
                 lfs_with_eval[member_name] = lf_with_eval
 
-            lfs_with_eval = collect_all_if(lfs_with_eval, eager)
+            lfs_with_eval = collect_all_if(lfs_with_eval, eager, **kwargs)
             for member_name, lf_with_eval in lfs_with_eval.items():
                 member_info = cls.members()[member_name]
 
@@ -714,7 +737,7 @@ class HospitalInvoiceData(dy.Collection):
 
         result = CollectionFilterResult(cls._init(results), failures)
         if eager:
-            return result.collect_all()
+            return result.collect_all(**kwargs)
         return result
 
     def join(
@@ -822,7 +845,7 @@ def collect_all(self, **kwargs: Any) -> Self:
             The same collection with all members collected once. Members annotated
             with :class:`~dataframely.DataFrame` are returned as DataFrames, while
             members annotated with :class:`~dataframely.LazyFrame` are returned as
-            "shallow-lazy" frames (obtained by calling ``.collect().lazy()``).
+            "shallow-lazy" frames (obtained by calling `.collect().lazy()`).
         """
         lazy_dict = self.to_dict()
         dfs = pl.collect_all(lazy_dict.values(), **kwargs)

dataframely/schema.py

@@ -501,6 +501,7 @@ def validate(
         *,
         cast: bool = False,
         eager: Literal[True] = True,
+        **kwargs: Any,
     ) -> DataFrame[Self]: ...
 
     @overload
@@ -512,6 +513,7 @@ def validate(
         *,
         cast: bool = False,
         eager: Literal[False],
+        **kwargs: Any,
     ) -> LazyFrame[Self]: ...
 
     @overload
@@ -523,6 +525,7 @@ def validate(
         *,
         cast: bool = False,
         eager: bool,
+        **kwargs: Any,
     ) -> DataFrame[Self] | LazyFrame[Self]: ...
 
     @classmethod
@@ -533,6 +536,7 @@ def validate(
         *,
         cast: bool = False,
         eager: bool = True,
+        **kwargs: Any,
     ) -> DataFrame[Self] | LazyFrame[Self]:
         """Validate that a data frame satisfies the schema.
 
@@ -554,6 +558,8 @@ def validate(
                     not surface *all* validation issues as the validation is aborted
                     once the first failure is encountered. Likewise, the reported
                     validation failure can be non-deterministic.
+            kwargs: Keyword arguments passed directly to :meth:`polars.LazyFrame.collect`
+                when `eager=True`.
 
         Returns:
             The input eager or lazy frame, wrapped in a generic version of the
@@ -574,7 +580,7 @@ def validate(
                 `eager=False`.
         """
         if eager:
-            out, failure = cls.filter(df, cast=cast, eager=True)
+            out, failure = cls.filter(df, cast=cast, eager=True, **kwargs)
             if len(failure) > 0:
                 counts = failure.counts()
                 raise ValidationError(
@@ -608,7 +614,7 @@ def validate(
 
     @classmethod
     def is_valid(
-        cls, df: pl.DataFrame | pl.LazyFrame, /, *, cast: bool = False
+        cls, df: pl.DataFrame | pl.LazyFrame, /, *, cast: bool = False, **kwargs: Any
     ) -> bool:
         """Check whether a data frame satisfies the schema.
 
@@ -627,6 +633,7 @@ def is_valid(
                 cast to the schema's defined data type before running validation. If set
                 to `False`, a wrong data type will result in a return value of
                 `False`.
+            kwargs: Keyword arguments passed directly to :meth:`polars.LazyFrame.collect`.
 
         Returns:
             Whether the provided dataframe can be validated with this schema.
@@ -646,12 +653,12 @@ def is_valid(
                 return (
                     lf.pipe(with_evaluation_rules, rules)
                     .select(all_rules(rules.keys()))
-                    .collect()
+                    .collect(**kwargs)
                     .item()
                 )
             # NOTE: We cannot simply return `True` here as, otherwise, we wouldn't
             #  validate the schema.
-            return lf.select(pl.lit(True)).collect().item()
+            return lf.select(pl.lit(True)).collect(**kwargs).item()
         except SchemaError:
             # If we encounter a schema error, we gracefully handle this as 'invalid'
             return False
@@ -699,6 +706,7 @@ def filter(
         *,
         cast: bool = False,
         eager: bool = True,
+        **kwargs: Any,
     ) -> FilterResult[Self] | LazyFilterResult[Self]:
         """Filter the data frame by the rules of this schema, returning `(valid,
         failures)`.
@@ -709,16 +717,16 @@ def filter(
         succeeds.
 
         Args:
-            df: The data frame to filter for valid rows.
-                The data frame is collected within this method, regardless of whether
-                a :class:`~polars.DataFrame` or :class:`~polars.LazyFrame` is passed.
-            cast:
-                Whether columns with a wrong data type in the input data frame are
+            df: The data frame to filter for valid rows. The data frame is collected
+                within this method, regardless of whether a :class:`~polars.DataFrame`
+                or :class:`~polars.LazyFrame` is passed.
+            cast: Whether columns with a wrong data type in the input data frame are
                 cast to the schema's defined data type if possible. Rows for which the
                 cast fails for any column are filtered out.
             eager: Whether the filter operation should be performed eagerly. If `False`, the
-                returned lazy frame will
-                fail to collect if the validation does not pass.
+                returned lazy frame will fail to collect if the validation does not pass.
+            kwargs: Keyword arguments passed directly to :meth:`polars.LazyFrame.collect`
+                when `eager=True`.
 
         Returns:
             A tuple of the validated rows in the input data frame (potentially
@@ -755,7 +763,7 @@ def filter(
         )
         if rules := cls._validation_rules(with_cast=cast):
             evaluated = lf.pipe(cls._with_evaluated_rules, rules).pipe(
-                collect_if, eager
+                collect_if, eager, **kwargs
             )
             filtered = evaluated.filter(pl.col(_COLUMN_VALID)).select(
                 cls.column_names()

tests/benches/test_collection.py

@@ -95,18 +95,24 @@ def one_to_at_least_one_reverse(self) -> pl.LazyFrame:
 
 
 @pytest.mark.benchmark(group="collection-filter-multi")
+@pytest.mark.parametrize("engine", ["in-memory", "streaming"])
 def test_multi_filter_validate(
-    benchmark: BenchmarkFixture, partitioned_dataset: dict[str, pl.DataFrame]
+    benchmark: BenchmarkFixture,
+    partitioned_dataset: dict[str, pl.DataFrame],
+    engine: str,
 ) -> None:
-    benchmark(MultiFilterCollection.validate, partitioned_dataset)
+    benchmark(MultiFilterCollection.validate, partitioned_dataset, engine=engine)
 
 
 @pytest.mark.benchmark(group="collection-filter-multi")
+@pytest.mark.parametrize("engine", ["in-memory", "streaming"])
 def test_multi_filter_filter(
-    benchmark: BenchmarkFixture, partitioned_dataset: dict[str, pl.DataFrame]
+    benchmark: BenchmarkFixture,
+    partitioned_dataset: dict[str, pl.DataFrame],
+    engine: str,
 ) -> None:
     def benchmark_fn() -> None:
-        _, failure = MultiFilterCollection.filter(partitioned_dataset)
+        _, failure = MultiFilterCollection.filter(partitioned_dataset, engine=engine)
         _ = [len(f) for f in failure.values()]
 
     benchmark(benchmark_fn)