1#!/usr/bin/env python32"""WP44 — OpenAlex public snapshot ingestion (works entity).34Streams every part file of the works snapshot listed in the manifest, verifies5content length and records a SHA-256 of the gzip bytes (immutable provenance),6parses each JSON line and writes:78 * ``cs/<partition>__<part>.jsonl.zst`` — the FULL OpenAlex record of every9 work matched by the CS filter (any topic in field 17 Computer Science,10 publication_year in [YEAR_MIN, YEAR_MAX]); XPAC works are kept with the11 ``is_xpac`` flag so that the corpus filter can be applied later.12 * ``slim/<partition>__<part>.parquet`` — a slim row for EVERY work in the13 snapshot (identifiers, year, type, primary topic, counts) used to resolve14 cited works outside computer science.1516A SQLite ledger makes the job resumable and keeps job-level provenance17(attempts, timings, hashes, counts, errors). Raw .gz files are deleted after a18successful parse unless --keep-raw is given (the snapshot is a public,19versioned release; the manifest and per-file hashes are kept).20"""21from __future__ import annotations2223import argparse24import datetime as dt25import gzip26import hashlib27import json28import multiprocessing as mp29import os30import sqlite331import subprocess32import sys33import time34import traceback3536import orjson37import pyarrow as pa38import pyarrow.parquet as pq39import zstandard as zstd4041S3_HTTP = "https://openalex.s3.amazonaws.com/"42CS_FIELD = "https://openalex.org/fields/17"43YEAR_MIN, YEAR_MAX = 2018, 20264445SLIM_SCHEMA = pa.schema([46 ("id", pa.int64()),47 ("doi", pa.string()),48 ("publication_year", pa.int32()),49 ("publication_date", pa.string()),50 ("type", pa.string()),51 ("cited_by_count", pa.int32()),52 ("primary_topic_id", pa.int32()),53 ("primary_subfield_id", pa.int32()),54 ("primary_field_id", pa.int32()),55 ("primary_domain_id", pa.int32()),56 ("cs_any_topic", pa.bool_()),57 ("is_xpac", pa.bool_()),58 ("is_oa", pa.bool_()),59 ("is_retracted", pa.bool_()),60 ("language", pa.string()),61 ("source_id", pa.int64()),62 ("n_refs", pa.int32()),63 ("n_authors", pa.int32()),64 ("has_abstract", pa.bool_()),65 ("fwci", pa.float32()),66])676869def _oa_int(url, prefix_len=None):70 """'https://openalex.org/W123' -> 123 (None-safe)."""71 if not url:72 return None73 tail = url.rsplit("/", 1)[-1]74 # strip leading letter(s): W, T, I, S...75 i = 076 while i < len(tail) and not tail[i].isdigit():77 i += 178 try:79 return int(tail[i:])80 except ValueError:81 return None828384def slim_row(w: dict) -> tuple:85 pt = w.get("primary_topic") or {}86 topics = w.get("topics") or []87 cs_any = any(((t.get("field") or {}).get("id") == CS_FIELD) for t in topics) or (88 (pt.get("field") or {}).get("id") == CS_FIELD)89 loc = w.get("primary_location") or {}90 src = (loc.get("source") or {}) if loc else {}91 oa = w.get("open_access") or {}92 doi = w.get("doi")93 if doi and doi.startswith("https://doi.org/"):94 doi = doi[16:]95 return (96 _oa_int(w.get("id")),97 doi.lower() if doi else None,98 w.get("publication_year"),99 w.get("publication_date"),100 w.get("type"),101 w.get("cited_by_count"),102 _oa_int(pt.get("id")),103 _oa_int((pt.get("subfield") or {}).get("id")),104 _oa_int((pt.get("field") or {}).get("id")),105 _oa_int((pt.get("domain") or {}).get("id")),106 bool(cs_any),107 bool(w.get("is_xpac")),108 bool(oa.get("is_oa")),109 bool(w.get("is_retracted")),110 w.get("language"),111 _oa_int(src.get("id")) if src else None,112 len(w.get("referenced_works") or []),113 w.get("authors_count") if w.get("authors_count") is not None else len(w.get("authorships") or []),114 w.get("abstract_inverted_index") is not None,115 w.get("fwci"),116 ), cs_any117118119def process_file(args):120 """Download, hash, parse and filter one part file. Returns a result dict."""121 url, expected_len, out_dir, keep_raw = args122 rel = url.replace("s3://openalex/data/jsonl/works/", "")123 name = rel.replace("/", "__") # updated_date=X__part_0000.gz124 http = url.replace("s3://openalex/", S3_HTTP)125 tmp_dir = os.path.join(out_dir, "tmp")126 cs_dir = os.path.join(out_dir, "cs")127 slim_dir = os.path.join(out_dir, "slim")128 for d in (tmp_dir, cs_dir, slim_dir):129 os.makedirs(d, exist_ok=True)130 raw_path = os.path.join(tmp_dir, name)131 t0 = time.time()132 res = {"url": url, "name": name, "expected_len": expected_len}133 try:134 # ---- download (curl, retries, resume) ----135 for attempt in range(1, 6):136 cmd = ["curl", "-sS", "-L", "--retry", "5", "--retry-delay", "3", "--retry-all-errors",137 "-C", "-", "-o", raw_path, http]138 p = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)139 if os.path.exists(raw_path) and os.path.getsize(raw_path) == expected_len:140 break141 time.sleep(5 * attempt)142 got = os.path.getsize(raw_path) if os.path.exists(raw_path) else -1143 if got != expected_len:144 raise RuntimeError(f"size mismatch: got {got}, expected {expected_len}")145 t_dl = time.time() - t0146 # ---- hash ----147 h = hashlib.sha256()148 with open(raw_path, "rb") as fh:149 for chunk in iter(lambda: fh.read(1 << 22), b""):150 h.update(chunk)151 sha = h.hexdigest()152 # ---- parse ----153 cols = [[] for _ in SLIM_SCHEMA]154 n_all = n_cs = n_bad = 0155 cs_tmp = os.path.join(cs_dir, name[:-3] + ".jsonl.zst.part")156 cctx = zstd.ZstdCompressor(level=6)157 with open(cs_tmp, "wb") as cs_fh, cctx.stream_writer(cs_fh) as cs_w, gzip.open(raw_path, "rb") as g:158 for line in g:159 if not line.strip():160 continue161 try:162 w = orjson.loads(line)163 except Exception:164 n_bad += 1165 continue166 n_all += 1167 row, cs_any = slim_row(w)168 for i, v in enumerate(row):169 cols[i].append(v)170 y = w.get("publication_year")171 if cs_any and y is not None and YEAR_MIN <= y <= YEAR_MAX:172 cs_w.write(line if line.endswith(b"\n") else line + b"\n")173 n_cs += 1174 os.replace(cs_tmp, cs_tmp[:-5])175 table = pa.Table.from_arrays([pa.array(c, type=f.type) for c, f in zip(cols, SLIM_SCHEMA)], schema=SLIM_SCHEMA)176 slim_tmp = os.path.join(slim_dir, name[:-3] + ".parquet.part")177 pq.write_table(table, slim_tmp, compression="zstd")178 os.replace(slim_tmp, slim_tmp[:-5])179 if not keep_raw:180 os.remove(raw_path)181 res.update(status="done", sha256=sha, n_records=n_all, n_cs=n_cs, n_bad=n_bad,182 t_download=round(t_dl, 1), t_total=round(time.time() - t0, 1), error=None)183 except Exception as e: # noqa: BLE001184 res.update(status="error", error=f"{e}\n{traceback.format_exc()[-800:]}", t_total=round(time.time() - t0, 1))185 return res186187188def main():189 ap = argparse.ArgumentParser()190 ap.add_argument("--out", required=True)191 ap.add_argument("--manifest", required=True, help="local copy of works manifest.json")192 ap.add_argument("--workers", type=int, default=6)193 ap.add_argument("--keep-raw", action="store_true")194 ap.add_argument("--limit", type=int, default=0, help="process only N files (smoke test)")195 ap.add_argument("--smallest-first", action="store_true")196 a = ap.parse_args()197 os.makedirs(a.out, exist_ok=True)198 man = json.load(open(a.manifest))199 files = [(f["url"], f["meta"]["content_length"], f["meta"].get("record_count")) for f in man["files"]]200 if a.smallest_first:201 files.sort(key=lambda x: x[1])202 if a.limit:203 files = files[: a.limit]204 db = sqlite3.connect(os.path.join(a.out, "ledger.sqlite"))205 db.execute("""CREATE TABLE IF NOT EXISTS files(url TEXT PRIMARY KEY, name TEXT, expected_len INTEGER,206 manifest_records INTEGER, status TEXT, sha256 TEXT, n_records INTEGER, n_cs INTEGER, n_bad INTEGER,207 t_download REAL, t_total REAL, attempts INTEGER DEFAULT 0, started TEXT, finished TEXT, error TEXT)""")208 db.execute("""CREATE TABLE IF NOT EXISTS runs(started TEXT, manifest_date TEXT, manifest_records INTEGER,209 manifest_bytes INTEGER, workers INTEGER, host TEXT, argv TEXT)""")210 db.execute("INSERT INTO runs VALUES(?,?,?,?,?,?,?)", (dt.datetime.now(dt.UTC).isoformat(), man.get("date"),211 man.get("record_count"), man.get("content_length"), a.workers, os.uname().nodename, " ".join(sys.argv)))212 for url, ln, rc in files:213 db.execute("INSERT OR IGNORE INTO files(url,name,expected_len,manifest_records,status) VALUES(?,?,?,?,'pending')",214 (url, url.replace("s3://openalex/data/jsonl/works/", "").replace("/", "__"), ln, rc))215 db.commit()216 done = {r[0] for r in db.execute("SELECT url FROM files WHERE status='done'")}217 todo = [(u, ln, a.out, a.keep_raw) for u, ln, _ in files if u not in done]218 total_bytes = sum(t[1] for t in todo)219 print(f"[{dt.datetime.now():%H:%M:%S}] manifest {man.get('date')}: {len(files)} files, {len(done)} done, "220 f"{len(todo)} to do ({total_bytes/1e9:.1f} GB), workers={a.workers}", flush=True)221 t0 = time.time(); done_bytes = 0; n_ok = n_err = 0222 with mp.Pool(a.workers, maxtasksperchild=50) as pool:223 for r in pool.imap_unordered(process_file, todo):224 now = dt.datetime.now(dt.UTC).isoformat()225 db.execute("UPDATE files SET attempts=attempts+1 WHERE url=?", (r["url"],))226 if r["status"] == "done":227 n_ok += 1; done_bytes += r["expected_len"]228 db.execute("""UPDATE files SET status='done', sha256=?, n_records=?, n_cs=?, n_bad=?, t_download=?,229 t_total=?, finished=?, error=NULL WHERE url=?""",230 (r["sha256"], r["n_records"], r["n_cs"], r["n_bad"], r["t_download"], r["t_total"], now, r["url"]))231 else:232 n_err += 1233 db.execute("UPDATE files SET status='error', error=?, finished=? WHERE url=?", (r["error"], now, r["url"]))234 db.commit()235 el = time.time() - t0236 rate = done_bytes / el if el else 0237 eta = (total_bytes - done_bytes) / rate / 60 if rate else float("nan")238 print(f"[{dt.datetime.now():%H:%M:%S}] {r['status']:5s} {r['name']:40s} rec={r.get('n_records')} "239 f"cs={r.get('n_cs')} | {n_ok} ok {n_err} err | {done_bytes/1e9:.1f}/{total_bytes/1e9:.1f} GB "240 f"{rate/1e6:.1f} MB/s ETA {eta:.0f} min", flush=True)241 print("finished", n_ok, "ok", n_err, "errors", flush=True)242243244if __name__ == "__main__":245 main()246