Skip to content

refactor!: Replace eager: bool with "lazy-in-lazy-out" #372

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Oliver Borchert (borchero) merged 4 commits into from
Jul 31, 2026
Merged

refactor!: Replace eager: bool with "lazy-in-lazy-out" #372

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0b811b2
refactor!: Replace `eager: bool` with "lazy-in-lazy-out"
borchero Jul 23, 2026
462477b
Merge branch 'main' into lazy-in-lazy-out
borchero Jul 30, 2026
f48e52a
Review
borchero Jul 30, 2026
d1670de
Merge branch 'no-group-by' into lazy-in-lazy-out
borchero Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 39 additions & 28 deletions dataframely/collection/collection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -354,7 +354,7 @@ def validate(
/,
*,
cast: bool = False,
eager: bool = True,
lazy: bool = False,
skip_member_validation: bool = False,
**kwargs: Any,
) -> Self:
Expand All @@ -366,13 +366,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.
Expand All @@ -382,15 +380,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=True`, a :class:`~polars.exceptions.ComputeError` is raised
upon collecting.
Comment thread
borchero marked this conversation as resolved.

Returns:
Expand All @@ -399,15 +397,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,
)
Expand Down Expand Up @@ -448,9 +447,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
Expand Down Expand Up @@ -546,7 +543,7 @@ def filter(
/,
*,
cast: bool = False,
eager: bool = True,
lazy: bool = False,
skip_member_validation: bool = False,
**kwargs: Any,
) -> CollectionFilterResult[Self]:
Expand All @@ -561,16 +558,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
Expand Down Expand Up @@ -602,6 +601,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
Expand All @@ -623,7 +623,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()

Expand All @@ -639,7 +643,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": (
Expand All @@ -649,7 +653,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.
Expand Down Expand Up @@ -677,7 +681,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]

Expand Down Expand Up @@ -742,9 +746,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,
Expand Down Expand Up @@ -1013,6 +1017,13 @@ def scan_parquet(cls, directory: str | Path, **kwargs: Any) -> Self:

# ----------------------------------- 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)
Expand Down
97 changes: 33 additions & 64 deletions dataframely/schema.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
import polars as pl

from ._base_schema import ORIGINAL_COLUMN_PREFIX, BaseSchema
from ._compat import pydantic, sa
from ._compat import _polars_version_tuple, pydantic, sa
from ._match_to_schema import match_to_schema
from ._native import format_rule_failures
from ._plugin import all_rules, all_rules_horizontal, all_rules_required
Expand Down Expand Up @@ -477,46 +477,31 @@ 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]: ...

@overload
@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,
df: pl.DataFrame | pl.LazyFrame,
/,
*,
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.
Expand All @@ -530,38 +515,35 @@ 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
input's data frame type to reflect schema adherence. Columns not defined
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
`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
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 the input data frame is lazy.
Comment thread
borchero marked this conversation as resolved.
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(
Expand Down Expand Up @@ -650,47 +632,32 @@ 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,
df: pl.DataFrame | pl.LazyFrame,
/,
*,
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
Expand All @@ -704,10 +671,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
Expand Down Expand Up @@ -744,7 +709,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()
Expand All @@ -763,8 +732,8 @@ def filter(
# Build the result objects
failure_info = FailureInfo(lf=failure_lf, rule_columns=list(rules.keys()))
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
Expand Down
2 changes: 1 addition & 1 deletion tests/benches/test_collection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading
Loading