code/canon_repack.py

"""Repack processed/canonical/skills/skills-{0..f}.parquet: sort by sha, 5,000-row row groups.

canon_build's final stage wrote parts 2..f as ~1,800 tiny row groups each (one per streamed join batch), so a
remote point lookup by sha touched thousands of ranges. Sorted + larger row groups -> min/max stats prune a lookup
to one row group. Same rows, same schema.

Per part (resumable, work/repack/<pre>.done): download -> split by sha[1] into 16 temp files (bounded memory) ->
sort each and append -> verify row count -> upload over the original -> delete local copies.
Usage: python code/canon_repack.py
"""
import subprocess, sys
from pathlib import Path
import pyarrow as pa, pyarrow.parquet as pq, pyarrow.compute as pc

BUCKET, PRE = "Mercity/SkillsStorage", "processed/canonical/skills"
W = Path("/root/skills-db/work/repack"); W.mkdir(parents=True, exist_ok=True)
RG = 5000


def hub(call):  # transfers run in a short-lived child: the hub client keeps transfer buffers after big files
    subprocess.run([sys.executable, "-c", "from dotenv import load_dotenv; load_dotenv('/root/skills-db/.env'); "
                    "from huggingface_hub import HfApi; api = HfApi(); " + call], check=True)


for h in "0123456789abcdef":
    if (W / f"{h}.done").exists(): continue
    src, out, tmp = W / f"in-{h}.parquet", W / f"skills-{h}.parquet", W / f"skills-{h}.parquet.tmp"
    if not src.exists():
        hub(f"api.download_bucket_files('{BUCKET}', [('{PRE}/skills-{h}.parquet', '{src}')])")
    pf = pq.ParquetFile(src); schema, n_in = pf.schema_arrow, pf.metadata.num_rows
    subs = {x: pq.ParquetWriter(W / f"sub-{x}.parquet", schema) for x in "0123456789abcdef"}
    for b in pf.iter_batches(batch_size=2000):
        key = pc.utf8_slice_codeunits(b.column("sha"), 1, 2)
        for x, w in subs.items():
            part = b.filter(pc.equal(key, x))
            if part.num_rows: w.write_batch(part)
        pa.default_memory_pool().release_unused()
    for w in subs.values(): w.close()
    del pf
    n_out = 0
    with pq.ParquetWriter(tmp, schema) as w:
        for x in "0123456789abcdef":
            t = pq.read_table(W / f"sub-{x}.parquet").sort_by("sha")
            w.write_table(t, row_group_size=RG); n_out += t.num_rows
            del t; (W / f"sub-{x}.parquet").unlink(); pa.default_memory_pool().release_unused()
    meta = pq.ParquetFile(tmp).metadata
    assert n_out == n_in == meta.num_rows, (h, n_in, n_out, meta.num_rows)
    tmp.rename(out)
    hub(f"api.batch_bucket_files('{BUCKET}', add=[('{out}', '{PRE}/skills-{h}.parquet')])")
    print(f"skills-{h}: {n_in:,} rows, {meta.num_row_groups} row groups (sorted by sha), uploaded", flush=True)
    src.unlink(); out.unlink(); (W / f"{h}.done").write_text("ok\n")
print("done", flush=True)