Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
50 changes: 44 additions & 6 deletions dataframely/collection/collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -386,6 +386,7 @@ def validate(
*,
cast: bool = False,
eager: bool = True,
skip_member_validation: bool = False,
Comment thread
borchero marked this conversation as resolved.
**kwargs: Any,
) -> Self:
"""Validate that a set of data frames satisfy the collection's invariants.
Expand All @@ -406,6 +407,11 @@ 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.
skip_member_validation: Whether to skip validating 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 be valid.
Comment thread
borchero marked this conversation as resolved.
kwargs: Keyword arguments passed directly to :meth:`polars.collect_all` and
:meth:`polars.LazyFrame.collect` when `eager=True`.

Expand All @@ -429,7 +435,13 @@ 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, **kwargs)
filtered, failures = cls.filter(
data,
cast=cast,
eager=True,
skip_member_validation=skip_member_validation,
**kwargs,
)
if any(len(failure) > 0 for failure in failures.values()):
errors: dict[str, str] = {}
for member, failure in failures.items():
Expand Down Expand Up @@ -460,7 +472,17 @@ def validate(
# efficiently as we cannot easily propagate error messages from different
# members anyways.
members: dict[str, pl.LazyFrame] = {
name: member.schema.validate(data[name].lazy(), cast=cast, eager=False)
name: (
(
member.schema.cast(data[name].lazy())
if cast
else data[name].lazy()
)
if skip_member_validation
else member.schema.validate(
data[name].lazy(), cast=cast, eager=False
)
)
for name, member in cls.members().items()
if name in data
}
Expand Down Expand Up @@ -556,6 +578,7 @@ def filter(
*,
cast: bool = False,
eager: bool = True,
skip_member_validation: bool = False,
**kwargs: Any,
) -> CollectionFilterResult[Self]:
"""Filter the members data frame by their schemas and the collection's filters.
Expand All @@ -572,6 +595,11 @@ 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.
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`.

Expand Down Expand Up @@ -615,10 +643,20 @@ class HospitalInvoiceData(dy.Collection):
if member.is_optional and member_name not in data:
continue

member_result, failures[member_name] = member.schema.filter(
data[member_name].lazy(), cast=cast, eager=eager, **kwargs
)
results[member_name] = member_result.lazy()
if skip_member_validation:
results[member_name] = (
member.schema.cast(data[member_name].lazy())
if cast
else data[member_name].lazy()
)
failures[member_name] = FailureInfo._create_empty(
member.schema, with_casting_rules=cast
)
else:
member_result, failures[member_name] = member.schema.filter(
data[member_name].lazy(), cast=cast, eager=eager, **kwargs
)
results[member_name] = member_result.lazy()

# Once we've done that, we can apply the filters on this collection. To this end,
# we iterate over all filters and store the filter results.
Expand Down
11 changes: 11 additions & 0 deletions dataframely/filter_result.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,17 @@ def __init__(
self._rule_columns = rule_columns
self.schema = schema

@classmethod
def _create_empty(cls, schema: type[S], with_casting_rules: bool) -> FailureInfo[S]:
rules = schema._validation_rules(with_cast=with_casting_rules)
lf = pl.LazyFrame(
schema={
**schema.to_polars_schema(), # type: ignore
**{rule: pl.Boolean for rule in rules},
}
)
return cls(lf=lf, rule_columns=list(rules.keys()), schema=schema)

@cached_property
def _df(self) -> pl.DataFrame:
return self._lf.collect()
Expand Down
108 changes: 108 additions & 0 deletions tests/collection/test_skip_member_validation.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
# Copyright (c) QuantCo 2025-2026
# SPDX-License-Identifier: BSD-3-Clause

import polars as pl
import polars.exceptions as plexc
import pytest
Comment thread
borchero marked this conversation as resolved.

import dataframely as dy


class FirstSchema(dy.Schema):
a = dy.Float64(min=5)


class SecondSchema(dy.Schema):
a = dy.String()


class Collection(dy.Collection):
first: dy.LazyFrame[FirstSchema]
second: dy.LazyFrame[SecondSchema]


@pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame])
def test_validate_skip_member_validation_eager(
df_type: type[pl.DataFrame] | type[pl.LazyFrame],
) -> None:
first = df_type({"a": [3, 4, 5]}) # NOTE: first two rows are violations
second = df_type({"a": ["1", "2", "3"]})

with pytest.raises(dy.exc.ValidationError):
Collection.validate({"first": first, "second": second}, cast=True) # type: ignore

Collection.validate(
{"first": first, "second": second}, # type: ignore
cast=True,
skip_member_validation=True,
)
Comment thread
borchero marked this conversation as resolved.


@pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame])
def test_validate_skip_member_validation_lazy(
df_type: type[pl.DataFrame] | type[pl.LazyFrame],
) -> None:
first = df_type({"a": [3, 4, 5]}) # NOTE: first two rows are violations
second = df_type({"a": ["1", "2", "3"]})

with pytest.raises(plexc.ComputeError):
Collection.validate(
{"first": first, "second": second}, # type: ignore
cast=True,
eager=False,
).collect_all()

Collection.validate(
{"first": first, "second": second}, # type: ignore
cast=True,
skip_member_validation=True,
eager=False,
).collect_all()
Comment thread
borchero marked this conversation as resolved.


@pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame])
def test_filter_skip_member_validation_eager(
df_type: type[pl.DataFrame] | type[pl.LazyFrame],
) -> None:
first = df_type({"a": [3, 4, 5]}) # NOTE: first two rows are violations
second = df_type({"a": ["1", "2", "3"]})

_, failure_info = Collection.filter(
{"first": first, "second": second}, # type: ignore
cast=True,
)
assert failure_info["first"].counts() == {"a|min": 2}
assert failure_info["second"].counts() == {}

_, failure_info = Collection.filter(
{"first": first, "second": second}, # type: ignore
cast=True,
skip_member_validation=True,
)
assert failure_info["first"].counts() == {}
assert failure_info["second"].counts() == {}
Comment thread
borchero marked this conversation as resolved.


@pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame])
def test_filter_skip_member_validation_lazy(
df_type: type[pl.DataFrame] | type[pl.LazyFrame],
) -> None:
first = df_type({"a": [3, 4, 5]}) # NOTE: first two rows are violations
second = df_type({"a": ["1", "2", "3"]})

_, failure_info = Collection.filter(
{"first": first, "second": second}, # type: ignore
cast=True,
eager=False,
)
assert failure_info["first"].counts() == {"a|min": 2}
assert failure_info["second"].counts() == {}

_, failure_info = Collection.filter(
{"first": first, "second": second}, # type: ignore
cast=True,
skip_member_validation=True,
eager=False,
)
assert failure_info["first"].counts() == {}
assert failure_info["second"].counts() == {}
Comment thread
borchero marked this conversation as resolved.
Loading