#!/usr/bin/env python3
"""Compact the per-part Parquet tables (1,898 files per table) into one Parquet file per table
(data/processed/compact/<table>.parquet, zstd, ~1M-row row groups) so that downstream DuckDB work does not
hit the macOS open-file limit. Also compacts the slim table of all OpenAlex works."""
import duckdb, os, sys, time
R = sys.argv[1]; suffix = sys.argv[2] if len(sys.argv) > 2 else ""   # e.g. "_p2" for the second snapshot pass
tables_dir = os.path.join(R, "data/processed/tables" + suffix); out = os.path.join(R, "data/processed/compact"); os.makedirs(out, exist_ok=True)
con = duckdb.connect(); con.execute("SET threads TO 24; SET memory_limit='70GB'; SET preserve_insertion_order=false;")
jobs = [(t + suffix, f"{tables_dir}/{t}/*.parquet") for t in ["works", "authorships", "work_topics", "work_keywords", "references", "abstracts", "locations", "grants", "counts_by_year"]]
if not suffix:
    jobs.append(("slim_all_works", f"{R}/data/raw/openalex/works/slim/*.parquet"))
for t, src in jobs:
    dst = f"{out}/{t}.parquet"
    if os.path.exists(dst):
        print("skip", t); continue
    t0 = time.time()
    con.execute(f"COPY (SELECT * FROM read_parquet('{src}')) TO '{dst}.part' (FORMAT PARQUET, COMPRESSION ZSTD, ROW_GROUP_SIZE 1000000)")
    os.replace(dst + ".part", dst)
    n = con.execute(f"SELECT count(*) FROM read_parquet('{dst}')").fetchone()[0]
    print(f"{t}: {n:,} rows, {os.path.getsize(dst)/1e9:.2f} GB, {time.time()-t0:.0f}s", flush=True)