matchlab 0.1.0.dev1__py3-none-any.whl

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.adapters import DuckDBAdapter  # noqa: PLC0415
+    from matchlab.eval import review  # noqa: PLC0415
+
+    adapter = DuckDBAdapter(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 adapter is not None and ":" not in target:
+            known = adapter.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,
+        adapter=adapter,
+        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 ExtractTransformError(MatchlabError):
+    """A source's extract/transform SQL is not valid."""
+
+
+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

matchlab/core/hash.py

@@ -0,0 +1,141 @@
+"""Utilities for hashing data and creating unique identifiers."""
+
+import hashlib
+from enum import StrEnum
+
+import polars as pl
+import polars.expr as plx
+import polars_hash as plh
+import pyarrow as pa
+
+HASH_FUNC = hashlib.sha256
+
+
+class HashMethod(StrEnum):
+    """Supported hash methods for row hashing."""
+
+    XXH3_128 = "xxh3_128"
+    SHA256 = "sha256"
+
+
+def _process_column_for_hashing(column_name: str, schema_type: pl.DataType) -> plx.Expr:
+    r"""Normalise `column_name` to a string column safe to hash.
+
+    Nulls become the sentinel `"\x00"`, so a null hashes differently from an empty
+    value. Binary columns hex-encode, structs JSON-encode and lists join on `,`,
+    before falling back to a plain string cast.
+    """
+    if isinstance(schema_type, pl.Binary):
+        return (
+            pl.col(column_name).fill_null("\x00").bin.encode("hex").alias(column_name)
+        )
+    elif isinstance(schema_type, pl.Struct):
+        return (
+            pl.col(column_name)
+            .struct.json_encode()
+            .fill_null("\x00")
+            .alias(column_name)
+        )
+    elif isinstance(schema_type, pl.List):
+        return pl.col(column_name).list.join(",").fill_null("\x00").alias(column_name)
+    else:
+        return pl.col(column_name).cast(pl.Utf8).fill_null("\x00").alias(column_name)
+
+
+def hash_rows(
+    df: pl.DataFrame, columns: list[str], method: HashMethod = HashMethod.XXH3_128
+) -> pl.Series:
+    """Hash each row of `df` over `columns`.
+
+    Each row's hash covers both column names and values, joined with the record and
+    unit separator symbols (`␞`, `␟`), so a column's name is part of what gets hashed
+    alongside its value.
+    """
+    expr_list = [
+        _process_column_for_hashing(column, df.schema[column]) for column in columns
+    ]
+    df_processed = df.with_columns(expr_list)
+
+    record_separator = "␞"
+    unit_separator = "␟"
+
+    str_concatenation: list[pl.Expr] = []
+    for c in columns:
+        str_concatenation.extend(
+            [
+                pl.lit(c),  # column name
+                pl.lit(unit_separator),
+                pl.col(c),  # column value
+                pl.lit(record_separator),
+            ]
+        )
+
+    if method == HashMethod.XXH3_128:
+        row_hashes = df_processed.select(
+            plh.concat_str(str_concatenation).nchash.xxh3_128().alias("row_hash")
+        )
+        return row_hashes["row_hash"]
+    elif method == HashMethod.SHA256:
+        row_hashes = df_processed.select(
+            plh.concat_str(str_concatenation)
+            .chash.sha2_256()
+            .str.decode("hex")
+            .alias("row_hash")
+        )
+        return row_hashes["row_hash"]
+    else:
+        raise ValueError(f"Unsupported hash method: {method}")
+
+
+def hash_arrow_table(
+    table: pa.Table,
+    method: HashMethod = HashMethod.XXH3_128,
+    as_sorted_list: list[str] | None = None,
+) -> bytes:
+    """Content-address `table` with a hash invariant to row and column order.
+
+    Pass `as_sorted_list` (2 or more column names, e.g. `["left_id", "right_id"]`) to
+    hash those columns as a sorted set instead of individually, so `(1, 2)` and
+    `(2, 1)` hash the same. This replaces the named columns with one `sorted_list`
+    column.
+
+    Combining `as_sorted_list` with a nullable column can null the whole row's hash
+    input, because Polars' `concat_list` returns null when any input is null.
+
+    Raises:
+        ValueError: If `as_sorted_list` names fewer than two columns, or a column
+            `table` doesn't have.
+    """
+    df = pl.from_arrow(table)
+
+    if df.height == 0:
+        return b"empty_table_hash"
+
+    if as_sorted_list:
+        if len(as_sorted_list) < 2:
+            raise ValueError(
+                "Lists passed to as_sorted_list must contain at least 2 column names"
+            )
+
+        missing_cols = [col for col in as_sorted_list if col not in df.columns]
+        if missing_cols:
+            raise ValueError(f"Columns not found in dataframe: {missing_cols}")
+
+        df = df.with_columns(
+            pl.concat_list(as_sorted_list).list.sort().alias("sorted_list")
+        ).drop(as_sorted_list)
+
+    columns: list[str] = sorted(df.columns)
+    df = df.select(columns)
+
+    # Explode list columns so each element sorts and hashes on its own. That's what
+    # makes the hash invariant to the order of elements within a list field.
+    for column in columns:
+        if isinstance(df.schema[column], pl.List):
+            df = df.explode(column, empty_as_null=True)
+
+    df = df.sort(by=columns)
+    row_hashes = hash_rows(df=df, columns=columns, method=method)
+    all_hashes: bytes = b"".join(row_hashes.sort().to_list())
+
+    return HASH_FUNC(all_hashes).digest()

matchlab/core/kinds.py

@@ -0,0 +1,12 @@
+"""The kinds of step a plan is made of."""
+
+from enum import StrEnum
+
+
+class StepKind(StrEnum):
+    """Enumeration of the kinds of step a plan is made of."""
+
+    SOURCE = "source"
+    TRANSFORM = "transform"
+    MODEL = "model"
+    RESOLVER = "resolver"