code/build_locator.py

"""Bundle locator index: which bucket parquet file holds each skill_uid's `skills` row and its `files` rows.

Crawl shards are written in crawl order, so a single bundle could be in any of ~3,500 files; this index turns
"rebuild one skill" into one lookup + one or two small reads (see rebuild_skill.py).

Output: processed/index/locator.parquet  (skill_uid, table ['skills'|'files'], path, n_rows) — sorted by skill_uid.
Resumable: per-batch results in work/locator/, done list in work/locator/done.txt.
Usage: python code/build_locator.py
"""
import sys, time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from bucket_query import connect, BUCKET
from huggingface_hub import HfApi
import pyarrow as pa, pyarrow.parquet as pq

W = Path("/root/skills-db/work/locator"); W.mkdir(parents=True, exist_ok=True)
DONE = W / "done.txt"
BASE = f"https://huggingface.co/buckets/{BUCKET}/resolve/"
PREFIXES = ["processed/gitskills/skills", "processed/gitskills/files", "processed/clawhub/skills.parquet",
            "processed/clawhub/files.parquet", "processed/clawhub/api_delta_skills.parquet",
            "processed/clawhub/api_delta_files.parquet"] + \
           [f"processed/github_delta/{r}/{t}" for r in ("pass1_sourcegraph", "pass2_new", "backfill_gitskills")
            for t in ("skills", "files")]
BATCH = 60

api = HfApi()
paths = []
for pre in PREFIXES:
    paths += sorted(f.path for f in api.list_bucket_tree(BUCKET, prefix=pre, recursive=True)
                    if getattr(f, "size", None) and f.path.endswith(".parquet"))
done = set(DONE.read_text().split()) if DONE.exists() else set()
todo = [p for p in paths if p not in done]
print(f"{len(paths)} table files, {len(todo)} to index", flush=True)

c = connect("300MB"); c.execute("SET threads=6"); c.execute(f"SET temp_directory='{W}/duck_tmp'")
for i in range(0, len(todo), BATCH):
    batch, t = todo[i:i + BATCH], time.time()
    lst = "[" + ",".join(f"'{BASE}{p}'" for p in batch) + "]"
    tbl = c.execute(f"""select skill_uid, src_file, count(*)::bigint n_rows
                        from read_parquet({lst}, filename='src_file') group by all""").arrow()
    tbl = tbl.read_all() if hasattr(tbl, "read_all") else tbl
    path = pa.array([f[len(BASE):] for f in tbl.column("src_file").to_pylist()], pa.string())
    kind = pa.array(["files" if p.split("/")[-2] == "files" or p.endswith("files.parquet") else "skills"
                     for p in path.to_pylist()], pa.string())  # crawl shards are named files-*.parquet in both dirs
    out = pa.table({"skill_uid": tbl.column("skill_uid"), "table": kind, "path": path, "n_rows": tbl.column("n_rows")})
    tmp = W / f"b-{i:05d}.parquet.tmp"; pq.write_table(out, tmp); tmp.rename(W / f"b-{i:05d}-{int(t)}.parquet")
    with open(DONE, "a") as f: f.write("\n".join(batch) + "\n")
    print(f"  {i + len(batch)}/{len(todo)} files, {out.num_rows:,} uids, {time.time() - t:.0f}s", flush=True)

final = W / "locator.parquet"
c.execute(f"SET preserve_insertion_order=false")
c.execute(f"copy (select * from read_parquet('{W}/b-*.parquet') order by skill_uid, \"table\", path) "
          f"to '{final}.tmp' (format parquet, row_group_size 100000)")
Path(f"{final}.tmp").rename(final)
n = pq.ParquetFile(final).metadata.num_rows
api.batch_bucket_files(BUCKET, add=[(str(final), "processed/index/locator.parquet")])
print(f"done: {n:,} rows -> processed/index/locator.parquet", flush=True)