matchlab 0.1.0.dev1__py3-none-any.whl

matchlab/recordstep.py

@@ -0,0 +1,181 @@
+"""RecordStep, the records a model matches over, and the verbs that build a plan.
+
+A `RecordStep` is a step whose artifact is a table of records carrying an `id`.
+`Source`, `Resolver` and `Transform` are the three kinds. On a `Source`, `id` is the
+leaf. On a `Resolver`, `id` is the entity root. Every `RecordStep` chains the same
+verbs (`select`, `clean`, `group`, `transform`, `dedupe`, `link`). A `Model` reads a
+`RecordStep` and yields edges, not records.
+
+`RecordStep` is not user-facing. Users hold a `Source`, a `Resolver` or a `Transform`
+and call verbs on it. Each concrete kind supplies only how it materialises
+(`_read_cache`) and which source rows it stands for (`_identifier_reads`).
+
+A record step does not store `identifiers()`, the `(id, source, key, leaf)` mapping a
+downstream resolver needs. Reshaping cannot change it, so it is read back from the
+source leaves and the upstream resolver output instead.
+"""
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+import polars as pl
+
+from matchlab.adapters import Adapter, Fingerprint
+from matchlab.core.dataframes import (
+    DataFrameClass,
+    DataFrameType,
+    qualify,
+    to_dataframe,
+)
+from matchlab.core.sql import SQLExpression
+from matchlab.steps import Step
+
+if TYPE_CHECKING:
+    # Each of these modules imports this one. Only for annotations, so importing the
+    # real classes (rather than typing them `Any`) costs nothing at runtime.
+    from matchlab.models import Model
+    from matchlab.models.dedupers.base import Deduper, DeduperSettings
+    from matchlab.models.linkers.base import Linker, LinkerSettings
+    from matchlab.resolvers import Resolver
+    from matchlab.sources import Source
+    from matchlab.transformers import Transform, Transformer
+
+# What one source's identifiers are read with, `(source_fp, source_name, resolver_fp)`,
+# the arguments of `Adapter.read_identifiers`. Deduplicating these is deduplicating the
+# queries, which is why they travel as a value rather than as a call.
+IdentifierRead = tuple[Fingerprint | None, str, Fingerprint | None]
+
+
+def build_record_step(
+    adapter: Adapter,
+    sources: "tuple[Source, ...]",
+    resolver: "Resolver | None",
+) -> pl.DataFrame:
+    """Assemble the `id` + qualified-columns table from stored extracts and identifiers.
+
+    Every source's extract is prefixed with the source name (`company` → `crn_company`)
+    and joined to its identifiers, so each row gains the `id` it belongs to. That is the
+    record's leaf when read directly, or the entity root when read through `resolver`.
+    Several sources are concatenated diagonally, each row carrying its own source's
+    columns and nulls for the rest. This is what both `Source` and `Resolver` read.
+    """
+    resolver_fp = resolver._fp if resolver is not None else None
+    identifiers = pl.concat(
+        [
+            adapter.read_identifiers(source._fp, source.name, resolver_fp)
+            for source in sources
+        ],
+        how="vertical",
+    )
+
+    extracts = [
+        adapter.read_source_extract(source._fp)
+        .select(pl.all().name.prefix(qualify(source.name)))
+        .with_columns(pl.lit(source.name).alias("source"))
+        .rename({source.qualified_key: "key"})
+        for source in sources
+    ]
+
+    return (
+        pl.concat(extracts, how="diagonal")
+        .join(identifiers.drop("leaf"), how="inner", on=("source", "key"))
+        .drop("source", "key")
+    )
+
+
+class RecordStep(Step):
+    """A step whose artifact is a table of records a model matches over."""
+
+    # -- RecordStep contract -------------------------------------------------------
+
+    @abstractmethod
+    def _read_cache(self, adapter: Adapter) -> pl.DataFrame:
+        """Return this record step's records from wherever they are materialised."""
+        ...
+
+    @property
+    @abstractmethod
+    def _identifier_reads(self) -> tuple[IdentifierRead, ...]:
+        """The `Adapter.read_identifiers` arguments this record step's data comes from.
+
+        One per source, naming what is read rather than reading it. Two record steps
+        built from the same source and resolver read identical rows, even if they then
+        reshape those rows differently. A downstream resolver wants the *set* of these
+        reads across every record step feeding it, so naming them lets it ask once
+        instead of once per record step.
+        """
+        ...
+
+    def identifiers(self, adapter: Adapter) -> pl.DataFrame:
+        """Return `(id, source, key, leaf)` for every record this record step reads.
+
+        `id` is the resolver's entity root when reading through one, otherwise the
+        source leaf. This is the upstream resolver output a downstream resolver needs to
+        carry every reachable leaf forward, including records no model matched.
+        """
+        return pl.concat(
+            [adapter.read_identifiers(*read) for read in self._identifier_reads],
+            how="vertical",
+        )
+
+    def data(self, return_type: DataFrameType = DataFrameType.POLARS) -> DataFrameClass:
+        """Return this record step's records, collecting the plan first if needed."""
+        if not self.is_collected:
+            self.collect()
+        return to_dataframe(self._read_cache(self._require_adapter()), return_type)
+
+    # -- verbs ------------------------------------------------------------------------
+
+    def transform(
+        self,
+        transformer: "Transformer | type[Transformer] | str",
+        transformer_settings: dict | None = None,
+    ) -> "Transform":
+        """Reshape this record step with a transformer."""
+        from matchlab.transformers import Transform  # noqa: PLC0415 - avoids a cycle
+
+        return Transform(self, transformer, transformer_settings)
+
+    def select(self, *columns: str) -> "Transform":
+        """Keep only the named columns, plus `id`."""
+        from matchlab.transformers import Select  # noqa: PLC0415 - avoids a cycle
+
+        return self.transform(Select(*columns))
+
+    def clean(self, cleaning: dict[str, SQLExpression]) -> "Transform":
+        """Derive columns with DuckDB SQL, keeping the rest."""
+        from matchlab.transformers import Clean  # noqa: PLC0415 - avoids a cycle
+
+        return self.transform(Clean(cleaning=cleaning))
+
+    def group(self, aggregates: dict[str, SQLExpression]) -> "Transform":
+        """Collapse each `id` to one row using aggregate SQL."""
+        from matchlab.transformers import Group  # noqa: PLC0415 - avoids a cycle
+
+        return self.transform(Group(aggregates=aggregates))
+
+    def dedupe(
+        self,
+        model_class: "type[Deduper] | str",
+        model_settings: "DeduperSettings | dict",
+    ) -> "Model":
+        """Deduplicate this record step."""
+        from matchlab.models import Model  # noqa: PLC0415 - avoids a cycle
+
+        return Model(left=self, model_class=model_class, model_settings=model_settings)
+
+    def link(
+        self,
+        other: "RecordStep",
+        model_class: "type[Linker] | str",
+        model_settings: "LinkerSettings | dict",
+    ) -> "Model":
+        """Link this record step to another. A `Source` is one, needing no wrapping."""
+        from matchlab.models import Model  # noqa: PLC0415 - avoids a cycle
+
+        return Model(
+            left=self,
+            right=other,
+            model_class=model_class,
+            model_settings=model_settings,
+        )

matchlab/resolvers/__init__.py

@@ -0,0 +1,10 @@
+"""Resolver methodologies and the Resolver plan node."""
+
+from matchlab.resolvers.components import Components, ComponentsSettings
+from matchlab.resolvers.resolvers import Resolver
+
+__all__ = (
+    "Components",
+    "ComponentsSettings",
+    "Resolver",
+)

matchlab/resolvers/base.py

@@ -0,0 +1,63 @@
+"""Base classes for resolver methodologies."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from enum import StrEnum
+from typing import ClassVar, Final
+
+import polars as pl
+import pyarrow as pa
+from pydantic import BaseModel, ConfigDict
+
+SCHEMA_CLUSTERS: Final[pa.Schema] = pa.schema(
+    [
+        ("parent_id", pa.uint64()),
+        ("child_id", pa.uint64()),
+    ]
+)
+"""A resolver methodology's cluster assignments take this shape."""
+
+
+class ResolverType(StrEnum):
+    """Enumeration of supported resolver methodology types.
+
+    This lives with `ResolverMethod`, not in `matchlab.specs`, because it is not part
+    of any spec. A `ResolverSpec` identifies a methodology only by its registered
+    `resolver_class`, so a `ResolverType` never reaches a fingerprint or a document.
+    It is just a class attribute methodologies use to declare which family they
+    belong to, and nothing more.
+    """
+
+    COMPONENTS = "components"
+
+
+class ResolverSettings(BaseModel, ABC):
+    """Base settings type for resolver methodologies.
+
+    Settings that point at one of the resolver's inputs refer to it by **position**.
+    `Resolver` accepts the `Model` object at the API and translates it. Nothing here
+    holds a plan node. A methodology stays a pure function of edges and numbers.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ResolverMethod(BaseModel, ABC):
+    """Base class for resolver methodologies."""
+
+    resolver_type: ClassVar[ResolverType]
+    settings: ResolverSettings
+
+    @abstractmethod
+    def compute_clusters(self, model_edges: Mapping[int, pl.DataFrame]) -> pl.DataFrame:
+        """Compute cluster assignments from model edges.
+
+        Args:
+            model_edges: Input position to that model's edges, conforming to
+                SCHEMA_MODEL_EDGES. Positions index the resolver's inputs, in the
+                order they were given, and are what per-model settings key by.
+
+        Returns:
+            A Polars DataFrame which conforms to SCHEMA_CLUSTERS
+        """
+        ...

matchlab/resolvers/components.py

@@ -0,0 +1,66 @@
+"""Connected-components resolver methodology."""
+
+from collections.abc import Mapping
+from typing import Annotated, ClassVar
+
+import polars as pl
+from pydantic import Field
+
+from matchlab.core.dsu import DisjointSet
+from matchlab.resolvers.base import (
+    SCHEMA_CLUSTERS,
+    ResolverMethod,
+    ResolverSettings,
+    ResolverType,
+)
+
+
+class ComponentsSettings(ResolverSettings):
+    """Settings for the Components resolver methodology."""
+
+    thresholds: dict[int, Annotated[float, Field(ge=0.0, le=1.0)]] = Field(
+        default_factory=dict,
+        description=(
+            "Minimum score for an edge to count, per input, keyed by the input's "
+            "position. Write these as `{model: 0.9}`. `Resolver` takes the model "
+            "object and works out the position."
+        ),
+    )
+
+
+class Components(ResolverMethod):
+    """Resolver methodology that computes connected components.
+
+    A threshold defaults to 0.0 if not set.
+    """
+
+    resolver_type: ClassVar[ResolverType] = ResolverType.COMPONENTS
+    settings: ComponentsSettings
+
+    def compute_clusters(  # noqa: D102
+        self, model_edges: Mapping[int, pl.DataFrame]
+    ) -> pl.DataFrame:
+        djs = DisjointSet[int]()
+
+        for position, edges in model_edges.items():
+            if edges.height == 0:
+                continue
+
+            threshold = self.settings.thresholds.get(position, 0.0)
+            filtered_edges = edges.filter(pl.col("score") >= threshold)
+
+            for left_id, right_id in filtered_edges.select(
+                "left_id", "right_id"
+            ).iter_rows():
+                djs.union(left_id, right_id)
+
+        rows: list[dict[str, int]] = []
+        for parent_id, component in enumerate(djs.get_components(), start=1):
+            rows.extend(
+                {"parent_id": parent_id, "child_id": node_id} for node_id in component
+            )
+
+        if not rows:
+            return pl.from_arrow(SCHEMA_CLUSTERS.empty_table())
+
+        return pl.DataFrame(rows).cast(pl.Schema(SCHEMA_CLUSTERS))