From c13108a9910e8ffef624ed8c33f5f5e58d59cb46 Mon Sep 17 00:00:00 2001 From: Oliver Borchert Date: Fri, 31 Jul 2026 12:58:06 +0200 Subject: [PATCH 1/3] refactor!: Remove group rules in favor of `over` expressions --- README.md | 4 +- dataframely/_base_schema.py | 18 +-- dataframely/_rule.py | 121 ++---------------- docs/guides/examples/real-world.ipynb | 15 +-- docs/guides/faq.md | 4 +- docs/guides/quickstart.md | 30 +---- skills/SKILL.md | 3 +- tests/benches/test_schema.py | 56 ++++---- tests/core_validation/test_rule_evaluation.py | 49 +------ tests/schema/test_filter.py | 8 +- tests/schema/test_matches.py | 42 +----- tests/schema/test_repr.py | 6 +- tests/schema/test_rule_implementation.py | 81 +----------- tests/schema/test_sample.py | 8 +- tests/schema/test_validate.py | 12 +- 15 files changed, 66 insertions(+), 391 deletions(-) diff --git a/README.md b/README.md index d41d562c..89014e6e 100644 --- a/README.md +++ b/README.md @@ -56,9 +56,9 @@ class HouseSchema(dy.Schema): ratio = pl.col("num_bathrooms") / pl.col("num_bedrooms") return (ratio >= 1 / 3) & (ratio <= 3) - @dy.rule(group_by=["zip_code"]) + @dy.rule() def minimum_zip_code_count(cls) -> pl.Expr: - return pl.len() >= 2 + return (pl.len() >= 2).over("zip_code") ``` ### Validating data against schema diff --git a/dataframely/_base_schema.py b/dataframely/_base_schema.py index e7ab585f..204c67d6 100644 --- a/dataframely/_base_schema.py +++ b/dataframely/_base_schema.py @@ -14,7 +14,7 @@ import polars as pl from ._native import arrow_c_schema -from ._rule import DtypeCastRule, GroupRule, Rule, RuleFactory +from ._rule import DtypeCastRule, Rule, RuleFactory from .columns import Column from .exc import ImplementationError @@ -142,21 +142,7 @@ def __new__( f"{len(common_names)} overlaps: {common_list}." ) - # 2) Check that the columns referenced in the group rules exist. - for rule_name, rule in rules.items(): - if isinstance(rule, GroupRule): - missing_columns = set(rule.group_columns) - set(result.columns) - if len(missing_columns) > 0: - missing_list = ", ".join( - sorted(f"'{col}'" for col in missing_columns) - ) - raise ImplementationError( - f"Group validation rule '{rule_name}' has been implemented " - f"incorrectly. It references {len(missing_columns)} columns " - f"which are not in the schema: {missing_list}." - ) - - # 3) Check that all members are non-pathological (i.e., user errors). + # 2) Check that all members are non-pathological (i.e., user errors). for attr, value in namespace.items(): if attr.startswith("__"): continue diff --git a/dataframely/_rule.py b/dataframely/_rule.py index 3480bb07..cda83918 100644 --- a/dataframely/_rule.py +++ b/dataframely/_rule.py @@ -4,9 +4,8 @@ from __future__ import annotations import sys -from collections import defaultdict from collections.abc import Callable -from typing import Any, Literal +from typing import Any import polars as pl @@ -54,73 +53,26 @@ class DtypeCastRule(Rule): """ -class GroupRule(Rule): - """Rule that is evaluated on a group of columns.""" - - def __init__( - self, expr: pl.Expr | Callable[[], pl.Expr], group_columns: list[str] - ) -> None: - super().__init__(expr) - self.group_columns = group_columns - - def matches(self, other: Rule) -> bool: - if not isinstance(other, GroupRule): - return False - return super().matches(other) and self.group_columns == other.group_columns - - def __repr__(self) -> str: - return f"{super().__repr__()} grouped by {self.group_columns}" - - # -------------------------------------- FACTORY ------------------------------------- # class RuleFactory: """Factory class for rules created within schemas.""" - def __init__( - self, - validation_fn: Callable[[Any], pl.Expr], - group_columns: list[str] | Literal["primary_key"] | None, - ) -> None: + def __init__(self, validation_fn: Callable[[Any], pl.Expr]) -> None: self.validation_fn = validation_fn - self.group_columns = group_columns @classmethod def from_rule(cls, rule: Rule) -> Self: """Create a rule factory from an existing rule.""" - if isinstance(rule, GroupRule): - return cls( - validation_fn=lambda _: rule.expr, - group_columns=rule.group_columns, - ) - return cls(validation_fn=lambda _: rule.expr, group_columns=None) + return cls(validation_fn=lambda _: rule.expr) def make(self, schema: Any) -> Rule: """Create a new rule from this factory.""" - group_columns: list[str] | None - if self.group_columns == "primary_key": - from dataframely.exc import ImplementationError - - group_columns = schema.primary_key() - if not group_columns: - raise ImplementationError( - "Rule uses `group_by='primary_key'` but the schema has no" - " primary key." - ) - else: - group_columns = self.group_columns - if group_columns is not None: - return GroupRule( - expr=lambda: self.validation_fn(schema), - group_columns=group_columns, - ) return Rule(expr=lambda: self.validation_fn(schema)) -def rule( - *, group_by: list[str] | Literal["primary_key"] | None = None -) -> Callable[[ValidationFunction], RuleFactory]: +def rule() -> Callable[[ValidationFunction], RuleFactory]: """Mark a function as a rule to evaluate during validation. The name of the function will be used as the name of the rule. The function should @@ -132,35 +84,21 @@ def rule( - Validation requires accessing multiple columns (e.g. if valid values of column A depend on the value in column B). - Validation must be performed on groups of rows (e.g. if a column A must not - contain any duplicate values among rows with the same value in column B). + contain any duplicate values among rows with the same value in column B). This + can be achieved with an `over` expression. In all other instances, column-level validation rules should be preferred as it aids readability and improves error messages. - Args: - group_by: An optional list of columns to group by for rules operating on groups - of rows. If this list is provided, the returned expression must return a - single boolean value, i.e. some kind of aggregation function must be used - (e.g. `sum`, `any`, ...). Pass ``"primary_key"`` to dynamically resolve to - the schema's primary key columns at class creation time. This is useful for - defining rules in mixin classes where the primary key is not known at - definition time. - Note: You'll need to explicitly handle `null` values in your columns when defining rules. By default, any rule that evaluates to `null` because one of the columns used in the rule is `null` is interpreted as `true`, i.e. the row is assumed to be valid. - - Attention: - The rule logic should return a static result. - Other implementations using arbitrary python logic works for filtering and - validation, but may lead to wrong results in Schema comparisons - and (de-)serialization. """ def decorator(validation_fn: ValidationFunction) -> RuleFactory: - return RuleFactory(validation_fn=validation_fn, group_columns=group_by) + return RuleFactory(validation_fn=validation_fn) return decorator @@ -183,26 +121,12 @@ def with_evaluation_rules(lf: pl.LazyFrame, rules: dict[str, Rule]) -> pl.LazyFr of the rule. For each rule, a value of `True` indicates successful validation while `False` indicates an issue. """ - # Rules must be distinguished into two types of rules: - # 1. Simple rules can simply be selected on the data frame (this includes rules - # that check whether dtype casts succeeded) - # 2. "Group" rules require a `group_by` and a subsequent join - simple_exprs = { - name: rule.expr - for name, rule in rules.items() - if not isinstance(rule, GroupRule) - } - group_rules = { - name: rule for name, rule in rules.items() if isinstance(rule, GroupRule) - } - - # Before we can select all of the simple expressions, we need to turn the - # group rules into something to use in a `select` statement as well. + exprs = {name: rule.expr for name, rule in rules.items()} result = ( # NOTE: A value of `null` always validates successfully as nullability should # already be checked via dedicated rules. - lf.pipe(_with_group_rules, group_rules).with_columns( - **{name: expr.fill_null(True) for name, expr in simple_exprs.items()}, + lf.with_columns( + **{name: expr.fill_null(True) for name, expr in exprs.items()}, ) ) @@ -226,28 +150,3 @@ def with_evaluation_rules(lf: pl.LazyFrame, rules: dict[str, Rule]) -> pl.LazyFr ) return result - - -def _with_group_rules(lf: pl.LazyFrame, rules: dict[str, GroupRule]) -> pl.LazyFrame: - # First, we partition the rules by group columns. This will minimize the number - # of `group_by` calls and joins to make. - grouped_rules: dict[frozenset[str], dict[str, pl.Expr]] = defaultdict(dict) - for name, rule in rules.items(): - # NOTE: `null` indicates validity, see note above. - grouped_rules[frozenset(rule.group_columns)][name] = rule.expr.fill_null(True) - - # Then, for each `group_by`, we apply the relevant rules and keep all the rule - # evaluations around - group_evaluations: dict[frozenset[str], pl.LazyFrame] = {} - for group_columns, group_rules in grouped_rules.items(): - # We group by the group columns and apply all expressions - group_evaluations[group_columns] = lf.group_by(group_columns).agg(**group_rules) - - # Eventually, we apply the rule evaluations onto the input data frame. For this, we - # "broadcast" the results within each group across rows in the same group. - result = lf - for group_columns, frame in group_evaluations.items(): - result = result.join( - frame, on=list(group_columns), nulls_equal=True, maintain_order="left" - ) - return result diff --git a/docs/guides/examples/real-world.ipynb b/docs/guides/examples/real-world.ipynb index 8cd1466d..f47e3788 100644 --- a/docs/guides/examples/real-world.ipynb +++ b/docs/guides/examples/real-world.ipynb @@ -329,16 +329,9 @@ " diagnosis_code = dy.String(primary_key=True, regex=r\"[A-Z][0-9]{2,4}\")\n", " is_main = dy.Bool(nullable=False)\n", "\n", - " @dy.rule(group_by=[\"invoice_id\"])\n", + " @dy.rule()\n", " def exactly_one_main_diagnosis(cls) -> pl.Expr:\n", - " return pl.col(\"is_main\").sum() == 1" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note how we can also define validation rules on groups of rows using `@dy.rule(group_by=[...])`.\n" + " return (pl.col(\"is_main\").sum() == 1).over(\"invoice_id\")" ] }, { @@ -382,9 +375,9 @@ " diagnosis_code = dy.String(primary_key=True, regex=r\"[A-Z][0-9]{2,4}\")\n", " is_main = dy.Bool(nullable=False)\n", "\n", - " @dy.rule(group_by=[\"invoice_id\"])\n", + " @dy.rule()\n", " def exactly_one_main_diagnosis(cls) -> pl.Expr:\n", - " return pl.col(\"is_main\").sum() == 1" + " return (pl.col(\"is_main\").sum() == 1).over(\"invoice_id\")" ] }, { diff --git a/docs/guides/faq.md b/docs/guides/faq.md index 83636fed..b800600f 100644 --- a/docs/guides/faq.md +++ b/docs/guides/faq.md @@ -19,10 +19,10 @@ class UserSchema(dy.Schema): username = dy.String(nullable=False) email = dy.String(nullable=True) # Must be unique, or null. - @dy.rule(group_by=["username"]) + @dy.rule() def unique_username(cls) -> pl.Expr: """Username, a non-nullable field, must be total unique.""" - return pl.len() == 1 + return (pl.len() == 1).over("username") @dy.rule() def unique_email_or_null(cls) -> pl.Expr: diff --git a/docs/guides/quickstart.md b/docs/guides/quickstart.md index 11477c5d..d498ad71 100644 --- a/docs/guides/quickstart.md +++ b/docs/guides/quickstart.md @@ -64,35 +64,7 @@ The decorator `@dy.rule()` "registers" the function as a rule using its name (i. The returned expression provides a boolean value for each row of the data which evaluates to `True` whenever the data are valid with respect to this rule. -## Group rules - -For defining even more complex rules, the `@dy.rule` decorator allows for a `group_by` -parameter: this allows to evaluate a rule across _rows_. -For our housing data, this allows us to specify, for example, that we want to observe at least two houses per zip code: - -```python -import dataframely as dy - - -class HouseSchema(dy.Schema): - zip_code = dy.String(nullable=False, min_length=3) - num_bedrooms = dy.UInt8(nullable=False) - num_bathrooms = dy.UInt8(nullable=False) - price = dy.Float64(nullable=False) - - @dy.rule() - def reasonable_bathroom_to_bedroom_ratio(cls) -> pl.Expr: - ratio = pl.col("num_bathrooms") / pl.col("num_bedrooms") - return (ratio >= 1 / 3) & (ratio <= 3) - - @dy.rule(group_by=["zip_code"]) - def minimum_zip_code_count(cls) -> pl.Expr: - return pl.len() >= 2 -``` - -When defining rules on groups, we have to take care to use some kind of "aggregate function" -in order to produce exactly one value per group: -in group rules, the "input" that the expression is evaluated on is a set of rows. +Note that, for more complex rules, that should evaluate across _rows_, you can simply use an `over` expression. ````{note} If you are using [`ruff`](https://docs.astral.sh/ruff/) to lint your code, you'll need to tell `ruff` to treat rules like classmethods. To this end, you can add the following to your `pyproject.toml`: diff --git a/skills/SKILL.md b/skills/SKILL.md index dcb9b58a..fd181fe7 100644 --- a/skills/SKILL.md +++ b/skills/SKILL.md @@ -67,8 +67,7 @@ about the column contents. return cls.col1.col > cls.col2.col ``` -- Use group rules (i.e. methods decorated with `@dy.rule(group_by=...)`) for cross-row constraints beyond primary key - checks. +- Use rules with an `over` expression for cross-row constraints beyond primary key checks. ### Referencing Columns diff --git a/tests/benches/test_schema.py b/tests/benches/test_schema.py index 02df0903..f1c680ae 100644 --- a/tests/benches/test_schema.py +++ b/tests/benches/test_schema.py @@ -108,67 +108,63 @@ def test_multi_primary_key_filter( ) -# ---------------------------------- SINGLE GROUP-BY --------------------------------- # +# ------------------------------------ SINGLE OVER ----------------------------------- # -class SingleGroupBySchema(dy.Schema): +class SingleOverSchema(dy.Schema): elevation = dy.UInt16(nullable=True) aspect = dy.UInt16(nullable=True) slope = dy.UInt8(nullable=True) - @dy.rule(group_by=["slope"]) + @dy.rule() def average_elevation_at_least_2500(cls) -> pl.Expr: - return pl.col("elevation").mean() > 2500 + return (pl.col("elevation").mean() > 2500).over("slope") -@pytest.mark.benchmark(group="schema-group-by-single") -def test_single_group_by_validate( +@pytest.mark.benchmark(group="schema-over-single") +def test_single_over_validate( benchmark: BenchmarkFixture, dataset: pl.DataFrame ) -> None: - benchmark(SingleGroupBySchema.validate, dataset) + benchmark(SingleOverSchema.validate, dataset) -@pytest.mark.benchmark(group="schema-group-by-single") -def test_single_group_by_filter( - benchmark: BenchmarkFixture, dataset: pl.DataFrame -) -> None: - benchmark(SingleGroupBySchema.filter, dataset) +@pytest.mark.benchmark(group="schema-over-single") +def test_single_over_filter(benchmark: BenchmarkFixture, dataset: pl.DataFrame) -> None: + benchmark(SingleOverSchema.filter, dataset) -# ---------------------------------- MULTI GROUP-BY ---------------------------------- # +# ------------------------------------ MULTI OVER ------------------------------------ # -class MultiGroupBySchema(dy.Schema): +class MultiOverSchema(dy.Schema): elevation = dy.UInt16(nullable=True) aspect = dy.UInt16(nullable=True) slope = dy.UInt8(nullable=True) - @dy.rule(group_by=["slope"]) + @dy.rule() def average_elevation_at_least_2500(cls) -> pl.Expr: - return pl.col("elevation").mean() > 2500 + return (pl.col("elevation").mean() > 2500).over("slope") - @dy.rule(group_by=["slope"]) + @dy.rule() def at_least_one_elevation_2500(cls) -> pl.Expr: - return (pl.col("elevation") > 2500).any() + return (pl.col("elevation") > 2500).any().over("slope") - @dy.rule(group_by=["aspect"]) + @dy.rule() def at_least_50_aspects(cls) -> pl.Expr: - return pl.len() > 50 + return (pl.len() > 50).over("aspect") - @dy.rule(group_by=["aspect", "slope"]) + @dy.rule() def some_useless_filter(cls) -> pl.Expr: - return pl.len() >= 1 + return (pl.len() >= 1).over("aspect", "slope") -@pytest.mark.benchmark(group="schema-group-by-multiple") -def test_multi_group_by_validate( +@pytest.mark.benchmark(group="schema-over-multiple") +def test_multi_over_validate( benchmark: BenchmarkFixture, dataset: pl.DataFrame ) -> None: - benchmark(MultiGroupBySchema.validate, dataset) + benchmark(MultiOverSchema.validate, dataset) -@pytest.mark.benchmark(group="schema-group-by-multiple") -def test_multi_group_by_filter( - benchmark: BenchmarkFixture, dataset: pl.DataFrame -) -> None: - benchmark(MultiGroupBySchema.filter, dataset) +@pytest.mark.benchmark(group="schema-over-multiple") +def test_multi_over_filter(benchmark: BenchmarkFixture, dataset: pl.DataFrame) -> None: + benchmark(MultiOverSchema.filter, dataset) diff --git a/tests/core_validation/test_rule_evaluation.py b/tests/core_validation/test_rule_evaluation.py index 625e0078..7da16194 100644 --- a/tests/core_validation/test_rule_evaluation.py +++ b/tests/core_validation/test_rule_evaluation.py @@ -4,7 +4,7 @@ import polars as pl from polars.testing import assert_frame_equal -from dataframely._rule import GroupRule, Rule +from dataframely._rule import Rule from dataframely.testing import evaluate_rules @@ -59,50 +59,3 @@ def test_cross_column_rule() -> None: expected = pl.LazyFrame({"primary_key": [False, False, True, True]}) assert_frame_equal(actual, expected) - - -def test_group_rule() -> None: - lf = pl.LazyFrame({"a": [1, 1, 2, 2, 3], "b": [1, 1, 1, 2, 1]}) - rules: dict[str, Rule] = { - "unique_b": GroupRule(pl.col("b").n_unique() == 1, group_columns=["a"]) - } - actual = evaluate_rules(lf, rules) - - expected = pl.LazyFrame({"unique_b": [True, True, False, False, True]}) - assert_frame_equal(actual, expected) - - -def test_simple_rule_and_group_rule() -> None: - lf = pl.LazyFrame({"a": [1, 1, 2, 2, 3], "b": [1, 1, 1, 2, 1]}) - rules: dict[str, Rule] = { - "b|max": Rule(pl.col("b") <= 1), - "unique_b": GroupRule(pl.col("b").n_unique() == 1, group_columns=["a"]), - } - actual = evaluate_rules(lf, rules) - - expected = pl.LazyFrame( - { - "b|max": [True, True, True, False, True], - "unique_b": [True, True, False, False, True], - } - ) - assert_frame_equal(actual, expected, check_column_order=False) - - -def test_multiple_group_rules() -> None: - lf = pl.LazyFrame({"a": [1, 1, 2, 2, 3], "b": [1, 1, 1, 2, 1]}) - rules: dict[str, Rule] = { - "unique_b": GroupRule(pl.col("b").n_unique() == 1, group_columns=["a"]), - "sum_b": GroupRule(pl.col("b").sum() >= 2, group_columns=["a"]), - "group_count": GroupRule(pl.len() >= 2, group_columns=["a", "b"]), - } - actual = evaluate_rules(lf, rules) - - expected = pl.LazyFrame( - { - "unique_b": [True, True, False, False, True], - "sum_b": [True, True, True, True, False], - "group_count": [True, True, False, False, False], - } - ) - assert_frame_equal(actual, expected) diff --git a/tests/schema/test_filter.py b/tests/schema/test_filter.py index 99557e79..62de947a 100644 --- a/tests/schema/test_filter.py +++ b/tests/schema/test_filter.py @@ -10,7 +10,7 @@ from polars.testing import assert_frame_equal import dataframely as dy -from dataframely._rule import GroupRule +from dataframely._rule import Rule from dataframely.exc import SchemaError from dataframely.filter_result import FilterResult from dataframely.random import Generator @@ -228,11 +228,7 @@ def test_filter_maintain_order(eager: bool) -> None: schema = create_schema( "test", {"a": dy.UInt16(), "b": dy.UInt8()}, - { - "at_least_fifty_per_b": GroupRule( - lambda: pl.len() >= 50, group_columns=["b"] - ) - }, + {"at_least_fifty_per_b": Rule(lambda: (pl.len() >= 50).over("b"))}, ) generator = Generator() df = pl.DataFrame( diff --git a/tests/schema/test_matches.py b/tests/schema/test_matches.py index 137e8f40..c8f272c2 100644 --- a/tests/schema/test_matches.py +++ b/tests/schema/test_matches.py @@ -5,7 +5,7 @@ import pytest import dataframely as dy -from dataframely._rule import GroupRule, Rule +from dataframely._rule import Rule from dataframely.testing import create_schema @@ -94,47 +94,7 @@ def test_reflexivity() -> None: ), False, ), - ( # equal group rules - create_schema( - "test1", - columns={"a": dy.Int16()}, - rules={ - "rule1": Rule(pl.col("a") > 0), - "rule2": GroupRule(pl.len() > 2, group_columns=["a"]), - }, - ), - create_schema( - "test2", - columns={"a": dy.Int16()}, - rules={ - "rule1": Rule(pl.col("a") > 0), - "rule2": GroupRule(pl.len() > 2, group_columns=["a"]), - }, - ), - True, - ), - ( # dfifferent group columns - create_schema( - "test1", - columns={"a": dy.Int16(), "b": dy.Int32()}, - rules={ - "rule2": GroupRule(pl.len() > 2, group_columns=["a"]), - }, - ), - create_schema( - "test2", - columns={"a": dy.Int16(), "b": dy.Int32()}, - rules={ - "rule2": GroupRule(pl.len() > 2, group_columns=["a", "b"]), - }, - ), - False, - ), ], ) def test_matches(lhs: type[dy.Schema], rhs: type[dy.Schema], expected: bool) -> None: assert lhs.matches(rhs) == expected - - -def test_group_rule_inequality_type_mismatch() -> None: - assert not GroupRule(pl.len() > 2, group_columns=["a"]).matches(Rule(pl.len() > 2)) diff --git a/tests/schema/test_repr.py b/tests/schema/test_repr.py index 3f49cd55..2fcc7a54 100644 --- a/tests/schema/test_repr.py +++ b/tests/schema/test_repr.py @@ -39,9 +39,9 @@ class SchemaWithRules(dy.Schema): def my_rule(cls) -> pl.Expr: return pl.col("a") < 100 - @dy.rule(group_by=["a"]) + @dy.rule() def my_group_rule(cls) -> pl.Expr: - return pl.col("a").sum() > 50 + return (pl.col("a").sum() > 50).over("a") def test_repr_with_rules() -> None: @@ -52,7 +52,7 @@ def test_repr_with_rules() -> None: - "b2": String(primary_key=True, regex='^[A-Z]{3}$') Rules: - "my_rule": [(col("a")) < (dyn int: 100)] - - "my_group_rule": [(col("a").sum()) > (dyn int: 50)] grouped by ['a'] + - "my_group_rule": [(col("a").sum()) > (dyn int: 50)].over([col("a")]) """ assert repr(SchemaWithRules) == textwrap.dedent(expected) diff --git a/tests/schema/test_rule_implementation.py b/tests/schema/test_rule_implementation.py index 74eedc97..2d73e01b 100644 --- a/tests/schema/test_rule_implementation.py +++ b/tests/schema/test_rule_implementation.py @@ -5,90 +5,11 @@ import pytest import dataframely as dy -from dataframely._rule import GroupRule, Rule +from dataframely._rule import Rule from dataframely.exc import ImplementationError from dataframely.testing import create_schema -def test_group_rule_group_by_error() -> None: - with pytest.raises( - ImplementationError, - match=( - r"Group validation rule 'b_greater_zero' has been implemented " - r"incorrectly\. It references 1 columns which are not in the schema" - ), - ): - create_schema( - "test", - columns={"a": dy.Integer(), "b": dy.Integer()}, - rules={ - "b_greater_zero": GroupRule( - (pl.col("b") > 0).all(), group_columns=["c"] - ) - }, - ) - - -def test_group_rule_primary_key_single() -> None: - class MySchema(dy.Schema): - a = dy.Int64(primary_key=True) - b = dy.Int64() - - @dy.rule(group_by="primary_key") - def b_positive(cls) -> pl.Expr: - return (pl.col("b") > 0).all() - - rules = MySchema._schema_validation_rules() - assert isinstance(rules["b_positive"], GroupRule) - assert rules["b_positive"].group_columns == ["a"] - - -def test_group_rule_primary_key_composite() -> None: - class MySchema(dy.Schema): - a = dy.Int64(primary_key=True) - b = dy.Int64(primary_key=True) - c = dy.Int64() - - @dy.rule(group_by="primary_key") - def c_positive(cls) -> pl.Expr: - return (pl.col("c") > 0).all() - - rules = MySchema._schema_validation_rules() - assert isinstance(rules["c_positive"], GroupRule) - assert sorted(rules["c_positive"].group_columns) == ["a", "b"] - - -def test_group_rule_primary_key_no_pk() -> None: - with pytest.raises( - ImplementationError, - match=r"group_by='primary_key'.*no primary key", - ): - - class MySchema(dy.Schema): - a = dy.Int64() - - @dy.rule(group_by="primary_key") - def a_positive(cls) -> pl.Expr: - return (pl.col("a") > 0).all() - - -def test_group_rule_primary_key_mixin() -> None: - class MyMixin: - id = dy.Int64(primary_key=True) - value = dy.Int64() - - @dy.rule(group_by="primary_key") - def value_positive(cls) -> pl.Expr: - return (pl.col("value") > 0).all() - - class MySchema(MyMixin, dy.Schema): - other_id = dy.Int64(primary_key=True) - - rules = MySchema._schema_validation_rules() - assert isinstance(rules["value_positive"], GroupRule) - assert rules["value_positive"].group_columns == ["id", "other_id"] - - def test_rule_column_overlap_error() -> None: with pytest.raises( ImplementationError, diff --git a/tests/schema/test_sample.py b/tests/schema/test_sample.py index 870f695c..f74dac7f 100644 --- a/tests/schema/test_sample.py +++ b/tests/schema/test_sample.py @@ -39,9 +39,9 @@ class ComplexSchema(dy.Schema): def a_greater_b(cls) -> pl.Expr: return pl.col("a") > pl.col("b") - @dy.rule(group_by=["a"]) + @dy.rule() def minimum_two_per_a(cls) -> pl.Expr: - return pl.len() >= 2 + return (pl.len() >= 2).over("a") class LimitedComplexSchema(dy.Schema): @@ -52,10 +52,10 @@ class LimitedComplexSchema(dy.Schema): def a_greater_b(cls) -> pl.Expr: return pl.col("a") > pl.col("b") - @dy.rule(group_by=["a"]) + @dy.rule() def minimum_two_per_a(cls) -> pl.Expr: # We cannot generate more than 768 rows with this rule - return pl.len() <= 3 + return (pl.len() <= 3).over("a") class OrderedSchema(dy.Schema): diff --git a/tests/schema/test_validate.py b/tests/schema/test_validate.py index 908d4e90..6f858ce7 100644 --- a/tests/schema/test_validate.py +++ b/tests/schema/test_validate.py @@ -9,7 +9,7 @@ from polars.testing import assert_frame_equal import dataframely as dy -from dataframely._rule import GroupRule +from dataframely._rule import Rule from dataframely.exc import SchemaError, ValidationError from dataframely.random import Generator from dataframely.testing import create_schema @@ -29,9 +29,9 @@ class MyComplexSchema(dy.Schema): def b_greater_a(cls) -> pl.Expr: return pl.col("b") > pl.col("a") - @dy.rule(group_by=["a"]) + @dy.rule() def b_unique_within_a(cls) -> pl.Expr: - return pl.col("b").n_unique() == 1 + return (pl.col("b").n_unique() == 1).over("a") class MyComplexSchemaWithLazyRules(dy.Schema): @@ -42,9 +42,9 @@ class MyComplexSchemaWithLazyRules(dy.Schema): def b_greater_a(cls) -> pl.Expr: return cls.b.col > cls.a.col - @dy.rule(group_by=["a"]) + @dy.rule() def b_unique_within_a(cls) -> pl.Expr: - return cls.b.col.n_unique() == SOME_CONSTANT_DEFINED_LATER + return (cls.b.col.n_unique() == SOME_CONSTANT_DEFINED_LATER).over("a") SOME_CONSTANT_DEFINED_LATER = 1 @@ -217,7 +217,7 @@ def test_validate_maintain_order() -> None: schema = create_schema( "test", {"a": dy.UInt16(), "b": dy.UInt8()}, - {"at_least_fifty_per_b": GroupRule(lambda: pl.len() >= 2, group_columns=["b"])}, + {"at_least_fifty_per_b": Rule(lambda: (pl.len() >= 2).over("b"))}, ) generator = Generator() df = pl.DataFrame( From 6a72de1aaaeda824f790ddbcd79f716a703d5442 Mon Sep 17 00:00:00 2001 From: Oliver Borchert Date: Fri, 31 Jul 2026 13:05:21 +0200 Subject: [PATCH 2/3] Add test --- tests/schema/test_validate.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/schema/test_validate.py b/tests/schema/test_validate.py index 6f858ce7..e514e770 100644 --- a/tests/schema/test_validate.py +++ b/tests/schema/test_validate.py @@ -213,6 +213,28 @@ def test_group_rule_on_nulls( assert not schema.is_valid(df, cast=True) +@pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame]) +@pytest.mark.parametrize("eager", [True, False]) +def test_dynamic_column_selection_unaffected_by_other_rules( + df_type: type[pl.DataFrame] | type[pl.LazyFrame], eager: bool +) -> None: + # Regression test for https://github.com/Quantco/dataframely/issues/332: a rule + # relying on dynamic column selection must not "see" the boolean columns produced + # by evaluating other rules, regardless of whether other (`over`) rules are present. + schema = create_schema( + "test", + {"x": dy.Bool()}, + rules={ + "no_boolean_column_is_true": Rule(~pl.any_horizontal(pl.col(pl.Boolean))) + }, + ) + + df = df_type({"x": [False]}) + result = _validate_and_collect(schema, df, eager=eager) + assert len(result) == 1 + assert schema.is_valid(df) + + def test_validate_maintain_order() -> None: schema = create_schema( "test", From 98db12944d92b64ee3770d04d32f431a4ee052d7 Mon Sep 17 00:00:00 2001 From: Oliver Borchert Date: Fri, 31 Jul 2026 13:12:11 +0200 Subject: [PATCH 3/3] Review --- docs/guides/quickstart.md | 2 +- tests/schema/test_validate.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/guides/quickstart.md b/docs/guides/quickstart.md index d498ad71..4d7e24b1 100644 --- a/docs/guides/quickstart.md +++ b/docs/guides/quickstart.md @@ -64,7 +64,7 @@ The decorator `@dy.rule()` "registers" the function as a rule using its name (i. The returned expression provides a boolean value for each row of the data which evaluates to `True` whenever the data are valid with respect to this rule. -Note that, for more complex rules, that should evaluate across _rows_, you can simply use an `over` expression. +Note that, for more complex rules that should evaluate across _rows_, you can simply use an `over` expression. ````{note} If you are using [`ruff`](https://docs.astral.sh/ruff/) to lint your code, you'll need to tell `ruff` to treat rules like classmethods. To this end, you can add the following to your `pyproject.toml`: diff --git a/tests/schema/test_validate.py b/tests/schema/test_validate.py index e514e770..e9649bf0 100644 --- a/tests/schema/test_validate.py +++ b/tests/schema/test_validate.py @@ -225,7 +225,8 @@ def test_dynamic_column_selection_unaffected_by_other_rules( "test", {"x": dy.Bool()}, rules={ - "no_boolean_column_is_true": Rule(~pl.any_horizontal(pl.col(pl.Boolean))) + "no_boolean_column_is_true": Rule(~pl.any_horizontal(pl.col(pl.Boolean))), + "some_over_rule": Rule(pl.len().over("x") >= 1), }, )