refactor!: Remove group rules in favor of over expressions

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

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

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

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\")"
    ]
   },
   {

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:

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`:

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