SPB Git

spb/wp44_uqo

Public
1.5 KB · 21 lines python
1#!/usr/bin/env python32"""Compact the per-part Parquet tables (1,898 files per table) into one Parquet file per table3(data/processed/compact/<table>.parquet, zstd, ~1M-row row groups) so that downstream DuckDB work does not4hit the macOS open-file limit. Also compacts the slim table of all OpenAlex works."""5import duckdb, os, sys, time6R = sys.argv[1]; suffix = sys.argv[2] if len(sys.argv) > 2 else ""   # e.g. "_p2" for the second snapshot pass7tables_dir = os.path.join(R, "data/processed/tables" + suffix); out = os.path.join(R, "data/processed/compact"); os.makedirs(out, exist_ok=True)8con = duckdb.connect(); con.execute("SET threads TO 24; SET memory_limit='" + os.environ.get("WP44_DUCKDB_MEM", "60GB") + "'; SET preserve_insertion_order=false;")9jobs = [(t + suffix, f"{tables_dir}/{t}/*.parquet") for t in ["works", "authorships", "work_topics", "work_keywords", "references", "abstracts", "locations", "grants", "counts_by_year"]]10if not suffix:11    jobs.append(("slim_all_works", f"{R}/data/raw/openalex/works/slim/*.parquet"))12for t, src in jobs:13    dst = f"{out}/{t}.parquet"14    if os.path.exists(dst):15        print("skip", t); continue16    t0 = time.time()17    con.execute(f"COPY (SELECT * FROM read_parquet('{src}')) TO '{dst}.part' (FORMAT PARQUET, COMPRESSION ZSTD, ROW_GROUP_SIZE 1000000)")18    os.replace(dst + ".part", dst)19    n = con.execute(f"SELECT count(*) FROM read_parquet('{dst}')").fetchone()[0]20    print(f"{t}: {n:,} rows, {os.path.getsize(dst)/1e9:.2f} GB, {time.time()-t0:.0f}s", flush=True)21
Files · main
Menu

Theme