Extarct utils

pipeline/utils/fuzzy_join.py

@@ -0,0 +1,179 @@
+import re
+import shutil
+import tempfile
+from concurrent.futures import ProcessPoolExecutor
+from os import cpu_count
+from pathlib import Path
+
+import polars as pl
+from thefuzz import fuzz
+from tqdm import tqdm
+
+_NUMBER_RE = re.compile(r'\d+')
+
+
+def _normalize(s: pl.Expr) -> pl.Expr:
+    return (
+        s.str.to_uppercase()
+        .str.replace_all(r'[,.\-]', ' ')
+        .str.replace_all(r'\s+', ' ')
+        .str.strip_chars()
+    )
+
+
+def fuzzy_join_on_postcode(
+    left: pl.LazyFrame,
+    right: pl.LazyFrame,
+    left_address_col: str,
+    right_address_col: str,
+    left_postcode_col: str,
+    right_postcode_col: str,
+) -> pl.LazyFrame:
+    """Fuzzy join two LazyFrames by matching addresses within postcode buckets.
+
+    Sinks each side to a temporary parquet file so the upstream pipeline
+    executes only once.  The matching phase collects just three narrow
+    columns (index, address, postcode) via projection pushdown, and the
+    final join reads the remaining columns lazily.
+
+    Returns a LazyFrame with all left and right columns.  Unmatched rows
+    have null right columns.
+    """
+
+    tmpdir = tempfile.mkdtemp(prefix='fuzzy_join_')
+    left_path = Path(tmpdir) / 'left.parquet'
+    right_path = Path(tmpdir) / 'right.parquet'
+
+    try:
+        # Materialise each side exactly once, with a row index, to temp parquet.
+        left.with_row_index('_left_idx').sink_parquet(left_path)
+        right.with_row_index('_right_idx').sink_parquet(right_path)
+
+        # Collect only the narrow columns needed for matching (projection pushdown).
+        left_match = (
+            pl.scan_parquet(left_path)
+            .select(
+                '_left_idx',
+                _normalize(pl.col(left_address_col)).alias('_left_address'),
+                pl.col(left_postcode_col).str.strip_chars().str.to_uppercase().alias('_left_postcode'),
+            )
+            .collect()
+        )
+
+        right_match = (
+            pl.scan_parquet(right_path)
+            .select(
+                '_right_idx',
+                _normalize(pl.col(right_address_col)).alias('_right_address'),
+                pl.col(right_postcode_col).str.strip_chars().str.to_uppercase().alias('_right_postcode'),
+            )
+            .unique(subset=['_right_address', '_right_postcode'], keep='first')
+            .collect()
+        )
+
+        # Group right side by postcode for fast lookup
+        right_by_postcode: dict[str, list[tuple[int, str]]] = {}
+        for idx, postcode, address in zip(
+            right_match['_right_idx'], right_match['_right_postcode'], right_match['_right_address']
+        ):
+            if postcode is not None:
+                right_by_postcode.setdefault(postcode, []).append((idx, address))
+
+        # Group left side by postcode
+        left_by_postcode: dict[str, list[tuple[int, str]]] = {}
+        for idx, postcode, address in zip(
+            left_match['_left_idx'], left_match['_left_postcode'], left_match['_left_address']
+        ):
+            if address is not None and postcode is not None:
+                left_by_postcode.setdefault(postcode, []).append((idx, address))
+
+        del left_match, right_match
+
+        # Build tasks for each postcode bucket
+        tasks = [
+            (left_entries, right_by_postcode[postcode])
+            for postcode, left_entries in left_by_postcode.items()
+            if postcode in right_by_postcode
+        ]
+
+        # Score all pairwise matches in parallel, then greedily assign from
+        # highest score downward so best pairs lock in first.
+        all_pairs: list[tuple[int, int, int]] = []  # (score, left_idx, right_idx)
+        with ProcessPoolExecutor(max_workers=cpu_count()) as executor:
+            for pairs in tqdm(
+                executor.map(_score_bucket, tasks, chunksize=64),
+                total=len(tasks),
+                desc='Fuzzy matching',
+            ):
+                all_pairs.extend(pairs)
+
+        del tasks, left_by_postcode, right_by_postcode
+
+        # Sort descending by score so best matches are assigned first
+        all_pairs.sort(key=lambda t: (t[0], -t[1]), reverse=True)
+
+        matches: list[tuple[int, int]] = []
+        matched_left: set[int] = set()
+        matched_right: set[int] = set()
+
+        for _score, left_idx, right_idx in all_pairs:
+            if left_idx in matched_left or right_idx in matched_right:
+                continue
+            matches.append((left_idx, right_idx))
+            matched_left.add(left_idx)
+            matched_right.add(right_idx)
+
+        del all_pairs, matched_left, matched_right
+
+        # Build a small mapping LazyFrame and join back to the cached parquets.
+        if matches:
+            mapping = pl.LazyFrame({
+                '_left_idx': pl.Series([m[0] for m in matches], dtype=pl.UInt32),
+                '_right_idx': pl.Series([m[1] for m in matches], dtype=pl.UInt32),
+            })
+        else:
+            mapping = pl.LazyFrame({
+                '_left_idx': pl.Series([], dtype=pl.UInt32),
+                '_right_idx': pl.Series([], dtype=pl.UInt32),
+            })
+
+        left_cached = pl.scan_parquet(left_path)
+        right_cached = pl.scan_parquet(right_path)
+
+        return (
+            left_cached
+            .join(mapping, on='_left_idx', how='left')
+            .join(right_cached, on='_right_idx', how='left')
+            .drop('_left_idx', '_right_idx')
+        )
+    except BaseException:
+        shutil.rmtree(tmpdir, ignore_errors=True)
+        raise
+
+
+def _numbers_compatible(a: str, b: str) -> bool:
+    """Check that numeric tokens (flat/house numbers) in the shorter set are a subset of the longer.
+
+    Returns False if one address has numbers and the other doesn't.
+    """
+    nums_a = set(_NUMBER_RE.findall(a))
+    nums_b = set(_NUMBER_RE.findall(b))
+    smaller, larger = (nums_a, nums_b) if len(nums_a) <= len(nums_b) else (nums_b, nums_a)
+    if not smaller and larger:
+        return False
+    return smaller.issubset(larger)
+
+
+def _score_bucket(
+    args: tuple[list[tuple[int, str]], list[tuple[int, str]], int],
+) -> list[tuple[int, int, int]]:
+    """Score all address pairs within a single postcode bucket."""
+    left_entries, right_entries = args
+    pairs = []
+    for left_row, left_address in left_entries:
+        for right_row, right_address in right_entries:
+            if not _numbers_compatible(left_address, right_address):
+                continue
+            score = fuzz.token_sort_ratio(left_address, right_address)
+            pairs.append((score, left_row, right_row))
+    return pairs