fix(pfs): families — exception-safe test restore, narrow load_families catch, non-empty activity guard, token-bucketed stem merge (refs #687)

src/pfs/families.py

@@ -8,10 +8,13 @@ DuckDB, no narwhals — it must import inside the ``llm`` container.
 
 from __future__ import annotations
 
+import logging
 import re
 from dataclasses import dataclass
 from typing import Any, Mapping, Sequence
 
+_logger = logging.getLogger(__name__)
+
 # HCPCS level II (letter + 4 digits) or CPT (5 digits). Validation against
 # the fee schedule happens at lookup time (pfs.valuation), not here.
 CODE_RE = re.compile(r"\b(?:[A-Z]\d{4}|\d{5})\b", re.IGNORECASE)
@@ -218,15 +221,28 @@ def derive_families(
         c: frozenset(e.value for e in elements.get(c, ()) if e.type == "activity")
         for c in codes
     }
-    for i, a in enumerate(codes):
+    # Inverted index token -> codes, so we only compare pairs that share at
-        for b in codes[i + 1 :]:
+    # least one stem token. Exact for the >= 0.5 threshold: any qualifying
-            if (
+    # pair has a non-empty intersection (jaccard > 0 implies shared tokens),
-                stems[a]
+    # so no comparison is skipped, but the O(n^2) all-pairs scan is.
-                and stems[b]
+    buckets: dict[str, list[str]] = {}
-                and _jaccard(stems[a], stems[b]) >= 0.5
+    for c, toks in stems.items():
-                and acts[a] == acts[b]
+        for t in toks:
-            ):
+            buckets.setdefault(t, []).append(c)
-                union(a, b)
+    seen: set[tuple[str, str]] = set()
+    for members in buckets.values():
+        for i, a in enumerate(members):
+            for b in members[i + 1 :]:
+                pair = (a, b) if a < b else (b, a)
+                if pair in seen:
+                    continue
+                seen.add(pair)
+                if (
+                    acts[a]
+                    and acts[a] == acts[b]
+                    and _jaccard(stems[a], stems[b]) >= 0.5
+                ):
+                    union(a, b)
 
     groups: dict[str, list[str]] = {}
     for c in codes:
@@ -285,8 +301,12 @@ def load_families(con: Any) -> dict[str, Family]:
 
     try:
         rows = read_families(con)
-    except Exception:  # table missing on an old replica
+    except Exception as exc:  # table/schema missing on an old replica
-        return {}
+        msg = str(exc)
+        if "code_family" in msg and ("Catalog" in msg or "does not exist" in msg):
+            _logger.info("pfs.code_family not present; no derived families: %s", exc)
+            return {}
+        raise
     out: dict[str, Family] = {}
     for r in rows:
         fam = out.get(r.key)

tests/pfs/test_families.py

@@ -2,7 +2,10 @@
 
 from __future__ import annotations
 
+import time
+
 import duckdb
+import pytest
 
 from pfs.codetables import (
     ElementRow,
@@ -25,6 +28,21 @@ from pfs.families import (
 )
 
 
+@pytest.fixture
+def restore_families():
+    """Snapshot ``pfs.families.FAMILIES`` and restore it in teardown, even
+    if the test body raises — a test that calls ``refresh_from`` must not
+    leave the live registry clobbered for the rest of the session."""
+    from pfs import families as mod
+
+    before = dict(mod.FAMILIES)
+    try:
+        yield mod
+    finally:
+        mod.FAMILIES.clear()
+        mod.FAMILIES.update(before)
+
+
 class TestFindCodes:
     def test_hcpcs_and_cpt(self):
         assert find_codes("Codes G0556 and 99490 apply; see g0557.") == (
@@ -197,11 +215,70 @@ class TestDerive:
         keys = {r.key for r in rows}
         assert len(keys) == 1 and next(iter(keys)).startswith("REM-MNTR")
 
+    def test_no_activity_elements_never_merges_on_stem_alone(self):
+        # Both codes have identical stem tokens (jaccard == 1.0, well past
+        # the 0.5 threshold) but neither has an activity element — an empty
+        # activity set must not compare equal-and-qualifying, or every
+        # activity-less code with a common description would merge with
+        # every other one. Word order differs so a stem-based merge would
+        # be visible as a shared key even though the codes never touch.
+        elements = {"11111": [], "22222": []}
+        descriptions = {"11111": "Foo bar widget", "22222": "Bar widget foo"}
+        rows = derive_families(elements, {}, descriptions)
+        keys = {r.code: r.key for r in rows}
+        assert keys["11111"] != keys["22222"]
+
+    def test_disjoint_stem_never_merges(self):
+        # 10000/10001 share enough stem tokens (and a matching activity) to
+        # merge; 20000's stem shares no token with either, so it must never
+        # even be compared, let alone merged.
+        elements = {
+            "10000": [_el("10000", "activity", "act-a")],
+            "10001": [_el("10001", "activity", "act-a")],
+            "20000": [_el("20000", "activity", "act-a")],
+        }
+        descriptions = {
+            "10000": "Alpha beta gamma",
+            "10001": "Alpha beta delta",
+            "20000": "Zeta eta theta",
+        }
+        rows = derive_families(elements, {}, descriptions)
+        keys = {r.code: r.key for r in rows}
+        assert keys["10000"] == keys["10001"]
+        assert keys["20000"] != keys["10000"]
+
+    def test_many_distinct_stems_is_fast(self):
+        # ~2,000 codes whose descriptions share no stem token with any
+        # other code's. The token-bucketed merge must not degrade to the
+        # O(n^2) all-pairs scan this guards against; kept generous (< 2s)
+        # so CI timing noise doesn't make it flaky.
+        def _word(i: int) -> str:
+            letters = []
+            n = i + 1
+            while n:
+                n, r = divmod(n - 1, 26)
+                letters.append(chr(97 + r))
+            return "".join(reversed(letters))
+
+        n = 2000
+        elements = {}
+        descriptions = {}
+        for i in range(n):
+            code = f"{10000 + i}"
+            elements[code] = [_el(code, "activity", "act")]
+            descriptions[code] = f"stem{_word(i)} term{_word(i + n)}"
+
+        start = time.perf_counter()
+        rows = derive_families(elements, {}, descriptions)
+        elapsed = time.perf_counter() - start
+
+        assert len(rows) == n
+        assert elapsed < 2.0
+
 
 class TestLoadAndRefresh:
-    def test_load_and_refresh_in_place(self):
+    def test_load_and_refresh_in_place(self, restore_families):
-        from pfs import families as mod
+        mod = restore_families
-
         con = duckdb.connect(":memory:")
         try:
             ensure_tables(con)
@@ -239,8 +316,29 @@ class TestLoadAndRefresh:
                 and set(mod.FAMILIES) == {"CCM"}
                 and "G2058" in mod.FAMILIES["CCM"].codes
             )
-            # restore for other tests
+        finally:
-            mod.FAMILIES.clear()
+            con.close()
-            mod.FAMILIES.update(before)
+
+    def test_schema_present_table_absent_returns_empty(self):
+        # pfs schema created (e.g. by an earlier ensure_tables call for a
+        # sibling table) but pfs.code_family itself never materialized —
+        # an old replica shape, not a real error.
+        con = duckdb.connect(":memory:")
+        try:
+            con.execute("CREATE SCHEMA IF NOT EXISTS pfs;")
+            assert load_families(con) == {}
+        finally:
+            con.close()
+
+    def test_wrong_shape_table_raises(self):
+        # A pfs.code_family table exists but not in the expected shape —
+        # this is a real bug, not an absent table, and must not be
+        # swallowed.
+        con = duckdb.connect(":memory:")
+        try:
+            con.execute("CREATE SCHEMA IF NOT EXISTS pfs;")
+            con.execute("CREATE TABLE pfs.code_family (only_col VARCHAR);")
+            with pytest.raises(Exception):
+                load_families(con)
         finally:
             con.close()