diff --git a/examples/online/.gitignore b/examples/online/.gitignore new file mode 100644 index 00000000..2355c543 --- /dev/null +++ b/examples/online/.gitignore @@ -0,0 +1,7 @@ +# Generated at run time — regenerated by the scripts, never committed. +warehouse.sqlite +model.json +.online_tmp/ +.mem_*.bin +.mem_*.json +__pycache__/ diff --git a/examples/online/README.md b/examples/online/README.md new file mode 100644 index 00000000..f0b5b908 --- /dev/null +++ b/examples/online/README.md @@ -0,0 +1,130 @@ +# matchlab vs pure Splink: expression, cost, and a single online match + +This example runs one pipeline two ways over ~100k entities scattered across four sources +as ~390k records (generated by the neighbouring `examples/companies/benchmark.py`). The +pipeline is the one from the Ministry of Justice's [*Running Splink in Production*](https://moj-analytical-services.github.io/splink/blog/2026/01/29/running-splink-in-production.html): +standardise each source, deduplicate each in isolation, link every pair with a shared +Splink model, and cluster once at the end. + +Both sides run that *identical* algorithm. `pipeline_matchlab.py` expresses it as a +matchlab plan. `pipeline_splink.py` hand-rolls it with Splink and a union-find, the way +`examples/companies/by_hand.py` is the fair opponent to `with_matchlab.py`. So the +comparison isolates one thing: what matchlab adds on top of the same work. + +Run `python benchmark.py` for the whole story in one table. Figures are from one run on +the dev container, so read them as orders of magnitude. + +## Expression + +matchlab maps onto the blog's shape one-to-one: `Source.view(cleaning=)` standardises, +`.dedupe()` deduplicates each source, a `SplinkLinker` per pair generates cross-dataset +edges, and one `Resolver` consolidates every edge and clusters once. Both pipelines +recover the ground truth exactly, 100,000 clusters with no precision or recall loss. + +## Time and memory + +| batch build | Splink | matchlab (memory) | matchlab (on disk) | +| --- | ---: | ---: | ---: | +| time | 5.3 s | 11.7 s | 12.7 s | +| peak memory | 1,675 MiB | 2,307 MiB | 2,228 MiB | + +matchlab takes about 2.2 times the wall clock at about 1.4 times the peak memory. Pointing +the store at a file rather than memory barely moves either, which the next section +explains. + +## Why matchlab is slower, and where the memory goes + +Both sides do the same matching, so the difference is not the matching. Timing the collect +by phase makes it concrete: + + 1.7 s read each source and content-hash every row into a leaf id, then store + 2.1 s dedupe each source, resolve, materialise + 10.1 s entity views, the six Splink links, and the final materialise + —— + 4.9 s of which: the six Splink predicts, shared with the hand-rolled baseline + +So about 5 s is the shared matching. matchlab's extra ~6 s is content-addressing and +writing: hashing every row to a leaf, minting content-addressed cluster ids, building the +complete `(root, leaf, key, source)` resolution and persisting it. That is the cache and +the queryable store being built, not the match being computed. + +The memory tells the same story from the other side. Putting the store on disk barely +lowers the peak, because the peak is the transient working set, Splink's own in-memory +engine on each of the six links plus the intermediate frames, not the persisted tables. +The store's own footprint is only the ~80 MiB gap between the two matchlab columns. A row +store helps online serving, below, not the batch build's memory. + +## DAG shape, and what the cost buys + +The blog shapes its DAG around compute cost and recovery. It slices cross-dataset linking +into pairwise pieces, clusters once at the end, and keeps training as a separate stage, +because each run is expensive and has to be restartable. Both pipelines here follow that +shape. matchlab makes it cheaper to live with. + +- Every step is a content-addressed artefact, so re-running an unchanged slice does no + work. Recovery and iteration are structural, not something you hand-build with + checkpoints and run identifiers. +- The shape is open. Spine, all-pairs, or coherent subsets are just which `.link()` calls + you make, so you tune the topology to the data. +- The backend is a choice. The blog already splits storage by access pattern, DuckDB for + the batch and a row store for the published lookup table teams join against, because + point selects are slow on a column store. matchlab makes that split explicit, which is + the next section. + +## A single online match + +Online serving is the blog's published lookup table: join a record's `(source, key)` to +get its entity, one record at a time. That is a point select, the opposite of the batch +scan Splink is built for. We time the same single lookup against matchlab's own DuckDB +store and against the identical resolution in SQLite. + +| single online match | latency | throughput | +| --- | ---: | ---: | +| DuckDB (matchlab store) | 6.13 ms | 163 /s | +| SQLite (row store) | 0.011 ms | 93,000 /s | + +**A single match in the pure Splink world costs about 6 ms.** Splink returns a frame and +keeps no lookup table, so serving one match means a point select against wherever you +persist the result. Its engine is DuckDB, and matchlab's store is that same DuckDB today, +so both land at ~6 ms. + +**A single match in a SQLite-backed matchlab costs about 0.011 ms**, roughly 500 times +faster. A row store answers a point select with a couple of B-tree reads, where a column +store pays a per-query scan overhead. This is why Matchbox proper is backed by PostgreSQL: +a column store for the batch build, a row store for online serving. + +## Is it worth it + +Clear-eyed, it depends on whether you look at the result again. + +- **Pure Splink** gives you the matcher, and a frame. You build the DAG, the checkpoints, + the storage, the lookup, and the caching around it, which is what `pipeline_splink.py` + is. +- **matchlab** gives you the same matcher wrapped in a DAG of content-addressed artefacts, + so caching, recovery, and iteration come for free, plus a materialised queryable + resolution, a `lookup_key` API, and a backend you pick per access pattern. + +So it is worth it if you iterate, re-resolve, or serve matches online. The ~2.2× batch +cost is paid once and buys near-free re-runs and a single online match that is hundreds of +times cheaper. It is not worth it if you resolve once and never look again, where the +multiplier buys nothing pure Splink does not already give you. + +## Caveats + +- Four sources, so linking every pair is six models. At the blog's scale (eight sources, + 28 pairs) you would group sources into coherent subsets instead. +- Synthetic data gives perfect quality, so the numbers show the two pipelines agreeing, + not a claim about linkage accuracy. +- We block on name only. The generator derives postcodes from the entity id, so blocking + on them explodes the candidate set, the blog's candidate-explosion warning made concrete. +- The SQLite figure is measured here, not just estimated. + +## Files + +| file | role | +| --- | --- | +| `data.py` | build the warehouse via the companies generator, plus hidden ground truth | +| `model.py` | shared standardisation and the train-once model artefact | +| `pipeline_matchlab.py` | the pipeline as a matchlab plan | +| `pipeline_splink.py` | the same pipeline hand-rolled with Splink and a union-find | +| `benchmark.py` | expression, batch time and memory, single-match latency | diff --git a/examples/online/benchmark.py b/examples/online/benchmark.py new file mode 100644 index 00000000..78d3e72c --- /dev/null +++ b/examples/online/benchmark.py @@ -0,0 +1,195 @@ +"""matchlab vs plain Splink: expression, batch cost, and a single online match. + +Both pipelines run the identical algorithm (dedupe each source, link every pair with the +shared Splink model, cluster once), so the comparison isolates what matchlab adds: a +plan, and a materialised store. We check they agree, time and memory-profile the batch +build, then time a single online lookup against matchlab's own DuckDB store versus the +same resolution in SQLite. + + python benchmark.py # the whole story, one table + python benchmark.py _one matchlab # run one pipeline once (used by memray) +""" + +from __future__ import annotations + +import json +import sqlite3 +import subprocess +import sys +import time +import timeit +import warnings +from pathlib import Path + +import polars as pl + +warnings.filterwarnings("ignore") +import logging # noqa: E402 + +logging.disable(logging.WARNING) + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) + +import data # noqa: E402 +import model # noqa: E402 +import pipeline_matchlab # noqa: E402 +import pipeline_splink # noqa: E402 + +from matchlab.adapters import DuckDBAdapter # noqa: E402 + +N_ROWS = 300_000 +REPEATS = 3 +N_PROBES = 2_000 + + +def build_matchlab(trained: dict, on_disk: bool = False) -> None: + """Collect the matchlab plan into a fresh store. Times/measures the build only. + + A fresh store each time, so content-addressed caching doesn't turn a repeat into a + no-op. On disk, the materialised tables live in the DuckDB file rather than in RAM. + """ + path = HERE / f".ml_{time.time_ns()}.duckdb" if on_disk else None + store = DuckDBAdapter(path) if path else DuckDBAdapter() + pipeline_matchlab.build_plan(trained).collect(store) + store.close() + if path: + path.unlink(missing_ok=True) + + +def matchlab_resolution(trained: dict) -> pl.DataFrame: + """The matchlab resolution, for the correctness check.""" + return pipeline_matchlab.build_plan(trained).collect(DuckDBAdapter()).resolution() + + +def _mixed(res: pl.DataFrame, truth: pl.DataFrame) -> int: + """Clusters that mix more than one true entity (0 means perfect).""" + return ( + res.join(truth, on=["source", "key"]) + .group_by("root") + .agg(pl.col("entity").n_unique().alias("n")) + .filter(pl.col("n") > 1) + .height + ) + + +def _peak_mib(which: str) -> float: + """Peak heap of one pipeline, measured by memray in a subprocess.""" + binf, jsonf = HERE / f".mem_{which}.bin", HERE / f".mem_{which}.json" + run = ["uv", "run", "memray"] + kw = {"check": True, "capture_output": True, "cwd": HERE} + subprocess.run( + [*run, "run", "-q", "-f", "-o", str(binf), __file__, "_one", which], **kw + ) + subprocess.run([*run, "stats", "--json", "-fo", str(jsonf), str(binf)], **kw) + peak = json.loads(jsonf.read_text())["metadata"]["peak_memory"] + binf.unlink(), jsonf.unlink() + return peak / 1024**2 + + +def _time(lookup, probes: list) -> float: + """Mean latency of a single lookup, in milliseconds.""" + lookup(*probes[0]) # warm up + start = time.perf_counter() + for source, key in probes: + lookup(source, key) + return (time.perf_counter() - start) / len(probes) * 1000 + + +def _one(which: str) -> None: + data.ensure_warehouse(N_ROWS) + trained = model.load_or_train(pipeline_splink.all_nodes()) + if which == "splink": + pipeline_splink.resolution(trained) + else: + build_matchlab(trained, on_disk=which == "matchlab_disk") + + +def main() -> None: + data.ensure_warehouse(N_ROWS) + truth = data.truth() + trained = model.load_or_train(pipeline_splink.all_nodes()) + print(f"{truth.height:,} records, {truth['entity'].n_unique():,} true entities\n") + + # Expression: both pipelines must agree with the ground truth. + for name, res in ( + ("matchlab", matchlab_resolution(trained)), + ("splink", pipeline_splink.resolution(trained)), + ): + good = res["root"].n_unique() == truth["entity"].n_unique() and not _mixed( + res, truth + ) + print(f"expression {name:<9} {'ok' if good else 'FAIL'}") + print() + + # Batch build: time (fresh store each repeat) and peak memory. matchlab is timed + # for the build only (the store is the product), both in memory and on disk. + def best(fn): + return min(timeit.repeat(fn, number=1, repeat=REPEATS)) + + ts = best(lambda: pipeline_splink.resolution(trained)) + tm = best(lambda: build_matchlab(trained)) + td = best(lambda: build_matchlab(trained, on_disk=True)) + ms, mm, md = _peak_mib("splink"), _peak_mib("matchlab"), _peak_mib("matchlab_disk") + print( + f"batch build {'Splink':>10}{'matchlab (mem)':>16}{'matchlab (disk)':>17}" + ) + print(f" time {ts:>8.1f} s{tm:>14.1f} s{td:>15.1f} s") + print(f" peak memory {ms:>6,.0f} MiB{mm:>12,.0f} MiB{md:>13,.0f} MiB\n") + + # A single online match, against matchlab's own DuckDB store and the same + # resolution in SQLite. No extra copies: the store is the one the plan built. + store_path = HERE / ".online_store.duckdb" + store_path.unlink(missing_ok=True) + store = DuckDBAdapter(store_path) + apex = pipeline_matchlab.build_plan(trained) + apex.collect(store) + fp = apex._fp + con = store.conn + + def duck(source, key): + root = con.execute( + "SELECT root FROM resolution WHERE fp=? AND source=? AND key=?", + [fp, source, key], + ).fetchone()[0] + return con.execute( + "SELECT source, key FROM resolution WHERE fp=? AND root=?", [fp, root] + ).fetchall() + + res = con.execute("SELECT root, key, source FROM resolution WHERE fp=?", [fp]).pl() + sqlite_path = HERE / ".online_store.sqlite" + sqlite_path.unlink(missing_ok=True) + scon = sqlite3.connect(sqlite_path) + scon.execute("CREATE TABLE resolution (root TEXT, key TEXT, source TEXT)") + scon.executemany( + "INSERT INTO resolution VALUES (?,?,?)", + res.select(pl.col("root").cast(pl.Utf8), "key", "source").iter_rows(), + ) + scon.execute("CREATE INDEX i1 ON resolution (source, key)") + scon.execute("CREATE INDEX i2 ON resolution (root)") + + def row(source, key): + root = scon.execute( + "SELECT root FROM resolution WHERE source=? AND key=?", [source, key] + ).fetchone()[0] + return scon.execute( + "SELECT source, key FROM resolution WHERE root=?", [root] + ).fetchall() + + probes = list(res.select("source", "key").sample(N_PROBES, seed=1).iter_rows()) + duck_ms, row_ms = _time(duck, probes), _time(row, probes) + print("single online match latency throughput") + print(f" DuckDB (matchlab store) {duck_ms:7.3f} ms {1000 / duck_ms:>8,.0f} /s") + print(f" SQLite (row store) {row_ms:7.3f} ms {1000 / row_ms:>8,.0f} /s") + + scon.close() + store.close() + store_path.unlink(missing_ok=True) + sqlite_path.unlink(missing_ok=True) + + +if __name__ == "__main__": + if len(sys.argv) >= 3 and sys.argv[1] == "_one": + _one(sys.argv[2]) + else: + main() diff --git a/examples/online/data.py b/examples/online/data.py new file mode 100644 index 00000000..18f82fc5 --- /dev/null +++ b/examples/online/data.py @@ -0,0 +1,71 @@ +"""The warehouse both pipelines read. + +Reuses the fast vectorised generator in ``examples/companies/benchmark.py`` (the +factory subsystem's per-row Python generation is minutes-slow at this size). ~100k +true entities become ~390k records across four sources, each keeping its own column +names. A hidden ``_truth`` table records the real entity behind every key, for scoring. + + python data.py 300_000 +""" + +from __future__ import annotations + +import sqlite3 +import sys +from pathlib import Path + +import polars as pl +from sqlalchemy import create_engine + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "companies")) +import benchmark as companies # noqa: E402 + +WAREHOUSE = Path(__file__).parent / "warehouse.sqlite" +# name -> (key column, name column, postcode column), as the companies example uses. +SOURCES: dict[str, tuple[str, str, str]] = dict(companies.COLUMNS) + + +def build_warehouse(n_rows: int = 300_000, path: Path = WAREHOUSE) -> Path: + """Generate and write the warehouse: one table per source, plus hidden truth.""" + frames = companies.generate(n_rows) + path.unlink(missing_ok=True) + with sqlite3.connect(path) as conn: + truth_rows: list[tuple[str, str, int]] = [] + for source, (key_col, name_col, postcode_col) in SOURCES.items(): + frame = frames[source] + conn.execute( + f"CREATE TABLE {source} " + f"({key_col} TEXT, {name_col} TEXT, {postcode_col} TEXT)" + ) + conn.executemany( + f"INSERT INTO {source} VALUES (?, ?, ?)", + frame.select("key", "name", "postcode").iter_rows(), + ) + conn.execute(f"CREATE INDEX idx_{source}_key ON {source} ({key_col})") + truth_rows.extend( + (source, key, entity) + for key, entity in frame.select("key", "_truth").iter_rows() + ) + conn.execute("CREATE TABLE _truth (source TEXT, key TEXT, entity INTEGER)") + conn.executemany("INSERT INTO _truth VALUES (?, ?, ?)", truth_rows) + return path + + +def ensure_warehouse(n_rows: int = 300_000, path: Path = WAREHOUSE) -> Path: + """Build the warehouse if missing, otherwise reuse it.""" + if not path.exists(): + build_warehouse(n_rows, path) + return path + + +def truth(path: Path = WAREHOUSE) -> pl.DataFrame: + """The hidden ground truth `(source, key, entity)`, for scoring only.""" + engine = create_engine(f"sqlite:///{path}") + return pl.read_database("SELECT source, key, entity FROM _truth", engine.connect()) + + +if __name__ == "__main__": + n = int(sys.argv[1].replace("_", "")) if len(sys.argv) > 1 else 300_000 + t = truth(build_warehouse(n)) + print(f"warehouse: {WAREHOUSE}") + print(f" {t.height:,} keys, {t['entity'].n_unique():,} true entities") diff --git a/examples/online/model.py b/examples/online/model.py new file mode 100644 index 00000000..0fa1d09d --- /dev/null +++ b/examples/online/model.py @@ -0,0 +1,68 @@ +"""The shared linkage model, trained once and reused. + +The blog keeps estimation out of the repeatable run: train rarely, persist a +``model.json``, reuse it for prediction. Both pipelines here share one such artefact, +and the same standardisation (cleaning to a common ``name`` and ``postcode``). +""" + +from __future__ import annotations + +import json +import warnings +from pathlib import Path + +import polars as pl +import splink.comparison_library as cl +from splink import DuckDBAPI, Linker, SettingsCreator, block_on + +MODEL_JSON = Path(__file__).parent / "model.json" + +# Standardisation: upper-case, strip whitespace, punctuation and company suffixes off +# the name, strip whitespace off the postcode. ``{0}`` is the raw (qualified) column. +CLEAN_NAME = ( + "regexp_replace(upper({0}), '\\s+|\\.|\\bLIMITED\\b|\\bLTD\\b|\\bPLC\\b', '', 'g')" +) +CLEAN_POSTCODE = "regexp_replace(upper({0}), '\\s+', '', 'g')" +THRESHOLD = 0.9 + + +def _settings(link_type: str, trained: dict | None = None) -> SettingsCreator: + """A Splink settings object, loading a trained artefact's m/u if given. + + Blocks on name only: this data's postcodes are low-cardinality (derived from the + entity id), so blocking on them explodes the candidate set. + """ + if trained is not None: + settings = SettingsCreator.from_path_or_dict(trained) + settings.link_type = link_type + return settings + return SettingsCreator( + link_type=link_type, + unique_id_column_name="id", + retain_matching_columns=False, + retain_intermediate_calculation_columns=False, + blocking_rules_to_generate_predictions=[block_on("name")], + comparisons=[ + cl.JaroWinklerAtThresholds("name", [0.9, 0.7]), + cl.ExactMatch("postcode"), + ], + ) + + +def load_or_train(nodes: pl.DataFrame) -> dict: + """Return the trained model dict, training and persisting it on a miss.""" + if MODEL_JSON.exists(): + return json.loads(MODEL_JSON.read_text()) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + linker = Linker(nodes.to_pandas(), _settings("dedupe_only"), db_api=DuckDBAPI()) + linker.training.estimate_probability_two_random_records_match( + [block_on("name")], recall=0.8 + ) + linker.training.estimate_u_using_random_sampling(max_pairs=1e6, seed=1) + linker.training.estimate_parameters_using_expectation_maximisation( + block_on("name") + ) + trained = linker.misc.save_model_to_json() + MODEL_JSON.write_text(json.dumps(trained)) + return trained diff --git a/examples/online/pipeline_matchlab.py b/examples/online/pipeline_matchlab.py new file mode 100644 index 00000000..a3a80884 --- /dev/null +++ b/examples/online/pipeline_matchlab.py @@ -0,0 +1,77 @@ +"""The blog's pipeline as a matchlab plan. + +Dedupe each dataset, link every pair (four sources, so six models, the blog's "small +number" middle ground), then consolidate all edges and cluster once. matchlab is an +opinionated implementation of exactly that shape: a DAG of content-addressed artefacts. +""" + +from __future__ import annotations + +from itertools import combinations + +from data import SOURCES, WAREHOUSE +from model import CLEAN_NAME, CLEAN_POSTCODE, THRESHOLD, _settings +from sqlalchemy import create_engine + +from matchlab import Source +from matchlab.locations import RelationalDBLocation +from matchlab.models.dedupers import NaiveDeduper +from matchlab.models.linkers import SplinkLinker +from matchlab.models.linkers.splinklinker import SplinkSettings + + +def build_plan(trained: dict, warehouse=WAREHOUSE): + """Build the lazy plan. Nothing runs until it is collected.""" + location = RelationalDBLocation(name="warehouse").set_client( + create_engine(f"sqlite:///{warehouse}") + ) + link_settings = SplinkSettings( + left_id="id", + right_id="id", + threshold=THRESHOLD, + linker_settings=_settings("link_only", trained), # reuse the artefact + linker_training_functions=[], # no per-link training + ) + + sources, nodes = {}, {} + for name, (key_col, name_col, postcode_col) in SOURCES.items(): + src = Source( + location=location, + name=name, + extract_transform=( + f"select {key_col}, {name_col}, {postcode_col} from {name}" + ), + key_field=key_col, + ) + sources[name] = src + clean_name = CLEAN_NAME.format(src.f(name_col)) + clean_postcode = CLEAN_POSTCODE.format(src.f(postcode_col)) + # Dedupe each source on its standardised name and postcode. + deduped = src.view(cleaning={"name": clean_name, "postcode": clean_postcode}) + nodes[name] = deduped + + dedupes = [ + v.dedupe(NaiveDeduper, {"unique_fields": ["name", "postcode"]}) + for v in nodes.values() + ] + resolver = dedupes[0].resolve(*dedupes[1:]) + + # Standardised entity nodes: one row per entity, read through the dedupe. + entity_nodes = {} + for name, src in sources.items(): + _, name_col, postcode_col = SOURCES[name] + entity_nodes[name] = resolver.view( + src, + cleaning={ + "name": f"any_value({CLEAN_NAME.format(src.f(name_col))})", + "postcode": f"any_value({CLEAN_POSTCODE.format(src.f(postcode_col))})", + }, + group=True, + ) + + # Link every pair (Splink), then consolidate all edges and cluster once. + links = [ + entity_nodes[a].link(entity_nodes[b], SplinkLinker, link_settings) + for a, b in combinations(entity_nodes, 2) + ] + return links[0].resolve(*links[1:]) diff --git a/examples/online/pipeline_splink.py b/examples/online/pipeline_splink.py new file mode 100644 index 00000000..5b758eff --- /dev/null +++ b/examples/online/pipeline_splink.py @@ -0,0 +1,118 @@ +"""The same pipeline, hand-rolled with Splink instead of matchlab. + +The fair opponent, in the spirit of `examples/companies/by_hand.py`: the *same* +algorithm the matchlab plan runs, written out by hand. Dedupe each source on the +standardised name and postcode, link every pair of deduped sources with the shared +Splink model, consolidate the edges, and cluster once with connected components. This is +also the blog's production shape: dedupe in isolation, pairwise transient links, cluster +once. + +The only difference from `pipeline_matchlab.py` is that this holds intermediates in +memory and returns a frame, where matchlab builds a plan and materialises a store. +""" + +from __future__ import annotations + +import warnings +from itertools import combinations + +import duckdb +import polars as pl +from data import SOURCES, WAREHOUSE +from model import CLEAN_NAME, CLEAN_POSTCODE, THRESHOLD, _settings +from splink import DuckDBAPI, Linker +from sqlalchemy import create_engine + + +def load_nodes(warehouse=WAREHOUSE) -> dict[str, pl.DataFrame]: + """Load and standardise every source into `(id, name, postcode)` record frames.""" + engine = create_engine(f"sqlite:///{warehouse}") + con = duckdb.connect(":memory:") + nodes: dict[str, pl.DataFrame] = {} + for source, (key_col, name_col, postcode_col) in SOURCES.items(): + raw = pl.read_database( + f"select {key_col} as key, {name_col} as nm, {postcode_col} as pc " + f"from {source}", + engine.connect(), + ) + con.register("raw", raw) + nodes[source] = con.execute( + f"select key as id, {CLEAN_NAME.format('nm')} as name, " + f"{CLEAN_POSTCODE.format('pc')} as postcode from raw" + ).pl() + con.unregister("raw") + con.close() + return nodes + + +def all_nodes(warehouse=WAREHOUSE) -> pl.DataFrame: + """Every source's standardised records in one frame, the model's training input.""" + return pl.concat(load_nodes(warehouse).values(), how="vertical") + + +def _dedupe( + nodes: dict[str, pl.DataFrame], +) -> tuple[dict[str, pl.DataFrame], pl.DataFrame]: + """Exact-dedupe each source on (name, postcode). Same rule as the NaiveDeduper. + + Returns one entity frame per source (`id`, name, postcode, one row per entity, id + globally unique) and a `(source, key, id)` map from record key to its entity. + """ + entities: dict[str, pl.DataFrame] = {} + key_to_entity: list[pl.DataFrame] = [] + for source, df in nodes.items(): + grouped = df.group_by("name", "postcode").agg(pl.col("id").alias("keys")) + grouped = grouped.with_row_index("i").with_columns( + (pl.lit(f"{source}:") + pl.col("i").cast(pl.Utf8)).alias("id") + ) + entities[source] = grouped.select("id", "name", "postcode") + key_to_entity.append( + grouped.explode("keys").select( + pl.lit(source).alias("source"), + pl.col("keys").alias("key"), + "id", + ) + ) + return entities, pl.concat(key_to_entity) + + +def _link(entities: dict[str, pl.DataFrame], trained: dict) -> list[tuple[str, str]]: + """Link every pair of sources with the shared Splink model; return entity edges.""" + edges: list[tuple[str, str]] = [] + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + for left, right in combinations(entities, 2): + linker = Linker( + [entities[left].to_pandas(), entities[right].to_pandas()], + _settings("link_only", trained), + input_table_aliases=[left, right], + db_api=DuckDBAPI(), + ) + pred = linker.inference.predict( + threshold_match_probability=THRESHOLD + ).as_pandas_dataframe() + edges.extend(zip(pred["id_l"], pred["id_r"], strict=True)) + return edges + + +def resolution(trained: dict, warehouse=WAREHOUSE) -> pl.DataFrame: + """Dedupe, link, cluster once; return `(source, key, root)`.""" + entities, key_to_entity = _dedupe(load_nodes(warehouse)) + edges = _link(entities, trained) + + # Connected components over the entity edges (union-find), then key -> root. + parent: dict[str, str] = { + e: e for df in entities.values() for e in df["id"].to_list() + } + + def find(x: str) -> str: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + for left, right in edges: + parent[find(left)] = find(right) + + roots = pl.DataFrame({"id": list(parent), "root": [find(e) for e in parent]}) + return key_to_entity.join(roots, on="id").select("source", "key", "root") diff --git a/uv.lock b/uv.lock index c2839b48..ed9a1317 100644 --- a/uv.lock +++ b/uv.lock @@ -521,9 +521,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/c6/dba32cab7e3a625b011aa5647486e2d28423a48845a2998c126dd69c85e1/greenlet-3.4.0-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:805bebb4945094acbab757d34d6e1098be6de8966009ab9ca54f06ff492def58", size = 285504, upload-time = "2026-04-08T15:52:14.071Z" }, { url = "https://files.pythonhosted.org/packages/54/f4/7cb5c2b1feb9a1f50e038be79980dfa969aa91979e5e3a18fdbcfad2c517/greenlet-3.4.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:439fc2f12b9b512d9dfa681c5afe5f6b3232c708d13e6f02c845e0d9f4c2d8c6", size = 605476, upload-time = "2026-04-08T16:24:37.064Z" }, { url = "https://files.pythonhosted.org/packages/d6/af/b66ab0b2f9a4c5a867c136bf66d9599f34f21a1bcca26a2884a29c450bd9/greenlet-3.4.0-cp311-cp311-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a70ed1cb0295bee1df57b63bf7f46b4e56a5c93709eea769c1fec1bb23a95875", size = 618336, upload-time = "2026-04-08T16:30:56.59Z" }, - { url = "https://files.pythonhosted.org/packages/6d/31/56c43d2b5de476f77d36ceeec436328533bff960a4cba9a07616e93063ab/greenlet-3.4.0-cp311-cp311-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8c5696c42e6bb5cfb7c6ff4453789081c66b9b91f061e5e9367fa15792644e76", size = 625045, upload-time = "2026-04-08T16:40:37.111Z" }, { url = "https://files.pythonhosted.org/packages/e5/5c/8c5633ece6ba611d64bf2770219a98dd439921d6424e4e8cf16b0ac74ea5/greenlet-3.4.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c660bce1940a1acae5f51f0a064f1bc785d07ea16efcb4bc708090afc4d69e83", size = 613515, upload-time = "2026-04-08T15:56:32.478Z" }, - { url = "https://files.pythonhosted.org/packages/80/ca/704d4e2c90acb8bdf7ae593f5cbc95f58e82de95cc540fb75631c1054533/greenlet-3.4.0-cp311-cp311-manylinux_2_39_riscv64.whl", hash = "sha256:89995ce5ddcd2896d89615116dd39b9703bfa0c07b583b85b89bf1b5d6eddf81", size = 419745, upload-time = "2026-04-08T16:43:04.022Z" }, { url = "https://files.pythonhosted.org/packages/a9/df/950d15bca0d90a0e7395eb777903060504cdb509b7b705631e8fb69ff415/greenlet-3.4.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:ee407d4d1ca9dc632265aee1c8732c4a2d60adff848057cdebfe5fe94eb2c8a2", size = 1574623, upload-time = "2026-04-08T16:26:18.596Z" }, { url = "https://files.pythonhosted.org/packages/1a/e7/0839afab829fcb7333c9ff6d80c040949510055d2d4d63251f0d1c7c804e/greenlet-3.4.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:956215d5e355fffa7c021d168728321fd4d31fd730ac609b1653b450f6a4bc71", size = 1639579, upload-time = "2026-04-08T15:57:29.231Z" }, { url = "https://files.pythonhosted.org/packages/d9/2b/b4482401e9bcaf9f5c97f67ead38db89c19520ff6d0d6699979c6efcc200/greenlet-3.4.0-cp311-cp311-win_amd64.whl", hash = "sha256:5cb614ace7c27571270354e9c9f696554d073f8aa9319079dcba466bbdead711", size = 238233, upload-time = "2026-04-08T17:02:54.286Z" }, @@ -531,9 +529,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/65/8b/3669ad3b3f247a791b2b4aceb3aa5a31f5f6817bf547e4e1ff712338145a/greenlet-3.4.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:1a54a921561dd9518d31d2d3db4d7f80e589083063ab4d3e2e950756ef809e1a", size = 286902, upload-time = "2026-04-08T15:52:12.138Z" }, { url = "https://files.pythonhosted.org/packages/38/3e/3c0e19b82900873e2d8469b590a6c4b3dfd2b316d0591f1c26b38a4879a5/greenlet-3.4.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:16dec271460a9a2b154e3b1c2fa1050ce6280878430320e85e08c166772e3f97", size = 606099, upload-time = "2026-04-08T16:24:38.408Z" }, { url = "https://files.pythonhosted.org/packages/b5/33/99fef65e7754fc76a4ed14794074c38c9ed3394a5bd129d7f61b705f3168/greenlet-3.4.0-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:90036ce224ed6fe75508c1907a77e4540176dcf0744473627785dd519c6f9996", size = 618837, upload-time = "2026-04-08T16:30:58.298Z" }, - { url = "https://files.pythonhosted.org/packages/44/57/eae2cac10421feae6c0987e3dc106c6d86262b1cb379e171b017aba893a6/greenlet-3.4.0-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6f0def07ec9a71d72315cf26c061aceee53b306c36ed38c35caba952ea1b319d", size = 624901, upload-time = "2026-04-08T16:40:38.981Z" }, { url = "https://files.pythonhosted.org/packages/36/f7/229f3aed6948faa20e0616a0b8568da22e365ede6a54d7d369058b128afd/greenlet-3.4.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a1c4f6b453006efb8310affb2d132832e9bbb4fc01ce6df6b70d810d38f1f6dc", size = 615062, upload-time = "2026-04-08T15:56:33.766Z" }, - { url = "https://files.pythonhosted.org/packages/6a/8a/0e73c9b94f31d1cc257fe79a0eff621674141cdae7d6d00f40de378a1e42/greenlet-3.4.0-cp312-cp312-manylinux_2_39_riscv64.whl", hash = "sha256:0e1254cf0cbaa17b04320c3a78575f29f3c161ef38f59c977108f19ffddaf077", size = 423927, upload-time = "2026-04-08T16:43:05.293Z" }, { url = "https://files.pythonhosted.org/packages/08/97/d988180011aa40135c46cd0d0cf01dd97f7162bae14139b4a3ef54889ba5/greenlet-3.4.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:9b2d9a138ffa0e306d0e2b72976d2fb10b97e690d40ab36a472acaab0838e2de", size = 1573511, upload-time = "2026-04-08T16:26:20.058Z" }, { url = "https://files.pythonhosted.org/packages/d4/0f/a5a26fe152fb3d12e6a474181f6e9848283504d0afd095f353d85726374b/greenlet-3.4.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8424683caf46eb0eb6f626cb95e008e8cc30d0cb675bdfa48200925c79b38a08", size = 1640396, upload-time = "2026-04-08T15:57:30.88Z" }, { url = "https://files.pythonhosted.org/packages/42/cf/bb2c32d9a100e36ee9f6e38fad6b1e082b8184010cb06259b49e1266ca01/greenlet-3.4.0-cp312-cp312-win_amd64.whl", hash = "sha256:a0a53fb071531d003b075c444014ff8f8b1a9898d36bb88abd9ac7b3524648a2", size = 238892, upload-time = "2026-04-08T17:03:10.094Z" }, @@ -541,9 +537,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7a/75/7e9cd1126a1e1f0cd67b0eda02e5221b28488d352684704a78ed505bd719/greenlet-3.4.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:43748988b097f9c6f09364f260741aa73c80747f63389824435c7a50bfdfd5c1", size = 285856, upload-time = "2026-04-08T15:52:45.82Z" }, { url = "https://files.pythonhosted.org/packages/9d/c4/3e2df392e5cb199527c4d9dbcaa75c14edcc394b45040f0189f649631e3c/greenlet-3.4.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5566e4e2cd7a880e8c27618e3eab20f3494452d12fd5129edef7b2f7aa9a36d1", size = 610208, upload-time = "2026-04-08T16:24:39.674Z" }, { url = "https://files.pythonhosted.org/packages/da/af/750cdfda1d1bd30a6c28080245be8d0346e669a98fdbae7f4102aa95fff3/greenlet-3.4.0-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1054c5a3c78e2ab599d452f23f7adafef55062a783a8e241d24f3b633ba6ff82", size = 621269, upload-time = "2026-04-08T16:30:59.767Z" }, - { url = "https://files.pythonhosted.org/packages/e0/93/c8c508d68ba93232784bbc1b5474d92371f2897dfc6bc281b419f2e0d492/greenlet-3.4.0-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:98eedd1803353daf1cd9ef23eef23eda5a4d22f99b1f998d273a8b78b70dd47f", size = 628455, upload-time = "2026-04-08T16:40:40.698Z" }, { url = "https://files.pythonhosted.org/packages/54/78/0cbc693622cd54ebe25207efbb3a0eb07c2639cb8594f6e3aaaa0bb077a8/greenlet-3.4.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f82cb6cddc27dd81c96b1506f4aa7def15070c3b2a67d4e46fd19016aacce6cf", size = 617549, upload-time = "2026-04-08T15:56:34.893Z" }, - { url = "https://files.pythonhosted.org/packages/7f/46/cfaaa0ade435a60550fd83d07dfd5c41f873a01da17ede5c4cade0b9bab8/greenlet-3.4.0-cp313-cp313-manylinux_2_39_riscv64.whl", hash = "sha256:b7857e2202aae67bc5725e0c1f6403c20a8ff46094ece015e7d474f5f7020b55", size = 426238, upload-time = "2026-04-08T16:43:06.865Z" }, { url = "https://files.pythonhosted.org/packages/ba/c0/8966767de01343c1ff47e8b855dc78e7d1a8ed2b7b9c83576a57e289f81d/greenlet-3.4.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:227a46251ecba4ff46ae742bc5ce95c91d5aceb4b02f885487aff269c127a729", size = 1575310, upload-time = "2026-04-08T16:26:21.671Z" }, { url = "https://files.pythonhosted.org/packages/b8/38/bcdc71ba05e9a5fda87f63ffc2abcd1f15693b659346df994a48c968003d/greenlet-3.4.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5b99e87be7eba788dd5b75ba1cde5639edffdec5f91fe0d734a249535ec3408c", size = 1640435, upload-time = "2026-04-08T15:57:32.572Z" }, { url = "https://files.pythonhosted.org/packages/a1/c2/19b664b7173b9e4ef5f77e8cef9f14c20ec7fce7920dc1ccd7afd955d093/greenlet-3.4.0-cp313-cp313-win_amd64.whl", hash = "sha256:849f8bc17acd6295fcb5de8e46d55cc0e52381c56eaf50a2afd258e97bc65940", size = 238760, upload-time = "2026-04-08T17:04:03.878Z" },