matchlab 0.1.0__py3-none-any.whl

matchlab/__init__.py

@@ -0,0 +1,99 @@
+"""The matchlab plan API.
+
+Build a plan from a source, chain verbs, then `collect()`:
+
+```python
+import matchlab as mb
+
+crn = mb.read_database("crn", sql=..., client=engine, key_field="pk")
+deduped = crn.clean(...).dedupe(model_class=mb.NaiveDeduper, ...).resolve()
+lookup = deduped.collect().get_lookup()
+```
+
+Nothing runs until `collect()`, and collecting again only does the work whose plan
+or inputs changed.
+
+Everything needed to write a plan is re-exported here, including built-in methodologies.
+Evaluation stays behind `matchlab.eval`, since a plan runs without it.
+"""
+
+from importlib.metadata import version
+
+from matchlab.document import PlanDocument, dump, load
+from matchlab.models import Model, add_model_class
+from matchlab.models.dedupers import NaiveDeduper
+from matchlab.models.dedupers.base import Deduper
+from matchlab.models.linkers import (
+    DeterministicLinker,
+    SplinkLinker,
+    WeightedDeterministicLinker,
+)
+from matchlab.models.linkers.base import Linker
+from matchlab.recordstep import RecordStep
+from matchlab.resolvers import Components, Resolver
+from matchlab.resolvers.base import ResolverMethod
+from matchlab.resolvers.resolvers import add_resolver_class
+from matchlab.resources import FromResources, Resource
+from matchlab.sources import (
+    DataFrame,
+    Location,
+    RelationalDB,
+    Source,
+    add_location_class,
+    read_database,
+    read_dataframe,
+)
+from matchlab.steps import Step
+from matchlab.stores import DuckDBStore
+from matchlab.stores.default import default_store, set_default_store
+from matchlab.transformers import (
+    Clean,
+    Explode,
+    Group,
+    Select,
+    Transform,
+    Transformer,
+    add_transformer_class,
+)
+
+__version__ = version("matchlab")
+
+__all__ = (
+    "Clean",
+    "Components",
+    "DataFrame",
+    "Deduper",
+    "DeterministicLinker",
+    "DuckDBStore",
+    "Explode",
+    "FromResources",
+    "Group",
+    "Linker",
+    "Location",
+    "Model",
+    "NaiveDeduper",
+    "PlanDocument",
+    "RecordStep",
+    "RelationalDB",
+    "Resolver",
+    "ResolverMethod",
+    "Resource",
+    "Select",
+    "Source",
+    "SplinkLinker",
+    "Step",
+    "Transform",
+    "Transformer",
+    "WeightedDeterministicLinker",
+    "__version__",
+    "add_location_class",
+    "add_model_class",
+    "add_resolver_class",
+    "add_transformer_class",
+    "default_store",
+    "dump",
+    "load",
+    "read_database",
+    "read_dataframe",
+    "set_default_store",
+)

matchlab/cli.py

@@ -0,0 +1,177 @@
+"""The `matchlab` command.
+
+Two commands. A terminal does two things better than a Python prompt. It tells you
+what version you have, and it gives you a full-screen reviewer.
+
+`review` takes a `module:attribute` naming a resolver in your code. With `--store`,
+it instead takes the label of a resolver output already published. The second form
+needs no plan and no warehouse. A stored resolver output knows which source artifacts
+it covers, and their extracts hold the values. Name more than one target and you
+review their merged clusters.
+"""
+
+import argparse
+import importlib
+import logging
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from matchlab.resolvers import Resolver
+
+
+def _load_target(target: str) -> "Resolver":
+    """Import `module:attribute` and return the resolver it names.
+
+    The attribute may be a `Resolver`, or a callable returning one. The callable form
+    is useful when building the plan needs a warehouse connection you'd rather not
+    open at import time. A `Resolver` isn't callable, so the two never blur.
+
+    Raises:
+        SystemExit: If the target can't be imported or isn't a resolver.
+    """
+    from matchlab.resolvers import Resolver  # noqa: PLC0415 - keeps startup cheap
+
+    if ":" not in target:
+        raise SystemExit(
+            f"Target '{target}' should look like 'module:attribute', for example "
+            "'pipeline:entities', naming a Resolver in pipeline.py."
+        )
+
+    module_name, _, attribute = target.partition(":")
+
+    # Lets a script in the working directory be named directly.
+    if str(Path.cwd()) not in sys.path:
+        sys.path.insert(0, str(Path.cwd()))
+
+    try:
+        module = importlib.import_module(module_name)
+    except ImportError as exc:
+        raise SystemExit(f"Could not import '{module_name}': {exc}") from exc
+
+    try:
+        obj: Any = getattr(module, attribute)
+    except AttributeError as exc:
+        raise SystemExit(f"'{module_name}' has no attribute '{attribute}'") from exc
+
+    if callable(obj):
+        obj = obj()
+
+    if not isinstance(obj, Resolver):
+        raise SystemExit(
+            f"'{target}' is a {type(obj).__name__}, not a Resolver. Review needs the "
+            "resolver whose clusters you want to look at."
+        )
+
+    return obj
+
+
+def _redirect_logging(log_file: str) -> None:
+    """Send all logging to a file, so it doesn't scribble over the UI."""
+    root = logging.getLogger()
+    for handler in root.handlers[:]:
+        root.removeHandler(handler)
+
+    handler = logging.FileHandler(log_file, mode="w")
+    handler.setFormatter(
+        logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s")
+    )
+    root.addHandler(handler)
+    root.setLevel(logging.INFO)
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="matchlab",
+        description="Build, run and evaluate entity resolution pipelines.",
+    )
+    commands = parser.add_subparsers(dest="command", required=True)
+
+    commands.add_parser("version", help="Show the installed matchlab version.")
+
+    review = commands.add_parser(
+        "review",
+        help="Review a resolver's clusters interactively.",
+        description=(
+            "Open the cluster reviewer. TARGET names a Resolver as "
+            "'module:attribute', such as 'pipeline:companies' for a resolver assigned "
+            "to `companies` in pipeline.py, or a function returning one, if building "
+            "the plan needs a connection you'd rather open on demand. "
+            "With --store, TARGET is instead the label of a resolver output published "
+            "in that store. Then no Python is imported. The values shown come from the "
+            "extracts cached when the sources were collected, so no warehouse "
+            "connection is needed. "
+            "Name several targets to review their merged clusters, so one session's "
+            "judgements score all of them against the same records."
+        ),
+    )
+    review.add_argument(
+        "targets",
+        metavar="TARGET",
+        nargs="+",
+        help="module:attribute, or a resolver output's label when --store is given",
+    )
+    review.add_argument(
+        "-n", "--samples", type=int, default=5, help="Clusters to queue (default: 5)."
+    )
+    review.add_argument(
+        "-t", "--tag", help="Tag every judgement from this session, for later scoring."
+    )
+    review.add_argument(
+        "--seed",
+        type=int,
+        help="Fix which clusters are drawn, so someone else can review the same ones.",
+    )
+    review.add_argument(
+        "-s",
+        "--store",
+        help="DuckDB store to read from and write judgements to "
+        "(default: the shared store in your cache directory).",
+    )
+    review.add_argument(
+        "--log", dest="log_file", help="Send log output to this file instead."
+    )
+    return parser
+
+
+def main(argv: list[str] | None = None) -> None:
+    """Run the `matchlab` command."""
+    args = _build_parser().parse_args(argv)
+
+    if args.command == "version":
+        from importlib.metadata import version  # noqa: PLC0415
+
+        print(version("matchlab"))
+        return
+
+    if args.log_file:
+        _redirect_logging(args.log_file)
+
+    from matchlab.eval import review  # noqa: PLC0415
+    from matchlab.stores import DuckDBStore  # noqa: PLC0415
+
+    store = DuckDBStore(args.store) if args.store else None
+
+    # With a store, a bare word is a label on a stored resolver output, so it needs no
+    # plan.
+    targets: list[Any] = []
+    for target in args.targets:
+        if store is not None and ":" not in target:
+            known = store.labels()
+            if target not in known:
+                raise SystemExit(
+                    f"No resolver output is labelled '{target}' in {args.store}. "
+                    f"Known labels: {', '.join(known) or 'none'}."
+                )
+            targets.append(target)
+        else:
+            targets.append(_load_target(target))
+
+    review(
+        targets if len(targets) > 1 else targets[0],
+        n=args.samples,
+        store=store,
+        tag=args.tag,
+        seed=args.seed,
+    )

matchlab/core/__init__.py

@@ -0,0 +1 @@
+"""Common functions, tools and objects used throughout matchlab."""

matchlab/core/dataframes.py

@@ -0,0 +1,254 @@
+"""Dataframe types, conversion, column naming and SQL reading for matchlab."""
+
+import re
+from collections.abc import Callable, Iterator
+from enum import StrEnum
+from typing import (
+    Any,
+    Literal,
+    TypeAlias,
+    overload,
+)
+
+import polars as pl
+from adbc_driver_manager.dbapi import Connection as AdbcConnection
+from pandas import DataFrame as PandasDataFrame
+from polars import DataFrame as PolarsDataFrame
+from pyarrow import Table as ArrowTable
+from sqlalchemy.engine import Engine
+
+from matchlab.core.sql import SQLQuery
+
+
+class DataFrameType(StrEnum):
+    """Enumeration of dataframe types to return from query."""
+
+    PANDAS = "pandas"
+    POLARS = "polars"
+    ARROW = "arrow"
+
+
+DataFrameClass: TypeAlias = ArrowTable | PandasDataFrame | PolarsDataFrame
+
+# A source's name prefixes every column it contributes (`crn` gives `crn_company`),
+# and those land in cleaning SQL that sqlglot parses. `crn-x_company` would parse as
+# subtraction and `crn.x_company` as table.column, so the name has to be usable at the
+# start of an unquoted identifier. Reserved words are fine — the name is only ever a
+# prefix, never a bare identifier.
+_IDENTIFIER = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
+
+
+def validate_col_prefix(prefix: str) -> None:
+    """Ensure string is safe to use as a prefix in a dataframe column."""
+    if not _IDENTIFIER.match(prefix):
+        raise ValueError(
+            f"'{prefix}' can't prefix a column name. It must start with a letter or"
+            "underscore and contain only letters, digits and underscores."
+        )
+
+
+def qualify(source: str, field: str = "") -> str:
+    """The column name a source's field carries once qualified.
+
+    Args:
+        source: The source's name.
+        field: The field to qualify. Omit it for the bare prefix, which is what
+            `pl.all().name.prefix()` takes.
+
+    Returns:
+        The qualified column name.
+    """
+    return f"{source}_{field}"
+
+
+@overload
+def to_dataframe(
+    data: PolarsDataFrame, return_type: Literal[DataFrameType.POLARS]
+) -> PolarsDataFrame: ...
+
+
+@overload
+def to_dataframe(
+    data: PolarsDataFrame, return_type: Literal[DataFrameType.PANDAS]
+) -> PandasDataFrame: ...
+
+
+@overload
+def to_dataframe(
+    data: PolarsDataFrame, return_type: Literal[DataFrameType.ARROW]
+) -> ArrowTable: ...
+
+
+def to_dataframe(
+    data: PolarsDataFrame, return_type: DataFrameType = DataFrameType.POLARS
+) -> DataFrameClass:
+    """Convert a Polars dataframe to the requested dataframe type.
+
+    Args:
+        data: The frame to convert.
+        return_type: The dataframe type to return. One of "arrow", "pandas",
+            or "polars".
+
+    Returns:
+        The data in the requested format.
+
+    Raises:
+        ValueError: If `return_type` isn't a member of `DataFrameType`.
+    """
+    match return_type:
+        case DataFrameType.POLARS:
+            return data
+        case DataFrameType.PANDAS:
+            return data.to_pandas()
+        case DataFrameType.ARROW:
+            return data.to_arrow()
+        case _:
+            raise ValueError(f"Return type {return_type} is invalid")
+
+
+@overload
+def sql_to_df(
+    stmt: SQLQuery,
+    connection: Engine | AdbcConnection,
+    return_type: Literal[DataFrameType.POLARS],
+    *,
+    return_batches: Literal[False] = False,
+    batch_size: int | None = None,
+    rename: dict[str, str] | Callable | None = None,
+    schema_overrides: dict[str, pl.DataType] | None = None,
+    execute_options: dict[str, Any] | None = None,
+) -> PolarsDataFrame: ...
+
+
+@overload
+def sql_to_df(
+    stmt: SQLQuery,
+    connection: Engine | AdbcConnection,
+    return_type: Literal[DataFrameType.PANDAS],
+    *,
+    return_batches: Literal[False] = False,
+    batch_size: int | None = None,
+    rename: dict[str, str] | Callable | None = None,
+    schema_overrides: dict[str, pl.DataType] | None = None,
+    execute_options: dict[str, Any] | None = None,
+) -> PandasDataFrame: ...
+
+
+@overload
+def sql_to_df(
+    stmt: SQLQuery,
+    connection: Engine | AdbcConnection,
+    return_type: Literal[DataFrameType.ARROW],
+    *,
+    return_batches: Literal[False] = False,
+    batch_size: int | None = None,
+    rename: dict[str, str] | Callable | None = None,
+    schema_overrides: dict[str, pl.DataType] | None = None,
+    execute_options: dict[str, Any] | None = None,
+) -> ArrowTable: ...
+
+
+@overload
+def sql_to_df(
+    stmt: SQLQuery,
+    connection: Engine | AdbcConnection,
+    return_type: Literal[DataFrameType.POLARS],
+    *,
+    return_batches: Literal[True],
+    batch_size: int | None = None,
+    rename: dict[str, str] | Callable | None = None,
+    schema_overrides: dict[str, pl.DataType] | None = None,
+    execute_options: dict[str, Any] | None = None,
+) -> Iterator[PolarsDataFrame]: ...
+
+
+@overload
+def sql_to_df(
+    stmt: SQLQuery,
+    connection: Engine | AdbcConnection,
+    return_type: Literal[DataFrameType.PANDAS],
+    *,
+    return_batches: Literal[True],
+    batch_size: int | None = None,
+    rename: dict[str, str] | Callable | None = None,
+    schema_overrides: dict[str, pl.DataType] | None = None,
+    execute_options: dict[str, Any] | None = None,
+) -> Iterator[PandasDataFrame]: ...
+
+
+@overload
+def sql_to_df(
+    stmt: SQLQuery,
+    connection: Engine | AdbcConnection,
+    return_type: Literal[DataFrameType.ARROW],
+    *,
+    return_batches: Literal[True],
+    batch_size: int | None = None,
+    rename: dict[str, str] | Callable | None = None,
+    schema_overrides: dict[str, pl.DataType] | None = None,
+    execute_options: dict[str, Any] | None = None,
+) -> Iterator[ArrowTable]: ...
+
+
+def sql_to_df(
+    stmt: SQLQuery,
+    connection: Engine | AdbcConnection,
+    return_type: DataFrameType = DataFrameType.PANDAS,
+    *,
+    return_batches: bool = False,
+    batch_size: int | None = None,
+    rename: dict[str, str] | Callable | None = None,
+    schema_overrides: dict[str, pl.DataType] | None = None,
+    execute_options: dict[str, Any] | None = None,
+) -> DataFrameClass | Iterator[DataFrameClass]:
+    """Execute a SQL string with Polars, returning the requested dataframe type.
+
+    Args:
+        stmt: A SQL string to execute.
+        connection: A SQLAlchemy Engine or ADBC connection.
+        return_type: The dataframe type to return. One of "arrow", "pandas",
+            or "polars".
+        return_batches: If True, return an iterator that yields each batch
+            separately. If False, return a single dataframe with all results.
+        batch_size: The size of each batch, when processing in batches.
+        rename: A dictionary mapping old column names to new column names, or
+            a callable that takes a dataframe and returns a dataframe with
+            renamed columns.
+        schema_overrides: A dictionary mapping column names to dtypes.
+        execute_options: Passed through to the underlying query execution
+            method as kwargs.
+
+    Returns:
+        If return_batches is False: A dataframe of the query results in the specified
+            format.
+        If return_batches is True: An iterator of dataframes in the specified format.
+
+    Raises:
+        ValueError: If `batch_size` and `return_batches` aren't set together.
+    """
+    if not batch_size and return_batches:
+        raise ValueError("A batch size must be specified if return_batches is True")
+
+    if batch_size and not return_batches:
+        raise ValueError("Cannot set a batch size if return_batches is False")
+
+    def _to_format(results: PolarsDataFrame) -> DataFrameClass:
+        """Rename the results' columns, then convert to the specified format."""
+        if rename:
+            results = results.rename(rename)
+
+        return to_dataframe(results, return_type)
+
+    res = pl.read_database(
+        query=stmt,
+        connection=connection,
+        iter_batches=bool(batch_size),
+        batch_size=batch_size,
+        schema_overrides=schema_overrides,
+        execute_options=execute_options,
+    )
+
+    if return_batches:
+        return (_to_format(batch) for batch in res)
+
+    return _to_format(res)

matchlab/core/dsu.py

@@ -0,0 +1,62 @@
+"""Functions to transform data between tabular and graph structures."""
+
+from collections import defaultdict
+from collections.abc import Hashable
+from typing import Generic, TypeVar
+
+T = TypeVar("T", bound=Hashable)
+
+
+class DisjointSet(Generic[T]):
+    """Disjoint set forest with "path compression" and "union by rank" heuristics.
+
+    This follows implementation from Cormen, Thomas H., et al. Introduction to
+    algorithms. MIT press, 2022
+    """
+
+    def __init__(self) -> None:
+        """Initialise the disjoint set."""
+        self.parent: dict[T, T] = {}
+        self.rank: dict[T, int] = {}
+
+    def _make_set(self, x: T) -> None:
+        """Create a new set with a single element x."""
+        self.parent[x] = x
+        self.rank[x] = 0
+
+    def add(self, x: T) -> None:
+        """Add a new element to the disjoint set."""
+        if x not in self.parent:
+            self._make_set(x)
+
+    def union(self, x: T, y: T) -> None:
+        """Merge the sets containing elements x and y."""
+        self._link(self._find(x), self._find(y))
+
+    def _link(self, x: T, y: T) -> None:
+        """Merge the sets containing elements x and y."""
+        if self.rank[x] > self.rank[y]:
+            self.parent[y] = x
+        else:
+            self.parent[x] = y
+            if self.rank[x] == self.rank[y]:
+                self.rank[y] += 1
+
+    def _find(self, x: T) -> T:
+        """Return the representative element of the set containing x."""
+        if x not in self.parent:
+            self._make_set(x)
+            return x
+
+        if x != self.parent[x]:
+            self.parent[x] = self._find(self.parent[x])
+
+        return self.parent[x]
+
+    def get_components(self) -> list[set[T]]:
+        """Return the connected components of the disjoint set."""
+        components = defaultdict(set)
+        for x in self.parent:
+            root = self._find(x)
+            components[root].add(x)
+        return list(components.values())

matchlab/core/exceptions.py

@@ -0,0 +1,40 @@
+"""Exceptions raised by matchlab."""
+
+from pyarrow import Schema
+
+
+class MatchlabError(Exception):
+    """An error occurred in matchlab."""
+
+    def __init__(self, message: str | None = None) -> None:
+        """Initialise the error, defaulting the message to the class docstring."""
+        super().__init__(message or self.__doc__)
+
+
+class SchemaMismatch(MatchlabError):
+    """An Arrow table did not have the expected schema."""
+
+    def __init__(self, expected: Schema, actual: Schema) -> None:
+        """Initialise with the schemas that failed to match."""
+        super().__init__(f"Schema mismatch. Expected:\n{expected}\nGot:\n{actual}")
+
+
+class SourceTableError(MatchlabError):
+    """A source's table could not be read from its location."""
+
+
+class StepNotFound(MatchlabError):
+    """A step could not be found."""
+
+    def __init__(self, message: str | None = None, name: str | None = None) -> None:
+        """Initialise with an explicit message, or a name to build one from."""
+        if message is None:
+            message = (
+                f"Step {name} not found." if name is not None else "Step not found."
+            )
+        super().__init__(message)
+        self.name = name
+
+
+class ResourceError(ValueError):
+    """A resource could not be named, resolved, or bound unambiguously."""