matchlab 0.1.0.dev1__py3-none-any.whl

matchlab/testkit/compare.py

@@ -0,0 +1,178 @@
+"""Scoring: turn results into the vocabulary, then compare them with the answer.
+
+The functions here convert a resolver's or model's output into `Cluster` objects, and
+diff two `Cluster` lists. None of that needs the planted answer, which is why they stay
+as plain functions rather than methods. `LinkedSources.diff_resolver_output()` and
+`.diff_model_edges()` do need the planted answer, so those are methods on the one
+object that holds it.
+
+Every comparison returns `(identical, report)`, where the report classifies each cluster
+as perfect, subset, superset, wrong or invalid. That breakdown is the point. "Eight
+clusters were subsets" tells you the matcher is too strict, in a way a number cannot.
+"""
+
+from collections import Counter
+
+import polars as pl
+
+from matchlab.core.dsu import DisjointSet
+from matchlab.testkit.entities import Cluster, EntityReference
+
+
+def resolver_output_to_clusters(resolver_output: pl.DataFrame) -> set[Cluster]:
+    """Convert a collected resolver output into entities comparable with truth.
+
+    This is the other half of measuring a plan against generated data. The testkit
+    plants known entities, the plan resolves records into clusters, and this turns
+    those clusters back into the same currency the truth is expressed in. Pair it with
+    [`LinkedSources.true_entity_subset()`][matchlab.testkit.linked.LinkedSources.true_entity_subset]
+    and [`diff_entities()`][matchlab.testkit.compare.diff_entities]:
+
+        identical, report = diff_entities(
+            expected=linked.true_entity_subset("crn", "cdms"),
+            actual=list(resolver_output_to_clusters(resolver_output)),
+        )
+
+    A `Cluster` compares by its keys and never by its ID (`Cluster.__eq__`), which is
+    what lets this work at all. The resolver output's `root` is a content-derived hash
+    minted at collect time, and has no counterpart in the testkit's synthetic ID space.
+    Only the `(source, key)` membership is comparable, and that is exactly what a
+    cluster asserts.
+
+    Args:
+        resolver_output: A table conforming to `SCHEMA_RESOLVER_OUTPUT`, the table
+            `Resolver.entities()` returns, with `root`, `key` and `source` columns.
+
+    Returns:
+        One Cluster per distinct `root`.
+
+    Raises:
+        ValueError: If the required columns are absent.
+    """
+    required = {"root", "key", "source"}
+    missing = required - set(resolver_output.columns)
+    if missing:
+        raise ValueError(
+            f"Fields {sorted(missing)} must be included in the resolver output and "
+            f"are missing. Available: {sorted(resolver_output.columns)}"
+        )
+
+    clusters = resolver_output.group_by("root").agg("source", "key")
+
+    entities: set[Cluster] = set()
+    for cluster in clusters.iter_rows(named=True):
+        keys: dict[str, set[str]] = {}
+        for source, key in zip(cluster["source"], cluster["key"], strict=True):
+            keys.setdefault(source, set()).add(str(key))
+        entities.add(
+            Cluster(
+                keys=EntityReference(
+                    {source: frozenset(group) for source, group in keys.items()}
+                )
+            )
+        )
+
+    return entities
+
+
+def scores_to_clusters(
+    scores: pl.DataFrame,
+    left_clusters: tuple[Cluster, ...],
+    right_clusters: tuple[Cluster, ...] | None = None,
+    threshold: float = 0.0,
+) -> tuple[Cluster, ...]:
+    """Merge clusters connected by a score at or above threshold.
+
+    Left and right clusters that share an edge scoring at least `threshold` merge into
+    one `Cluster`. With no `right_clusters`, this merges within `left_clusters` alone
+    (deduplication).
+
+    Args:
+        scores: A `left_id`/`right_id`/`score` edge table.
+        left_clusters: Clusters the model read as its left input.
+        right_clusters: Clusters read as the right input, for a linker. `None` for a
+            dedupe.
+        threshold: Score at or above which an edge counts as a match.
+
+    Returns:
+        One merged `Cluster` per connected component.
+    """
+    left_lookup = {entity.id: entity for entity in left_clusters}
+    if right_clusters is not None:
+        right_lookup = {entity.id: entity for entity in right_clusters}
+    else:
+        right_lookup = left_lookup
+
+    djs = DisjointSet[Cluster]()
+
+    # Add ALL entities to the disjoint set
+    for entity in left_clusters:
+        djs.add(entity)
+    if right_clusters is not None:
+        for entity in right_clusters:
+            djs.add(entity)
+
+    # Add edges to the disjoint set
+    for record in scores.to_dicts():
+        if record["score"] >= threshold:
+            djs.union(
+                left_lookup[record["left_id"]],
+                right_lookup[record["right_id"]],
+            )
+
+    components: set[set[Cluster]] = djs.get_components()
+
+    entities: list[Cluster] = []
+    for component in components:
+        merged: Cluster = sum(component)
+        entities.append(merged)
+
+    return tuple(entities)
+
+
+def diff_entities(expected: list[Cluster], actual: list[Cluster]) -> tuple[bool, dict]:
+    """Compare two lists of Cluster against each other, with a diff report.
+
+    Args:
+        expected: The expected Cluster list.
+        actual: The actual Cluster list.
+
+    Returns:
+        `(identical, report)`. `identical` is `True` if the two lists match exactly.
+        `report` counts how each actual cluster relates to the expected ones:
+
+        - `perfect`: matches an expected cluster exactly.
+        - `subset`: is a subset of an expected cluster.
+        - `superset`: is a superset of an expected cluster.
+        - `wrong`: does not overlap any expected cluster.
+        - `invalid`: contains keys absent from every expected cluster.
+    """
+    expected_set, actual_set = set(expected), set(actual)
+    if expected_set == actual_set:
+        return True, {}
+
+    all_expected = sum(expected_set)
+    perfect_matches = expected_set & actual_set
+    remaining_actual = actual_set - perfect_matches
+
+    counter = Counter(
+        {
+            "perfect": len(perfect_matches),
+            "subset": 0,
+            "superset": 0,
+            "wrong": 0,
+            "invalid": 0,
+        }
+    )
+
+    for a in remaining_actual:
+        if any(a in e for e in expected_set):
+            counter["subset"] += 1
+        elif a not in all_expected:
+            counter["invalid"] += 1
+        elif any(e in a for e in expected_set):
+            counter["superset"] += 1
+        else:
+            counter["wrong"] += 1
+
+    return False, dict(counter)

matchlab/testkit/entities.py

@@ -0,0 +1,290 @@
+"""The vocabulary: what a planted answer is, and what a claimed answer is.
+
+Two types exist, and the difference between them matters:
+
+* `TrueEntity` is the *answer*: one planted real-world thing, identified by the values
+  it was generated from, spanning every source it landed in.
+* `Cluster` is *an answer*: a set of records claimed to be one entity, identified by
+  that membership alone.
+
+`TrueEntity.cluster()` projects the first onto the second, so the two become
+comparable. Only membership can be compared this way. A resolver output's IDs are
+minted at collect time and have no counterpart in generated data.
+"""
+
+from random import getrandbits
+from typing import Any, Self
+
+from frozendict import frozendict
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class EntityReference(frozendict):
+    """Reference to an entity's presence in specific sources.
+
+    Maps source names to sets of primary keys.
+    """
+
+    def __init__(
+        self,
+        mapping: dict[str, frozenset[str]] | None = None,
+    ) -> None:
+        """Initialise the EntityReference."""
+        super().__init__({} if mapping is None else mapping)
+
+    def __add__(self, other: "EntityReference") -> "EntityReference":
+        """Merge two EntityReferences by unioning keys for each source."""
+        if not isinstance(other, EntityReference):
+            return NotImplemented
+
+        return EntityReference(
+            {
+                k: self.get(k, frozenset()) | other.get(k, frozenset())
+                for k in self.keys() | other.keys()
+            }
+        )
+
+    def __le__(self, other: Self) -> bool:
+        """Test if self is a subset of other."""
+        if not isinstance(other, EntityReference):
+            return NotImplemented
+
+        return all(name in other and self[name] <= other[name] for name in self)
+
+
+class EntityIDMixin:
+    """Shared ID behaviour for entity classes.
+
+    Provides integer conversion and comparison operators, so entities sort by `id`.
+    """
+
+    id: int
+
+    def __int__(self) -> int:
+        """Allow converting an entity to an integer by returning its ID."""
+        return self.id
+
+    def __lt__(self, other: Self | int) -> bool:
+        """Compare based on ID for sorting operations."""
+        if hasattr(other, "id"):
+            return self.id < other.id
+        if isinstance(other, int):
+            return self.id < other
+        return NotImplemented
+
+    def __gt__(self, other: Self | int) -> bool:
+        """Compare based on ID for sorting operations."""
+        if hasattr(other, "id"):
+            return self.id > other.id
+        if isinstance(other, int):
+            return self.id > other
+        return NotImplemented
+
+    def __le__(self, other: Self | int) -> bool:
+        """Compare based on ID for sorting operations."""
+        if hasattr(other, "id"):
+            return self.id <= other.id
+        if isinstance(other, int):
+            return self.id <= other
+        return NotImplemented
+
+    def __ge__(self, other: Self | int) -> bool:
+        """Compare based on ID for sorting operations."""
+        if hasattr(other, "id"):
+            return self.id >= other.id
+        if isinstance(other, int):
+            return self.id >= other
+        return NotImplemented
+
+
+class SourceKeyMixin:
+    """Shared source-key behaviour for entity classes.
+
+    Provides `get_keys()`, for reading back the keys held for one source.
+    """
+
+    keys: EntityReference
+
+    def get_keys(self, name: str) -> set[str]:
+        """Get keys for a specific source.
+
+        Args:
+            name: Name of the source.
+
+        Returns:
+            Set of keys, empty if the source is not found.
+        """
+        return set(self.keys.get(name, frozenset()))
+
+
+class Cluster(BaseModel, EntityIDMixin, SourceKeyMixin):
+    """A set of records claimed to be one entity, an *answer*.
+
+    This is the unit of comparison. Both sides of `diff_entities` are clusters: the
+    expected side projected from `TrueEntity.cluster()`, and the actual side read back
+    off a model or a resolver's output. It carries membership and nothing else, because
+    membership is the only thing the two sides can agree on. A resolver output's IDs
+    are minted at collect time and have no counterpart in generated data.
+
+    Equality and hashing are over `keys` alone. `id` is carried for bookkeeping and is
+    deliberately ignored when comparing. Contrast `TrueEntity`, which is the answer key
+    rather than an answer.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)
+
+    id: int = Field(default_factory=lambda: getrandbits(63))  # 64 gives OverflowError
+    keys: EntityReference
+
+    def __add__(self, other: "Cluster") -> "Cluster":
+        """Combine two Cluster objects by combining the keys."""
+        if other is None:
+            return self
+        if not isinstance(other, Cluster):
+            return NotImplemented
+        return Cluster(keys=self.keys + other.keys)
+
+    def __radd__(self, other: object) -> "Cluster":
+        """Handle sum() by treating 0 as an empty Cluster."""
+        if other == 0:  # sum() starts with 0
+            return self
+        return NotImplemented
+
+    def __sub__(self, other: "Cluster") -> dict[str, frozenset[str]]:
+        """Return keys in self that aren't in other, by source.
+
+        Used to diff two Cluster objects.
+        """
+        if not isinstance(other, Cluster):
+            return NotImplemented
+
+        diff = {}
+        for name, our_keys in self.keys.items():
+            their_keys = other.keys.get(name, frozenset())
+            if remaining := our_keys - their_keys:
+                diff[name] = remaining
+
+        return diff
+
+    def __rsub__(self, other: "Cluster") -> dict[str, frozenset[str]]:
+        """Support reverse subtraction."""
+        if not isinstance(other, Cluster):
+            return NotImplemented
+        return other - self
+
+    def __eq__(self, other: object) -> bool:
+        """Compare based on keys."""
+        if not isinstance(other, Cluster):
+            return NotImplemented
+        return self.keys == other.keys
+
+    def __contains__(self, other: "Cluster") -> bool:
+        """Check if this entity contains all keys from other entity."""
+        return other.keys <= self.keys
+
+    def __hash__(self) -> int:
+        """Hash based on EntityReference which is itself hashable."""
+        return hash(self.keys)
+
+    def is_subset_of_true_entity(self, source_entity: "TrueEntity") -> bool:
+        """Check if this Cluster's references are a subset of a TrueEntity's."""
+        return self.keys <= source_entity.keys
+
+    def similarity_ratio(self, other: "Cluster") -> float:
+        """Return ratio of shared keys to total keys across all sources."""
+        total_keys = 0
+        shared_keys = 0
+
+        # Get all source names
+        all_sources = set(self.keys.keys()) | set(other.keys.keys())
+
+        for name in all_sources:
+            our_keys = self.keys.get(name, frozenset())
+            their_keys = other.keys.get(name, frozenset())
+
+            total_keys += len(our_keys | their_keys)
+            shared_keys += len(our_keys & their_keys)
+
+        return shared_keys / total_keys if total_keys > 0 else 0.0
+
+
+class TrueEntity(BaseModel, EntityIDMixin, SourceKeyMixin):
+    """One planted real-world thing, the *answer key*.
+
+    This is what the generator started from. It holds `base_values`, the feature values
+    its rows were derived from, and accumulates the keys it landed under in every source
+    it appears in. Equality is over `base_values`. Two entities are the same thing if
+    they were generated from the same values.
+
+    It spans every source, so it is not directly comparable with a result. Project it
+    onto the sources under test with `cluster()` to get something that is.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    id: int = Field(default_factory=lambda: getrandbits(63))  # 64 gives OverflowError
+    base_values: dict[str, Any] = Field(description="Feature name -> base value")
+    keys: EntityReference = Field(
+        description="Source to keys mapping",
+        default=EntityReference(mapping=frozenset()),
+    )
+    total_unique_variations: int = Field(default=0)
+
+    def __eq__(self, other: object) -> bool:
+        """Equal if base values are shared, or integer ID matches."""
+        if isinstance(other, TrueEntity):
+            return self.base_values == other.base_values
+        if isinstance(other, int):
+            return self.id == other
+        return NotImplemented
+
+    def __hash__(self) -> int:
+        """Hash based on sorted base values."""
+        return hash(tuple(sorted(self.base_values.items())))
+
+    def add_source_reference(self, name: str, keys: list[str]) -> None:
+        """Add or update a source reference.
+
+        Args:
+            name: Source name.
+            keys: List of primary keys for this source.
+        """
+        mapping = dict(self.keys)
+        mapping[name] = frozenset(keys)
+        self.keys = EntityReference(mapping)
+
+    def cluster(self, *names: str) -> Cluster | None:
+        """Project this true entity onto the given sources, making it comparable.
+
+        Comparing equality of `Cluster` sets is a simpler, more reliable test than
+        checking whether `Cluster` objects are subsets of `TrueEntity` objects. This
+        method is what makes that comparison possible:
+
+        ```python
+        actual: set[Cluster] = ...
+        expected: set[Cluster] = {
+            s.cluster("source1", "source2") for s in true_entities
+        }
+
+        is_identical = expected == actual
+        missing = expected - actual
+        extra = actual - expected
+        ```
+
+        Args:
+            *names: Names of sources to include in the Cluster.
+
+        Returns:
+            A `Cluster` with only the named sources' keys, or `None` if this entity has
+            no keys in any of them.
+        """
+        filtered = {
+            name: self.keys.get(name)
+            for name in names
+            if self.keys.get(name) is not None
+        }
+
+        if len(filtered) == 0:
+            return None
+
+        return Cluster(keys=EntityReference(filtered))

matchlab/testkit/features.py

@@ -0,0 +1,146 @@
+"""What to generate: feature declarations and the variations applied to them.
+
+Pure declaration. Nothing here generates or compares anything. `FeatureConfig`
+describes one column's values. Its `variations` are what let one true entity produce
+more than one form of the same value, which `_generate.generate_rows` turns into rows.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Generic, Self, TypeVar
+
+import polars as pl
+from adbc_driver_manager.dbapi import Connection as AdbcConnection
+from faker import Faker
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+from sqlalchemy import Engine
+
+T = TypeVar("T")
+
+
+class VariationRule(BaseModel, Generic[T], ABC):
+    """Abstract base class for variation rules."""
+
+    model_config = ConfigDict(frozen=True)
+
+    @property
+    @abstractmethod
+    def type(self) -> type[T]:
+        """Python type this rule can be applied to."""
+        ...
+
+    @abstractmethod
+    def apply(self, value: T) -> T:
+        """Apply the variation to a value."""
+        ...
+
+
+class SuffixRule(VariationRule[str]):
+    """Add a suffix to a value."""
+
+    suffix: str
+
+    @property
+    def type(self) -> type[str]:  # noqa: D102
+        return str
+
+    def apply(self, value: str) -> str:  # noqa: D102
+        return f"{value}{self.suffix}"
+
+
+class ReplaceRule(VariationRule[str]):
+    """Replace occurrences of a string with another."""
+
+    old: str
+    new: str
+
+    @property
+    def type(self) -> type[str]:  # noqa: D102
+        return str
+
+    def apply(self, value: str) -> str:  # noqa: D102
+        return value.replace(self.old, self.new)
+
+
+def infer_data_type(base: str, parameters: tuple | None) -> pl.DataType:
+    """Infer the Polars data type a Faker configuration produces.
+
+    Args:
+        base: Faker generator type.
+        parameters: Parameters for the generator.
+
+    Returns:
+        The Polars data type of the generated values.
+    """
+    generator = Faker()
+    value_generator = getattr(generator, base)
+    parameters = {} if not parameters else dict(parameters)
+    examples = [value_generator(**dict(parameters)) for _ in range(5)]
+    series = pl.Series(examples)
+    return series.dtype
+
+
+class FeatureConfig(BaseModel):
+    """Configuration for generating a feature with variations."""
+
+    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
+
+    name: str
+    base_generator: str
+    parameters: tuple | None = Field(
+        default=None,
+        description=(
+            "Parameters for the generator. A tuple of tuples passed to the generator."
+        ),
+    )
+    unique: bool = Field(
+        default=True,
+        description=(
+            "Whether the generator enforces uniqueness in the generated data. "
+            "For example, using unique=True with the 'boolean' generator will error "
+            "if more than two values are generated."
+        ),
+    )
+    drop_base: bool = Field(
+        default=False, description="Whether the base case is dropped."
+    )
+    variations: tuple[VariationRule, ...] = Field(default_factory=tuple)
+    datatype: pl.DataType = Field(
+        default_factory=lambda data: infer_data_type(
+            data["base_generator"], data["parameters"]
+        )
+    )
+
+    def add_variations(self, *rule: VariationRule) -> "FeatureConfig":
+        """Add a variation rule to the feature."""
+        return FeatureConfig(
+            name=self.name,
+            base_generator=self.base_generator,
+            parameters=self.parameters,
+            unique=self.unique,
+            drop_base=self.drop_base,
+            variations=self.variations + tuple(rule),
+        )
+
+    @field_validator("name", mode="after")
+    @classmethod
+    def protected_names(cls: type[Self], value: str) -> str:
+        """Ensure name is not a reserved keyword."""
+        if value in {"id", "key"}:
+            raise ValueError("Feature name cannot be 'id' or 'key'.")
+        return value
+
+
+class SourceParameters(BaseModel):
+    """Configuration for generating a source."""
+
+    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True, extra="forbid")
+
+    features: tuple[FeatureConfig, ...] = Field(default_factory=tuple)
+    name: str
+    # Defaults to None rather than an engine, so that every caller ends up with its own
+    # warehouse. A `Field(default=create_engine(...))` is evaluated once at class
+    # definition, which silently shared one in-memory SQLite database across the whole
+    # process, so two sources built without an explicit engine could not be joined.
+    engine: Engine | AdbcConnection | None = Field(default=None)
+    n_true_entities: int | None = Field(default=None)
+    repetition: int = Field(default=0)