matchlab 0.1.0.dev1__py3-none-any.whl

matchlab/steps.py

@@ -0,0 +1,284 @@
+"""The lazy plan node.
+
+A `Step` holds references to its **inputs only** (`upstream`). There is no registry,
+no parent pointer, and no downstream list. "The DAG" is whatever is reachable
+upstream from the node you hold, and lineage operations are pure functions of a root
+node (`matchlab.lineage`).
+
+Nothing is computed until `collect()`. Collection walks the plan upstream-first and
+runs only the steps whose artifact is not already stored.
+
+**Fingerprints identify artifacts.** A step's fingerprint combines its kind, its
+spec, and its inputs' fingerprints. Everything downstream of a source can therefore
+derive it from the *plan alone*, before any work happens, and `collect`
+can skip a cached step without running it. Sources are the exception. Raw data enters
+there, so a source's spec key includes a content hash of the data it read.
+Constructing a fresh `Source` therefore re-reads the warehouse (the documented way to
+refresh), while an existing `Source` object memoises its read.
+"""
+
+import json
+from abc import ABC, abstractmethod
+from typing import ClassVar, Self
+
+from platformdirs import user_cache_path
+from pydantic import BaseModel
+
+from matchlab import lineage
+from matchlab.adapters import Adapter, DuckDBAdapter, Fingerprint
+from matchlab.core.hash import HASH_FUNC
+from matchlab.core.kinds import StepKind
+from matchlab.lineage import StepStatus
+from matchlab.progress import report
+
+CACHE_DIR = user_cache_path("matchlab")
+
+_DEFAULT_ADAPTER: Adapter | None = None
+
+
+def set_default_adapter(adapter: Adapter | None) -> None:
+    """Set the adapter used by `collect()` when none is passed. `None` resets it."""
+    global _DEFAULT_ADAPTER  # module-level default, by design
+    _DEFAULT_ADAPTER = adapter
+
+
+def default_adapter() -> Adapter:
+    """Return the default adapter, creating a DuckDB store in the cache dir if unset."""
+    global _DEFAULT_ADAPTER  # module-level default, by design
+    if _DEFAULT_ADAPTER is None:
+        CACHE_DIR.mkdir(parents=True, exist_ok=True)
+        _DEFAULT_ADAPTER = DuckDBAdapter(CACHE_DIR / "store.duckdb")
+    return _DEFAULT_ADAPTER
+
+
+class Step(ABC):
+    """A node in a lazy plan."""
+
+    kind: ClassVar[StepKind]
+
+    def __init__(self, upstream: "tuple[Step, ...]" = ()) -> None:
+        """Initialise a plan node with its direct inputs.
+
+        Steps have no names. They are identified by **position**, where they fall in
+        `lineage.walk`, which is the order `collect` runs them in and the order
+        `PlanDocument` lists them in. `step 7` in a log, `[7]` in `draw()`, and
+        `steps[7]` in a document are therefore the same node.
+
+        A position is not stored here, because it is not a property of the step. It
+        belongs to the walk it came from, and the same step numbers differently in
+        `walk(deduped)` and `walk(companies)`. Whoever does the walking passes it to
+        whoever needs it, `collect` to its reporter, `draw` to its own renderer.
+
+        Finding a result later is a separate matter, and a separate act.
+        `Resolver.publish` points a **label** at a resolver's output. `Source` is the
+        one step with a name, and it means something else again. A source's name is
+        part of its output, prefixing every column it contributes and tagging its rows.
+        """
+        self.upstream = tuple(upstream)
+        self._fp: Fingerprint | None = None
+        self._adapter: Adapter | None = None
+
+    def __repr__(self) -> str:
+        """Return a short representation showing kind and collection state."""
+        return f"<{type(self).__name__} {'collected' if self.is_collected else 'lazy'}>"
+
+    def __str__(self) -> str:
+        """How this step appears in a drawing."""
+        return self.kind
+
+    @property
+    def is_collected(self) -> bool:
+        """Whether this step has been materialised."""
+        return self._fp is not None
+
+    # -- plan identity ----------------------------------------------------------------
+
+    @property
+    @abstractmethod
+    def spec(self) -> BaseModel:
+        """This step's settings, as a serialisable model.
+
+        One model per step kind, in `matchlab.specs`. It must carry everything this
+        step's output depends on and nothing else. That is the invariant `_spec_key`
+        rests on, and the one to check when adding a setting. Omit something that
+        changes the output and collect will hand back a stale artifact without
+        re-running (see `_fingerprint`).
+
+        Specs describe a step's own settings, not its inputs'. Edges live on
+        `upstream`, and `_fingerprint` already folds in the parents' fingerprints.
+        """
+        ...
+
+    def _spec_key(self) -> bytes:
+        """Bytes identifying this step's spec.
+
+        `Source` extends this with a content hash of the data it read, so that the
+        fingerprint tracks the warehouse. For every other kind the spec is the
+        whole story.
+        """
+        return json.dumps(self.spec.model_dump(mode="json"), sort_keys=True).encode()
+
+    def _fingerprint(self) -> Fingerprint:
+        """Address this step by kind, spec and inputs.
+
+        The key is derived from the *plan*, not from the step's output. That is what
+        makes it computable before the step runs, which is what lets `_ensure` skip
+        work. An output digest would only be knowable once the work was already done.
+        Sources are the exception. They fold a hash of the data they read into
+        `_spec_key`, because no spec reveals that the warehouse moved.
+
+        The trade-off is that the key can disagree with the bytes in both directions.
+
+        * the spec omits something that changes the output, a **stale hit**, because
+          `_ensure` never runs the step and reads the old artifact. `SplinkLinker`
+          without an explicit seed is a live example. The only defence is that every
+          `_spec_key` covers everything its step's output depends on. Treat that as
+          the invariant to protect when adding a step or a setting.
+        * the spec includes something that doesn't change the output, a **spurious
+          miss**, re-running the step and everything below it. No spec does this
+          today, and the way to keep it that way is to record settings only. A spec
+          that described an *input* would, since identity already arrives via the
+          parent fingerprint. A setting that must point at an input points at its
+          position, which is not redundant. It decides which input the setting
+          applies to.
+
+        There is also no early cutoff. A `Transform` whose SQL is reformatted but
+        semantically unchanged invalidates the whole subtree beneath it.
+
+        TODO(fingerprints): split this into an action key (plan-derived, as now)
+        mapping to an output digest (content-derived, recorded by the adapter on
+        write), and build a child's key from its parents' output digests rather than
+        their action keys. That buys early cutoff, storage dedup and rename tolerance.
+        It does **not** fix stale hits. An output digest governs propagation, never
+        admission, so a step whose action key is a hit still never runs and its digest
+        is never consulted. Costs an indirection table in the adapter contract, a hash
+        of every artifact on write, and the loss of a work-set knowable before running.
+        See PLAN.md, "Known limitations", for the full ledger.
+        """
+        parts: list[bytes] = [self.kind.encode(), self._spec_key()]
+        for parent in self.upstream:
+            if parent._fp is None:  # collect orders upstream first
+                raise RuntimeError(
+                    f"A {parent.kind} input of this {self.kind} has no fingerprint yet."
+                )
+            parts.append(parent._fp)
+        return HASH_FUNC(b"|".join(parts)).digest()
+
+    # -- execution --------------------------------------------------------------------
+
+    @abstractmethod
+    def _execute(self, adapter: Adapter, fp: Fingerprint) -> None:
+        """Compute this step and persist its artifact under `fp`."""
+        ...
+
+    def _ensure(self, adapter: Adapter) -> StepStatus:
+        """Materialise this step unless its artifact is already stored.
+
+        This classifies the outcome but reports none of it. `collect` holds the walk,
+        and therefore each step's position, the thing a record has to quote to be
+        findable in the plan. Reporting belongs there, with the single
+        `matchlab.progress.Progress` that owns what is drawn and what is logged. That is
+        also what keeps the two from ever disagreeing. This keeps to the work.
+
+        Returns:
+            What it took. `DONE` if this call computed the step, `CACHED` if the
+            artifact was already stored.
+        """
+        self._adapter = adapter
+
+        if self._fp is not None and adapter.has(self._fp):
+            return StepStatus.CACHED
+
+        fp = self._fingerprint()
+
+        if adapter.has(fp):  # cache hit, skip the work entirely
+            self._fp = fp
+            return StepStatus.CACHED
+
+        self._execute(adapter, fp)
+        self._fp = fp
+        return StepStatus.DONE
+
+    # -- public API -------------------------------------------------------------------
+
+    def collect(
+        self, adapter: Adapter | None = None, interactive: bool | None = None
+    ) -> Self:
+        """Materialise this step and everything it depends on.
+
+        Steps whose artifact is already stored are skipped without being run, so
+        re-collecting after adding a downstream step only does the new work.
+
+        Args:
+            adapter: Where to read and write artifacts. Defaults to the module-level
+                adapter (a DuckDB store in the user cache directory).
+            interactive: Whether someone is watching. `None`, the default, takes a
+                terminal or a notebook as a yes. When they are, the plan is drawn as a
+                live tree redrawn in place, and not logged. The tree on screen is the
+                key those `[step N]` records need, and it stays there. When they are
+                not, the plan is logged instead. See `matchlab.progress`.
+
+        Returns:
+            This step, now collected.
+        """
+        adapter = adapter or default_adapter()
+        # One walk, used for both. It fixes each step's position, so a `step 7` the
+        # reporter logs is the node it drew as `[7]`.
+        steps = lineage.walk(self)
+        with report(self, steps, interactive, adapter) as reporter:
+            for step in steps:
+                reporter.begin(step)
+                try:
+                    status = step._ensure(adapter)
+                except Exception:
+                    reporter.end(step, StepStatus.FAILED)
+                    raise
+                reporter.end(step, status)
+        return self
+
+    def lineage(self) -> list["Step"]:
+        """Return this step and all its inputs, upstream-first."""
+        return lineage.walk(self)
+
+    def draw(self) -> str:
+        """Render this step's sub-plan as a tree."""
+        return lineage.draw(self)
+
+    def fingerprints(self) -> set[Fingerprint]:
+        """Address every artifact this plan is made of, its own and its inputs'.
+
+        Which artifacts belong to a plan is the plan's own business, so this is where
+        a store gets told: `adapter.prune(keep=plan.fingerprints())` hands storage a set
+        of addresses it already understands, rather than a graph it would have to learn
+        to walk.
+
+        Returns:
+            One fingerprint per step in `lineage()`. A set, because two steps in one
+            plan can address the same artifact. Identical specs over identical
+            inputs is the same bytes, and it is stored once.
+
+        Raises:
+            RuntimeError: If any step has not been collected. An uncollected plan names
+                no artifacts at all, so answering with a smaller set would quietly tell
+                a caller that less is worth keeping than they think.
+        """
+        return {step._collected()[1] for step in self.lineage()}
+
+    # -- helpers for subclasses -------------------------------------------------------
+
+    def _collected(self) -> tuple[Adapter, Fingerprint]:
+        """Return the adapter this step was collected into, and its fingerprint.
+
+        Both together, because everything that reads a stored artifact needs both and
+        neither exists before collection. Returning the pair is what lets a caller
+        use the fingerprint without re-checking that it is there.
+        """
+        if self._adapter is None or self._fp is None:
+            raise RuntimeError(
+                f"This {self.kind} has not been collected. Call collect() first."
+            )
+        return self._adapter, self._fp
+
+    def _require_adapter(self) -> Adapter:
+        """Return the adapter this step was collected into."""
+        return self._collected()[0]

matchlab/testkit/__init__.py

@@ -0,0 +1,104 @@
+"""Generate data whose answer you already know, then score matching against it.
+
+This package makes entity resolution *measurable*. It plants entities you already know
+the answer for, so a test can assert precision and recall against truth rather than
+eyeball a fixture.
+
+The package has one job, done in four steps: **describe** the data (`features`),
+**generate** it and plant the answer (`sources`, `linked`), **match** it (`matchers`),
+then **score** the result (`compare`). `LinkedSources` is the handle for all of it. It
+planted the entities, so nothing downstream has to be told the answer.
+
+## Which do I use?
+
+| I want… | Use |
+|---|---|
+| data whose answer I know | `linked_sources_factory()` → `LinkedSources` |
+| one source only | `source_factory()` → `GeneratedSource` |
+| rows I wrote by hand | `source_from_tuple()` |
+| a plan to test | `linked.dedupe(...)` / `linked.link(..., through=...)` |
+| to score a collected plan | `linked.diff_resolver_output(resolver_output, *sources)` |
+| to score a methodology's edges | `linked.diff_model_edges(edges, left=...)` |
+| different columns | `FeatureConfig` + `SuffixRule` / `ReplaceRule` |
+| a matcher that already knows the answer | `PerfectDeduper` / `PerfectLinker` |
+| to build or bend the answer it matches on | `AnswerKey.from_sources(...)` |
+| to compose a comparison by hand | `resolver_output_to_clusters` + `diff_entities` |
+
+`dedupe()` and `link()` wire a perfect matcher up for you, so the last two rows are
+only for when you want a matcher that is deliberately wrong in a particular way.
+
+A testkit exposes what it knows about the *fixture*. Anything you do to the *plan* goes
+through the node it wraps: `source.source.dedupe(...)`, `model.model.resolve()`. There
+are no forwarding shortcuts, so it is always clear which of the two you are holding.
+
+```python
+from matchlab.testkit import linked_sources_factory
+
+linked = linked_sources_factory(
+    n_true_entities=10, engine=warehouse
+).write_to_location()
+resolver_output = linked.link("crn", "cdms").model.resolve().collect().entities()
+
+identical, report = linked.diff_resolver_output(resolver_output, "crn", "cdms")
+assert identical, report
+```
+
+## Two layers, two questions
+
+Ground truth is expressed twice, in different ID spaces. Neither subsumes the other, and
+picking the wrong one is the easiest mistake to make here:
+
+| | Synthetic-ID layer | Value-keyed layer |
+|---|---|---|
+| Asks | does this *methodology* work? | does the *plan* carry grouping forward? |
+| Built from | `GeneratedModel.predicted_clusters` | `AnswerKey` + `Perfect*` |
+| Scored with | `diff_model_edges` | `diff_resolver_output` |
+| Needs a warehouse | no, the record step read is mocked | yes, and a real `collect()` |
+| Example | `test/methodologies/` | `test/plan/test_ground_truth.py` |
+
+The split exists because matchlab derives record identity by content-hashing rows at
+collect time. IDs are unknowable when a fixture is built, so anything asserted against a
+*collected* plan has to be keyed by row values instead. See `matchers` for the full
+reasoning. `dedupe()`/`link()` build both, so either is available from one call.
+
+Names not re-exported below are still importable from their module, but they are
+plumbing rather than entry points.
+"""
+
+from matchlab.testkit.compare import diff_entities, resolver_output_to_clusters
+from matchlab.testkit.entities import Cluster, TrueEntity
+from matchlab.testkit.features import (
+    FeatureConfig,
+    ReplaceRule,
+    SourceParameters,
+    SuffixRule,
+)
+from matchlab.testkit.linked import LinkedSources, linked_sources_factory
+from matchlab.testkit.matchers import AnswerKey, PerfectDeduper, PerfectLinker
+from matchlab.testkit.models import GeneratedModel
+from matchlab.testkit.sources import GeneratedSource, source_factory, source_from_tuple
+
+__all__ = (
+    # generate data with a known answer
+    "linked_sources_factory",
+    "LinkedSources",
+    "source_factory",
+    "source_from_tuple",
+    "GeneratedSource",
+    "GeneratedModel",
+    # match, when you want something other than a perfect matcher
+    "PerfectDeduper",
+    "PerfectLinker",
+    "AnswerKey",
+    # describe what to generate
+    "FeatureConfig",
+    "SuffixRule",
+    "ReplaceRule",
+    "SourceParameters",
+    # score results against the answer
+    "diff_entities",
+    "resolver_output_to_clusters",
+    # the vocabulary you get back
+    "TrueEntity",
+    "Cluster",
+)

matchlab/testkit/_generate.py

@@ -0,0 +1,240 @@
+"""Row and entity generation.
+
+Internal. Reach for `source_factory` or `linked_sources_factory` instead. These are the
+pieces they are built from, and they take arguments that only make sense once the
+callers have worked them out.
+"""
+
+from functools import cache
+from itertools import product
+from math import prod
+
+import polars as pl
+import pyarrow as pa
+from faker import Faker
+
+from matchlab.testkit.entities import EntityReference, TrueEntity
+from matchlab.testkit.features import FeatureConfig
+
+
+@cache
+def generate_entities(
+    generator: Faker,
+    features: tuple[FeatureConfig, ...],
+    n: int,
+) -> tuple[TrueEntity]:
+    """Generate base entities with their ground truth values from the generator."""
+    entities = []
+    for _ in range(n):
+        base_values = {}
+        for feature in features:
+            generator_func = generator.unique if feature.unique else generator
+            value_generator = getattr(generator_func, feature.base_generator)
+            parameters = {} if not feature.parameters else dict(feature.parameters)
+
+            value = value_generator(**parameters)
+            # Explicitly cast lists to tuples to ensure they are hashable
+            if isinstance(value, list):
+                value = tuple(value)
+            base_values[feature.name] = value
+
+        entities.append(TrueEntity(base_values=base_values, keys=EntityReference()))
+    return tuple(entities)
+
+
+def generate_rows(
+    generator: Faker,
+    selected_entities: tuple[TrueEntity, ...],
+    features: tuple[FeatureConfig, ...],
+    repetition: int,
+) -> tuple[dict[str, list], dict[int, list[str]], dict[int, list[str]]]:
+    """Generate raw data rows, plus the maps that connect three different IDs.
+
+    Every row carries three IDs, each answering a different question. Confusing them is
+    the mistake this docstring exists to prevent.
+
+    * `key`: the row's own identifier. Like a source's primary key, unique within the
+      source but not across entities.
+    * `id`: matchlab's identifier for the row's *content*. Rows with identical feature
+      values share one `id`, whichever entity produced them.
+    * `entity`: the ID of the `TrueEntity` that generated the row. This is the planted
+      answer, and the one truth downstream tests check their result against.
+
+    `entity` groups rows by who generated them, and `id` groups rows by what content
+    they hold. The two groupings diverge whenever an entity produces more than one
+    variation: rows from the same entity can carry different `id`s, and (via
+    `repetition`) the same `id` can appear on more than one row.
+
+    Returns:
+        - raw_data: column arrays, ready to build a DataFrame from
+        - entity_keys: `TrueEntity` ID -> the keys it produced
+        - id_keys: content `id` -> the keys sharing that content
+
+    Worked example, two entities, each producing two variations twice over:
+
+    | id | key | company_name |
+    |----|-----|--------------|
+    | 1  | a   | alpha co     |
+    | 2  | b   | alpha ltd    |
+    | 1  | c   | alpha co     |  # same content as row a
+    | 2  | d   | alpha ltd    |  # same content as row b
+    | 3  | e   | beta co      |
+    | 4  | f   | beta ltd     |
+    | 3  | g   | beta co      |  # same content as row e
+    | 4  | h   | beta ltd     |  # same content as row f
+
+    As raw data:
+
+    ```python
+    raw_data = {
+        "id": [1, 2, 1, 2, 3, 4, 3, 4],
+        "key": ["a", "b", "c", "d", "e", "f", "g", "h"],
+        "company_name": [
+            "alpha co",
+            "alpha ltd",
+            "alpha co",
+            "alpha ltd",
+            "beta co",
+            "beta ltd",
+            "beta co",
+            "beta ltd",
+        ],
+    }
+    ```
+
+    Grouped by entity, the two variations sit together:
+
+    ```python
+    entity_keys = {
+        1: ["a", "b", "c", "d"],  # every key entity 1 produced
+        2: ["e", "f", "g", "h"],  # every key entity 2 produced
+    }
+    ```
+
+    Grouped by content, they split apart again:
+
+    ```python
+    id_keys = {
+        1: ["a", "c"],  # both rows read "alpha co"
+        2: ["b", "d"],  # both rows read "alpha ltd"
+        3: ["e", "g"],  # both rows read "beta co"
+        4: ["f", "h"],  # both rows read "beta ltd"
+    }
+    ```
+    """
+    raw_data = {"key": [], "id": []}
+    for feature in features:
+        raw_data[feature.name] = []
+
+    # Track entity locations and row identities
+    entity_keys = {entity.id: [] for entity in selected_entities}
+    id_keys = {}
+    value_to_id = {}
+
+    def add_row(entity_id: int, values: tuple) -> None:
+        """Add a row of data, handling IDs and keys."""
+        key = str(generator.uuid4())
+        entity_keys[entity_id].append(key)
+
+        if values not in value_to_id:
+            mb_id = generator.random_number(digits=16)
+            value_to_id[values] = mb_id
+            id_keys[mb_id] = []
+
+        row_id = value_to_id[values]
+        id_keys[row_id].append(key)
+
+        raw_data["key"].append(key)
+        raw_data["id"].append(row_id)
+        for feature, value in zip(features, values, strict=True):
+            raw_data[feature.name].append(value)
+
+    for entity in selected_entities:
+        # Collect all possible values, per feature.
+        possible_values = []
+        for feature in features:
+            base = entity.base_values[feature.name]
+
+            for rule in feature.variations:
+                if not isinstance(base, rule.type):
+                    raise TypeError(
+                        f"{rule.__class__.__name__} requires {rule.type} "
+                        f"but {feature.name} generates {type(base)}"
+                    )
+
+            variations = []
+            # Apply all variations as long as they change the value
+            for v in (rule.apply(base) for rule in feature.variations):
+                if v != base:
+                    variations.append(v)
+
+            values = variations if feature.drop_base else variations + [base]
+            possible_values.append(values or [base])
+
+        if (num_variations := prod(len(values) for values in possible_values)) > 100:
+            raise RuntimeError(
+                f"Entity {entity.id} would generate {num_variations:,} variations, "
+                "which exceeds the maximum of 100. This would cause a Cartesian "
+                "explosion."
+            )
+
+        for values in product(*possible_values):
+            for _ in range(repetition + 1):
+                add_row(entity.id, values)
+
+    return raw_data, entity_keys, id_keys
+
+
+@cache
+def generate_source(
+    generator: Faker,
+    n_true_entities: int,
+    features: tuple[FeatureConfig, ...],
+    repetition: int,
+    seed_entities: tuple[TrueEntity, ...] | None = None,
+) -> tuple[pa.Table, dict[int, set[str]], dict[int, set[str]]]:
+    """Generate raw data as PyArrow tables with entity tracking.
+
+    Note: this is `@cache`d but also *mutates* `total_unique_variations` on the shared
+    `TrueEntity` objects it is given, so a cache hit skips that mutation. Harmless
+    today because the first call for a given key does the writing, but it means the
+    function is not free of side effects and the cache cannot simply be dropped.
+
+    Returns:
+        - data: PyArrow table with generated data
+        - entity_keys: TrueEntity ID -> list of keys mapping
+        - id_keys: Unique row ID -> list of keys mapping for identical rows
+    """
+    # Select or generate entities
+    if seed_entities is None:
+        selected_entities = generate_entities(generator, features, n_true_entities)
+    else:
+        selected_entities = generator.random_elements(
+            elements=seed_entities,
+            unique=True,
+            length=min(n_true_entities, len(seed_entities)),
+        )
+
+    # Generate initial data
+    raw_data, entity_keys, id_keys = generate_rows(
+        generator=generator,
+        selected_entities=selected_entities,
+        features=features,
+        repetition=repetition,
+    )
+
+    # Create DataFrame
+    df = pl.DataFrame(raw_data)
+
+    # Update variation counts
+    for entity in selected_entities:
+        if entity.id in entity_keys:
+            # Count unique row IDs this entity appears in
+            entity_rows = df.filter(pl.col("key").is_in(list(entity_keys[entity.id])))
+            entity.total_unique_variations = entity_rows["id"].n_unique()
+
+    return (
+        df.to_arrow(),
+        entity_keys,
+        id_keys,
+    )