From 0b811b2ad109c62c999cbf4d1f85bc091d5a2e79 Mon Sep 17 00:00:00 2001 From: Oliver Borchert Date: Thu, 23 Jul 2026 14:31:34 +0200 Subject: [PATCH] refactor!: Replace `eager: bool` with "lazy-in-lazy-out" --- dataframely/collection/collection.py | 67 ++++++----- dataframely/schema.py | 104 +++++++----------- tests/benches/test_collection.py | 2 +- tests/benches/test_schema.py | 8 +- tests/collection/test_dataframe_members.py | 29 +++++ tests/collection/test_filter_validate.py | 36 +++--- .../collection/test_skip_member_validation.py | 8 +- tests/schema/test_filter.py | 64 +++++------ tests/schema/test_validate.py | 94 ++++++++-------- 9 files changed, 210 insertions(+), 202 deletions(-) diff --git a/dataframely/collection/collection.py b/dataframely/collection/collection.py index cf59836e..167dec06 100644 --- a/dataframely/collection/collection.py +++ b/dataframely/collection/collection.py @@ -395,7 +395,7 @@ def validate( /, *, cast: bool = False, - eager: bool = True, + lazy: bool = False, skip_member_validation: bool = False, **kwargs: Any, ) -> Self: @@ -407,13 +407,11 @@ def validate( 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. - eager: Whether the validation should be performed eagerly. If `True`, this - method raises a validation error and the returned collection contains - "shallow" lazy frames, i.e., lazy frames by simply calling - :meth:`~polars.DataFrame.lazy` on the validated data frame. If - `False`, this method only raises a `ValueError` if `data` does - not contain data for all required members. The returned collection - contains "true" lazy frames that will be validated upon calling + lazy: Whether the validation should be performed lazily. If `False`, this + method raises a validation error. If `True`, this method only raises a + `ValueError` if (1) `data` does not contain data for all required + members or (2) the collection defines any of its members as eager. + Validation will then be performed when calling :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. @@ -423,15 +421,15 @@ def validate( validated. This option is particularly useful in performance-critical scenarios where the members are known to be valid. kwargs: Keyword arguments passed directly to :meth:`polars.collect_all` and - :meth:`polars.LazyFrame.collect` when `eager=True`. + :meth:`polars.LazyFrame.collect` when `lazy=False`. Raises: ValueError: If an insufficient set of input data frames is provided, i.e. if any required member of this collection is missing in the input. - ValidationError: If `eager=True` and any of the input data frames does not + ValidationError: If `lazy=False` and any of the input data frames does not satisfy its schema definition or the filters on this collection result in the removal of at least one row across any of the input data frames. - If `eager=False`, a :class:`~polars.exceptions.ComputeError` is raised + If `lazy=False`, a :class:`~polars.exceptions.ComputeError` is raised upon collecting. Returns: @@ -440,15 +438,16 @@ def validate( collection did not remove rows from any member. The input order of each member is maintained. """ + cls._validate_lazy_param(lazy) cls._validate_input_keys(data) - if eager: + if not lazy: # 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, + lazy=False, skip_member_validation=skip_member_validation, **kwargs, ) @@ -489,9 +488,7 @@ def validate( else data[name].lazy() ) if skip_member_validation - else member.schema.validate( - data[name].lazy(), cast=cast, eager=False - ) + else member.schema.validate(data[name].lazy(), cast=cast, lazy=True) ) for name, member in cls.members().items() if name in data @@ -587,7 +584,7 @@ def filter( /, *, cast: bool = False, - eager: bool = True, + lazy: bool = False, skip_member_validation: bool = False, **kwargs: Any, ) -> CollectionFilterResult[Self]: @@ -602,16 +599,18 @@ def filter( :class:`~polars.LazyFrame`. cast: Whether columns with a wrong data type in the member data frame are cast to their schemas' defined data types if possible. - 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. + lazy: Whether the filter operation should be performed lazily. Note that, + before polars v1.43.0, eager filtering provided significant speedups due + to https://github.com/pola-rs/polars/pull/24129. As of polars v1.43.0, + lazy filtering is equally fast, provided that `POLARS_ALLOW_NESTED_CSPE=1` + is set. skip_member_validation: Whether to skip filtering individual members and only apply the collection filters. **Use this option with caution** as it requires the caller to ensure that the individual members have been validated. This option is particularly useful in performance-critical scenarios where the members are known to already be valid. kwargs: Keyword arguments passed directly to :meth:`polars.collect_all` and - :meth:`polars.LazyFrame.collect` when `eager=True`. + :meth:`polars.LazyFrame.collect` when `lazy=False`. Returns: A named tuple with fields `result` and `failure`. The `result` field @@ -643,6 +642,7 @@ class HospitalInvoiceData(dy.Collection): failed_df = failure.invoice.invalid() print(failed_df) """ + cls._validate_lazy_param(lazy) cls._validate_input_keys(data) # First, we iterate over all members in this collection and filter them @@ -664,7 +664,11 @@ class HospitalInvoiceData(dy.Collection): ) else: member_result, failures[member_name] = member.schema.filter( - data[member_name].lazy(), cast=cast, eager=eager, **kwargs + data[member_name].lazy() + if lazy + else data[member_name].lazy().collect(**kwargs), + cast=cast, + **kwargs, ) results[member_name] = member_result.lazy() @@ -680,7 +684,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, **kwargs) + keep = collect_all_if(keep, not lazy, **kwargs) drop: dict[str, pl.LazyFrame] = { f"{failure_propagating_member}|failure_propagation": ( @@ -690,7 +694,7 @@ class HospitalInvoiceData(dy.Collection): ) for failure_propagating_member in failure_propagating_members } - drop = collect_all_if(drop, eager, **kwargs) + drop = collect_all_if(drop, not lazy, **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. @@ -718,7 +722,7 @@ class HospitalInvoiceData(dy.Collection): lfs_with_eval[member_name] = lf_with_eval - lfs_with_eval = collect_all_if(lfs_with_eval, eager, **kwargs) + lfs_with_eval = collect_all_if(lfs_with_eval, not lazy, **kwargs) for member_name, lf_with_eval in lfs_with_eval.items(): member_info = cls.members()[member_name] @@ -784,9 +788,9 @@ class HospitalInvoiceData(dy.Collection): ) result = CollectionFilterResult(cls._init(results), failures) - if eager: - return result.collect_all(**kwargs) - return result + if lazy: + return result + return result.collect_all(**kwargs) def join( self, @@ -1421,6 +1425,13 @@ def _requires_validation_for_reading_parquets( # ----------------------------------- UTILITIES ---------------------------------- # + @classmethod + def _validate_lazy_param(cls, lazy: bool, /) -> None: + if lazy and any(not member.is_lazy for member in cls.members().values()): + raise ValueError( + "Cannot use `lazy=True` on a collection with eager members." + ) + @classmethod def _validate_input_keys(cls, data: Mapping[str, FrameType], /) -> None: actual = set(data) diff --git a/dataframely/schema.py b/dataframely/schema.py index dab8d151..23185c54 100644 --- a/dataframely/schema.py +++ b/dataframely/schema.py @@ -17,7 +17,14 @@ from polars._typing import FileSource from ._base_schema import ORIGINAL_COLUMN_PREFIX, BaseSchema -from ._compat import PartitionSchemeOrSinkDirectory, deltalake, pa, pydantic, sa +from ._compat import ( + PartitionSchemeOrSinkDirectory, + _polars_version_tuple, + deltalake, + pa, + pydantic, + sa, +) from ._deprecation import deprecated from ._match_to_schema import match_to_schema from ._native import format_rule_failures @@ -497,11 +504,10 @@ def _sampling_overrides(cls) -> dict[str, pl.Expr]: @classmethod def validate( cls, - df: pl.DataFrame | pl.LazyFrame, + df: pl.DataFrame, /, *, cast: bool = False, - eager: Literal[True] = True, **kwargs: Any, ) -> DataFrame[Self]: ... @@ -509,15 +515,13 @@ def validate( @classmethod def validate( cls, - df: pl.DataFrame | pl.LazyFrame, + df: pl.LazyFrame, /, *, cast: bool = False, - eager: Literal[False], **kwargs: Any, ) -> LazyFrame[Self]: ... - @overload @classmethod def validate( cls, @@ -525,18 +529,6 @@ def validate( /, *, cast: bool = False, - eager: bool, - **kwargs: Any, - ) -> DataFrame[Self] | LazyFrame[Self]: ... - - @classmethod - def validate( - cls, - df: pl.DataFrame | pl.LazyFrame, - /, - *, - cast: bool = False, - eager: bool = True, **kwargs: Any, ) -> DataFrame[Self] | LazyFrame[Self]: """Validate that a data frame satisfies the schema. @@ -550,17 +542,8 @@ def validate( df: The data frame to validate. cast: Whether columns with a wrong data type in the input data frame are cast to the schema's defined data type if possible. - eager: Whether the validation should be performed eagerly and this method - should raise upon failure. If `False`, the returned lazy frame will - fail to collect if the validation does not pass. - - Note: - If running on the streaming engine, lazy validation will potentially - 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`. + when the input data frame is eager. Returns: The input eager or lazy frame, wrapped in a generic version of the @@ -568,20 +551,26 @@ def validate( in the schema are removed from the output. This operation is guaranteed to maintain input ordering of rows. + Note: + If running on the streaming engine, lazy validation will potentially 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. + Raises: - SchemaError: If `eager=True` and the input data frame misses columns or + SchemaError: If the input data frame is eager and it misses columns or `cast=False` and any data type mismatches the definition in this - schema. Only raised upon collection if `eager=False`. - ValidationError: If `eager=True` and in any rule in the schema is - violated, i.e. the data does not pass the validation. When - `eager=False`, a :class:`~polars.exceptions.ComputeError` is raised + Only raised upon collection if the input data frame is lazy. + ValidationError: If the input data frame is eager and any rule in the schema + is violated, i.e. the data does not pass the validation. When the input + data frame is lazy, a :class:`~polars.exceptions.ComputeError` is raised upon collecting. - InvalidOperationError: If `eager=True`, `cast=True`, and the cast fails - for any value in the data. Only raised upon collection if - `eager=False`. + InvalidOperationError: If the input data frame is eager, `cast=True`, and + the cast fails for any value in the data. Only raised upon collection + if the input data frame is lazy. """ - if eager: - out, failure = cls.filter(df, cast=cast, eager=True, **kwargs) + if isinstance(df, pl.DataFrame): + out, failure = cls.filter(df, cast=cast, **kwargs) if len(failure) > 0: counts = failure.counts() raise ValidationError( @@ -670,25 +659,22 @@ def is_valid( @classmethod def filter( cls, - df: pl.DataFrame | pl.LazyFrame, + df: pl.DataFrame, /, *, cast: bool = False, - eager: Literal[True] = True, ) -> FilterResult[Self]: ... @overload @classmethod def filter( cls, - df: pl.DataFrame | pl.LazyFrame, + df: pl.LazyFrame, /, *, cast: bool = False, - eager: Literal[False], ) -> LazyFilterResult[Self]: ... - @overload @classmethod def filter( cls, @@ -696,21 +682,9 @@ def filter( /, *, cast: bool = False, - eager: bool, - ) -> FilterResult[Self] | LazyFilterResult[Self]: ... - - @classmethod - def filter( - cls, - df: pl.DataFrame | pl.LazyFrame, - /, - *, - 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)`. + """Filter the data frame by the rules of this schema. This method can be thought of as a "soft alternative" to :meth:`validate`. While :meth:`validate` raises an exception when a row does not adhere to the @@ -724,10 +698,8 @@ def filter( 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. - kwargs: Keyword arguments passed directly to :meth:`polars.LazyFrame.collect` - when `eager=True`. + kwargs: Keyword arguments passed directly to :meth:`polars.collect_all` + if the input data frame is eager. Returns: A tuple of the validated rows in the input data frame (potentially @@ -764,7 +736,11 @@ def filter( ) if rules := cls._validation_rules(with_cast=cast): evaluated = lf.pipe(cls._with_evaluated_rules, rules).pipe( - collect_if, eager, **kwargs + collect_if, + # NOTE: Polars 1.43.0 fixes a bug related to CSPE which allows us to leverage + # the `collect_all` below to perform the entire filtering efficiently. + isinstance(df, pl.DataFrame) and _polars_version_tuple < (1, 43), + **kwargs, ) filtered = evaluated.filter(pl.col(_COLUMN_VALID)).select( cls.column_names() @@ -785,8 +761,8 @@ def filter( lf=failure_lf, rule_columns=list(rules.keys()), schema=cls ) result = LazyFilterResult(filtered, failure_info) # type: ignore - if eager: - return result.collect_all() + if isinstance(df, pl.DataFrame): + return result.collect_all(**kwargs) return result @classmethod @@ -1379,7 +1355,7 @@ def _validate_if_needed( if cls._requires_validation_for_reading_parquet( deserialized_schema, validation, source=source ): - return df.pipe(cls.validate, cast=True, eager=isinstance(df, pl.DataFrame)) + return df.pipe(cls.validate, cast=True) return cls.cast(df) diff --git a/tests/benches/test_collection.py b/tests/benches/test_collection.py index eb5acb3d..f0d7e734 100644 --- a/tests/benches/test_collection.py +++ b/tests/benches/test_collection.py @@ -75,7 +75,7 @@ def test_single_filter_filter_lazy( benchmark: BenchmarkFixture, partitioned_dataset: dict[str, pl.DataFrame] ) -> None: def benchmark_fn() -> None: - result = SingleFilterCollection.filter(partitioned_dataset, eager=False) + result = SingleFilterCollection.filter(partitioned_dataset, lazy=True) result.collect_all() benchmark(benchmark_fn) diff --git a/tests/benches/test_schema.py b/tests/benches/test_schema.py index 02df0903..9b911ae5 100644 --- a/tests/benches/test_schema.py +++ b/tests/benches/test_schema.py @@ -33,12 +33,14 @@ class RowWiseValidationSchema(dy.Schema): @pytest.mark.benchmark(group="schema-row-wise") -@pytest.mark.parametrize("eager", [True, False]) +@pytest.mark.parametrize("lazy", [True, False]) def test_row_wise_validate( - benchmark: BenchmarkFixture, dataset: pl.DataFrame, eager: bool + benchmark: BenchmarkFixture, dataset: pl.DataFrame, lazy: bool ) -> None: + df = dataset.lazy() if lazy else dataset + def benchmark_fn() -> None: - RowWiseValidationSchema.validate(dataset, eager=eager).lazy().collect() + RowWiseValidationSchema.validate(df).lazy().collect() benchmark(benchmark_fn) diff --git a/tests/collection/test_dataframe_members.py b/tests/collection/test_dataframe_members.py index 7ccc5398..f7622256 100644 --- a/tests/collection/test_dataframe_members.py +++ b/tests/collection/test_dataframe_members.py @@ -125,3 +125,32 @@ def test_member_access_returns_correct_type( collection = collection_cls.validate(valid_data) for name, expected_type in expected_types.items(): assert isinstance(getattr(collection, name), expected_type) + + +# ------------------------------------------------------------------------------------ # +# LAZY FILTERING RESTRICTIONS # +# ------------------------------------------------------------------------------------ # + + +@pytest.mark.parametrize("collection_cls", [EagerCollection, MixedCollection]) +def test_filter_lazy_with_eager_members_raises( + collection_cls: type[dy.Collection], + valid_data: dict[str, pl.DataFrame], +) -> None: + with pytest.raises( + ValueError, + match=r"Cannot use `lazy=True` on a collection with eager members\.", + ): + collection_cls.filter(valid_data, lazy=True) + + +@pytest.mark.parametrize("collection_cls", [EagerCollection, MixedCollection]) +def test_validate_lazy_with_eager_members_raises( + collection_cls: type[dy.Collection], + valid_data: dict[str, pl.DataFrame], +) -> None: + with pytest.raises( + ValueError, + match=r"Cannot use `lazy=True` on a collection with eager members\.", + ): + collection_cls.validate(valid_data, lazy=True) diff --git a/tests/collection/test_filter_validate.py b/tests/collection/test_filter_validate.py index 6dd49ee6..19a37933 100644 --- a/tests/collection/test_filter_validate.py +++ b/tests/collection/test_filter_validate.py @@ -86,17 +86,17 @@ def data_with_filter_with_rule_violation() -> tuple[pl.LazyFrame, pl.LazyFrame]: # -------------------------------------- FILTER -------------------------------------- # -@pytest.mark.parametrize("eager", [True, False]) +@pytest.mark.parametrize("lazy", [False, True]) def test_filter_without_filter_without_rule_violation( data_without_filter_without_rule_violation: tuple[pl.LazyFrame, pl.LazyFrame], - eager: bool, + lazy: bool, ) -> None: out, failure = SimpleCollection.filter( { "first": data_without_filter_without_rule_violation[0], "second": data_without_filter_without_rule_violation[1], }, - eager=eager, + lazy=lazy, ) assert isinstance(out, SimpleCollection) @@ -106,17 +106,17 @@ def test_filter_without_filter_without_rule_violation( assert len(failure["second"]) == 0 -@pytest.mark.parametrize("eager", [True, False]) +@pytest.mark.parametrize("lazy", [False, True]) def test_filter_without_filter_with_rule_violation( data_without_filter_with_rule_violation: tuple[pl.LazyFrame, pl.LazyFrame], - eager: bool, + lazy: bool, ) -> None: out, failure = SimpleCollection.filter( { "first": data_without_filter_with_rule_violation[0], "second": data_without_filter_with_rule_violation[1], }, - eager=eager, + lazy=lazy, ) assert isinstance(out, SimpleCollection) @@ -126,17 +126,17 @@ def test_filter_without_filter_with_rule_violation( assert failure["second"].counts() == {"b|min": 1} -@pytest.mark.parametrize("eager", [True, False]) +@pytest.mark.parametrize("lazy", [False, True]) def test_filter_with_filter_without_rule_violation( data_with_filter_without_rule_violation: tuple[pl.LazyFrame, pl.LazyFrame], - eager: bool, + lazy: bool, ) -> None: out, failure = MyCollection.filter( { "first": data_with_filter_without_rule_violation[0], "second": data_with_filter_without_rule_violation[1], }, - eager=eager, + lazy=lazy, ) assert isinstance(out, MyCollection) @@ -152,17 +152,17 @@ def test_filter_with_filter_without_rule_violation( } -@pytest.mark.parametrize("eager", [True, False]) +@pytest.mark.parametrize("lazy", [False, True]) def test_filter_with_filter_with_rule_violation( data_with_filter_with_rule_violation: tuple[pl.LazyFrame, pl.LazyFrame], - eager: bool, + lazy: bool, ) -> None: out, failure = MyCollection.filter( { "first": data_with_filter_with_rule_violation[0], "second": data_with_filter_with_rule_violation[1], }, - eager=eager, + lazy=lazy, ) assert isinstance(out, MyCollection) @@ -175,17 +175,17 @@ def test_filter_with_filter_with_rule_violation( # -------------------------------- VALIDATE WITH DATA -------------------------------- # -@pytest.mark.parametrize("eager", [True, False]) +@pytest.mark.parametrize("lazy", [False, True]) def test_validate_without_filter_without_rule_violation( data_without_filter_without_rule_violation: tuple[pl.LazyFrame, pl.LazyFrame], - eager: bool, + lazy: bool, ) -> None: data = { "first": data_without_filter_without_rule_violation[0], "second": data_without_filter_without_rule_violation[1], } assert SimpleCollection.is_valid(data) - out = SimpleCollection.validate(data, eager=eager) + out = SimpleCollection.validate(data, lazy=lazy) assert isinstance(out, SimpleCollection) assert_frame_equal(out.first, data_without_filter_without_rule_violation[0]) @@ -235,7 +235,7 @@ def test_validate_without_filter_with_rule_violation_lazy( } assert not SimpleCollection.is_valid(data) - validated = SimpleCollection.validate(data, eager=False) + validated = SimpleCollection.validate(data, lazy=True) with pytest.raises(plexc.ComputeError): validated.collect_all() @@ -268,7 +268,7 @@ def test_validate_with_filter_without_rule_violation_lazy( } assert not MyCollection.is_valid(data) - validated = MyCollection.validate(data, eager=False) + validated = MyCollection.validate(data, lazy=True) with pytest.raises(plexc.ComputeError): validated.collect_all() @@ -300,7 +300,7 @@ def test_validate_with_filter_with_rule_violation_lazy( } assert not MyCollection.is_valid(data) - validated = MyCollection.validate(data, eager=False) + validated = MyCollection.validate(data, lazy=True) with pytest.raises(plexc.ComputeError): validated.collect_all() diff --git a/tests/collection/test_skip_member_validation.py b/tests/collection/test_skip_member_validation.py index 1777b231..24005acd 100644 --- a/tests/collection/test_skip_member_validation.py +++ b/tests/collection/test_skip_member_validation.py @@ -49,14 +49,14 @@ def test_validate_skip_member_validation_lazy( Collection.validate( {"first": first, "second": second}, # type: ignore cast=True, - eager=False, + lazy=True, ).collect_all() Collection.validate( {"first": first, "second": second}, # type: ignore cast=True, skip_member_validation=True, - eager=False, + lazy=True, ).collect_all() @@ -93,7 +93,7 @@ def test_filter_skip_member_validation_lazy( _, failure_info = Collection.filter( {"first": first, "second": second}, # type: ignore cast=True, - eager=False, + lazy=True, ) assert failure_info["first"].counts() == {"a|min": 2} assert failure_info["second"].counts() == {} @@ -102,7 +102,7 @@ def test_filter_skip_member_validation_lazy( {"first": first, "second": second}, # type: ignore cast=True, skip_member_validation=True, - eager=False, + lazy=True, ) assert failure_info["first"].counts() == {} assert failure_info["second"].counts() == {} diff --git a/tests/schema/test_filter.py b/tests/schema/test_filter.py index 99557e79..8835af4d 100644 --- a/tests/schema/test_filter.py +++ b/tests/schema/test_filter.py @@ -26,9 +26,12 @@ class MySchema(dy.Schema): def _filter_and_collect( - schema: type[S], df: pl.DataFrame | pl.LazyFrame, *, cast: bool = False, eager: bool + schema: type[S], df: pl.DataFrame | pl.LazyFrame, *, cast: bool = False ) -> FilterResult[S]: - result = schema.filter(df, cast=cast, eager=eager) + # Whether filtering is performed eagerly is now determined by the input type: an + # eager data frame yields an eager result, a lazy frame yields a lazy result. + eager = isinstance(df, pl.DataFrame) + result = schema.filter(df, cast=cast) assert isinstance(result.result, pl.DataFrame if eager else pl.LazyFrame) if not eager: return result.collect_all() # type: ignore @@ -39,7 +42,6 @@ def _filter_and_collect( @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) @pytest.mark.parametrize( ("schema", "expected_columns"), [ @@ -52,11 +54,10 @@ def test_filter_extra_columns( df_type: type[pl.DataFrame] | type[pl.LazyFrame], schema: dict[str, DataTypeClass], expected_columns: list[str] | None, - eager: bool, ) -> None: df = df_type(schema=schema) try: - filtered, _ = _filter_and_collect(MySchema, df, eager=eager) + filtered, _ = _filter_and_collect(MySchema, df) assert expected_columns is not None assert set(filtered.columns) == set(expected_columns) except SchemaError: @@ -66,7 +67,6 @@ def test_filter_extra_columns( @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) @pytest.mark.parametrize( ("schema", "cast", "success"), [ @@ -79,11 +79,10 @@ def test_filter_dtypes( schema: dict[str, DataTypeClass], cast: bool, success: bool, - eager: bool, ) -> None: df = df_type(schema=schema) try: - _filter_and_collect(MySchema, df, cast=cast, eager=eager) + _filter_and_collect(MySchema, df, cast=cast) assert success except SchemaError: assert not success @@ -95,7 +94,6 @@ def test_filter_dtypes( @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) @pytest.mark.parametrize( ("data_a", "data_b", "failure_mask", "counts", "cooccurrence_counts"), [ @@ -121,7 +119,6 @@ def test_filter_dtypes( ) def test_filter_failure( df_type: type[pl.DataFrame] | type[pl.LazyFrame], - eager: bool, data_a: list[int], data_b: list[str | None], failure_mask: list[bool], @@ -129,7 +126,7 @@ def test_filter_failure( cooccurrence_counts: dict[frozenset[str], int], ) -> None: df = df_type({"a": data_a, "b": data_b}) - df_valid, failures = _filter_and_collect(MySchema, df, eager=eager) + df_valid, failures = _filter_and_collect(MySchema, df) assert_frame_equal(df.filter(pl.Series(failure_mask)).lazy().collect(), df_valid) assert validation_mask(df, failures).to_list() == failure_mask assert len(failures) == (len(failure_mask) - sum(failure_mask)) @@ -138,13 +135,12 @@ def test_filter_failure( @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_filter_no_rules( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: schema = create_schema("test", {"a": dy.Int64(nullable=True)}) df = df_type({"a": [1, 2, 3]}) - df_valid, failures = _filter_and_collect(schema, df, eager=eager) + df_valid, failures = _filter_and_collect(schema, df) assert_frame_equal(df.lazy().collect(), df_valid) assert len(failures) == 0 assert failures.counts() == {} @@ -152,13 +148,12 @@ def test_filter_no_rules( @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_filter_with_rule_all_valid( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: schema = create_schema("test", {"a": dy.String(min_length=3)}) df = df_type({"a": ["foo", "foobar"]}) - df_valid, failures = _filter_and_collect(schema, df, eager=eager) + df_valid, failures = _filter_and_collect(schema, df) assert_frame_equal(df.lazy().collect(), df_valid) assert len(failures) == 0 assert failures.counts() == {} @@ -166,9 +161,8 @@ def test_filter_with_rule_all_valid( @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_filter_cast( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: data = { # validation: [true, true, false, false, false, false] @@ -177,7 +171,7 @@ def test_filter_cast( "b": [20, 2000, None, 30, 3000, 50], } df = df_type(data) - df_valid, failures = _filter_and_collect(MySchema, df, cast=True, eager=eager) + df_valid, failures = _filter_and_collect(MySchema, df, cast=True) assert df_valid.collect_schema().names() == MySchema.column_names() assert len(failures) == 5 assert failures.counts() == { @@ -195,26 +189,30 @@ def test_filter_cast( } -@pytest.mark.parametrize("eager", [True, False]) -def test_filter_nondeterministic_lazyframe(eager: bool) -> None: +@pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) +def test_filter_nondeterministic_lazyframe( + df_type: type[pl.DataFrame] | type[pl.LazyFrame], +) -> None: n = 10_000 - lf = pl.LazyFrame( + df = df_type( { "a": range(n), "b": [random.choice(["foo", "foobar"]) for _ in range(n)], } ).select(pl.all().shuffle()) - filtered, _ = _filter_and_collect(MySchema, lf, eager=eager) + filtered, _ = _filter_and_collect(MySchema, df) assert filtered.select(pl.col("b").n_unique()).item() == 1 -@pytest.mark.parametrize("eager", [True, False]) -def test_filter_failure_info_original_dtype(eager: bool) -> None: +@pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) +def test_filter_failure_info_original_dtype( + df_type: type[pl.DataFrame] | type[pl.LazyFrame], +) -> None: schema = create_schema("test", {"a": dy.UInt8()}) - lf = pl.LazyFrame({"a": [100, 200, 300]}, schema={"a": pl.Int64}) + df = df_type({"a": [100, 200, 300]}, schema={"a": pl.Int64}) - out, failures = _filter_and_collect(schema, lf, cast=True, eager=eager) + out, failures = _filter_and_collect(schema, df, cast=True) assert len(out) == 2 assert out.schema.dtypes() == [pl.UInt8] @@ -223,8 +221,7 @@ def test_filter_failure_info_original_dtype(eager: bool) -> None: assert failures.invalid().dtypes == [pl.Int64] -@pytest.mark.parametrize("eager", [True, False]) -def test_filter_maintain_order(eager: bool) -> None: +def test_filter_maintain_order() -> None: schema = create_schema( "test", {"a": dy.UInt16(), "b": dy.UInt8()}, @@ -241,19 +238,18 @@ def test_filter_maintain_order(eager: bool) -> None: "b": generator.sample_int(10_000, min=0, max=255), } ) - out, _ = _filter_and_collect(schema, df, cast=True, eager=eager) + out, _ = _filter_and_collect(schema, df, cast=True) assert out.get_column("a").is_sorted() -@pytest.mark.parametrize("eager", [True, False]) -def test_filter_details(eager: bool) -> None: +def test_filter_details() -> None: df = pl.DataFrame( { "a": [2, 2], "b": ["bar", "foobar"], } ) - _, fails = _filter_and_collect(MySchema, df, cast=True, eager=eager) + _, fails = _filter_and_collect(MySchema, df, cast=True) assert fails.details().to_dicts() == [ { diff --git a/tests/schema/test_validate.py b/tests/schema/test_validate.py index 908d4e90..bd144e85 100644 --- a/tests/schema/test_validate.py +++ b/tests/schema/test_validate.py @@ -55,10 +55,11 @@ def _validate_and_collect( df: pl.DataFrame | pl.LazyFrame, *, cast: bool = False, - eager: bool, ) -> pl.DataFrame: - result = schema.validate(df, cast=cast, eager=eager) - if eager: + # Whether validation is performed eagerly is now determined by the input type: an + # eager data frame raises within `validate`, a lazy frame raises upon collection. + result = schema.validate(df, cast=cast) + if isinstance(df, pl.DataFrame): assert isinstance(result, pl.DataFrame) return result else: @@ -70,37 +71,34 @@ def _validate_and_collect( @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_missing_columns( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: df = df_type({"a": [1], "b": [""]}) with pytest.raises(SchemaError, match=r"1 missing columns for schema 'MySchema'"): - _validate_and_collect(MySchema, df, eager=eager) + _validate_and_collect(MySchema, df) assert not MySchema.is_valid(df) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_invalid_dtype( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: df = df_type({"a": [1], "b": [1], "c": [1]}) with pytest.raises( SchemaError, match=r"2 columns with invalid dtype for schema 'MySchema'", ): - _validate_and_collect(MySchema, df, eager=eager) + _validate_and_collect(MySchema, df) assert not MySchema.is_valid(df) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_invalid_dtype_cast( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: df = df_type({"a": [1], "b": [1], "c": [1]}) - actual = _validate_and_collect(MySchema, df, cast=True, eager=eager) + actual = _validate_and_collect(MySchema, df, cast=True) expected = pl.DataFrame({"a": [1], "b": ["1"], "c": ["1"]}) assert_frame_equal(actual, expected) assert MySchema.is_valid(df, cast=True) @@ -110,84 +108,83 @@ def test_invalid_dtype_cast( @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_invalid_column_contents( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: + eager = df_type is pl.DataFrame df = df_type({"a": [1, 2, 3], "b": ["x", "longtext", None], "c": ["1", None, "3"]}) with pytest.raises( ValidationError if eager else plexc.ComputeError, match=r"2 rules failed validation" if eager else None, ): - _validate_and_collect(MySchema, df, eager=eager) + _validate_and_collect(MySchema, df) assert not MySchema.is_valid(df) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_invalid_primary_key( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: + eager = df_type is pl.DataFrame df = df_type({"a": [1, 1], "b": ["x", "y"], "c": ["1", "2"]}) with pytest.raises( ValidationError if eager else plexc.ComputeError, match=r"1 rules failed validation", ): - _validate_and_collect(MySchema, df, eager=eager) + _validate_and_collect(MySchema, df) assert not MySchema.is_valid(df) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_invalid_primary_key_with_examples( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: + eager = df_type is pl.DataFrame df = df_type({"a": [1, 1], "b": ["x", "y"], "c": ["1", "2"]}) with dy.Config(max_failure_examples=5): with pytest.raises( ValidationError if eager else plexc.ComputeError, match=r"'primary_key' failed for 2 rows; examples: \[\{'a': 1\}\]", ): - _validate_and_collect(MySchema, df, eager=eager) + _validate_and_collect(MySchema, df) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_invalid_column_contents_with_examples( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: + eager = df_type is pl.DataFrame df = df_type({"a": [1, 2, 3], "b": ["x", "longtext", None], "c": ["1", None, "3"]}) with dy.Config(max_failure_examples=5): with pytest.raises( ValidationError if eager else plexc.ComputeError, match=r"examples: \[\{'b': 'longtext'\}\]" if eager else r"examples: \[", ): - _validate_and_collect(MySchema, df, eager=eager) + _validate_and_collect(MySchema, df) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_violated_custom_rule( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: + eager = df_type is pl.DataFrame df = df_type({"a": [1, 1, 2, 3, 3], "b": [2, 2, 2, 4, 5]}) with pytest.raises( ValidationError if eager else plexc.ComputeError, match=r"2 rules failed validation" if eager else None, ): - _validate_and_collect(MyComplexSchema, df, eager=eager) + _validate_and_collect(MyComplexSchema, df) assert not MyComplexSchema.is_valid(df) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_success_multi_row_strip_cast( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: df = df_type( {"a": [1, 2, 3], "b": ["x", "y", "z"], "c": [1, None, None], "d": [1, 2, 3]} ) - actual = _validate_and_collect(MySchema, df, cast=True, eager=eager) + actual = _validate_and_collect(MySchema, df, cast=True) expected = pl.DataFrame( {"a": [1, 2, 3], "b": ["x", "y", "z"], "c": ["1", None, None]} ) @@ -197,19 +194,18 @@ def test_success_multi_row_strip_cast( @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) @pytest.mark.parametrize("schema", [MyComplexSchema, MyComplexSchemaWithLazyRules]) -@pytest.mark.parametrize("eager", [True, False]) def test_group_rule_on_nulls( df_type: type[pl.DataFrame] | type[pl.LazyFrame], schema: type[MyComplexSchema] | type[MyComplexSchemaWithLazyRules], - eager: bool, ) -> None: + eager = df_type is pl.DataFrame # The schema is violated because we have multiple "b" values for the same "a" value df = df_type({"a": [None, None], "b": [1, 2]}) with pytest.raises( ValidationError if eager else plexc.ComputeError, match=r"1 rules failed validation", ): - _validate_and_collect(schema, df, cast=True, eager=eager) + _validate_and_collect(schema, df, cast=True) assert not schema.is_valid(df, cast=True) @@ -250,9 +246,8 @@ class NullableUniqueSchema(dy.Schema): @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_unique_constraint_valid( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: df = df_type( { @@ -261,63 +256,62 @@ def test_unique_constraint_valid( "name": ["A", "B", "C"], } ) - result = _validate_and_collect(UniqueSchema, df, eager=eager) + result = _validate_and_collect(UniqueSchema, df) assert len(result) == 3 assert UniqueSchema.is_valid(df) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_unique_constraint_invalid( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: + eager = df_type is pl.DataFrame df = df_type({"id": [1, 2], "email": ["a@x.com", "a@x.com"], "name": ["A", "B"]}) with pytest.raises( ValidationError if eager else plexc.ComputeError, match=r"1 rules failed validation", ): - _validate_and_collect(UniqueSchema, df, eager=eager) + _validate_and_collect(UniqueSchema, df) assert not UniqueSchema.is_valid(df) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_multiple_unique_columns_valid( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: df = df_type({"a": [1, 2, 3], "b": ["x", "y", "z"]}) - result = _validate_and_collect(MultiUniqueSchema, df, eager=eager) + result = _validate_and_collect(MultiUniqueSchema, df) assert len(result) == 3 assert MultiUniqueSchema.is_valid(df) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_multiple_unique_columns_one_invalid( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: + eager = df_type is pl.DataFrame # a is unique, b has duplicates df = df_type({"a": [1, 2, 3], "b": ["x", "x", "z"]}) with pytest.raises( ValidationError if eager else plexc.ComputeError, match=r"1 rules failed validation", ): - _validate_and_collect(MultiUniqueSchema, df, eager=eager) + _validate_and_collect(MultiUniqueSchema, df) assert not MultiUniqueSchema.is_valid(df) @pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) -@pytest.mark.parametrize("eager", [True, False]) def test_multiple_unique_columns_both_invalid( - df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool + df_type: type[pl.DataFrame] | type[pl.LazyFrame], ) -> None: + eager = df_type is pl.DataFrame # Both columns have duplicates df = df_type({"a": [1, 1, 3], "b": ["x", "y", "y"]}) with pytest.raises( ValidationError if eager else plexc.ComputeError, match=r"2 rules failed validation" if eager else None, ): - _validate_and_collect(MultiUniqueSchema, df, eager=eager) + _validate_and_collect(MultiUniqueSchema, df) assert not MultiUniqueSchema.is_valid(df) @@ -330,7 +324,7 @@ def test_lazy_validation_scan_parquet_length(tmp_path: Path) -> None: # Act height = ( pl.scan_parquet(tmp_path / "test.parquet") - .pipe(schema.validate, eager=False) + .pipe(schema.validate) .select(pl.len()) .collect() .item() @@ -346,7 +340,7 @@ def test_lazy_validation_scan_parquet_length(tmp_path: Path) -> None: def test_lazy_validate_does_not_block_streaming_engine() -> None: schema = create_schema("test", {"a": dy.Int64(), "b": dy.Int64()}) lf = pl.LazyFrame({"a": [1, 2, 3], "b": [2, 3, 4]}).lazy() - out = schema.validate(lf, eager=False) + out = schema.validate(lf) graph = out.show_graph(engine="streaming", plan_stage="physical", raw_output=True) assert graph is not None assert "in-memory-map" not in graph