code/minhash_cluster.py

"""Step 3: near-duplicate clustering of the canonical skills (MinHash + LSH), bounded RAM on a 3.8 GB box.

  sig     per skill: word 5-gram shingles (unigrams for very short texts) -> 128 multiply-shift MinHash values,
          stored in a disk memmap (N x 128 uint32) + an index table (row -> sha); 2 worker processes
  lsh     16 bands x 8 rows -> (band, bucket) per skill -> DuckDB groups buckets with >1 member on disk
  union   every bucket member is verified against the bucket's first member (estimated Jaccard >= 0.8 from
          signature agreement) and merged with union-find
  out     clusters.parquet: sha, cluster_id (sha of representative), cluster_size, is_rep
          representative = most July copies, then most stars, then most complete bundle
  upload  processed/canonical/clusters.parquet + stats/near_dup_summary.json

Usage: python code/minhash_cluster.py [sig|lsh|union|out|upload|all]     input: work/canon/final/part-*.parquet
"""
import json, os, re, sys, time, zlib
from pathlib import Path
import numpy as np
import pyarrow as pa, pyarrow.parquet as pq

ROOT = Path(__file__).resolve().parent.parent
os.environ.setdefault("HOME", "/root")
W = ROOT / os.environ.get("MH_DIR", "work/minhash"); W.mkdir(parents=True, exist_ok=True)
SRC = os.environ.get("MH_SRC", str(ROOT / "work/canon/text/*.parquet"))   # distinct texts (sha, skill_md) = canonical set
META = os.environ.get("MH_META", str(ROOT / "work/canon/canon/canonical_meta-*.parquet"))   # ranking fields per sha
K, BANDS, ROWS, SHINGLE, THRESH = 128, 16, 8, 5, 0.8
rng = np.random.default_rng(20261002)
A = (rng.integers(1, 2**63, K, dtype=np.uint64) | np.uint64(1))   # odd multipliers (multiply-shift hashing)
B = rng.integers(0, 2**63, K, dtype=np.uint64)
TOK = re.compile(r"\w+")
P = np.uint64(1099511628211)


def signature(text):
    toks = TOK.findall(text.lower())
    if not toks: return None
    ids = np.fromiter((zlib.crc32(t.encode()) for t in toks), dtype=np.uint64, count=len(toks))
    if len(ids) >= SHINGLE:   # rolling polynomial hash of word 5-grams
        h = np.zeros(len(ids) - SHINGLE + 1, dtype=np.uint64)
        for j in range(SHINGLE): h = h * P + ids[j:len(ids) - SHINGLE + 1 + j]
    else:
        h = ids
    h = np.unique(h)[:4000]   # bottom-k: the 4,000 smallest shingle hashes (consistent sample across documents)
    with np.errstate(over="ignore"):
        v = (h[:, None] * A[None, :] + B[None, :]) >> np.uint64(32)
    return v.min(axis=0).astype(np.uint32)


def _sig_part(args):
    path, out = args
    if Path(out).exists(): return out
    shas, sigs = [], []
    for bt in pq.ParquetFile(path).iter_batches(batch_size=2000, columns=["sha", "skill_md"]):
        for r in bt.to_pylist():
            s = signature(r["skill_md"] or "")
            if s is None: continue
            shas.append(r["sha"]); sigs.append(s)
    np.save(out + ".tmp.npy", np.stack(sigs) if sigs else np.zeros((0, K), np.uint32))
    pq.write_table(pa.table({"sha": shas}), out + ".sha.parquet")
    os.replace(out + ".tmp.npy", out)
    return out


def sig():
    import glob
    from multiprocessing import Pool
    parts = sorted(glob.glob(SRC))
    jobs = [(p, str(W / f"sig-{Path(p).stem}.npy")) for p in parts]
    with Pool(2) as pool:
        for o in pool.imap_unordered(_sig_part, jobs): print("  signed", Path(o).name, flush=True)
    # concatenate into one memmap + index
    n = sum(np.load(o, mmap_mode="r").shape[0] for _, o in jobs)
    mm = np.lib.format.open_memmap(W / "sigs.npy", mode="w+", dtype=np.uint32, shape=(n, K)); i = 0; shas = []
    for _, o in jobs:
        a = np.load(o, mmap_mode="r"); mm[i:i + len(a)] = a; i += len(a)
        shas += pq.read_table(o + ".sha.parquet").column("sha").to_pylist()
    mm.flush(); del mm
    pq.write_table(pa.table({"row": np.arange(n, dtype=np.int64), "sha": shas}), W / "index.parquet")
    for _, o in jobs: os.remove(o); os.remove(o + ".sha.parquet")
    print("sig done", n, flush=True)


def lsh():
    """Per band: hash the band's 8 values, argsort by bucket hash, emit runs of >1 equal hashes as candidate groups.
    Pure numpy (~120 MB per band) — DuckDB's list() aggregation over all 41M (band, bucket) rows ran out of memory."""
    sigs = np.load(W / "sigs.npy", mmap_mode="r"); n = len(sigs)
    schema = pa.schema([("band", pa.int8()), ("bucket", pa.int64()), ("members", pa.list_(pa.int64()))])
    w = pq.ParquetWriter(W / "buckets.parquet.tmp", schema, compression="zstd"); groups = members_total = 0
    for b in range(BANDS):
        band = np.asarray(sigs[:, b * ROWS:(b + 1) * ROWS]).astype(np.uint64)
        h = np.zeros(n, dtype=np.uint64)
        with np.errstate(over="ignore"):
            for j in range(ROWS): h = h * P + band[:, j]
        del band
        order = np.argsort(h, kind="stable"); hs = h[order]
        cut = np.flatnonzero(np.diff(hs) != 0) + 1
        starts = np.concatenate(([0], cut)); ends = np.concatenate((cut, [n]))
        big = np.flatnonzero(ends - starts > 1)
        for k in range(0, len(big), 50000):
            sel = big[k:k + 50000]
            mem = [order[starts[i]:ends[i]].astype(np.int64) for i in sel]
            w.write_table(pa.table({"band": np.full(len(sel), b, np.int8), "bucket": hs[starts[sel]].view(np.int64),
                                    "members": pa.array(mem, type=pa.list_(pa.int64()))}, schema=schema))
            groups += len(sel); members_total += sum(len(m) for m in mem)
        print(f"  band {b}: {len(big):,} candidate buckets", flush=True)
        del h, order, hs
    w.close(); os.replace(W / "buckets.parquet.tmp", W / "buckets.parquet")
    (W / "bands.parquet").unlink(missing_ok=True)
    print("lsh done", (groups, members_total), flush=True)


def union():
    sigs = np.load(W / "sigs.npy", mmap_mode="r"); n = len(sigs)
    parent = np.arange(n, dtype=np.int64)

    def find(x):
        root = x
        while parent[root] != root: root = parent[root]
        while parent[x] != root: parent[x], x = root, parent[x]
        return root
    merged = checked = 0
    for bt in pq.ParquetFile(W / "buckets.parquet").iter_batches(batch_size=20000, columns=["members"]):
        for members in bt.column("members").to_pylist():
            head = members[0]; hs = np.asarray(sigs[head]); rest = np.asarray(members[1:], dtype=np.int64)
            for k in range(0, len(rest), 4096):
                chunk = rest[k:k + 4096]
                sim = (np.asarray(sigs[np.sort(chunk)]) == hs).mean(axis=1)   # memmap fancy-index wants sorted rows
                for m, s_ in zip(np.sort(chunk), sim):
                    checked += 1
                    if s_ >= THRESH:
                        a, b = find(head), find(int(m))
                        if a != b: parent[max(a, b)] = min(a, b); merged += 1
    roots = np.fromiter((find(i) for i in range(n)), dtype=np.int64, count=n)
    np.save(W / "roots.npy", roots)
    print(f"union done: checked {checked:,} pairs, merged {merged:,}, clusters {len(np.unique(roots)):,} of {n:,}", flush=True)


def out():
    """Cluster table without DuckDB windows: Arrow join for ranking fields, numpy lexsort to pick each cluster's
    representative (most July copies, then most stars, then most complete), np.unique for sizes."""
    import glob
    roots = np.load(W / "roots.npy")
    idx = pq.read_table(W / "index.parquet")                          # row, sha
    meta = pa.concat_tables([pq.read_table(p, columns=["sha", "n_copies_jul", "canon_stars", "completeness"])
                             for p in sorted(glob.glob(META))])
    j = idx.join(meta, keys="sha", join_type="left outer").sort_by("row")
    copies = np.nan_to_num(j["n_copies_jul"].to_numpy(zero_copy_only=False).astype(np.float64), nan=0)
    stars = np.nan_to_num(j["canon_stars"].to_numpy(zero_copy_only=False).astype(np.float64), nan=-1)
    compl = np.nan_to_num(j["completeness"].to_numpy(zero_copy_only=False).astype(np.float64), nan=0)
    order = np.lexsort((-compl, -stars, -copies, roots))             # last key = primary
    first = np.ones(len(order), bool); first[1:] = roots[order][1:] != roots[order][:-1]
    rep_of_root = {}
    rep_rows = order[first]
    rep_of_root = dict(zip(roots[rep_rows].tolist(), rep_rows.tolist()))
    uniq, counts = np.unique(roots, return_counts=True); size_of = dict(zip(uniq.tolist(), counts.tolist()))
    import pyarrow.compute as pc
    rep_idx = np.fromiter((rep_of_root[r] for r in roots.tolist()), dtype=np.int64, count=len(roots))
    sizes = np.fromiter((size_of[r] for r in roots.tolist()), dtype=np.int64, count=len(roots))
    tbl = pa.table({"sha": j["sha"], "cluster_id": pc.take(j["sha"], pa.array(rep_idx)),   # Arrow-native: no Python strings
                    "cluster_size": sizes, "is_rep": rep_idx == np.arange(len(roots))})
    pq.write_table(tbl, W / "clusters.parquet", compression="zstd")
    summ = {"skills": len(roots), "near_dup_clusters": len(uniq), "singletons": int((counts == 1).sum()),
            "largest_cluster": int(counts.max()), "clusters_ge_10": int((counts >= 10).sum()),
            "skills_in_multi_clusters": int(counts[counts > 1].sum()),
            "params": {"num_perm": K, "bands": BANDS, "rows": ROWS, "shingle_words": SHINGLE, "jaccard_threshold": THRESH}}
    (W / "near_dup_summary.json").write_text(json.dumps(summ, indent=1)); print("out done", json.dumps(summ), flush=True)


def upload():
    from dotenv import load_dotenv; load_dotenv(ROOT / ".env")
    from huggingface_hub import HfApi
    HfApi().batch_bucket_files("Mercity/SkillsStorage", add=[(str(W / "clusters.parquet"), "processed/canonical/clusters.parquet"),
                                                            (str(W / "near_dup_summary.json"), "stats/near_dup_summary.json")])
    print("upload done", flush=True)


if __name__ == "__main__":
    stage = sys.argv[1] if len(sys.argv) > 1 else "all"
    gate = ROOT / "work/canon/select.ok"   # needs the complete text store (scan) + canonical meta (select); not the final table
    while "MH_SRC" not in os.environ and not gate.exists():
        print("waiting for canon select stage...", flush=True); time.sleep(120)
    # NOTE: sig reads the text store, which canon keeps until minhash is done; delete work/canon/text afterwards
    for s in (["sig", "lsh", "union", "out", "upload"] if stage == "all" else [stage]):
        if (W / f"{s}.ok").exists(): print(f"{s}: already done", flush=True); continue
        t = time.time(); globals()[s](); (W / f"{s}.ok").write_text("ok\n"); print(f"{s}: {time.time() - t:.0f}s", flush=True)