matchlab 0.1.0.dev1__py3-none-any.whl

matchlab/models/linkers/deterministic.py

@@ -0,0 +1,229 @@
+"""A linking methodology based on a deterministic set of conditions."""
+
+import json
+
+import duckdb
+import polars as pl
+from pydantic import Field, field_validator
+
+from matchlab.core.logging import logger
+from matchlab.core.sql import SQLCondition, SQLQuery
+from matchlab.models import comparison
+from matchlab.models.linkers.base import Linker, LinkerSettings
+
+
+class DeterministicSettings(LinkerSettings):
+    """Settings for `DeterministicLinker`."""
+
+    comparisons: list[SQLCondition] | list[list[SQLCondition]] = Field(
+        description="""
+            Match conditions, in DuckDB SQL. Qualify every column with `l` or `r`.
+
+            A flat list applies every condition in parallel, unioned with OR logic:
+
+                [
+                    "l.company_number = r.company_number",
+                    "l.name = r.name",
+                ]
+
+            A nested list runs sequential rounds instead. Conditions within a round
+            use OR logic. After a round, matched records leave the pool before the
+            next round runs:
+
+                [
+                    [
+                        "l.company_number = r.company_number",
+                        "l.name = r.name",
+                    ],
+                    [
+                        "l.name_normalised = r.name_normalised",
+                        "l.website = r.website",
+                    ],
+                ]
+
+            Supports any DuckDB SQL expression, not just equality.
+        """,
+    )
+
+    @field_validator("comparisons", mode="before")
+    @classmethod
+    def validate_comparison(
+        cls, value: SQLCondition | list[SQLCondition] | list[list[SQLCondition]]
+    ) -> list[list[SQLCondition]]:
+        """Normalise to a list of rounds, and validate each comparison string."""
+        if isinstance(value, str):
+            return [[comparison(value, dialect="duckdb")]]
+        if not value:
+            raise ValueError("comparisons cannot be empty")
+        if all(isinstance(v, str) for v in value):
+            return [[comparison(v, dialect="duckdb") for v in value]]
+        if all(isinstance(v, list) for v in value):
+            for round_idx, round_comparisons in enumerate(value):
+                if not round_comparisons:
+                    raise ValueError(f"Round {round_idx} cannot be empty")
+                if not all(isinstance(c, str) for c in round_comparisons):
+                    raise ValueError(f"Round {round_idx} must contain only strings")
+            return [[comparison(c, dialect="duckdb") for c in r] for r in value]
+        raise ValueError(
+            "comparisons must be a string, list of strings, or list of lists"
+        )
+
+
+class DeterministicLinker(Linker):
+    """A deterministic linker that links based on a set of boolean conditions.
+
+    Uses DuckDB as the SQL backend, enabling rich SQL operations while maintaining
+    a Polars DataFrame interface. Supports both parallel matching (single round)
+    and sequential matching (multiple rounds where matched records are removed
+    after each round).
+    """
+
+    settings: DeterministicSettings
+
+    def prepare(self, left: pl.DataFrame, right: pl.DataFrame) -> None:
+        """No preparation needed."""
+        pass
+
+    def link(self, left: pl.DataFrame, right: pl.DataFrame) -> pl.DataFrame:
+        """Link the left and right dataframes.
+
+        If comparisons is a flat list, applies all comparisons in parallel.
+        If comparisons is a nested list, applies each round sequentially,
+        removing matched records from the pool after each round.
+        """
+        con: duckdb.DuckDBPyConnection = duckdb.connect(":memory:")
+        try:
+            all_matches: list[pl.DataFrame] = []
+            remaining_left, remaining_right = left, right
+
+            for round_num, round_comparisons in enumerate(
+                self.settings.comparisons, start=1
+            ):
+                if remaining_left.is_empty() or remaining_right.is_empty():
+                    logger.info(f"Round {round_num}: Skipping - no records remaining")
+                    break
+
+                logger.info(
+                    f"Round {round_num}: {len(remaining_left):,} left × "
+                    f"{len(remaining_right):,} right"
+                )
+
+                matches = self._link_round(
+                    con, remaining_left, remaining_right, round_comparisons, round_num
+                )
+
+                logger.info(f"Round {round_num}: Found {len(matches):,} matches")
+
+                if not matches.is_empty():
+                    all_matches.append(matches)
+                    matched_left = matches.select("left_id").unique()
+                    matched_right = matches.select("right_id").unique()
+                    remaining_left = remaining_left.join(
+                        matched_left,
+                        left_on=self.settings.left_id,
+                        right_on="left_id",
+                        how="anti",
+                    )
+                    remaining_right = remaining_right.join(
+                        matched_right,
+                        left_on=self.settings.right_id,
+                        right_on="right_id",
+                        how="anti",
+                    )
+
+            return self._finalise_results(all_matches)
+        finally:
+            con.close()
+
+    def _link_round(
+        self,
+        con: duckdb.DuckDBPyConnection,
+        left: pl.DataFrame,
+        right: pl.DataFrame,
+        comparisons: list[SQLCondition],
+        round_num: int,
+    ) -> pl.DataFrame:
+        """Apply all comparisons in a round using OR logic via DuckDB."""
+        con.register("left_df", left)
+        con.register("right_df", right)
+
+        subqueries: list[SQLQuery] = []
+        for condition in comparisons:
+            subquery: SQLQuery = f"""
+                SELECT
+                    l.{self.settings.left_id} AS left_id,
+                    r.{self.settings.right_id} AS right_id,
+                    1.0 AS score
+                FROM left_df l
+                INNER JOIN right_df r
+                    ON {condition}
+            """
+            subqueries.append(subquery)
+
+        query: SQLQuery = f"""
+            SELECT DISTINCT *
+            FROM ({" UNION ALL ".join(subqueries)})
+        """
+
+        max_est: int = self._get_max_cardinality(con, query)
+        logger.info(f"Round {round_num}: Estimated max cardinality: {max_est:,}")
+
+        return con.execute(query).pl()
+
+    def _get_max_cardinality(
+        self, con: duckdb.DuckDBPyConnection, query: SQLQuery
+    ) -> int:
+        """Get max cardinality estimate from DuckDB plan, or -1 if unavailable."""
+        explain = con.execute(
+            f"PRAGMA explain_output = 'all'; EXPLAIN (FORMAT json) {query}"
+        ).fetchall()
+        plans = {k: json.loads(v) for k, v in dict(explain).items()}
+
+        estimates: list[dict] = []
+        for plan_name, plan_type in [
+            ("physical_plan", "physical"),
+            ("logical_opt", "optimised"),
+        ]:
+            if plan := plans.get(plan_name):
+                estimates.extend(self._traverse_plan(plan[0], plan_type))
+
+        # Debug: log full tree breakdown
+        logger.debug("Plan breakdown:")
+        for plan_type in ["physical", "optimised"]:
+            plan_ests = [e for e in estimates if e["plan_type"] == plan_type]
+            if plan_ests:
+                logger.debug(f"  {plan_type.capitalize()} plan:")
+                for est in plan_ests:
+                    indent = "    " + "  " * est["depth"]
+                    logger.debug(f"{indent}{est['node_name']}: {est['cardinality']:,}")
+
+        # Return max for info logging
+        return max(
+            (e["cardinality"] for e in estimates if e["cardinality"] > 0),
+            default=-1,
+        )
+
+    def _traverse_plan(self, node: dict, plan_type: str, depth: int = 0) -> list[dict]:
+        """Recursively collect cardinality estimates with metadata from plan tree."""
+        estimates: list[dict] = []
+        cardinality = node.get("extra_info", {}).get("Estimated Cardinality")
+        if cardinality is not None:
+            estimates.append(
+                {
+                    "cardinality": int(cardinality),
+                    "node_name": node.get("name", "UNKNOWN"),
+                    "depth": depth,
+                    "plan_type": plan_type,
+                }
+            )
+        for child in node.get("children", []):
+            estimates.extend(self._traverse_plan(child, plan_type, depth + 1))
+        return estimates
+
+    def _finalise_results(self, all_matches: list[pl.DataFrame]) -> pl.DataFrame:
+        """Combine matches from all rounds and ensure correct schema."""
+        if all_matches:
+            return pl.concat(all_matches).with_columns(pl.col("score").cast(pl.Float32))
+        return pl.DataFrame({"left_id": [], "right_id": [], "score": []}).with_columns(
+            pl.col("score").cast(pl.Float32)
+        )

matchlab/models/linkers/splinklinker.py

@@ -0,0 +1,262 @@
+"""A linking methodology leveraging Splink."""
+
+import inspect
+import json
+from typing import Any
+
+import polars as pl
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    field_serializer,
+    field_validator,
+    model_validator,
+)
+from pydantic.functional_serializers import SerializationInfo
+from splink import DuckDBAPI, SettingsCreator
+from splink import Linker as SplinkLibLinkerClass
+from splink.internals.linker_components.training import LinkerTraining
+
+from matchlab.core.logging import logger
+from matchlab.models.linkers.base import Linker, LinkerSettings
+
+
+class SplinkLinkerFunction(BaseModel):
+    """A method of splink.Linker.training used to train the linker."""
+
+    function: str
+    arguments: dict[str, Any]
+
+    @model_validator(mode="after")
+    def validate_function_and_arguments(self) -> "SplinkLinkerFunction":
+        """Ensure the function and arguments are valid."""
+        if not hasattr(LinkerTraining, self.function):
+            raise ValueError(
+                f"Function {self.function} not found as method of Splink Linker class"
+            )
+
+        splink_linker_func = getattr(LinkerTraining, self.function)
+        splink_linker_func_param_set = set(
+            inspect.signature(splink_linker_func).parameters.keys()
+        )
+        current_func_param_set = set(self.arguments.keys())
+
+        if not current_func_param_set <= splink_linker_func_param_set:
+            raise ValueError(
+                f"Function {self.function} given incorrect arguments: "
+                f"{current_func_param_set.difference(splink_linker_func_param_set)}. "
+                "Consider referring back to the Splink documentation: "
+                "https://moj-analytical-services.github.io/splink/linker.html"
+            )
+
+        return self
+
+
+class SplinkSettings(LinkerSettings):
+    """A data class to enforce the Splink linker's settings dictionary shape.
+
+    !!! warning "Pass a seed to any sampling training function"
+        Step fingerprints are derived from configuration, so two collections that
+        share these settings share a cache entry. Training functions that sample —
+        `estimate_u_using_random_sampling`, for one — are non-deterministic unless
+        given a `seed` in their `arguments`, and without one the first result is
+        cached and silently reused. See `Step._fingerprint`.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    linker_training_functions: list[SplinkLinkerFunction] = Field(
+        description="""
+            A list of dictionaries where keys are the names of methods for
+            splink.Linker.training and values are dictionaries encoding the arguments of
+            those methods. Each function will be run in the order supplied.
+
+            Example:
+            
+                >>> linker_training_functions=[
+                ...     {
+                ...         "function": "estimate_probability_two_random_records_match",
+                ...         "arguments": {
+                ...             "deterministic_matching_rules": \"""
+                ...                 l.company_name = r.company_name
+                ...             \""",
+                ...             "recall": 0.7,
+                ...         },
+                ...     },
+                ...     {
+                ...         "function": "estimate_u_using_random_sampling",
+                ...         "arguments": {"max_pairs": 1e6},
+                ...     }
+                ... ]
+            
+        """
+    )
+    linker_settings: SettingsCreator = Field(
+        description="""
+            A valid Splink SettingsCreator.
+
+            See Splink's documentation for a full description of available settings.
+            https://moj-analytical-services.github.io/splink/api_docs/settings_dict_guide.html
+
+            * link_type must be set to "link_only"
+            * unique_id_name is overridden to the value of left_id and right_id,
+                which must match
+
+            Example:
+
+                >>> from splink import SettingsCreator, block_on
+                ... import splink.comparison_library as cl
+                ... import splink.comparison_template_library as ctl
+                ... 
+                ... splink_settings = SettingsCreator(
+                ...     retain_matching_columns=False,
+                ...     retain_intermediate_calculation_columns=False,
+                ...     blocking_rules_to_generate_predictions=[
+                ...         block_on("company_name"),
+                ...         block_on("postcode"),
+                ...     ],
+                ...     comparisons=[
+                ...         cl.jaro_winkler_at_thresholds(
+                ...             "company_name", 
+                ...             [0.9, 0.6], 
+                ...             term_frequency_adjustments=True
+                ...         ),
+                ...         ctl.postcode_comparison("postcode"), 
+                ...     ]
+                ... )         
+        """
+    )
+    threshold: float | None = Field(
+        default=None,
+        description="""
+            The score above which matches will be kept.
+
+            None is used to indicate no threshold.
+            
+            Inclusive, so a value of 1 will keep only exact matches across all 
+            comparisons.
+        """,
+        gt=0,
+        le=1,
+    )
+
+    @model_validator(mode="after")
+    def check_link_only(self) -> "SplinkSettings":
+        """Ensure link_type is set to "link_only"."""
+        if self.linker_settings.link_type != "link_only":
+            raise ValueError('link_type must be set to "link_only"')
+        return self
+
+    @model_validator(mode="after")
+    def add_enforced_settings(self) -> "SplinkSettings":
+        """Ensure ID is the only field we link on."""
+        self.linker_settings.unique_id_column_name = self.left_id
+        return self
+
+    @field_validator("linker_settings", mode="before")
+    @classmethod
+    def load_linker_settings(cls, value: str | SettingsCreator) -> SettingsCreator:
+        """Load serialised settings into SettingsCreator."""
+        if isinstance(value, str):
+            value = SettingsCreator.from_path_or_dict(json.loads(value))
+        return value
+
+    @field_serializer("linker_settings")
+    def serialise_settings(
+        self, value: SettingsCreator, info: SerializationInfo
+    ) -> str:
+        """Convert Splink settings to string."""
+        return json.dumps(value.create_settings_dict("duckdb"))
+
+
+class SplinkLinker(Linker):
+    """A linker that leverages Bayesian record linkage using Splink."""
+
+    settings: SplinkSettings
+
+    _linker: SplinkLibLinkerClass
+    _id_dtype_l: pl.DataType
+    _id_dtype_r: pl.DataType
+
+    def prepare(self, left: pl.DataFrame, right: pl.DataFrame) -> None:
+        """Build the Splink linker over left and right, and run its training functions.
+
+        Runs each function in `settings.linker_training_functions`, in order, against
+        the built linker. `link()` then just predicts against the trained state.
+        """
+        if (set(left.columns) != set(right.columns)) or not left.dtypes == right.dtypes:
+            raise ValueError(
+                "SplinkLinker requires input data to be conformant, meaning they "
+                "share the same column names and data formats."
+            )
+
+        self._id_dtype_l = left[self.settings.left_id].dtype
+        self._id_dtype_r = right[self.settings.right_id].dtype
+
+        # Convert to pandas for Splink compatibility
+        left_pd = left.with_columns(
+            pl.col(self.settings.left_id).cast(pl.String)
+        ).to_pandas()
+        right_pd = right.with_columns(
+            pl.col(self.settings.right_id).cast(pl.String)
+        ).to_pandas()
+
+        self._linker = SplinkLibLinkerClass(
+            input_table_or_tables=[left_pd, right_pd],
+            input_table_aliases=["l", "r"],
+            settings=self.settings.linker_settings,
+            db_api=DuckDBAPI(),
+        )
+
+        for func in self.settings.linker_training_functions:
+            proc_func = getattr(self._linker.training, func.function)
+            proc_func(**func.arguments)
+
+    def link(
+        self, left: pl.DataFrame = None, right: pl.DataFrame = None
+    ) -> pl.DataFrame:
+        """Predict match scores using the linker trained in `prepare()`.
+
+        `left`/`right` are accepted only to satisfy the `Linker` contract. The data
+        was already fixed when `prepare()` built the underlying Splink linker, so
+        passing values here logs a warning and has no effect.
+        """
+        if left is not None or right is not None:
+            logger.warning(
+                "Left and right data are declared in .prepare() for SplinkLinker. "
+                "These values will be ignored"
+            )
+
+        res = self._linker.inference.predict(
+            threshold_match_probability=self.settings.threshold
+        )
+
+        return (
+            res.as_duckdbpyrelation()
+            .pl()
+            .lazy()
+            .select(
+                [
+                    f"{self.settings.left_id}_l",
+                    f"{self.settings.right_id}_r",
+                    "match_probability",
+                ]
+            )
+            .with_columns(
+                [
+                    pl.col(f"{self.settings.left_id}_l")
+                    .cast(self._id_dtype_l)
+                    .alias("left_id"),
+                    pl.col(f"{self.settings.right_id}_r")
+                    .cast(self._id_dtype_r)
+                    .alias("right_id"),
+                    pl.col("match_probability").cast(pl.Float32).alias("score"),
+                ]
+            )
+            .select(["left_id", "right_id", "score"])
+            # Multiple blocking rules can lead to multiple matches
+            .group_by(["left_id", "right_id"])
+            .agg(pl.col("score").max())
+            .collect()
+        )

matchlab/models/linkers/weighteddeterministic.py

@@ -0,0 +1,143 @@
+"""A linking methodology that applies different weights to field comparisons."""
+
+import duckdb
+import polars as pl
+from pydantic import BaseModel, Field, field_validator
+
+from matchlab.core.sql import SQLCondition
+from matchlab.models import comparison
+from matchlab.models.linkers.base import Linker, LinkerSettings
+
+
+class WeightedComparison(BaseModel):
+    """A comparison condition, and the weight it contributes to a pair's score."""
+
+    comparison: SQLCondition = Field(
+        description="""
+            A valid ON clause comparing fields between the left and the right data.
+
+            Qualify every column with `l` or `r`. For example:
+
+            "l.company_name = r.company_name"
+        """
+    )
+    weight: float = Field(
+        description="""
+            A weight to give this comparison. Use 1 for all comparisons to give
+            uniform weight to each.
+        """
+    )
+
+    @field_validator("comparison")
+    @classmethod
+    def validate_comparison(cls, v: SQLCondition) -> SQLCondition:
+        """Validate the comparison string."""
+        comp_val = comparison(v, dialect="duckdb")
+        return comp_val
+
+
+class WeightedDeterministicSettings(LinkerSettings):
+    """Settings for `WeightedDeterministicLinker`.
+
+    Example:
+        >>> {
+        ...     weighted_comparisons: [
+        ...         ("l.company_name = r.company_name", 0.7),
+        ...         ("l.postcode = r.postcode", 0.7),
+        ...         ("l.company_id = r.company_id", 1),
+        ...     ],
+        ...     threshold: 0.8,
+        ... }
+    """
+
+    weighted_comparisons: list[WeightedComparison] = Field(
+        description="A list of tuples in the form of a comparison, and a weight."
+    )
+    threshold: float = Field(
+        description="""
+            The score above which matches will be kept. 
+            
+            Inclusive, so a value of 1 will keep only exact matches across all 
+            comparisons.
+        """,
+        ge=0,
+        le=1,
+    )
+
+
+class WeightedDeterministicLinker(Linker):
+    """Scores a pair by the weighted share of comparisons it matches on."""
+
+    settings: WeightedDeterministicSettings
+
+    _id_dtype_l: pl.DataType
+    _id_dtype_r: pl.DataType
+
+    def prepare(self, left: pl.DataFrame, right: pl.DataFrame) -> None:
+        """No preparation needed."""
+        pass
+
+    def link(self, left: pl.DataFrame, right: pl.DataFrame) -> pl.DataFrame:
+        """Score each pair by summed matching weight over total weight.
+
+        Keeps only pairs scoring at or above `settings.threshold`.
+        """
+        self._id_dtype_l = left[self.settings.left_id].dtype
+        self._id_dtype_r = right[self.settings.right_id].dtype
+
+        # Used below but ruff can't detect
+        left_df = left.clone()  # noqa: F841
+        right_df = right.clone()  # noqa: F841
+
+        match_subquery = []
+        weights = []
+
+        for weighted_comparison in self.settings.weighted_comparisons:
+            match_subquery.append(
+                f"""
+                    select distinct on (list_sort([raw.left_id, raw.right_id]))
+                        raw.left_id,
+                        raw.right_id,
+                        1.0 * {weighted_comparison.weight} as score
+                    from (
+                        select
+                            l.{self.settings.left_id} as left_id,
+                            r.{self.settings.right_id} as right_id,
+                        from
+                            left_df l
+                        inner join right_df r on
+                            {weighted_comparison.comparison}
+                    ) raw
+                """
+            )
+            weights.append(weighted_comparison.weight)
+
+        match_subquery = " union all ".join(match_subquery)
+        total_weight = sum(weights)
+
+        sql = f"""
+            select
+                matches.left_id,
+                matches.right_id,
+                sum(matches.score) / {total_weight} as score
+            from
+                ({match_subquery}) matches
+            group by
+                matches.left_id,
+                matches.right_id
+            having
+                sum(matches.score) / 
+                    {total_weight} >= {self.settings.threshold};
+        """
+
+        return (
+            duckdb.sql(sql)
+            .pl()
+            .with_columns(
+                [
+                    pl.col("left_id").cast(self._id_dtype_l),
+                    pl.col("right_id").cast(self._id_dtype_r),
+                    pl.col("score").cast(pl.Float32),
+                ]
+            )
+        )