feat: Allow semantic comparison of collections (#62)

Co-authored-by: Oliver Borchert <[email protected]>

Author: AndreasAlbertQC

dataframely/_base_schema.py

@@ -149,9 +149,9 @@ def _get_metadata(source: dict[str, Any]) -> Metadata:
     def __repr__(cls) -> str:
         parts = [f'[Schema "{cls.__name__}"]']
         parts.append(textwrap.indent("Columns:", prefix=" " * 2))
-        for name, col in cls.columns().items():  # type: ignore
+        for name, col in cls.columns().items():
             parts.append(textwrap.indent(f'- "{name}": {col!r}', prefix=" " * 4))
-        if validation_rules := cls._schema_validation_rules():  # type: ignore
+        if validation_rules := cls._schema_validation_rules():
             parts.append(textwrap.indent("Rules:", prefix=" " * 2))
             for name, rule in validation_rules.items():
                 parts.append(textwrap.indent(f'- "{name}": {rule!r}', prefix=" " * 4))

dataframely/_filter.py

@@ -34,6 +34,12 @@ def filter() -> Callable[[Callable[[C], pl.LazyFrame]], Filter[C]]:
     Attention:
         Make sure to provide unique combinations of the primary keys or the filters
         might introduce duplicate rows.
+
+    Attention:
+        The filter logic should return a lazy frame with a static computational graph.
+        Other implementations using arbitrary python logic works for filtering and
+        validation, but may lead to wrong results in Collection comparisons
+        and (de-)serialization.
     """
 
     def decorator(validation_fn: Callable[[C], pl.LazyFrame]) -> Filter[C]:

dataframely/_rule.py

@@ -106,6 +106,12 @@ def rule(*, group_by: list[str] | None = None) -> Callable[[ValidationFunction],
         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) -> Rule:

dataframely/collection.py

@@ -4,6 +4,7 @@
 import warnings
 from abc import ABC
 from collections.abc import Mapping, Sequence
+from dataclasses import asdict
 from pathlib import Path
 from typing import Any, Self, cast
 
@@ -225,6 +226,65 @@ def sample(
         # 3) Eventually, we initialize the final collection and return
         return cls.validate(members)
 
+    @classmethod
+    def matches(cls, other: type["Collection"]) -> bool:
+        """Check whether this collection semantically matches another.
+
+        Args:
+            other: The collection to compare with.
+
+        Returns:
+            Whether the two collections are semantically equal.
+
+        Attention:
+            For custom filters, reliable comparison results are only guaranteed
+            if the filter always returns a static polars expression.
+            Otherwise, this function may falsely indicate a match.
+        """
+
+        def _members_match() -> bool:
+            members_lhs = cls.members()
+            members_rhs = other.members()
+
+            # Member names must match
+            if members_lhs.keys() != members_rhs.keys():
+                return False
+
+            # Member attributes must match
+            for name in members_lhs:
+                lhs = asdict(members_lhs[name])
+                rhs = asdict(members_rhs[name])
+                for attr in lhs.keys() | rhs.keys():
+                    if attr == "schema":
+                        if not lhs[attr].matches(rhs[attr]):
+                            return False
+                    else:
+                        if lhs[attr] != rhs[attr]:
+                            return False
+            return True
+
+        def _filters_match() -> bool:
+            filters_lhs = cls._filters()
+            filters_rhs = other._filters()
+
+            # Filter names must match
+            if filters_lhs.keys() != filters_rhs.keys():
+                return False
+
+            # Computational graph of filter logic must match
+            # Evaluate on empty dataframes
+            empty_left = cls.create_empty()
+            empty_right = other.create_empty()
+
+            for name in filters_lhs:
+                lhs = filters_lhs[name].logic(empty_left)
+                rhs = filters_rhs[name].logic(empty_right)
+                if lhs.serialize(format="json") != rhs.serialize(format="json"):
+                    return False
+            return True
+
+        return _members_match() and _filters_match()
+
     @classmethod
     def _preprocess_sample(
         cls, sample: dict[str, Any], index: int, generator: Generator

tests/collection/test_matches.py

@@ -0,0 +1,119 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+import polars as pl
+
+import dataframely as dy
+
+
+def test_collection_matches_itself() -> None:
+    """Collections should match themselves."""
+
+    class MySchema(dy.Schema):
+        foo = dy.Integer()
+
+    # First collection has one member
+    class MyCollection1(dy.Collection):
+        x: dy.LazyFrame[MySchema]
+
+    assert MyCollection1.matches(MyCollection1)
+
+
+def test_collection_matches_different_members() -> None:
+    """Collections should count as different if they have members with different
+    names."""
+
+    class MySchema(dy.Schema):
+        foo = dy.Integer()
+
+    class MyCollection1(dy.Collection):
+        x: dy.LazyFrame[MySchema]
+
+    class MyCollection2(dy.Collection):
+        y: dy.LazyFrame[MySchema]
+
+    # Should not match
+    assert not MyCollection1.matches(MyCollection2)
+
+
+def test_collection_matches_different_schemas() -> None:
+    """Collections should count as different if their members have different schemas."""
+
+    class MyIntSchema(dy.Schema):
+        foo = dy.Integer()
+
+    class MyStringSchema(dy.Schema):
+        foo = dy.String()
+
+    assert not MyIntSchema.matches(MyStringSchema), (
+        "Test schemas must not match for test setup to make sense"
+    )
+
+    # Collections have the same member name
+    # but mismatching schemas
+    class MyCollection1(dy.Collection):
+        x: dy.LazyFrame[MyIntSchema]
+
+    class MyCollection2(dy.Collection):
+        x: dy.LazyFrame[MyStringSchema]
+
+    # Should not match
+    assert not MyCollection1.matches(MyCollection2)
+
+
+def test_collection_matches_different_filter_names() -> None:
+    """Collections should count as different if they have the same members but different
+    names."""
+
+    class MyIntSchema(dy.Schema):
+        foo = dy.Integer(primary_key=True)
+
+    class MyCollection1(dy.Collection):
+        x: dy.LazyFrame[MyIntSchema]
+
+    class MyCollection2(MyCollection1):
+        @dy.filter()
+        def test_filter(self) -> pl.LazyFrame:
+            return dy.filter_relationship_one_to_one(self.x, self.x, ["foo"])
+
+    # Should not match
+    assert not MyCollection1.matches(MyCollection2)
+
+
+def test_collection_matches_different_filter_logc() -> None:
+    """Collections should count as different if they have the same members but different
+    filter logic."""
+
+    class MyIntSchema(dy.Schema):
+        foo = dy.Integer(primary_key=True)
+
+    class BaseCollection(dy.Collection):
+        x: dy.LazyFrame[MyIntSchema]
+
+    class MyCollection1(BaseCollection):
+        @dy.filter()
+        def test_filter(self) -> pl.LazyFrame:
+            return dy.filter_relationship_one_to_one(self.x, self.x, ["foo"])
+
+    class MyCollection2(BaseCollection):
+        @dy.filter()
+        def test_filter(self) -> pl.LazyFrame:
+            return dy.filter_relationship_one_to_at_least_one(self.x, self.x, ["foo"])
+
+    assert not MyCollection1.matches(MyCollection2)
+
+
+def test_collection_matches_different_optional() -> None:
+    """Collections should count as different if their members differ in whether they are
+    optional or not."""
+
+    class FooSchema(dy.Schema):
+        x = dy.Integer()
+
+    class MandatoryFooCollection(dy.Collection):
+        foo: dy.LazyFrame[FooSchema]
+
+    class OptionalFooCollection(dy.Collection):
+        foo: dy.LazyFrame[FooSchema] | None
+
+    assert not MandatoryFooCollection.matches(OptionalFooCollection)