matchlab 0.1.0.dev1__py3-none-any.whl

matchlab/testkit/matchers.py

@@ -0,0 +1,233 @@
+"""Matchers that already know the answer.
+
+A perfect matcher makes no mistakes, so pointing one at generated data leaves the *plan*
+as the only variable. That is what makes an end-to-end assertion mean something. Build
+them through `LinkedSources.dedupe()` / `.link()`. They are public here only so the
+registry can name them.
+
+The subtlety is how a matcher identifies a record. Pre-generating edges between record
+IDs does not work, because matchlab derives cluster IDs by content-hashing rows at
+collect time. The IDs a model actually receives are unknowable when a fixture is built,
+which is why an `AnswerKey` is keyed by the row's *values* instead. At match time the
+matcher reads those columns out of the record step it was handed, looks up each row's
+true entity, and emits edges between whatever IDs are actually present. That makes it
+independent of how identity is assigned, and lets it survive any future change to ID
+minting.
+"""
+
+from collections.abc import Iterable
+from itertools import combinations
+
+import polars as pl
+from pydantic import BaseModel, ConfigDict
+
+from matchlab.core.hash import HASH_FUNC
+from matchlab.core.schemas import SCHEMA_MODEL_EDGES
+from matchlab.models import add_model_class
+from matchlab.models.dedupers.base import Deduper, DeduperSettings
+from matchlab.models.linkers.base import Linker, LinkerSettings
+from matchlab.testkit.entities import TrueEntity
+from matchlab.testkit.sources import GeneratedSource
+
+# Registered answer keys, by content hash. Only ever added to. A testkit process is
+# short-lived and the keys are small.
+_ANSWERS: dict[str, "AnswerKey"] = {}
+
+
+class AnswerKey(BaseModel):
+    """The sheet a perfect matcher looks answers up in: row values → true entity.
+
+    Three things happen to one of these, in order:
+
+    * **build**: `AnswerKey.from_sources()` derives it from generated data. Pass a
+      subset of the true entities to make a matcher that knows only part of the answer.
+    * **store**: `.register()` puts it in a process-local registry and returns a
+      content-addressed ID. A matcher's settings carry that ID rather than the table
+      itself, because settings are JSON-serialised into a step's fingerprint and a
+      lookup table is not. Hashing the content keeps the fingerprint honest. A different
+      answer produces a different ID, and so a different model artifact.
+    * **use**: `.dedupe_edges()` / `.link_edges()` are what `PerfectDeduper` and
+      `PerfectLinker` call at match time, against whatever record step they were
+      handed.
+
+    `groups` maps a tuple of column values to a true-entity ID. `columns` names the
+    columns to read, in the same order they appear in the record step the model is
+    given (source-qualified). Linkers carry a second set for the right-hand record
+    step. Both sides map into the same entity-ID space, which is what lets them be
+    joined.
+    """
+
+    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
+
+    left_columns: tuple[str, ...]
+    left_groups: dict[tuple, int]
+    right_columns: tuple[str, ...] | None = None
+    right_groups: dict[tuple, int] | None = None
+    score: float = 1.0
+
+    @classmethod
+    def from_sources(
+        cls,
+        left: GeneratedSource,
+        true_entities: Iterable[TrueEntity],
+        right: GeneratedSource | None = None,
+        score: float = 1.0,
+    ) -> "AnswerKey":
+        """Derive the lookup from generated sources.
+
+        Each generated row is mapped to the true entity that owns its key, then keyed by
+        its feature values under the names the model will see them by, source-qualified,
+        because that is how they arrive in the record step a matcher is handed.
+
+        Args:
+            left: The generated source the matcher reads as its left input.
+            right: Its right input, for a linker. `None` for a deduper.
+            true_entities: The planted entities to answer for. Pass a subset to build a
+                matcher that knows only part of the answer.
+            score: The score to emit on every edge.
+        """
+        entities = tuple(true_entities)
+
+        def side(source: GeneratedSource) -> tuple[tuple[str, ...], dict[tuple, int]]:
+            key_to_entity: dict[str, int] = {
+                str(key): entity.id
+                for entity in entities
+                for key in entity.get_keys(source.source.name)
+            }
+            # From the testkit, not the plan node. The node would have to read the
+            # warehouse to list its columns, and this runs before any collect.
+            # `field_names` also covers a source set by hand, where `features` is None.
+            features = source.field_names
+            groups: dict[tuple, int] = {}
+            for row in pl.DataFrame(source.data).iter_rows(named=True):
+                entity_id = key_to_entity.get(str(row["key"]))
+                if entity_id is not None:
+                    groups[tuple(row[feature] for feature in features)] = entity_id
+            columns = tuple(f"{source.source.name}_{feature}" for feature in features)
+            return columns, groups
+
+        left_columns, left_groups = side(left)
+        right_columns, right_groups = side(right) if right is not None else (None, None)
+        return cls(
+            left_columns=left_columns,
+            left_groups=left_groups,
+            right_columns=right_columns,
+            right_groups=right_groups,
+            score=score,
+        )
+
+    def _assign(
+        self, data: pl.DataFrame, columns: tuple[str, ...], groups: dict[tuple, int]
+    ) -> dict[int, list[int]]:
+        """Bucket the row IDs by the true entity their values belong to."""
+        missing = [column for column in columns if column not in data.columns]
+        if missing:
+            raise KeyError(
+                f"Ground-truth columns absent from the model input: {missing}. "
+                f"Available: {sorted(data.columns)}"
+            )
+
+        buckets: dict[int, list[int]] = {}
+        for row in data.select("id", *columns).iter_rows(named=True):
+            entity = groups.get(tuple(row[column] for column in columns))
+            if entity is not None:
+                buckets.setdefault(entity, []).append(row["id"])
+        return buckets
+
+    def _edges(self, pairs: list[tuple[int, int]]) -> pl.DataFrame:
+        if not pairs:
+            return pl.DataFrame(schema=pl.Schema(SCHEMA_MODEL_EDGES))
+        left, right = zip(*pairs, strict=True)
+        return pl.DataFrame(
+            {
+                "left_id": list(left),
+                "right_id": list(right),
+                "score": [self.score] * len(pairs),
+            },
+            schema=pl.Schema(SCHEMA_MODEL_EDGES),
+        )
+
+    def dedupe_edges(self, data: pl.DataFrame) -> pl.DataFrame:
+        """All within-entity pairs among the record step's records."""
+        pairs: list[tuple[int, int]] = []
+        for ids in self._assign(data, self.left_columns, self.left_groups).values():
+            unique = sorted(set(ids))
+            pairs.extend(combinations(unique, 2))
+        return self._edges(pairs)
+
+    def link_edges(self, left: pl.DataFrame, right: pl.DataFrame) -> pl.DataFrame:
+        """All cross-side pairs whose records share a true entity."""
+        left_buckets = self._assign(left, self.left_columns, self.left_groups)
+        right_buckets = self._assign(
+            right, self.right_columns or (), self.right_groups or {}
+        )
+        pairs = [
+            (left_id, right_id)
+            for entity, left_ids in left_buckets.items()
+            for left_id in sorted(set(left_ids))
+            for right_id in sorted(set(right_buckets.get(entity, [])))
+        ]
+        return self._edges(pairs)
+
+    def register(self) -> str:
+        """Store this key, returning a content-addressed ID for it.
+
+        Settings are JSON-serialised into a step's fingerprint and a lookup table is
+        not, so the key itself cannot live there. Hashing its content keeps the
+        fingerprint honest. A different answer key produces a different ID, and
+        therefore a different model artifact.
+        """
+        payload = repr(
+            (
+                self.left_columns,
+                sorted(map(repr, self.left_groups.items())),
+                self.right_columns,
+                sorted(map(repr, (self.right_groups or {}).items())),
+                self.score,
+            )
+        ).encode()
+        key_id = HASH_FUNC(payload).hexdigest()
+        _ANSWERS[key_id] = self
+        return key_id
+
+
+class PerfectDeduperSettings(DeduperSettings):
+    """Settings for a deduper that matches from registered ground truth."""
+
+    truth_id: str
+
+
+class PerfectDeduper(Deduper):
+    """A perfect deduper. Emits every within-entity pair it is given."""
+
+    settings: PerfectDeduperSettings
+
+    def prepare(self, data: pl.DataFrame) -> None:
+        """No preparation required."""
+
+    def dedupe(self, data: pl.DataFrame) -> pl.DataFrame:
+        """Emit edges between records sharing a true entity."""
+        return _ANSWERS[self.settings.truth_id].dedupe_edges(data)
+
+
+class PerfectLinkerSettings(LinkerSettings):
+    """Settings for a linker that matches from registered ground truth."""
+
+    truth_id: str
+
+
+class PerfectLinker(Linker):
+    """A perfect linker. Emits every cross-side pair sharing a true entity."""
+
+    settings: PerfectLinkerSettings
+
+    def prepare(self, left: pl.DataFrame, right: pl.DataFrame) -> None:
+        """No preparation required."""
+
+    def link(self, left: pl.DataFrame, right: pl.DataFrame) -> pl.DataFrame:
+        """Emit edges between records sharing a true entity."""
+        return _ANSWERS[self.settings.truth_id].link_edges(left, right)
+
+
+add_model_class(PerfectDeduper)
+add_model_class(PerfectLinker)

matchlab/testkit/models.py

@@ -0,0 +1,165 @@
+"""The generated model, and the expectation it is asserted against.
+
+`GeneratedModel.scores` and `.predicted_clusters` express the expected result in the
+testkit's *own* entity ID space, which is what `LinkedSources.diff_model_edges` compares
+a methodology against without running a plan. The model itself matches on row values
+(see `matchers`), so a collected plan reproduces the same answer over the
+content-derived IDs the record step actually carries.
+
+Build one with `LinkedSources.dedupe()` or `.link()`. They know the answer, so nothing
+has to be threaded through by hand.
+"""
+
+import numpy as np
+import polars as pl
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from matchlab.core.schemas import SCHEMA_MODEL_EDGES
+from matchlab.models.models import Model
+from matchlab.testkit.compare import scores_to_clusters
+from matchlab.testkit.entities import Cluster, TrueEntity
+from matchlab.testkit.sources import GeneratedSource
+
+
+def generate_entity_scores(
+    left_entities: frozenset[Cluster],
+    right_entities: frozenset[Cluster] | None,
+    true_entities: frozenset[TrueEntity],
+    score_range: tuple[float, float] = (0.8, 1.0),
+    seed: int = 42,
+) -> pl.DataFrame:
+    """Generate scores that recover the planted entity relationships.
+
+    Maps each `Cluster` to the `TrueEntity` it is a subset of, then emits a score for
+    every same-entity pair. The result is a fully connected, perfectly correct score
+    graph. This function does not generate partial, wrong, or missing matches.
+
+    Args:
+        left_entities: Cluster objects from the left input.
+        right_entities: Cluster objects from the right input. `None` to deduplicate
+            `left_entities` instead.
+        true_entities: The planted true entities to score against.
+        score_range: Assign each matching pair a random score in this range.
+        seed: Random seed for reproducibility.
+
+    Returns:
+        A Polars DataFrame with `left_id`, `right_id`, and `score` columns.
+
+    Raises:
+        ValueError: If `score_range` is not increasing within [0, 1], or a `Cluster` is
+            a subset of more than one true entity.
+    """
+    # Validate inputs
+    if not (0 <= score_range[0] <= score_range[1] <= 1):
+        raise ValueError("Scores must be increasing values between 0 and 1")
+
+    # Handle deduplication case
+    if right_entities is None:
+        right_entities = left_entities
+
+    # Create mapping of Cluster -> TrueEntity
+    entity_mapping: dict[Cluster, TrueEntity] = {}
+
+    def _map_entity(entity: Cluster) -> None:
+        matching_sources = [
+            source
+            for source in true_entities
+            if entity.is_subset_of_true_entity(source)
+        ]
+        if len(matching_sources) > 1:
+            raise ValueError(
+                f"Cluster with ID {entity.id} is a subset of multiple "
+                f"true entities. This violates the uniqueness constraint."
+            )
+        if matching_sources:
+            entity_mapping[entity] = matching_sources[0]
+
+    # Map all entities, checking constraints
+    for entity in left_entities:
+        _map_entity(entity)
+    if right_entities is not left_entities:
+        for entity in right_entities:
+            _map_entity(entity)
+
+    # Group by TrueEntity
+    source_groups: dict[TrueEntity, tuple[set[Cluster], set[Cluster]]] = {}
+    for entity, source in entity_mapping.items():
+        if source not in source_groups:
+            source_groups[source] = (set(), set())
+        # Add to left or right group based on which input set it came from
+        if entity in left_entities:
+            source_groups[source][0].add(entity)
+        if entity in right_entities:  # Note: could be in both for deduplication
+            source_groups[source][1].add(entity)
+
+    # Generate score edges for each group
+    edges = []
+    rng = np.random.default_rng(seed=seed)
+
+    for left_group, right_group in source_groups.values():
+        # Skip empty groups
+        if not left_group or not right_group:
+            continue
+
+        # Generate all pairs within this group
+        for left_entity in left_group:
+            for right_entity in right_group:
+                # Deduplication counts each pair once, so keep left_id < right_id.
+                if (
+                    right_entities == left_entities
+                    and left_entity.id >= right_entity.id
+                ):
+                    continue
+
+                # Generate random score in range
+                score = np.float32(rng.uniform(score_range[0], score_range[1]))
+                edges.append((left_entity.id, right_entity.id, score))
+
+    # If no edges were generated, return empty table with correct schema
+    if not edges:
+        return pl.DataFrame(schema=pl.Schema(SCHEMA_MODEL_EDGES))
+
+    return pl.DataFrame(edges, orient="row", schema=pl.Schema(SCHEMA_MODEL_EDGES))
+
+
+class GeneratedModel(BaseModel):
+    """A generated model: the plan node, its inputs, and the answer expected of it.
+
+    As with `GeneratedSource`, this exposes what it knows about the fixture and nothing
+    else. To run the plan, go through `.model`, e.g. `model.model.resolve()`.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)
+
+    # The plan node. Run it with `model.resolve()`.
+    model: Model
+    # The generated source it reads as its left input. Its rows and starting partition
+    # are reachable from here rather than copied onto this object.
+    left_source: GeneratedSource
+    # Its right input, for a linker. `None` for a deduper.
+    right_source: GeneratedSource | None = None
+    # The edges a perfect matcher should emit, in the sources' own entity ID space.
+    scores: pl.DataFrame = Field(frozen=True)
+
+    _predicted_clusters: tuple[Cluster, ...]
+
+    @property
+    def predicted_clusters(self) -> tuple[Cluster, ...]:
+        """What the model's scores imply, to compare against the planted answer."""
+        return self._predicted_clusters
+
+    @model_validator(mode="after")
+    def _derive_predicted_clusters(self) -> "GeneratedModel":
+        """Resolve the expected scores into the clusters they imply."""
+        self._predicted_clusters = scores_to_clusters(
+            scores=self.scores,
+            left_clusters=self.left_source.input_clusters,
+            right_clusters=(
+                self.right_source.input_clusters
+                if self.right_source is not None
+                else None
+            ),
+            threshold=0,
+        )
+
+        return self