matchlab 0.1.0.dev1__py3-none-any.whl

matchlab/sources.py

@@ -0,0 +1,291 @@
+"""Source — the leaf of a plan.
+
+A source reads rows from a warehouse and content-addresses them. It takes no inputs,
+so it is where raw data, and therefore non-determinism, enters a plan. Its spec key
+includes a hash of the data it read, which is what makes a freshly constructed `Source`
+pick up warehouse changes while an existing object memoises its read.
+
+**The extract/transform is the whole declaration.** Every column it returns is part of
+the record, so every column except the key contributes to that record's identity. Two
+rows are the same leaf when they agree on all of them. A column
+you do not want to affect identity is a column you should not select, and a type you
+want pinned is a `cast` in the SELECT.
+"""
+
+import tempfile
+from collections.abc import Callable, Generator, Iterable
+from pathlib import Path
+from typing import ClassVar
+
+import polars as pl
+import pyarrow.parquet as pq
+
+from matchlab.adapters import Adapter, Fingerprint
+from matchlab.core.dataframes import (
+    DataFrameClass,
+    DataFrameType,
+    qualify,
+    validate_col_prefix,
+)
+from matchlab.core.hash import HashMethod, hash_arrow_table, hash_rows
+from matchlab.core.kinds import StepKind
+from matchlab.core.logging import logger
+from matchlab.core.resolver_output import leaf_id
+from matchlab.core.sql import SQLQuery
+from matchlab.locations import Location
+from matchlab.recordstep import IdentifierRead, RecordStep, build_record_step
+from matchlab.specs import SourceSpec
+
+
+class Source(RecordStep):
+    """A warehouse table, extracted and content-addressed.
+
+    A source is a `RecordStep`. Read directly, its records are its extract, with `id`
+    set to each row's content-addressed leaf, so a model can match over it with no
+    intervening step.
+    """
+
+    kind: ClassVar[StepKind] = StepKind.SOURCE
+
+    def __init__(
+        self,
+        location: Location,
+        name: str,
+        extract_transform: SQLQuery,
+        key_field: str,
+        validate_etl: bool = True,
+    ) -> None:
+        """Define a source.
+
+        Args:
+            location: Where the data lives.
+            name: The source's name within the plan. Prefixes every column this
+                source contributes, so it must be usable in a SQL identifier.
+            extract_transform: SQL producing the rows to index. Every column it
+                returns other than the key contributes to record identity.
+            key_field: The name of the unique identifier field. Read as a string
+                whatever the warehouse returns it as.
+            validate_etl: Validate the extract/transform SQL up front.
+
+        Raises:
+            ValueError: If the name could not prefix a SQL identifier, or if
+                `key_field` is not a column name.
+        """
+        validate_col_prefix(name)
+
+        # A source's name is not a way of finding it later, the way a published
+        # resolver output's is. It prefixes every column this source contributes and
+        # tags its rows in a resolver's output, so it is part of the output. That is
+        # why it is required, lives on the source itself, and is hashed into the
+        # fingerprint.
+        self.name = name
+        super().__init__()
+
+        if not isinstance(key_field, str):
+            raise ValueError(
+                f"key_field must be a column name, got {type(key_field).__name__}"
+            )
+
+        if validate_etl:
+            location.validate_extract_transform(extract_transform)
+
+        self.location = location
+        self.extract_transform = extract_transform
+        self.key_field = key_field
+
+        # Memoised warehouse read: (extract, hashes). Populated on first fingerprint.
+        self._read: tuple[pl.DataFrame, pl.DataFrame] | None = None
+
+    # -- spec -------------------------------------------------------------------------
+
+    def __str__(self) -> str:
+        """A source is drawn with its name, since it has one."""
+        return f"{self.kind} '{self.name}'"
+
+    @property
+    def spec(self) -> SourceSpec:
+        """The serialisable spec for this source."""
+        return SourceSpec(
+            name=self.name,
+            extract_transform=self.extract_transform,
+            key_field=self.key_field,
+        )
+
+    # -- column naming ----------------------------------------------------------------
+    #
+    # A record step built over several sources qualifies every column with the source
+    # it came from, so `company` from `crn` becomes `crn_company`. These say what a
+    # column will be called once that has happened, which is how you refer to it in a
+    # cleaning expression or a model's settings, before anything has been collected.
+
+    def f(self, fields: str | Iterable[str]) -> str | list[str]:
+        """Prefix one or more field names with this source's name."""
+        if isinstance(fields, str):
+            return qualify(self.name, fields)
+        return [qualify(self.name, field) for field in fields]
+
+    @property
+    def qualified_key(self) -> str:
+        """This source's key field, prefixed with the source name."""
+        return qualify(self.name, self.key_field)
+
+    @property
+    def index_fields(self) -> list[str]:
+        """Every column the extract returns except the key, in sorted order.
+
+        Read from the warehouse rather than declared, so it cannot drift from the SQL.
+        """
+        extract, _ = self._read_warehouse()
+        return sorted(column for column in extract.columns if column != self.key_field)
+
+    # -- warehouse access -------------------------------------------------------------
+
+    def fetch(
+        self,
+        qualify_names: bool = False,
+        batch_size: int | None = None,
+        return_type: DataFrameType = DataFrameType.POLARS,
+        keys: list[str] | None = None,
+    ) -> Generator[DataFrameClass, None, None]:
+        """Apply the extract/transform and yield the resulting rows in batches."""
+        rename: Callable[[str], str] | None = None
+        if qualify_names:
+
+            def rename(column: str) -> str:
+                return qualify(self.name, column)
+
+        # Keys are always strings, whatever the warehouse returns them as. Every
+        # other column is read as the warehouse types it. Pin one with a `cast` in
+        # the extract/transform if you need it stable across drivers.
+        selection = (self.key_field, keys) if keys else None
+        yield from self.location.execute(
+            extract_transform=self.extract_transform,
+            schema_overrides={self.key_field: pl.String},
+            rename=rename,
+            batch_size=batch_size,
+            return_type=return_type,
+            **({"keys": selection} if selection else {}),
+        )
+
+    def sample(
+        self, n: int = 100, return_type: DataFrameType = DataFrameType.POLARS
+    ) -> DataFrameClass:
+        """Peek at the first `n` rows without collecting."""
+        return next(self.fetch(batch_size=n, return_type=return_type))
+
+    def _read_warehouse(self) -> tuple[pl.DataFrame, pl.DataFrame]:
+        """Read the source and content-address every row. Memoised.
+
+        Every column except the key contributes to the row hash, so two rows are the
+        same leaf exactly when the extract returns identical values for both. Identity
+        is content, not the key. See `matchlab.core.resolver_output.leaf_id` for why
+        that is the correct anchor for evaluation rather than a storage trick.
+
+        Returns:
+            `(extract, hashes)`, the raw rows, and a `hash → keys` index in which
+            byte-identical rows share a hash, so records a judge cannot tell apart share
+            a leaf.
+        """
+        if self._read is not None:
+            return self._read
+
+        logger.info("Reading from the warehouse", prefix=self.name)
+        key = self.key_field
+
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / f"{self.name}.parquet"
+            writer = None
+            for batch in self.fetch(return_type=DataFrameType.ARROW):
+                if writer is None:
+                    writer = pq.ParquetWriter(path, schema=batch.schema)
+                writer.write_table(batch)
+            if writer is None:
+                raise ValueError(f"Source '{self.name}' returned no rows")
+            writer.close()
+
+            frames: list[pl.DataFrame] = []
+            hashed: list[pl.DataFrame] = []
+            parquet = pq.ParquetFile(path)
+            for group in range(parquet.num_row_groups):
+                batch = pl.from_arrow(parquet.read_row_group(group))
+                if batch[key].is_null().any():
+                    raise ValueError(f"Source '{self.name}' has null keys")
+                frames.append(batch)
+                index = sorted(column for column in batch.columns if column != key)
+                if not index:
+                    raise ValueError(
+                        f"Source '{self.name}' returns only its key field. The "
+                        "extract/transform must select at least one other column."
+                    )
+                # XXH3_128 rather than SHA-256. `leaf_id`
+                # slices the row hash to 8 bytes, so a cryptographic digest's strength
+                # is discarded before it identifies anything, and the 128 bits that
+                # survive into `_spec_key` are ample for telling rows of one table
+                # apart. It is ~7x faster, on a path every collect pays, because
+                # `_spec_key` re-reads and re-hashes the warehouse even when the
+                # artifact is cached.
+                row_hashes = hash_rows(
+                    df=batch, columns=index, method=HashMethod.XXH3_128
+                )
+                hashed.append(
+                    batch.select(pl.col(key).alias("keys")).with_columns(
+                        row_hashes.alias("hash")
+                    )
+                )
+
+        extract = pl.concat(frames)
+        hashes = pl.concat(hashed).group_by("hash").agg(pl.col("keys"))
+        self._read = (extract, hashes)
+        return self._read
+
+    def leaves(self) -> pl.DataFrame:
+        """Return `(key, leaf)`, each source key mapped to its leaf cluster."""
+        _, hashes = self._read_warehouse()
+        return hashes.explode("keys", empty_as_null=True).select(
+            pl.col("keys").alias("key"),
+            leaf_id(pl.col("hash")).alias("leaf"),
+        )
+
+    # -- Step contract ----------------------------------------------------------------
+
+    def _spec_key(self) -> bytes:
+        """The spec plus a content hash of the data read.
+
+        Including the data hash is what lets a plan detect that the warehouse changed.
+        A new `Source` object re-reads and gets a different key, invalidating
+        everything downstream of it.
+
+        Hashing `hashes` alone is sufficient because every non-key column feeds the row
+        hash, and `hashes` records which keys carry each one. A change to any selected
+        column therefore moves the fingerprint. This was not true while a separate list
+        of index fields could be narrower than the extract. A column outside it changed
+        the stored extract without touching the fingerprint, so the source cache-hit,
+        never re-stored, and downstream record steps kept reading the stale value.
+        """
+        _, hashes = self._read_warehouse()
+        return super()._spec_key() + hash_arrow_table(hashes.to_arrow())
+
+    def _execute(self, adapter: Adapter, fp: Fingerprint) -> None:
+        extract, _ = self._read_warehouse()
+        adapter.store_source(
+            fp=fp,
+            key_field=self.key_field,
+            extract=extract,
+            leaves=self.leaves(),
+        )
+
+    # -- RecordStep contract -------------------------------------------------------
+
+    def _read_cache(self, adapter: Adapter) -> pl.DataFrame:
+        """Return this source's records with each row's leaf as `id`.
+
+        Derived from the extract and leaves stored by `_execute`, a cheap join rather
+        than a separately cached artifact. Reading a source on its own means `id` is
+        the leaf, so there is no resolver in the read.
+        """
+        return build_record_step(adapter, (self,), None)
+
+    @property
+    def _identifier_reads(self) -> tuple[IdentifierRead, ...]:
+        """Read directly, so `id` is the leaf and no resolver enters the read."""
+        return ((self._fp, self.name, None),)

matchlab/specs.py

@@ -0,0 +1,129 @@
+"""Step specs — what a fingerprint hashes.
+
+One model per step kind, and `Step._spec_key` hashes it to address that step's artifact.
+What belongs in a spec is therefore decided entirely by the fingerprint invariant.
+**It must carry everything the step's output depends on, and nothing else.** Omit
+something that changes the output, and `collect` hands back a stale artifact without
+re-running. Include something that doesn't, and the step re-runs for nothing, along
+with everything below it.
+
+That rules out two things. **Edges**, because `_fingerprint` already folds in the
+parents' fingerprints. Describing an input here would record it twice, and a rename
+would then invalidate a subtree producing identical bytes. **Names** are also excluded,
+with one exception: `SourceSpec.name` is in the output it prefixes. A setting that must
+point at an input points at its position instead (`ComponentsSettings.thresholds`).
+
+A spec is therefore not the serialisation format. Rebuilding a plan needs the edges
+that identifying a step must exclude. `matchlab.document` carries them alongside it.
+"""
+
+import textwrap
+from enum import StrEnum
+from typing import Any
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+)
+
+from matchlab.core.sql import SQLQuery
+
+
+class ModelType(StrEnum):
+    """Whether a model is a deduper or a linker."""
+
+    LINKER = "linker"
+    DEDUPER = "deduper"
+
+
+class SourceSpec(BaseModel):
+    """Specification of a source: what it extracts, and what keys it.
+
+    There is no separate list of indexed fields. The extract/transform is the single
+    declaration of what a source *is*. Every column it returns is part of the record,
+    and therefore part of that record's identity. A column you do not want to affect
+    identity is a column you should not select.
+
+    **Not where the rows came from.** A source's key already folds in a content hash of
+    what it actually read, so anything a location could change about the output is
+    caught by the data itself.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    name: str = Field(
+        description=(
+            "The source's name within the plan. Part of the spec because it is "
+            "part of the output: it prefixes every column this source contributes "
+            "and tags its rows in a resolver's output."
+        )
+    )
+    extract_transform: SQLQuery = Field(
+        description=(
+            "Logic to extract and transform data from the source. "
+            "Language is location dependent."
+        )
+    )
+    key_field: str = Field(
+        description=textwrap.dedent("""
+            The name of the key field. This is the source's key for unique
+            entities, such as a primary key in a relational database.
+
+            Keys are always read as strings, whatever the warehouse returns.
+
+            For example, if the source describes companies, it may have used
+            a Companies House number as its key.
+
+            This key is ALWAYS correct. It should be something generated and
+            owned by the source being indexed.
+
+            For example, your organisation's CRM ID is a key field within the CRM.
+
+            A CRM ID entered by hand in another dataset shouldn't be used
+            as a key field.
+        """),
+    )
+
+
+class TransformSpec(BaseModel):
+    """Specification of a transform, naming the transformer that reshapes the step."""
+
+    model_config = ConfigDict(frozen=True)
+
+    transformer_class: str = Field(
+        description="The registered name of the Transformer subclass."
+    )
+    transformer_settings: dict[str, Any] = Field(
+        description="That transformer's configuration, dumped to JSON."
+    )
+
+
+class ModelSpec(BaseModel):
+    """Specification of a deduper or linker."""
+
+    model_config = ConfigDict(frozen=True)
+
+    model_type: ModelType = Field(description="Whether this is a deduper or a linker.")
+    model_class: str = Field(
+        description="The registered name of the Deduper or Linker subclass."
+    )
+    model_settings: dict[str, Any] = Field(
+        description="That class's settings, dumped to JSON."
+    )
+
+
+class ResolverSpec(BaseModel):
+    """Specification of a resolver over one or more models."""
+
+    model_config = ConfigDict(frozen=True)
+
+    resolver_class: str = Field(
+        description="The registered name of the ResolverMethod subclass."
+    )
+    resolver_settings: dict[str, Any] = Field(
+        description=(
+            "That class's settings, dumped to JSON. Settings that point at one of the "
+            "resolver's inputs key by the input's position, so no name appears here."
+        )
+    )