Compute document centroids during the vector scan

haiku_rag_slim/haiku/rag/doctor.py

@@ -243,43 +243,36 @@ class _DuplicateFamily(BaseModel):
 
 
 def _duplicate_families(
-    doc_vectors: dict[str, np.ndarray], cfg: DuplicateDetectionConfig
+    doc_ids: list[str],
+    centroids: np.ndarray,
+    counts: np.ndarray,
+    cfg: DuplicateDetectionConfig,
 ) -> list[_DuplicateFamily]:
     """Cluster documents whose embedding centroids are nearly identical.
 
-    One unit centroid per document, pairwise cosine, then connected components
+    ``centroids`` holds one summed (unnormalized) centroid per document and
-    of pairs above ``similarity_threshold``. One family per component, each
+    ``counts`` its embedded-chunk count. Documents below the small-document
-    carrying every member's highest cosine to another member of the family.
+    floor are dropped; the rest are normalized and clustered by union-find over
+    pairwise cosine above ``similarity_threshold``. One family per component,
+    each carrying every member's highest cosine to another member.
     """
-    # Drop unembedded (zero) vectors and documents below the small-document
+    centroids = np.asarray(centroids, dtype=np.float32)
-    # floor; reduce each remaining document to a unit centroid.
+    counts = np.asarray(counts)
-    centroids_by_doc: dict[str, np.ndarray] = {}
+    norms = np.linalg.norm(centroids, axis=1)
-    sizes: dict[str, int] = {}
+    eligible = np.nonzero((counts >= cfg.min_chunks) & (norms > 0))[0]
-    for doc_id, matrix in doc_vectors.items():
+    if eligible.size < 2:
-        m = np.asarray(matrix, dtype=np.float32)
-        if m.ndim != 2 or m.shape[0] == 0:
-            continue
-        m = m[np.linalg.norm(m, axis=1) > 0]
-        if m.shape[0] < cfg.min_chunks:
-            continue
-        centroid = m.mean(axis=0)
-        norm = np.linalg.norm(centroid)
-        if norm == 0:
-            continue
-        centroids_by_doc[doc_id] = centroid / norm
-        sizes[doc_id] = int(m.shape[0])
-    if len(centroids_by_doc) < 2:
         return []
-
+    unit = centroids[eligible] / norms[eligible][:, None]
-    order = sorted(centroids_by_doc)
+    ids = [doc_ids[i] for i in eligible]
-    centroids = np.array([centroids_by_doc[d] for d in order])
+    sizes = {doc_ids[i]: int(counts[i]) for i in eligible}
+    n = len(ids)
 
     # Pairwise cosine, block-wise to avoid a full D×D matrix at once. Each row
     # only compares against higher-indexed documents (upper triangle). Cluster
     # with union-find and keep only each document's best similarity to a twin —
     # a self-similar corpus forms one clique, so storing every pair would be
     # O(D²) objects.
-    parent = list(range(len(order)))
+    parent = list(range(n))
 
     def find(x: int) -> int:
         while parent[x] != x:
@@ -287,11 +280,11 @@ def _duplicate_families(
             x = parent[x]
         return x
 
-    best = np.zeros(len(order), dtype=np.float32)
+    best = np.zeros(n, dtype=np.float32)
     linked = False
     block = 512
-    for start in range(0, len(order), block):
+    for start in range(0, n, block):
-        sims = centroids[start : start + block] @ centroids.T
+        sims = unit[start : start + block] @ unit.T
         for row in range(sims.shape[0]):
             gi = start + row
             cols = (
@@ -310,21 +303,21 @@ def _duplicate_families(
         return []
 
     components: dict[int, list[int]] = {}
-    for idx in range(len(order)):
+    for idx in range(n):
         components.setdefault(find(idx), []).append(idx)
 
     families: list[_DuplicateFamily] = []
     for indices in components.values():
         if len(indices) < 2:
             continue
-        members = sorted(order[i] for i in indices)
+        members = sorted(ids[i] for i in indices)
         # Largest document (most chunks) is the one to keep; smallest id on a tie.
         keep = min(members, key=lambda d: (-sizes[d], d))
         families.append(
             _DuplicateFamily(
                 members=members,
                 keep=keep,
-                similarity={order[i]: round(float(best[i]), 3) for i in indices},
+                similarity={ids[i]: round(float(best[i]), 3) for i in indices},
                 sizes={d: sizes[d] for d in members},
             )
         )
@@ -375,13 +368,15 @@ def _write_duplicates_out(
 
 
 def _check_duplicate_documents(
-    doc_vectors: dict[str, np.ndarray],
+    doc_ids: list[str],
+    centroids: np.ndarray,
+    counts: np.ndarray,
     uri_by_doc: Mapping[str, str | None],
     title_by_doc: Mapping[str, str | None],
     cfg: DuplicateDetectionConfig,
     yaml_path: Path | None = None,
 ) -> CheckResult:
-    families = _duplicate_families(doc_vectors, cfg)
+    families = _duplicate_families(doc_ids, centroids, counts, cfg)
 
     def label(doc_id: str) -> str:
         return uri_by_doc.get(doc_id) or title_by_doc.get(doc_id) or doc_id
@@ -618,43 +613,59 @@ async def run_db_checks(
             )
         )
 
-    ids = arrow.column("id").to_pylist()
     # Reshape the Arrow fixed-size-list child buffer directly into an (N, dim)
     # float32 matrix. Going through to_pylist() would box N*dim Python floats
     # (tens of GB and most of the wall-clock on large corpora); the stored
     # vectors are already float32, so this keeps the layout and the dtype.
     vec_col = arrow.column("vector").combine_chunks()
     vectors = vec_col.values.to_numpy(zero_copy_only=False).reshape(-1, actual_dim)
-    zero_ids: list[str] = []
+    embedded = vectors.any(axis=1) if vectors.size else np.zeros(0, dtype=bool)
-    if vectors.size:
+
-        zero_ids = [ids[i] for i in np.nonzero(~vectors.any(axis=1))[0]]
+    # Unembedded (all-zero) chunks: report a count and a few sampled ids without
+    # materializing every chunk id.
+    zero_rows = np.nonzero(~embedded)[0]
+    zero_count = int(zero_rows.size)
+    id_col = arrow.column("id")
+    zero_sample = [id_col[int(i)].as_py() for i in zero_rows[:_SAMPLE_LIMIT]]
+    if zero_count > _SAMPLE_LIMIT:
+        zero_sample.append(f"... (+{zero_count - _SAMPLE_LIMIT} more)")
     results.append(
         CheckResult(
             name="unembedded_chunks",
-            severity=Severity.WARN if zero_ids else Severity.OK,
+            severity=Severity.WARN if zero_count else Severity.OK,
             message=(
-                f"{len(zero_ids)} chunk(s) have an all-zero (unembedded) vector."
+                f"{zero_count} chunk(s) have an all-zero (unembedded) vector."
-                if zero_ids
+                if zero_count
                 else "All chunks are embedded."
             ),
-            remediation="haiku-rag rebuild --embed-only" if zero_ids else None,
+            remediation="haiku-rag rebuild --embed-only" if zero_count else None,
-            details=_sample(zero_ids),
+            details=zero_sample,
         )
     )
 
-    # Near-identical documents (centroid cosine), grouped from the same vector
+    # Near-identical documents (centroid cosine). Reduce each document's chunk
-    # scan rather than a second pass.
+    # vectors to one summed centroid during the scan: dictionary-encode the
-    chunk_doc_ids_ordered = arrow.column("document_id").to_pylist()
+    # document ids into integer codes, then sum each document's embedded rows in
-    indices_by_doc: dict[str, list[int]] = {}
+    # a single pass per document — no second full copy of the vector matrix.
-    for index, doc_id in enumerate(chunk_doc_ids_ordered):
+    encoded = arrow.column("document_id").combine_chunks().dictionary_encode()
-        indices_by_doc.setdefault(doc_id, []).append(index)
+    doc_ids = encoded.dictionary.to_pylist()
-    doc_vectors = {doc_id: vectors[idx] for doc_id, idx in indices_by_doc.items()}
+    codes = encoded.indices.to_numpy(zero_copy_only=False)
-    # Per-doc fancy indexing has copied every vector; release the full matrix so
+    centroids = np.zeros((len(doc_ids), actual_dim), dtype=np.float32)
-    # both copies are not resident during duplicate detection.
+    counts = np.zeros(len(doc_ids), dtype=np.int64)
+    order = np.argsort(codes, kind="stable")
+    bounds = np.searchsorted(codes, np.arange(len(doc_ids) + 1), sorter=order)
+    for d in range(len(doc_ids)):
+        rows = order[bounds[d] : bounds[d + 1]]
+        rows = rows[embedded[rows]]
+        counts[d] = rows.size
+        if rows.size:
+            centroids[d] = vectors[rows].sum(axis=0)
     del vectors
     results.append(
         _check_duplicate_documents(
-            doc_vectors,
+            doc_ids,
+            centroids,
+            counts,
             uri_by_doc,
             title_by_doc,
             config.doctor.duplicates,

tests/test_doctor.py

@@ -944,22 +944,28 @@ async def test_probe_endpoint_connection_error():
 # --- Duplicate-document detection ----------------------------------------
 
 
-def _docs(spec: dict[str, list[int]], dim: int = 8) -> dict[str, np.ndarray]:
+def _centroids(
-    """Build per-document chunk matrices from one-hot indices.
+    spec: dict[str, list[int]], dim: int = 8
+) -> tuple[list[str], np.ndarray, np.ndarray]:
+    """Summed one-hot centroids + chunk counts per document, as
+    ``_duplicate_families`` consumes them.
 
     Orthogonal one-hot chunks make the centroid cosine of two documents equal to
     ``shared / sqrt(len(a) * len(b))``: identical documents score 1.0, fully
     distinct documents score 0.0.
     """
     eye = np.eye(dim)
-    return {
+    doc_ids = list(spec)
-        doc: np.array([eye[i] for i in idxs], dtype=float) for doc, idxs in spec.items()
+    centroids = np.array(
-    }
+        [eye[idxs].sum(axis=0) for idxs in spec.values()], dtype=np.float32
+    )
+    counts = np.array([len(idxs) for idxs in spec.values()], dtype=np.int64)
+    return doc_ids, centroids, counts
 
 
 def test_duplicate_families_identical_docs_flagged():
     families = _duplicate_families(
-        _docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3]}), DuplicateDetectionConfig()
+        *_centroids({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3]}), DuplicateDetectionConfig()
     )
     assert len(families) == 1
     assert set(families[0].members) == {"a", "b"}
@@ -970,21 +976,22 @@ def test_duplicate_families_append_only_not_flagged():
     # A is fully contained in the larger B, but their centroids diverge
     # (cosine sqrt(3/6) ~= 0.71), so it stays below the similarity cutoff.
     families = _duplicate_families(
-        _docs({"a": [0, 1, 2], "b": [0, 1, 2, 3, 4, 5]}), DuplicateDetectionConfig()
+        *_centroids({"a": [0, 1, 2], "b": [0, 1, 2, 3, 4, 5]}),
+        DuplicateDetectionConfig(),
     )
     assert families == []
 
 
 def test_duplicate_families_distinct_docs_none():
     families = _duplicate_families(
-        _docs({"a": [0, 1, 2], "b": [3, 4, 5]}), DuplicateDetectionConfig()
+        *_centroids({"a": [0, 1, 2], "b": [3, 4, 5]}), DuplicateDetectionConfig()
     )
     assert families == []
 
 
 def test_duplicate_families_three_way_one_family():
     families = _duplicate_families(
-        _docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [0, 1, 2, 3]}),
+        *_centroids({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [0, 1, 2, 3]}),
         DuplicateDetectionConfig(),
     )
     assert len(families) == 1
@@ -997,7 +1004,7 @@ def test_duplicate_families_clique_single_family():
     # A self-similar corpus (all identical) is one clique. Union-find collapses
     # it to a single family without materializing every pair.
     spec = {chr(ord("a") + k): [0, 1, 2, 3] for k in range(8)}
-    families = _duplicate_families(_docs(spec), DuplicateDetectionConfig())
+    families = _duplicate_families(*_centroids(spec), DuplicateDetectionConfig())
     assert len(families) == 1
     assert set(families[0].members) == set(spec)
     assert all(s == pytest.approx(1.0) for s in families[0].similarity.values())
@@ -1006,7 +1013,7 @@ def test_duplicate_families_clique_single_family():
 def test_duplicate_families_tiny_docs_ignored():
     # min_chunks = 3 excludes the one-chunk documents.
     families = _duplicate_families(
-        _docs({"a": [0], "b": [0]}), DuplicateDetectionConfig()
+        *_centroids({"a": [0], "b": [0]}), DuplicateDetectionConfig()
     )
     assert families == []
 
@@ -1014,9 +1021,9 @@ def test_duplicate_families_tiny_docs_ignored():
 def test_duplicate_families_threshold_is_configurable():
     # Share 3 of 4 chunks each -> centroid cosine 0.75.
     spec = {"a": [0, 1, 2, 3], "b": [0, 1, 2, 4]}
-    assert _duplicate_families(_docs(spec), DuplicateDetectionConfig()) == []
+    assert _duplicate_families(*_centroids(spec), DuplicateDetectionConfig()) == []
     flagged = _duplicate_families(
-        _docs(spec), DuplicateDetectionConfig(similarity_threshold=0.7)
+        *_centroids(spec), DuplicateDetectionConfig(similarity_threshold=0.7)
     )
     assert len(flagged) == 1
     assert set(flagged[0].members) == {"a", "b"}
@@ -1029,9 +1036,10 @@ def test_duplicate_documents_report_truncates_summary():
         idx = [3 * k, 3 * k + 1, 3 * k + 2]
         spec[f"a{k}"] = idx
         spec[f"b{k}"] = list(idx)
-    docs = _docs(spec, dim=3 * pairs)
     uris = {d: f"file:///srv/shared/library/docs/{d}.pdf" for d in spec}
-    result = _check_duplicate_documents(docs, uris, {}, DuplicateDetectionConfig())
+    result = _check_duplicate_documents(
+        *_centroids(spec, dim=3 * pairs), uris, {}, DuplicateDetectionConfig()
+    )
     assert result.severity is Severity.WARN
     # The summary message still reports the full total.
     assert f"{pairs} group(s)" in result.message
@@ -1042,10 +1050,14 @@ def test_duplicate_documents_report_truncates_summary():
 
 
 def test_duplicate_documents_report_factors_common_path():
-    docs = _docs({"a": [0, 1, 2], "b": [0, 1, 2]}, dim=3)
     base = "file:///srv/shared/library/docs/"
     uris = {"a": base + "alpha.pdf", "b": base + "beta.pdf"}
-    result = _check_duplicate_documents(docs, uris, {}, DuplicateDetectionConfig())
+    result = _check_duplicate_documents(
+        *_centroids({"a": [0, 1, 2], "b": [0, 1, 2]}, dim=3),
+        uris,
+        {},
+        DuplicateDetectionConfig(),
+    )
     assert f"common path: {base}" in result.details
     member_lines = [d for d in result.details if d.lstrip().startswith("#")]
     assert {d.strip() for d in member_lines} == {"#1 alpha.pdf", "#2 beta.pdf"}
@@ -1054,11 +1066,11 @@ def test_duplicate_documents_report_factors_common_path():
 
 def test_duplicate_documents_writes_yaml(tmp_path):
     # a,b identical (a 4-chunk duplicate); c distinct and excluded.
-    docs = _docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [4, 5, 6]}, dim=8)
+    spec = {"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [4, 5, 6]}
     uris = {"a": "file:///x/a.pdf", "b": "file:///x/b.pdf", "c": "file:///x/c.pdf"}
     out = tmp_path / "dups.yaml"
     _check_duplicate_documents(
-        docs, uris, {}, DuplicateDetectionConfig(), yaml_path=out
+        *_centroids(spec, dim=8), uris, {}, DuplicateDetectionConfig(), yaml_path=out
     )
     data = yaml.safe_load(out.read_text())
     assert len(data["groups"]) == 1
@@ -1076,10 +1088,14 @@ def test_duplicate_documents_writes_yaml(tmp_path):
 
 
 def test_duplicate_documents_writes_empty_yaml_when_none(tmp_path):
-    docs = _docs({"a": [0, 1, 2], "b": [3, 4, 5]}, dim=6)  # distinct
+    spec = {"a": [0, 1, 2], "b": [3, 4, 5]}  # distinct
     out = tmp_path / "dups.yaml"
     _check_duplicate_documents(
-        docs, {"a": "u", "b": "v"}, {}, DuplicateDetectionConfig(), yaml_path=out
+        *_centroids(spec, dim=6),
+        {"a": "u", "b": "v"},
+        {},
+        DuplicateDetectionConfig(),
+        yaml_path=out,
     )
     assert yaml.safe_load(out.read_text()) == {"groups": []}