matchlab 0.1.0.dev1__py3-none-any.whl

matchlab/testkit/sources.py

@@ -0,0 +1,327 @@
+"""One generated source: a real plan node, its rows, and the partition they imply."""
+
+from collections.abc import Callable
+from functools import cache, wraps
+from typing import Any, ParamSpec, Self, TypeVar
+
+import polars as pl
+import pyarrow as pa
+from faker import Faker
+from pydantic import BaseModel, ConfigDict, Field
+from sqlalchemy import Engine, create_engine
+from sqlglot import cast, select
+from sqlglot.expressions import column
+
+from matchlab.locations import RelationalDBLocation
+from matchlab.sources import Source
+from matchlab.testkit._generate import generate_entities, generate_source
+from matchlab.testkit.entities import Cluster, EntityReference, TrueEntity
+from matchlab.testkit.features import FeatureConfig
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+def make_features_hashable(func: Callable[P, R]) -> Callable[P, R]:
+    """Let callers configure `source_factory` with plain dicts.
+
+    Converts each dict to a `FeatureConfig` before the wrapped function runs, so
+    `source_factory` stays hashable, which `@cache` needs, while its callers don't have
+    to build `FeatureConfig` objects by hand.
+    """
+
+    @wraps(func)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        # Handle features in first positional arg
+        if args and args[0] is not None:
+            if isinstance(args[0][0], dict):
+                args = (tuple(FeatureConfig(**d) for d in args[0]),) + args[1:]
+            else:
+                args = (tuple(args[0]),) + args[1:]
+
+        # Handle features in kwargs
+        if "features" in kwargs and kwargs["features"] is not None:
+            if isinstance(kwargs["features"][0], dict):
+                kwargs["features"] = tuple(
+                    FeatureConfig(**d) for d in kwargs["features"]
+                )
+            else:
+                kwargs["features"] = tuple(kwargs["features"])
+
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+class GeneratedSource(BaseModel):
+    """Generated rows, and the `Source` plan node that reads them.
+
+    Sources that share entities are what make linking testable. Use
+    `linked_sources_factory` in `linked` instead when you need that.
+
+    This exposes what it knows about the *fixture*: the rows, the features they were
+    generated from, and the partition they imply. Anything you do to the *plan* goes
+    through `.source`, which is the node itself: `source.source.clean(...)`,
+    `source.source.name`, `source.source.spec`. There is deliberately no shortcut. A
+    testkit that forwarded those would be indistinguishable from the node it wraps, and
+    knowing which one you are holding is the whole point.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    source: Source = Field(
+        description="The Source object containing the spec and convenience methods."
+    )
+    features: tuple[FeatureConfig, ...] | None = Field(
+        description=(
+            "The features used to generate the data. "
+            "If None, the source data was not generated, but set manually."
+        ),
+        default=None,
+    )
+    data: pa.Table = Field(
+        description="Data corresponding to the output of queries, without leaf data."
+    )
+    input_clusters: tuple[Cluster, ...] = Field(
+        description=(
+            "The partition a matcher starts from: generated rows grouped by identical "
+            "content. Not the answer, that is LinkedSources.true_entities."
+        )
+    )
+
+    @property
+    def field_names(self) -> list[str]:
+        """The non-key columns of this testkit's data, in generation order.
+
+        Taken from the testkit rather than from `Source.index_fields`, which would have
+        to read the warehouse. These names are needed before a plan is even built.
+        Falls back to the data's own columns when the source was set manually rather
+        than generated from features.
+        """
+        if self.features is not None:
+            return [feature.name for feature in self.features]
+        return [
+            column for column in self.data.column_names if column not in ("key", "id")
+        ]
+
+    def write_to_location(self, client: Any | None = None) -> Self:  # noqa: ANN401
+        """Write the data to the source's location.
+
+        Args:
+            client: Write through this client instead of the source's own, rebuilding
+                the source's location around it. Testkits are generated against an
+                in-memory SQLite engine by default, so a test that wants a real
+                warehouse repoints one here.
+        """
+        if client:
+            self.source.location = type(self.source.location)(
+                name=self.source.location.name, client=client
+            )
+
+        pl.from_arrow(self.data).write_database(
+            table_name=self.source.name,
+            connection=self.source.location.client,
+            if_table_exists="replace",
+        )
+
+        return self
+
+
+def values_of(
+    entity: TrueEntity | Cluster, sources: dict[str, GeneratedSource]
+) -> dict[str, dict[str, list[str]]]:
+    """Every unique value this entity took, per source and per feature.
+
+    Args:
+        entity: The entity whose rows to look up.
+        sources: The generated sources to look in, by name.
+
+    Returns:
+        `{source: {feature: sorted unique values}}`. Each source may vary the same base
+        value differently, so they are kept apart.
+
+    Raises:
+        ValueError: If the entity references a source that was not supplied.
+    """
+    values: dict[str, dict[str, list[str]]] = {}
+
+    for source_name, keys in entity.keys.items():
+        source = sources.get(source_name)
+        if source is None:
+            raise ValueError(f"Source testkit not found: {source_name}")
+
+        rows = pl.DataFrame(source.data).filter(pl.col("key").is_in(list(keys)))
+        values[source_name] = {
+            feature.name: sorted(rows[feature.name].unique().to_list())
+            for feature in source.features
+        }
+
+    return values
+
+
+@make_features_hashable
+@cache
+def source_factory(
+    features: list[FeatureConfig] | list[dict] | None = None,
+    name: str | None = None,
+    location_name: str = "dbname",
+    engine: Engine | None = None,
+    n_true_entities: int = 10,
+    repetition: int = 0,
+    seed: int = 42,
+) -> GeneratedSource:
+    """Generate a complete source testkit from configured features.
+
+    Sources created with the factory system can only use a RelationalDBLocation,
+    and the data at that location will be stored in a single table.
+
+    Args:
+        features: List of FeatureConfig objects or dictionaries to use for generating
+            the source data. If None, defaults to a set of common features.
+        name: Name of the source. If None, a unique name is generated. This will be
+            used as the name of the table in the RelationalDBLocation, but also in
+            the str for the source.
+        location_name: Name of the location for the source.
+        engine: SQLAlchemy engine to use for the source's RelationalDBLocation. If
+            None, an in-memory SQLite engine is created.
+        n_true_entities: Number of true entities to generate. Defaults to 10.
+        repetition: Number of times to repeat the generated data. Defaults to 0.
+        seed: Random seed for reproducibility. Defaults to 42.
+    """
+    generator = Faker()
+    generator.seed_instance(seed)
+
+    if features is None:
+        features = (
+            FeatureConfig(
+                name="company_name",
+                base_generator="company",
+            ),
+            FeatureConfig(
+                name="crn",
+                base_generator="bothify",
+                parameters=(("text", "???-###-???-###"),),
+            ),
+        )
+
+    if name is None:
+        name = generator.unique.word()
+
+    if engine is None:
+        engine = create_engine("sqlite:///:memory:")
+
+    # Generate base entities
+    base_entities = generate_entities(
+        generator=generator,
+        features=features,
+        n=n_true_entities,
+    )
+
+    # Generate data using the base entities
+    data, entity_keys, row_keys = generate_source(
+        generator=generator,
+        n_true_entities=n_true_entities,
+        features=features,
+        repetition=repetition,
+        seed_entities=base_entities,
+    )
+
+    # Create true entities with references
+    true_entities = []
+    for entity in base_entities:
+        keys = entity_keys.get(entity.id, [])
+        if keys:
+            entity.add_source_reference(name, keys)
+            true_entities.append(entity)
+
+    # Create Cluster objects from row_keys
+    clusters = [
+        Cluster(
+            id=row_id,
+            keys=EntityReference({name: frozenset(keys)}),
+        )
+        for row_id, keys in row_keys.items()
+    ]
+
+    # Create source config
+    source = Source(
+        location=RelationalDBLocation(name=location_name, client=engine),
+        name=name,
+        extract_transform=select(
+            cast(column("key"), "string").as_("key"),
+            *[column(feature.name) for feature in features],
+        )
+        .from_(name)
+        .sql(),
+        key_field="key",
+    )
+
+    return GeneratedSource(
+        source=source,
+        features=features,
+        data=data,
+        input_clusters=tuple(sorted(clusters)),
+    )
+
+
+def source_from_tuple(
+    data_tuple: tuple[dict[str, Any], ...],
+    data_keys: tuple[Any],
+    name: str | None = None,
+    location_name: str = "dbname",
+    engine: Engine | None = None,
+    seed: int = 42,
+) -> GeneratedSource:
+    """Generate a complete source testkit from dummy data."""
+    generator = Faker()
+    generator.seed_instance(seed)
+
+    if name is None:
+        name = generator.unique.word()
+
+    if engine is None:
+        engine = create_engine("sqlite:///:memory:")
+
+    base_entities = tuple(TrueEntity(base_values=row) for row in data_tuple)
+
+    # Create true entities with references
+    true_entities: list[TrueEntity] = []
+    for entity, key in zip(base_entities, data_keys, strict=True):
+        entity.add_source_reference(name, [key])
+        true_entities.append(entity)
+    # A list, not a set: these are zipped against `data_keys` below, so they must stay
+    # in row order. A set comprehension paired each key with an arbitrary entity.
+    entity_ids = [entity.id for entity in true_entities]
+
+    # Create Cluster objects from row_keys
+    clusters = [
+        Cluster(
+            id=entity_id,
+            keys=EntityReference({name: frozenset([key])}),
+        )
+        for key, entity_id in zip(data_keys, entity_ids, strict=True)
+    ]
+
+    # Create source config
+    source = Source(
+        location=RelationalDBLocation(name=location_name, client=engine),
+        name=name,
+        extract_transform=select(
+            cast(column("key"), "string").as_("key"),
+            *[column(field) for field in data_tuple[0]],
+        )
+        .from_(name)
+        .sql(),
+        key_field="key",
+    )
+
+    raw_data = pa.Table.from_pylist(list(data_tuple))
+    raw_keys = pa.array(data_keys)
+
+    data = raw_data.append_column("id", [entity_ids]).append_column("key", raw_keys)
+
+    return GeneratedSource(
+        source=source,
+        data=data,
+        input_clusters=tuple(sorted(clusters)),
+    )

matchlab/transformers/__init__.py

@@ -0,0 +1,24 @@
+"""Pluggable, serialisable reshapings of a record step.
+
+`Select` projects, `Clean` derives, `Group` collapses each `id` to one row, and
+`Explode` cross-joins each `id`'s columns back out to every combination, each one job.
+Register a custom one with `add_transformer_class`, as `add_model_class` does a
+deduper.
+"""
+
+from matchlab.transformers.base import Transformer
+from matchlab.transformers.clean import Clean
+from matchlab.transformers.explode import Explode
+from matchlab.transformers.group import Group
+from matchlab.transformers.select import Select
+from matchlab.transformers.transform import Transform, add_transformer_class
+
+__all__ = (
+    "Clean",
+    "Explode",
+    "Group",
+    "Select",
+    "Transform",
+    "Transformer",
+    "add_transformer_class",
+)

matchlab/transformers/base.py

@@ -0,0 +1,41 @@
+"""Base class for transformer methodologies, and the DuckDB query runner some share.
+
+A `Transformer` is a pure function of a record step, called by a `Transform` step's
+`apply()` on every collect. Both input and output must carry `id`, the grouping every
+downstream model and resolver reads.
+
+Transformers are declarative and serialisable. Their fields are the configuration a
+`Transform` folds into its spec and cache key, so keep them to plain data, with no
+callables, exactly as a `Deduper`'s settings are.
+"""
+
+from abc import ABC, abstractmethod
+
+import duckdb
+import polars as pl
+from pydantic import BaseModel, ConfigDict
+
+from matchlab.core.sql import SQLQuery
+
+
+class Transformer(BaseModel, ABC):
+    """Base contract every transformer implements.
+
+    Concrete transformers (`Select`, `Clean`, `Group`, `Explode`) carry their
+    configuration as flat fields, so `MyTransformer(...)` reads naturally, and
+    `model_dump(mode="json")` is the whole of its serialisation.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    @abstractmethod
+    def apply(self, data: pl.DataFrame) -> pl.DataFrame:
+        """Return `data` reshaped. Both the input and the output carry `id`."""
+        ...
+
+
+def run_sql(query: SQLQuery, data: pl.DataFrame) -> pl.DataFrame:
+    """Execute one DuckDB query against `data`, registered as `data`."""
+    with duckdb.connect(":memory:") as connection:
+        connection.register("data", data)
+        return connection.execute(query).pl()

matchlab/transformers/clean.py

@@ -0,0 +1,48 @@
+"""Clean derives columns. Add or replace named columns, keep the rest."""
+
+import polars as pl
+from pydantic import Field, field_validator
+from sqlglot import expressions, parse_one
+from sqlglot import select as sqlglot_select
+
+from matchlab.core.sql import SQLExpression
+from matchlab.transformers.base import Transformer, run_sql
+
+
+class Clean(Transformer):
+    """Derive columns with DuckDB SQL, without dropping unrelated fields.
+
+    Each entry maps an output column name to a SQL expression over the record step's
+    columns
+    (the source-qualified names, `crn_company`). An alias that names an existing column
+    replaces it. A new alias is added. Every other column, and `id`, passes through
+    untouched. Dropping is `Select`'s job, not `Clean`'s.
+    """
+
+    cleaning: dict[str, SQLExpression] = Field(
+        description="Output column name to a DuckDB SQL expression over the step."
+    )
+
+    @field_validator("cleaning")
+    @classmethod
+    def _non_empty(cls, cleaning: dict[str, SQLExpression]) -> dict[str, SQLExpression]:
+        """A clean that derives nothing is meaningless. 'No cleaning' is no `Clean`."""
+        if not cleaning:
+            raise ValueError("Clean needs at least one expression to derive.")
+        return cleaning
+
+    def apply(self, data: pl.DataFrame) -> pl.DataFrame:
+        """Add or replace the named columns, keeping the rest.
+
+        Invalid SQL raises at build time, naming the expression.
+        """
+        collisions = [alias for alias in self.cleaning if alias in data.columns]
+        star = expressions.Star(except_=[expressions.column(a) for a in collisions])
+        projection: list[expressions.Expression] = [star]
+        for alias, sql in self.cleaning.items():
+            projection.append(
+                expressions.alias_(parse_one(sql, dialect="duckdb"), alias)
+            )
+
+        query = sqlglot_select(*projection, dialect="duckdb").from_("data")
+        return run_sql(query.sql(dialect="duckdb"), data)

matchlab/transformers/explode.py

@@ -0,0 +1,67 @@
+"""Explode changes granularity. Cross-join each `id`'s columns to every combination."""
+
+from functools import reduce
+from operator import mul
+
+import polars as pl
+from pydantic import Field, field_validator
+
+from matchlab.transformers.base import Transformer
+
+
+class Explode(Transformer):
+    """Give every combination of each `id`'s non-null column values its own row.
+
+    `Group` collapses each `id` to one row. `Explode` goes the other way, giving the
+    true cross product a `list` aggregate can't express. Every non-null value in one
+    column pairs with every non-null value in every other column, for the same `id`.
+    Values are deduplicated per column first, so a value repeated across records
+    counts once. A column that is all null for one `id` still keeps that row, with
+    `null` there, instead of dropping it.
+
+    An `id` with several populated columns can generate many rows fast, so
+    `max_combinations` caps it and raises, naming the `id`, rather than silently
+    building an enormous table.
+    """
+
+    max_combinations: int = Field(
+        default=100,
+        description="Reject a group whose cross product would exceed this many rows.",
+    )
+
+    @field_validator("max_combinations")
+    @classmethod
+    def _positive(cls, max_combinations: int) -> int:
+        """Zero or fewer combinations is a way to reject every group, not a cap."""
+        if max_combinations <= 0:
+            raise ValueError("max_combinations must be positive.")
+        return max_combinations
+
+    def apply(self, data: pl.DataFrame) -> pl.DataFrame:
+        """Cross-join each `id`'s columns to every combination of their values."""
+        columns = [column for column in data.columns if column != "id"]
+        if not columns:
+            return data.select("id").unique(maintain_order=True)
+
+        grouped = data.group_by("id", maintain_order=True).agg(
+            [pl.col(column).drop_nulls().unique() for column in columns]
+        )
+
+        # Length 1 stands in for an empty list too, since exploding one keeps the row.
+        sizes = [pl.col(column).list.len().clip(lower_bound=1) for column in columns]
+        counts = grouped.select("id", reduce(mul, sizes).alias("combinations"))
+        overflow = counts.filter(pl.col("combinations") > self.max_combinations)
+        if overflow.height:
+            row = overflow.row(0, named=True)
+            raise ValueError(
+                f"id {row['id']} would generate {row['combinations']:,} "
+                f"combinations, which exceeds the maximum of "
+                f"{self.max_combinations}. This would cause a Cartesian explosion."
+            )
+
+        # One column at a time compounds into the true cross product. Exploding
+        # several columns in one call zips them in parallel instead.
+        exploded = grouped
+        for column in columns:
+            exploded = exploded.explode(column, empty_as_null=True)
+        return exploded

matchlab/transformers/group.py

@@ -0,0 +1,56 @@
+"""Group changes granularity. Collapse each `id` to one row."""
+
+import polars as pl
+from pydantic import Field, field_validator
+from sqlglot import expressions, parse_one
+from sqlglot import select as sqlglot_select
+
+from matchlab.core.sql import SQLExpression
+from matchlab.transformers.base import Transformer, run_sql
+
+
+class Group(Transformer):
+    """Collapse each `id` to a single row, combining columns with aggregate SQL.
+
+    Reading several records per entity (through a resolver) gives a model more evidence,
+    but a comparison needs one row. `Group` says how each column combines. `id` is the
+    grouping key, and every named expression is a DuckDB aggregate, for example
+    `any_value(crn_company)` where records agree or `list(distinct crn_town)` where they
+    don't. Only `id` and the named aggregates survive, since grouping has to decide how
+    every column collapses.
+    """
+
+    aggregates: dict[str, SQLExpression] = Field(
+        description="Output column name to a DuckDB aggregate expression."
+    )
+
+    @field_validator("aggregates")
+    @classmethod
+    def _non_empty(
+        cls, aggregates: dict[str, SQLExpression]
+    ) -> dict[str, SQLExpression]:
+        """Grouping collapses rows, so each column must say how it combines."""
+        if not aggregates:
+            raise ValueError(
+                "Group needs aggregate expressions: grouping collapses several records "
+                "into one row, so each column has to say how it combines. Pass "
+                "aggregates such as {'name': 'any_value(crn_company)'}."
+            )
+        return aggregates
+
+    def apply(self, data: pl.DataFrame) -> pl.DataFrame:
+        """Collapse each `id` to one row using the aggregate expressions.
+
+        A non-aggregate expression raises DuckDB's own error, naming the column.
+        """
+        identifier = expressions.alias_(expressions.column("id"), "id")
+        projection: list[expressions.Expression] = [identifier]
+        for alias, sql in self.aggregates.items():
+            projection.append(
+                expressions.alias_(parse_one(sql, dialect="duckdb"), alias)
+            )
+
+        query = (
+            sqlglot_select(*projection, dialect="duckdb").from_("data").group_by("id")
+        )
+        return run_sql(query.sql(dialect="duckdb"), data)

matchlab/transformers/select.py

@@ -0,0 +1,31 @@
+"""Select projects. Keep the named columns, drop the rest."""
+
+import polars as pl
+from pydantic import Field
+
+from matchlab.transformers.base import Transformer
+
+
+class Select(Transformer):
+    """Keep only the named columns, plus `id`.
+
+    Projection is `Select`'s one job. Naming a column keeps it, and every column you do
+    not name is dropped. `id` survives whether or not you name it, since it is the
+    grouping a model matches on. Columns are the source-qualified names (`crn_company`).
+    """
+
+    columns: tuple[str, ...] = Field(
+        default=(),
+        description="The columns to keep, by their current (qualified) names.",
+    )
+
+    def __init__(self, *columns: str, **data: object) -> None:
+        """Take the columns to keep positionally, e.g. `Select('crn_company')`."""
+        if columns:
+            data["columns"] = columns
+        super().__init__(**data)
+
+    def apply(self, data: pl.DataFrame) -> pl.DataFrame:
+        """Project to `id` and the named columns, in that order."""
+        kept = [column for column in self.columns if column != "id"]
+        return data.select("id", *kept)