matchlab 0.1.0.dev1__py3-none-any.whl

matchlab/eval/samples.py

@@ -0,0 +1,415 @@
+"""Client-side helpers for retrieving and preparing evaluation samples.
+
+Everything here takes a **Resolver**: one you are holding, the label one was
+published under, or a sequence of either. The sequence form is what makes two
+methodologies comparable. Sampling across several resolvers unions their components,
+so one round of judging covers all of them and the scores answer the same question.
+"""
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, TypeAlias
+
+import polars as pl
+from pydantic import BaseModel
+
+from matchlab.core.dataframes import qualify
+from matchlab.core.dsu import DisjointSet
+from matchlab.core.exceptions import SourceTableError
+from matchlab.core.resolver_output import root_id
+from matchlab.eval.judgements import Judgement
+from matchlab.eval.metrics import PrecisionRecall, precision_recall
+
+if TYPE_CHECKING:
+    from matchlab.adapters import Adapter, Fingerprint
+    from matchlab.resolvers import Resolver
+else:
+    Adapter = Any
+    Fingerprint = Any
+    Resolver = Any
+
+ResolverRef: TypeAlias = "Resolver | str"
+"""One resolver to read: a live `Resolver`, or the label one was published under."""
+
+Reading: TypeAlias = "tuple[Adapter, Fingerprint]"
+"""A located resolver: the store holding it, and its fingerprint."""
+
+
+class EvaluationFieldMetadata(BaseModel):
+    """Metadata describing one field shown to a reviewer during evaluation."""
+
+    display_name: str
+    source_columns: list[str]
+
+
+class EvaluationItem(BaseModel):
+    """One cluster shown to a reviewer, with its leaves' source data alongside it.
+
+    `records` holds the leaf IDs and their source-qualified data columns. For example:
+
+    | leaf | src_a_first | src_a_last | src_b_first | src_b_last |
+    |------|-------------|------------|-------------|------------|
+    | 1    | Thomas      | Bayes      |             |            |
+    | 2    | Tommy       | B          |             |            |
+    | 12   |             |            | Tom         | Bayes      |
+
+    `fields` maps each display name to the source-qualified columns holding it, so a
+    reviewer sees one column per field rather than one per source. For example:
+
+    ```python
+    EvaluationFieldMetadata(
+        display_name="first", source_columns=["src_a_first", "src_b_first"]
+    )
+    ```
+    """
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    leaves: list[int]
+    records: pl.DataFrame
+    fields: list[EvaluationFieldMetadata]
+
+    def get_unique_record_groups(self) -> list[list[int]]:
+        """Group identical records by leaf ID.
+
+        Returns:
+            List of groups, where each group is a list of leaf IDs
+            that have identical values across all data fields.
+            Example: [[1, 3], [2], [4, 5, 6]] means records 1 & 3 are identical.
+        """
+        data_cols = [col for field in self.fields for col in field.source_columns]
+        grouped = self.records.group_by(data_cols, maintain_order=True).agg(
+            pl.col("leaf")
+        )
+        return [group for group in grouped["leaf"]]
+
+
+def create_judgement(
+    item: EvaluationItem,
+    assignments: dict[int, str],
+    tag: str | None = None,
+) -> Judgement:
+    """Build the Judgement for one reviewed item from its column assignments.
+
+    Args:
+        item: The reviewed cluster.
+        assignments: Group letter chosen for each unique record group, keyed by that
+            group's index in `item.get_unique_record_groups()`.
+        tag: Tag to record on the judgement.
+
+    Returns:
+        A Judgement with one endorsed group per distinct letter used in
+        `assignments`. A leaf whose group was never assigned is missing from
+        `endorsed`, which `Judgement`'s own validation rejects, so an incomplete
+        `assignments` raises rather than producing a partial judgement.
+    """
+    groups: dict[str, list[int]] = {}
+    unique_record_groups = item.get_unique_record_groups()
+
+    for col_idx, group in assignments.items():
+        leaf_ids = unique_record_groups[col_idx]
+        groups.setdefault(group, []).extend(leaf_ids)
+
+    endorsed = [sorted(set(leaf_ids)) for leaf_ids in groups.values()]
+    return Judgement(shown=item.leaves, endorsed=endorsed, tag=tag)
+
+
+def create_evaluation_item(
+    df: pl.DataFrame,
+    source_fields: list[tuple[str, list[str]]],
+    leaves: list[int],
+) -> EvaluationItem:
+    """Build an EvaluationItem, grouping each cluster's columns by field across sources.
+
+    Args:
+        df: The cluster's rows, with source-qualified data columns.
+        source_fields: `(prefix, qualified columns)` per source. The columns come
+            from the fetched data rather than from the sources, which would have to
+            re-read the warehouse just to list their names.
+        leaves: The leaf IDs in this cluster.
+    """
+    data_cols = [c for c in df.columns if c not in ["root", "leaf", "key"]]
+
+    # The same field in two sources shares a display name, which is what lines them
+    # up for review.
+    field_to_columns: dict[str, list[str]] = {}
+
+    for prefix, columns in source_fields:
+        for column in columns:
+            if column in data_cols:
+                field = column.removeprefix(prefix)
+                field_to_columns.setdefault(field, []).append(column)
+
+    fields: list[EvaluationFieldMetadata] = []
+    for field_name, source_columns in field_to_columns.items():
+        fields.append(
+            EvaluationFieldMetadata(
+                display_name=field_name, source_columns=source_columns
+            )
+        )
+
+    records = df.select(["leaf"] + data_cols)
+
+    return EvaluationItem(leaves=leaves, records=records, fields=fields)
+
+
+def _many(resolver: "ResolverRef | Sequence[ResolverRef]") -> bool:
+    """Whether several resolvers were asked for.
+
+    A `str` is itself a `Sequence`. A label names one resolver, not a pile of
+    one-character ones, so it is excluded before the sequence check, not after.
+    """
+    return not isinstance(resolver, str) and isinstance(resolver, Sequence)
+
+
+def _locate(
+    resolver: "ResolverRef", adapter: "Adapter | None"
+) -> tuple["Adapter", bytes]:
+    """Turn a resolver or a label into the store and fingerprint to read.
+
+    The two forms differ only in how the fingerprint is found. A live resolver
+    already knows its own fingerprint, while a label needs a lookup in a store.
+    Everything downstream wants the same pair, so this is the only place the
+    distinction exists.
+
+    Raises:
+        SourceTableError: If nothing is published under that label.
+    """
+    from matchlab.steps import default_adapter  # noqa: PLC0415 - avoids a cycle
+
+    if isinstance(resolver, str):
+        adapter = adapter or default_adapter()
+        fingerprint = adapter.find(resolver)
+        if fingerprint is None:
+            known = ", ".join(adapter.labels()) or "none"
+            raise SourceTableError(
+                f"No resolver is published under the label '{resolver}'. "
+                f"Known labels: {known}."
+            )
+        return adapter, fingerprint
+
+    if not resolver.is_collected:
+        resolver.collect(adapter)
+    collected_in, fingerprint = resolver._collected()
+    return adapter or collected_in, fingerprint
+
+
+def _readings(
+    resolver: "ResolverRef | Sequence[ResolverRef]", adapter: "Adapter | None"
+) -> list[Reading]:
+    """Locate every resolver asked for, in the order given."""
+    if not _many(resolver):
+        return [_locate(resolver, adapter)]
+    if not resolver:
+        raise ValueError("At least one resolver must be given.")
+    return [_locate(one, adapter) for one in resolver]
+
+
+def _sources_of(readings: list[Reading]) -> dict[str, Reading]:
+    """Map each source name to the store and fingerprint its rows come from.
+
+    Raises:
+        SourceTableError: If two resolvers cover the same source name with different
+            artifacts. Names repeat across generations of a source, so agreeing on the
+            name is not agreeing on the data, and comparing methodologies over different
+            data is not a comparison.
+    """
+    located: dict[str, Reading] = {}
+    for store, fp in readings:
+        for name, source_fp in store.resolver_output_sources(fp).items():
+            seen = located.setdefault(name, (store, source_fp))
+            if seen[1] != source_fp:
+                raise SourceTableError(
+                    f"These resolvers disagree about source '{name}': one covers "
+                    f"{seen[1].hex()[:8]}, another {source_fp.hex()[:8]}. They are "
+                    "built over different data, so their clusters cannot be compared. "
+                    "Re-collect them over the same sources."
+                )
+    return located
+
+
+def _merged_resolver_output(readings: list[Reading]) -> pl.DataFrame:
+    """Union several resolvers' output into one `(root, leaf, key, source)` table.
+
+    Two records land in the same merged component if either resolver put them
+    together. This is the right sample for a bake-off. Every cluster where the
+    methodologies could disagree ends up on screen, so one judgement settles it for
+    both, and neither resolver gets to pick the clusters it is scored on.
+
+    Merged roots are minted with `root_id`, the same content-addressed function a
+    resolver uses for its own roots, so two people running the same comparison arrive
+    at identical root IDs for identical components. Nothing persists these merged roots.
+    `store_judgement` re-mints its own root from the leaves it is given, so a merged
+    root only ever lives as far as the reviewer.
+    """
+    frames = [store.read_resolver(fp) for store, fp in readings]
+
+    components = DisjointSet[int]()
+    for frame in frames:
+        for leaves in frame.group_by("root").agg("leaf")["leaf"].to_list():
+            components.add(leaves[0])
+            for leaf in leaves[1:]:
+                components.union(leaves[0], leaf)
+
+    # `root_id` is invariant to leaf order but the caller does the sorting, and it is
+    # vectorised because there are as many clusters here as there are entities.
+    merged = (
+        pl.DataFrame(
+            {"leaf": [sorted(component) for component in components.get_components()]},
+            schema={"leaf": pl.List(pl.UInt64)},
+        )
+        .with_columns(root_id(pl.col("leaf")).alias("root"))
+        # A component always holds at least one leaf, so the empty-list case this
+        # settles cannot arise. Pinning it keeps the polars 2.0 default change quiet.
+        .explode("leaf", empty_as_null=False)
+    )
+
+    records = pl.concat(
+        [frame.select("leaf", "key", "source") for frame in frames]
+    ).unique()
+    return merged.join(records, on="leaf").select("root", "leaf", "key", "source")
+
+
+def _sample_clusters(
+    resolver_output: pl.DataFrame, n: int, seed: int | None
+) -> pl.DataFrame:
+    """Take up to `n` whole clusters from merge-forwarded Resolver output in memory.
+
+    The in-memory twin of `Adapter.sample`, used for merged Resolver output from
+    several readings. No store holds that merged output, since it exists only for
+    this comparison.
+    """
+    roots = resolver_output["root"].unique()
+    if n < roots.len():
+        roots = roots.sample(n=n, seed=seed, shuffle=True)
+    return resolver_output.filter(pl.col("root").is_in(roots.to_list()))
+
+
+def get_samples(
+    n: int,
+    resolver: "ResolverRef | Sequence[ResolverRef]",
+    adapter: "Adapter | None" = None,
+    seed: int | None = None,
+) -> dict[int, EvaluationItem]:
+    """Retrieve sampled clusters enriched with source data, as EvaluationItems.
+
+    Record values come from the extract stored when each source was collected, not
+    from a fresh warehouse read. This works offline, and shows the data the matching
+    actually saw.
+
+    Args:
+        n: Number of clusters to sample.
+        resolver: The resolver to sample from, collected first if it isn't
+            already, or the label one was published under. The label form needs no
+            plan, because a stored resolver's output records which source artifacts
+            it covers. Pass several and the sample is drawn from their merged
+            components, so one round of judging scores all of them against the same
+            clusters.
+        adapter: Where to read from. Defaults to the resolver's, else the module
+            default.
+        seed: Fixes which clusters come back. The same store, `n` and seed give the
+            same sample, which is how two people review the same clusters.
+
+    Returns:
+        Dictionary of cluster ID to EvaluationItems describing the cluster.
+
+    Raises:
+        SourceTableError: If nothing is published under `resolver`, if a source the
+            resolver's output covers isn't in the store, or if several resolvers
+            disagree about a source.
+        ValueError: If `resolver` is an empty sequence.
+    """
+    readings = _readings(resolver, adapter)
+
+    if len(readings) == 1:
+        store, resolver_fp = readings[0]
+        samples = store.sample(resolver_fp, n, seed)
+    else:
+        samples = _sample_clusters(_merged_resolver_output(readings), n, seed)
+
+    if not len(samples):
+        return {}
+
+    sources = _sources_of(readings)
+    results_by_source: list[pl.DataFrame] = []
+    source_fields: list[tuple[str, list[str]]] = []
+
+    for source_step in samples["source"].unique():
+        located = sources.get(source_step)
+        if located is None:
+            raise SourceTableError(
+                f"This resolver references source '{source_step}', which is not "
+                "in the store. Re-collect the plan to repopulate it."
+            )
+        source_store, source_fp = located
+
+        samples_by_source = samples.filter(pl.col("source") == source_step)
+        rows, qualified_key = source_store.read_source_records(
+            source_fp, source_step, samples_by_source["key"]
+        )
+        values = [column for column in rows.columns if column != qualified_key]
+
+        samples_and_source = samples_by_source.join(
+            rows, left_on="key", right_on=qualified_key
+        )
+        source_fields.append((qualify(source_step), values))
+        results_by_source.append(samples_and_source[["root", "leaf", "key"] + values])
+
+    if not results_by_source:
+        return {}
+
+    all_results: pl.DataFrame = pl.concat(results_by_source, how="diagonal")
+
+    results_by_root: dict[int, EvaluationItem] = {}
+    for root in all_results["root"].unique():
+        cluster_df = all_results.filter(pl.col("root") == root).drop("root")
+        leaves = cluster_df.select("leaf").to_series().unique().to_list()
+        evaluation_item = create_evaluation_item(cluster_df, source_fields, leaves)
+        results_by_root[root] = evaluation_item
+
+    return results_by_root
+
+
+class EvalData:
+    """Caches a store's judgements, and scores resolvers against them."""
+
+    def __init__(self, adapter: Adapter, tag: str | None = None) -> None:
+        """Load judgement and expansion data used to compute evaluation metrics.
+
+        Args:
+            adapter: The storage adapter holding judgements, for example the one a
+                resolver was collected into.
+            tag: Optional tag to filter judgements by.
+        """
+        self.adapter = adapter
+        self.tag = tag
+        self.judgements, self.expansion = adapter.read_eval_data(tag)
+
+    def precision_recall(
+        self, resolver: "ResolverRef | Sequence[ResolverRef]"
+    ) -> PrecisionRecall | list[PrecisionRecall]:
+        """Score one or more resolvers against these judgements.
+
+        Only pairs present in every resolver's output *and* in the judgements are
+        compared, so scoring several at once is the fair way to rank them. Each is
+        measured over the same records, and none is flattered by clusters the others
+        never saw. Scoring them one at a time gives each its own comparison set, and
+        those numbers do not line up.
+
+        Args:
+            resolver: A resolver, the label one was published under, or a sequence of
+                either.
+
+        Returns:
+            One `(precision, recall)` pair, or a list of them in the order given if a
+            sequence was passed.
+        """
+        readings = _readings(resolver, self.adapter)
+        scores = precision_recall(
+            [
+                store.read_resolver(fp).select("root", "leaf").unique()
+                for store, fp in readings
+            ],
+            self.judgements,
+            self.expansion,
+        )
+        return scores if _many(resolver) else scores[0]

matchlab/eval/tui/__init__.py

@@ -0,0 +1,66 @@
+"""An interactive reviewer for resolved clusters.
+
+`review()` opens a terminal app showing one cluster at a time: the records that were
+grouped together, laid out so you can see what matched. You paint them into groups and
+each decision is stored as a `Judgement`, which `EvalData.precision_recall` then scores
+a resolver against.
+
+`app` is imported inside `review()` rather than here, so that importing `matchlab.eval`
+for its scoring helpers doesn't pay for Textual's import.
+"""
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from matchlab.adapters import Adapter
+    from matchlab.eval.samples import ResolverRef
+
+
+def review(
+    resolver: "ResolverRef | Sequence[ResolverRef]",
+    n: int = 5,
+    adapter: "Adapter | None" = None,
+    tag: str | None = None,
+    seed: int | None = None,
+    show_help: bool = True,
+) -> None:
+    """Review resolved clusters interactively, recording judgements.
+
+    Takes either a live `Resolver`, collected first if it isn't already, or the
+    *name* of a resolver already published in a store. The second form needs no plan
+    and no warehouse. Record values come from the extract cached when each source was
+    collected, which is the data the matching actually saw.
+
+    ```python
+    review(companies, tag="session-1")  # from a plan
+    review("companies", adapter=DuckDBAdapter("run.duckdb"))  # from a store alone
+    review([naive, splink], tag="bakeoff")  # judge both at once
+    ```
+
+    Args:
+        resolver: The resolver whose clusters to review, or the name one was
+            published under. Pass several and you review their merged components, so
+            one session's judgements score every one of them.
+        n: How many clusters to hold in the queue at once.
+        adapter: Where judgements are stored, and where samples are drawn from.
+            Defaults to the module-level adapter.
+        tag: Tags every judgement made in this session, so a later
+            `EvalData(adapter, tag=...)` can score against just these.
+        seed: Fixes which clusters this session draws, so someone else with the same
+            store can review the same ones.
+        show_help: Show the key bindings on start.
+    """
+    from matchlab.eval.tui.app import EntityResolutionApp  # noqa: PLC0415 - lazy
+
+    EntityResolutionApp(
+        resolver=resolver,
+        num_samples=n,
+        adapter=adapter,
+        session_tag=tag,
+        seed=seed,
+        show_help=show_help,
+    ).run()
+
+
+__all__ = ["review"]