matchlab 0.1.0.dev1__py3-none-any.whl

matchlab/document.py

@@ -0,0 +1,329 @@
+"""Serialise a plan to a portable document, and rebuild it somewhere else.
+
+A `PlanDocument` is a **derived view** of a plan, not a source format. Humans write
+Python, and a document is what you hand to another environment so it can run the same
+plan. It is dumped, transferred and loaded, never hand-authored. That is why nodes
+refer to each other by position rather than by any human-facing name.
+
+Nodes come out in `lineage.walk` order, so every input index is smaller than the index
+of the step that consumes it. Positions also preserve structural sharing. A record step
+feeding two models is one node referenced twice. Nesting each step's inputs inside it
+would have inlined the whole subtree twice over instead.
+
+**Edges live here, settings live in the spec.** That split is the point of this
+module. A spec is hashed into a fingerprint, so it must carry everything that changes
+the step's output and nothing else. A serialised plan must also carry the edges, which
+a fingerprint already covers by folding in its parents' fingerprints. Putting both in
+one model would have forced a choice: a spec that lied about identity, where a rename
+invalidates a subtree that produces identical bytes, or a document that could not be
+rebuilt.
+
+Edges are not the only thing that lands on this side of the split. A source's
+`LocationRef` says which location class to build and under what name, so that `load`
+can attach a client. Neither fact changes a byte the source produces, which is why a
+`SourceSpec` records nothing about where its rows came from.
+
+**What a document cannot carry:**
+
+* *Clients.* A document names the locations its sources read and stops there. `load`
+  takes the clients and attaches them by name, so credentials never enter a document.
+* *Code.* `model_class` and `resolver_class` are registry names, so the target
+  environment must have the same classes registered (`add_model_class`). A document is
+  portable across environments, not across codebases.
+* *Labels.* Publishing is something you do to a *result*, using `Resolver.publish`. It
+  is not part of the plan, so the receiving environment collects and then publishes
+  under whatever label it wants. The only names in a document are sources' names,
+  which are part of their output.
+* *Data, or any hash of it.* A source's fingerprint folds in a content hash of what it
+  actually read, and that hash is derived on load from the target warehouse. Same rows
+  give the same fingerprints, so the target store hits cache instead of recomputing.
+  Different rows give different fingerprints, so it re-runs. Both are the intended
+  behaviour.
+"""
+
+from collections.abc import Mapping
+from typing import Any, TypeVar, cast
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from matchlab.core.kinds import StepKind
+from matchlab.lineage import walk
+from matchlab.locations import build_location
+from matchlab.models import Model
+from matchlab.recordstep import RecordStep
+from matchlab.resolvers import Resolver
+from matchlab.sources import Source
+from matchlab.specs import (
+    ModelSpec,
+    ResolverSpec,
+    SourceSpec,
+    TransformSpec,
+)
+from matchlab.steps import Step
+from matchlab.transformers import Transform
+
+StepSpec = SourceSpec | TransformSpec | ModelSpec | ResolverSpec
+
+SpecT = TypeVar("SpecT", bound=BaseModel)
+StepT = TypeVar("StepT", bound=Step)
+
+_SPEC_TYPES: dict[StepKind, type[BaseModel]] = {
+    StepKind.SOURCE: SourceSpec,
+    StepKind.TRANSFORM: TransformSpec,
+    StepKind.MODEL: ModelSpec,
+    StepKind.RESOLVER: ResolverSpec,
+}
+
+
+class LocationRef(BaseModel):
+    """Everything needed to rebuild a source's location."""
+
+    model_config = ConfigDict(frozen=True)
+
+    location_class: str = Field(
+        description="The registered name of the Location subclass to build."
+    )
+    name: str = Field(
+        description="The key `load` looks up in `clients` to find this location's."
+    )
+
+
+class StepNode(BaseModel):
+    """One step: what kind it is, what it specifies, and what it reads."""
+
+    model_config = ConfigDict(frozen=True)
+
+    kind: StepKind = Field(description="Which kind of step this is.")
+    spec: StepSpec = Field(
+        description="This step's own settings, exactly what its fingerprint hashes."
+    )
+    inputs: tuple[int, ...] = Field(
+        default=(),
+        description=(
+            "Positions of this step's inputs, in order. Always smaller than this "
+            "step's own position, and order matters: it is the order the fingerprint "
+            "folds parents in, and the order a linker's left and right arrive in."
+        ),
+    )
+    location: LocationRef | None = Field(
+        default=None,
+        description=(
+            "For a source, how to rebuild the location it reads. Here rather than in "
+            "the spec because it describes reconstruction rather than output. See "
+            "`LocationRef`. Every other kind of step leaves it unset."
+        ),
+    )
+
+    @model_validator(mode="after")
+    def _check_location(self) -> "StepNode":
+        """A source names a location. Nothing else has one to name."""
+        if self.kind is StepKind.SOURCE and self.location is None:
+            raise ValueError("A source step must name the location it reads.")
+        if self.kind is not StepKind.SOURCE and self.location is not None:
+            raise ValueError(f"A {self.kind} step must not name a location.")
+        return self
+
+    @model_validator(mode="before")
+    @classmethod
+    def _spec_from_kind(cls, data: Any) -> Any:  # noqa: ANN401 - raw pydantic input
+        """Parse `spec` as the type `kind` names, rather than guessing.
+
+        The specs are structurally close enough that a plain union mis-assigns. A
+        `TransformSpec` and a `ModelSpec` overlap once optional settings default, so a
+        union could bind the wrong one. `kind` is the discriminator the document holds.
+        """
+        if not isinstance(data, dict):
+            return data
+        # `kind` is still whatever the input carried, a plain string, from JSON. It
+        # finds its entry anyway, because a `StrEnum` member hashes and compares as
+        # its value.
+        kind, spec = data.get("kind"), data.get("spec")
+        if isinstance(spec, dict) and kind in _SPEC_TYPES:
+            return {**data, "spec": _SPEC_TYPES[kind].model_validate(spec)}
+        return data
+
+
+class PlanDocument(BaseModel):
+    """A whole plan: its steps in topological order, and the edges between them."""
+
+    model_config = ConfigDict(frozen=True)
+
+    steps: tuple[StepNode, ...] = Field(
+        description=(
+            "Every step reachable from the plan's apex, inputs before consumers."
+        )
+    )
+
+    @model_validator(mode="after")
+    def _check_edges(self) -> "PlanDocument":
+        """Reject a document whose edges do not point backwards at real steps."""
+        for position, node in enumerate(self.steps):
+            for target in node.inputs:
+                if not 0 <= target < position:
+                    raise ValueError(
+                        f"Step {position} ({node.kind}) reads step {target}, which is "
+                        "not an earlier step. Inputs must precede their consumers."
+                    )
+        return self
+
+
+def dump(root: Step) -> PlanDocument:
+    """Describe `root` and everything it reads as a portable document.
+
+    Args:
+        root: The apex of the plan. Only what is reachable upstream of it is included,
+            exactly as `collect()` would run.
+
+    Returns:
+        The document. It holds no client, no credentials and no data.
+    """
+    ordered = walk(root)
+    position = {id(step): index for index, step in enumerate(ordered)}
+    return PlanDocument(
+        steps=tuple(
+            _node(step, tuple(position[id(parent)] for parent in step.upstream))
+            for step in ordered
+        )
+    )
+
+
+def _node(step: Step, inputs: tuple[int, ...]) -> StepNode:
+    """Describe one step, checking its spec is the one its kind implies."""
+    spec_type = _SPEC_TYPES.get(step.kind)
+    if spec_type is None or not isinstance(step.spec, spec_type):
+        raise ValueError(
+            f"A {step.kind} step has no matching spec type. A new kind of step "
+            "needs an entry in `_SPEC_TYPES`."
+        )
+    return StepNode(
+        kind=step.kind,
+        spec=cast(StepSpec, step.spec),
+        inputs=inputs,
+        location=(
+            LocationRef(
+                location_class=type(step.location).__name__, name=step.location.name
+            )
+            if isinstance(step, Source)
+            else None
+        ),
+    )
+
+
+def load(document: PlanDocument, clients: Mapping[str, Any]) -> Step:
+    """Rebuild a plan from a document, returning its apex.
+
+    Nothing is collected. This reconstructs the same lazy plan, so the returned step
+    fingerprints identically to the one that was dumped, given the same data.
+
+    Args:
+        document: A dumped plan.
+        clients: Location name to client, for every location the document's sources
+            name. A `RelationalDBLocation` takes a SQLAlchemy engine or an ADBC
+            connection.
+
+    Returns:
+        The plan's apex, the last step in the document.
+
+    Raises:
+        ValueError: If the document is empty, names a location with no client or an
+            unregistered location class, or wires a step to an input of the wrong kind.
+    """
+    if not document.steps:
+        raise ValueError("A plan document needs at least one step")
+
+    built: list[Step] = []
+    for position, node in enumerate(document.steps):
+        built.append(_rebuild(node, position, [built[i] for i in node.inputs], clients))
+    return built[-1]
+
+
+def _rebuild(
+    node: StepNode, position: int, inputs: list[Step], clients: Mapping[str, Any]
+) -> Step:
+    """Rebuild one step from its node and its already-rebuilt inputs."""
+    match node.kind:
+        case StepKind.SOURCE:
+            spec = _expect(node, position, SourceSpec)
+            reference = cast(LocationRef, node.location)
+            if reference.name not in clients:
+                raise ValueError(
+                    f"Source '{spec.name}' reads location '{reference.name}', which "
+                    f"has no client. Pass one in `clients`: "
+                    f"{{'{reference.name}': engine}}."
+                )
+            return Source(
+                location=build_location(
+                    reference.location_class,
+                    name=reference.name,
+                    client=clients[reference.name],
+                ),
+                name=spec.name,
+                extract_transform=spec.extract_transform,
+                key_field=spec.key_field,
+            )
+
+        case StepKind.TRANSFORM:
+            spec = _expect(node, position, TransformSpec)
+            record_steps = _all_of(node, position, inputs, RecordStep)
+            if len(record_steps) != 1:
+                raise ValueError(
+                    f"Step {position} (transform) reads {len(record_steps)} record "
+                    "steps; a transform reshapes exactly one."
+                )
+            return Transform(
+                record_steps[0],
+                transformer=spec.transformer_class,
+                transformer_settings=spec.transformer_settings,
+            )
+
+        case StepKind.MODEL:
+            spec = _expect(node, position, ModelSpec)
+            record_steps = _all_of(node, position, inputs, RecordStep)
+            if not 1 <= len(record_steps) <= 2:
+                raise ValueError(
+                    f"Step {position} (model) reads {len(record_steps)} record "
+                    "steps; a deduper reads one and a linker two."
+                )
+            return Model(
+                left=record_steps[0],
+                right=record_steps[1] if len(record_steps) > 1 else None,
+                model_class=spec.model_class,
+                model_settings=spec.model_settings,
+            )
+
+        case StepKind.RESOLVER:
+            spec = _expect(node, position, ResolverSpec)
+            return Resolver(
+                *_all_of(node, position, inputs, Model),
+                resolver_class=spec.resolver_class,
+                resolver_settings=spec.resolver_settings,
+            )
+
+        case _:  # every member of StepKind is handled above
+            raise ValueError(
+                f"No rebuild for a {node.kind} step. A new kind of step needs an arm "
+                "here as well as an entry in `_SPEC_TYPES`."
+            )
+
+
+def _expect(node: StepNode, position: int, spec_type: type[SpecT]) -> SpecT:
+    """Narrow a node's spec to the type its kind implies."""
+    if not isinstance(node.spec, spec_type):
+        raise ValueError(  # the validator parses spec by kind
+            f"Step {position} ({node.kind}) has a "
+            f"{type(node.spec).__name__}, expected {spec_type.__name__}."
+        )
+    return node.spec
+
+
+def _all_of(
+    node: StepNode, position: int, inputs: list[Step], step_type: type[StepT]
+) -> tuple[StepT, ...]:
+    """Assert every input is the kind this step can read, and return them in order."""
+    checked = tuple(step for step in inputs if isinstance(step, step_type))
+    if len(checked) != len(inputs):
+        raise ValueError(
+            f"Step {position} ({node.kind}) must read only "
+            f"{step_type.__name__.lower()}s, got {[step.kind for step in inputs]}."
+        )
+    return checked

matchlab/eval/__init__.py

@@ -0,0 +1,11 @@
+"""Public evaluation helpers."""
+
+from matchlab.eval.samples import EvalData, create_judgement, get_samples
+from matchlab.eval.tui import review
+
+__all__ = [
+    "EvalData",
+    "create_judgement",
+    "get_samples",
+    "review",
+]

matchlab/eval/judgements.py

@@ -0,0 +1,41 @@
+"""Representation of human judgements."""
+
+from itertools import chain
+from typing import Self
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class Judgement(BaseModel):
+    """A user's decision on how one reviewed cluster's leaves split into entities."""
+
+    tag: str | None = None
+    shown: list[int] = Field(
+        description="Leaf IDs of the whole cluster shown to the user"
+    )
+    endorsed: list[list[int]] = Field(
+        description="Leaf IDs the user grouped together, one list per endorsed entity"
+    )
+
+    @model_validator(mode="after")
+    def check_no_duplicates(self) -> Self:
+        """Reject leaf IDs repeated within `shown`, or within `endorsed`."""
+        if len(self.shown) != len(set(self.shown)):
+            raise ValueError("One or more leaf IDs were repeated in the shown data")
+        concat_ids = list(chain(*self.endorsed))
+        if len(concat_ids) != len(set(concat_ids)):
+            raise ValueError(
+                "One or more leaf IDs were repeated in the endorsement data"
+            )
+
+        return self
+
+    @model_validator(mode="after")
+    def check_consistency(self) -> Self:
+        """Check that `endorsed`'s leaves are exactly the ones in `shown`."""
+        all_shown = set(self.shown)
+        all_endorsed = set(chain(*self.endorsed))
+        if all_shown != all_endorsed:
+            raise ValueError("Inconsistent leaf IDs between shown and endorsed")
+
+        return self

matchlab/eval/metrics.py

@@ -0,0 +1,244 @@
+"""Score a resolver's clusters against stored human judgements."""
+
+from itertools import chain, combinations
+from typing import TypeAlias
+
+import polars as pl
+
+Pair: TypeAlias = tuple[int, int]
+Pairs: TypeAlias = set[Pair]
+
+PrecisionRecall: TypeAlias = tuple[float, float]
+
+
+def precision_recall(
+    models_root_leaf: list[pl.DataFrame],
+    judgements: pl.DataFrame,
+    expansion: pl.DataFrame,
+) -> list[PrecisionRecall]:
+    """Compute precision and recall for each model's clusters against judgements.
+
+    The function turns both model clusters and judgements into pairs of leaves, then
+    compares them:
+
+    - Convert each model's clusters and the judgements into pairwise connections
+      between leaves. For judgements, this includes pairs the user was shown but
+      rejected, not just the ones they endorsed. Each pair's net score is the number
+      of times it was endorsed, minus the number of times it was rejected.
+    - Keep only pairs whose leaves appear in every model and in the judgements, so
+      each model is compared over the same data.
+    - Drop a pair from both model and validation pairs if it was rejected as often as
+      it was endorsed. If it was rejected more often than endorsed, drop it from
+      validation pairs only, and keep it in model pairs.
+    - Compute precision and recall for each model against the remaining validation
+      pairs.
+
+    This ignores user IDs, so judgements are not yet weighted by reviewer.
+
+    Args:
+        models_root_leaf: One `(root, leaf)` table per model. Each must be
+            merge-forward, including every leaf reachable from the model's inputs,
+            even ones no model touched, which appear as their own root.
+        judgements: Table following `SCHEMA_JUDGEMENTS`.
+        expansion: Table following `SCHEMA_CLUSTER_EXPANSION`.
+
+    Returns:
+        One `(precision, recall)` pair per model, in the order given.
+    """
+    leaves_per_set: list[set[int]] = []  # one entry for each model, one for judgements
+    pairs_per_model: list[Pairs] = []
+
+    if not len(judgements):
+        raise ValueError("Judgements data cannot be empty.")
+
+    # Process models and judgements
+    for root_leaf in models_root_leaf:
+        if not len(root_leaf):
+            raise ValueError("Model data cannot be empty.")
+        leaves_per_set.append(set(root_leaf["leaf"].to_list()))
+        clusters = (
+            root_leaf.group_by("root")
+            .agg(pl.col("leaf").alias("leaves"))
+            .select("leaves")
+            .to_series()
+            .to_list()
+        )
+        model_pairs = set()
+        for c in clusters:
+            model_pairs.update(combinations(sorted(c), r=2))
+        pairs_per_model.append(model_pairs)
+
+    validation_pairs, validation_net_count, validation_leaves = process_judgements(
+        judgements, expansion
+    )
+    leaves_per_set.append(validation_leaves)
+
+    # Filter pairs to the leaves shared by every model and the judgements. For
+    # example, if model1 has (1,2), (1,3), (2,3), model2 has (1,10), (2,20), and
+    # judgements have (1), (2,3), only (1,2) is kept.
+    shared_leaves = set.intersection(*leaves_per_set)
+
+    for i, model_pairs in enumerate(pairs_per_model):
+        pairs_per_model[i] = {
+            (a, b)
+            for a, b in model_pairs
+            if (
+                a in shared_leaves
+                and b in shared_leaves
+                # Remove neutrally-judged pairs from the model
+                and validation_net_count[(a, b)] != 0
+            )
+        }
+
+    validation_pairs = {
+        (a, b)
+        for a, b in validation_pairs
+        if a in shared_leaves
+        and b in shared_leaves
+        # Remove all neutrally or negatively judged pairs
+        and validation_net_count[(a, b)] > 0
+    }
+
+    if not validation_pairs:
+        raise ValueError("Validation data has no pairs to evaluate.")
+
+    # Compute PR scores for each model
+    pr_scores: list[PrecisionRecall] = []
+    for i, model_pairs in enumerate(pairs_per_model):
+        true_positive_pairs = model_pairs & validation_pairs
+        if not model_pairs:
+            raise ValueError(f"Model at index {i} has no pairs to evaluate.")
+        pr_scores.append(
+            (
+                len(true_positive_pairs) / len(model_pairs),
+                len(true_positive_pairs) / len(validation_pairs),
+            )
+        )
+
+    return pr_scores
+
+
+def process_judgements(
+    judgements: pl.DataFrame, expansion: pl.DataFrame
+) -> tuple[Pairs, dict[Pair, float], set[int]]:
+    """Turn judgements into leaf pairs, a net score per pair, and the leaves shown.
+
+    Expanding a cluster's leaves into every sorted pair is not quite enough on its
+    own. When a user is shown (123) and endorses (12), they also reject pairs (1,3)
+    and (2,3), not just endorse (1,2). Each pair's net score sums +1 for every
+    endorsement and -1 for every rejection implied this way.
+
+    This function assumes well-formed input:
+
+    - Every shown cluster ID has a matching row in `expansion`.
+    - Every endorsed cluster ID has a matching row in `expansion`, unless it is a
+      single leaf.
+    - No cluster is split across judgement rows unless every other part is expanded
+      too. For example, if (123)->(12) appears, (123)->(3) must appear too.
+
+    Args:
+        judgements: Table following `SCHEMA_JUDGEMENTS`.
+        expansion: Table following `SCHEMA_CLUSTER_EXPANSION`.
+
+    Returns:
+        A tuple of:
+
+        - Every pair implied by an endorsement or a rejection.
+        - Each pair's net score, where positive means endorsed more than rejected.
+        - Every leaf ID shown to a user.
+    """
+    expanded_judgements = (
+        judgements.join(expansion, left_on="shown", right_on="root")
+        .rename({"leaves": "shown_leaves"})
+        .join(
+            expansion, left_on="endorsed", right_on="root", how="left"
+        )  # left join as singleton leaves won't be expanded
+        .rename({"leaves": "endorsed_leaves"})
+        # Assume a missing expansion means a singleton leaf
+        .with_columns(
+            pl.when(pl.col("endorsed_leaves").is_null())
+            .then(
+                pl.col("endorsed").map_elements(
+                    lambda x: [x], return_dtype=pl.List(pl.UInt64)
+                )
+            )
+            .otherwise(pl.col("endorsed_leaves"))
+            .alias("endorsed_leaves")
+        )
+    )
+
+    all_leaves = set(
+        chain.from_iterable(expanded_judgements["endorsed_leaves"].to_list())
+    )
+
+    if len(expanded_judgements) != len(judgements):
+        raise ValueError("Malformed judgements / expansion data received.")
+
+    validation_net_count: dict[Pair, float] = {}
+
+    for row in expanded_judgements.rows(named=True):
+        shown = row["shown_leaves"]
+        endorsed = row["endorsed_leaves"]
+
+        positive_pairs = set(combinations(sorted(endorsed), r=2))
+        potential_pairs = set(combinations(sorted(shown), r=2))
+        negative_pairs = potential_pairs - positive_pairs
+
+        # -- WORKED EXAMPLE --
+        # A user is shown cluster (1234) and endorses three groups: (1), (23), (4).
+        # This makes pair (2,3) good, and (1,2), (1,3), (1,4), (2,4), (3,4) bad.
+        # Target result: +1 for (2,3), -1 for every bad pair.
+
+        # The data does not arrive as one row holding the whole judgement. It arrives
+        # as one row per endorsed group, and these rows can be mixed with other users'
+        # judgements:
+        #   Row A: shown (1234), endorsed (1)
+        #   Row B: shown (1234), endorsed (23)
+        #   Row C: shown (1234), endorsed (4)
+        # Each row below is scored on its own. The partial scores still add up to the
+        # right total, whatever order the rows arrive in.
+
+        # Row A: shown (1234), endorsed (1)
+        # - Potential pairs: (1,2) (1,3) (1,4) (2,3) (2,4) (3,4).
+        # - Endorsed pairs: none. A single item forms no pair.
+        # - Every potential pair is tentatively rejected, weight -1/4, because this
+        #   endorsed group holds 1 of the 4 shown items.
+        # - Running scores: all pairs = -0.25.
+
+        # Row B: shown (1234), endorsed (23)
+        # - Endorsed pair (2,3) gets +1, plus +2/4 to cancel the rejection weight it
+        #   took in rows A and C, the non-endorsed share of this group.
+        #   Total for (2,3): +1 + 0.5 = +1.5.
+        # - Remaining rejected pairs take a further -2/4 = -0.5.
+        # - Running scores: (2,3) = -0.25 + 1.5 = +1.25, others = -0.25 - 0.5 = -0.75.
+
+        # Row C: shown (1234), endorsed (4)
+        # - Final -1/4 rejection weight lands on every potential pair.
+        # - Final scores: (2,3) = +1.25 - 0.25 = +1.0, others = -0.75 - 0.25 = -1.0.
+
+        # Result: (2,3) nets +1, matching the endorsement, and every rejected pair
+        # nets -1, whatever order the rows arrive in or however they interleave with
+        # other judgements, as long as no cluster is split without expanding its
+        # other part too.
+
+        # Subtract negative weight for all shown pairs not endorsed on this row
+        negative_adjustment = len(endorsed) / len(shown)
+        validation_net_count.update(
+            {
+                p: validation_net_count.get(p, 0) - negative_adjustment
+                for p in negative_pairs
+            }
+        )
+
+        # Add 1 plus positive weight for all pairs endorsed on this row
+        positive_adjustment = 1 + ((len(shown) - len(endorsed)) / len(shown))
+        validation_net_count.update(
+            {
+                p: validation_net_count.get(p, 0) + positive_adjustment
+                for p in positive_pairs
+            }
+        )
+
+    all_pairs = set(validation_net_count.keys())
+
+    return all_pairs, validation_net_count, all_leaves