diff --git a/dataframely/_polars.py b/dataframely/_polars.py index cd58edb..e8c83ce 100644 --- a/dataframely/_polars.py +++ b/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 diff --git a/dataframely/collection/collection.py b/dataframely/collection/collection.py index 7583465..447909b 100644 --- a/dataframely/collection/collection.py +++ b/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) diff --git a/dataframely/schema.py b/dataframely/schema.py index 726aa6d..9d2e04e 100644 --- a/dataframely/schema.py +++ b/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() diff --git a/tests/benches/test_collection.py b/tests/benches/test_collection.py index 0bf0cb5..a28ca58 100644 --- a/tests/benches/test_collection.py +++ b/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)