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/"42ID_SET = None # optional: set of OpenAlex work ids (ints) to extract instead of the CS filter (second pass)43CS_FIELD = "https://openalex.org/fields/17"44YEAR_MIN, YEAR_MAX = 2018, 20264546SLIM_SCHEMA = pa.schema([47 ("id", pa.int64()),48 ("doi", pa.string()),49 ("publication_year", pa.int32()),50 ("publication_date", pa.string()),51 ("type", pa.string()),52 ("cited_by_count", pa.int32()),53 ("primary_topic_id", pa.int32()),54 ("primary_subfield_id", pa.int32()),55 ("primary_field_id", pa.int32()),56 ("primary_domain_id", pa.int32()),57 ("cs_any_topic", pa.bool_()),58 ("is_xpac", pa.bool_()),59 ("is_oa", pa.bool_()),60 ("is_retracted", pa.bool_()),61 ("language", pa.string()),62 ("source_id", pa.int64()),63 ("n_refs", pa.int32()),64 ("n_authors", pa.int32()),65 ("has_abstract", pa.bool_()),66 ("fwci", pa.float32()),67])686970def _oa_int(url, prefix_len=None):71 """'https://openalex.org/W123' -> 123 (None-safe)."""72 if not url:73 return None74 tail = url.rsplit("/", 1)[-1]75 # strip leading letter(s): W, T, I, S...76 i = 077 while i < len(tail) and not tail[i].isdigit():78 i += 179 try:80 return int(tail[i:])81 except ValueError:82 return None838485def slim_row(w: dict) -> tuple:86 pt = w.get("primary_topic") or {}87 topics = w.get("topics") or []88 cs_any = any(((t.get("field") or {}).get("id") == CS_FIELD) for t in topics) or (89 (pt.get("field") or {}).get("id") == CS_FIELD)90 loc = w.get("primary_location") or {}91 src = (loc.get("source") or {}) if loc else {}92 oa = w.get("open_access") or {}93 doi = w.get("doi")94 if doi and doi.startswith("https://doi.org/"):95 doi = doi[16:]96 return (97 _oa_int(w.get("id")),98 doi.lower() if doi else None,99 w.get("publication_year"),100 w.get("publication_date"),101 w.get("type"),102 w.get("cited_by_count"),103 _oa_int(pt.get("id")),104 _oa_int((pt.get("subfield") or {}).get("id")),105 _oa_int((pt.get("field") or {}).get("id")),106 _oa_int((pt.get("domain") or {}).get("id")),107 bool(cs_any),108 bool(w.get("is_xpac")),109 bool(oa.get("is_oa")),110 bool(w.get("is_retracted")),111 w.get("language"),112 _oa_int(src.get("id")) if src else None,113 len(w.get("referenced_works") or []),114 w.get("authors_count") if w.get("authors_count") is not None else len(w.get("authorships") or []),115 w.get("abstract_inverted_index") is not None,116 w.get("fwci"),117 ), cs_any118119120def _load_ids(path):121 with open(path) as fh:122 next(fh)123 return {int(x.strip().rsplit("W", 1)[-1]) for x in fh if x.strip()}124125126def _init_worker(id_list):127 """Pool initializer: macOS uses the spawn start method, so module globals must be rebuilt in each worker."""128 global ID_SET129 ID_SET = _load_ids(id_list) if id_list else None130131132def process_file(args):133 """Download, hash, parse and filter one part file. Returns a result dict."""134 url, expected_len, out_dir, keep_raw, mode = args135 rel = url.replace("s3://openalex/data/jsonl/works/", "")136 name = rel.replace("/", "__") # updated_date=X__part_0000.gz137 http = url.replace("s3://openalex/", S3_HTTP)138 tmp_dir = os.path.join(out_dir, "tmp")139 cs_dir = os.path.join(out_dir, "cs")140 slim_dir = os.path.join(out_dir, "slim")141 for d in (tmp_dir, cs_dir, slim_dir):142 os.makedirs(d, exist_ok=True)143 raw_path = os.path.join(tmp_dir, name)144 t0 = time.time()145 res = {"url": url, "name": name, "expected_len": expected_len}146 try:147 # ---- download (curl, retries, resume) ----148 for attempt in range(1, 6):149 cmd = ["curl", "-sS", "-L", "--retry", "5", "--retry-delay", "3", "--retry-all-errors",150 "-C", "-", "-o", raw_path, http]151 p = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)152 if os.path.exists(raw_path) and os.path.getsize(raw_path) == expected_len:153 break154 time.sleep(5 * attempt)155 got = os.path.getsize(raw_path) if os.path.exists(raw_path) else -1156 if got != expected_len:157 raise RuntimeError(f"size mismatch: got {got}, expected {expected_len}")158 t_dl = time.time() - t0159 # ---- hash ----160 h = hashlib.sha256()161 with open(raw_path, "rb") as fh:162 for chunk in iter(lambda: fh.read(1 << 22), b""):163 h.update(chunk)164 sha = h.hexdigest()165 # ---- parse ----166 cols = [[] for _ in SLIM_SCHEMA]167 n_all = n_cs = n_bad = 0168 cs_tmp = os.path.join(cs_dir, name[:-3] + ".jsonl.zst.part")169 cctx = zstd.ZstdCompressor(level=6)170 with open(cs_tmp, "wb") as cs_fh, cctx.stream_writer(cs_fh) as cs_w, gzip.open(raw_path, "rb") as g:171 for line in g:172 if not line.strip():173 continue174 try:175 w = orjson.loads(line)176 except Exception:177 n_bad += 1178 continue179 n_all += 1180 if mode == "idset":181 if _oa_int(w.get("id")) in ID_SET:182 cs_w.write(line if line.endswith(b"\n") else line + b"\n")183 n_cs += 1184 continue185 row, cs_any = slim_row(w)186 for i, v in enumerate(row):187 cols[i].append(v)188 y = w.get("publication_year")189 if cs_any and y is not None and YEAR_MIN <= y <= YEAR_MAX:190 cs_w.write(line if line.endswith(b"\n") else line + b"\n")191 n_cs += 1192 os.replace(cs_tmp, cs_tmp[:-5])193 if mode == "idset":194 cols = None195 table = None if cols is None else pa.Table.from_arrays([pa.array(c, type=f.type) for c, f in zip(cols, SLIM_SCHEMA)], schema=SLIM_SCHEMA)196 if table is not None:197 slim_tmp = os.path.join(slim_dir, name[:-3] + ".parquet.part")198 pq.write_table(table, slim_tmp, compression="zstd")199 os.replace(slim_tmp, slim_tmp[:-5])200 if not keep_raw:201 os.remove(raw_path)202 res.update(status="done", sha256=sha, n_records=n_all, n_cs=n_cs, n_bad=n_bad,203 t_download=round(t_dl, 1), t_total=round(time.time() - t0, 1), error=None)204 except Exception as e: # noqa: BLE001205 res.update(status="error", error=f"{e}\n{traceback.format_exc()[-800:]}", t_total=round(time.time() - t0, 1))206 return res207208209def main():210 ap = argparse.ArgumentParser()211 ap.add_argument("--out", required=True)212 ap.add_argument("--manifest", required=True, help="local copy of works manifest.json")213 ap.add_argument("--workers", type=int, default=6)214 ap.add_argument("--keep-raw", action="store_true")215 ap.add_argument("--limit", type=int, default=0, help="process only N files (smoke test)")216 ap.add_argument("--smallest-first", action="store_true")217 ap.add_argument("--id-list", default=None, help="CSV with a header and one OpenAlex work id per line: extract exactly these works (second pass)")218 a = ap.parse_args()219 global ID_SET220 mode = "cs"221 if a.id_list:222 ID_SET = _load_ids(a.id_list)223 mode = "idset"224 print(f"id-set mode: {len(ID_SET)} ids", flush=True)225 os.makedirs(a.out, exist_ok=True)226 man = json.load(open(a.manifest))227 files = [(f["url"], f["meta"]["content_length"], f["meta"].get("record_count")) for f in man["files"]]228 if a.smallest_first:229 files.sort(key=lambda x: x[1])230 if a.limit:231 files = files[: a.limit]232 db = sqlite3.connect(os.path.join(a.out, "ledger.sqlite"))233 db.execute("""CREATE TABLE IF NOT EXISTS files(url TEXT PRIMARY KEY, name TEXT, expected_len INTEGER,234 manifest_records INTEGER, status TEXT, sha256 TEXT, n_records INTEGER, n_cs INTEGER, n_bad INTEGER,235 t_download REAL, t_total REAL, attempts INTEGER DEFAULT 0, started TEXT, finished TEXT, error TEXT)""")236 db.execute("""CREATE TABLE IF NOT EXISTS runs(started TEXT, manifest_date TEXT, manifest_records INTEGER,237 manifest_bytes INTEGER, workers INTEGER, host TEXT, argv TEXT)""")238 db.execute("INSERT INTO runs VALUES(?,?,?,?,?,?,?)", (dt.datetime.now(dt.UTC).isoformat(), man.get("date"),239 man.get("record_count"), man.get("content_length"), a.workers, os.uname().nodename, " ".join(sys.argv)))240 for url, ln, rc in files:241 db.execute("INSERT OR IGNORE INTO files(url,name,expected_len,manifest_records,status) VALUES(?,?,?,?,'pending')",242 (url, url.replace("s3://openalex/data/jsonl/works/", "").replace("/", "__"), ln, rc))243 db.commit()244 done = {r[0] for r in db.execute("SELECT url FROM files WHERE status='done'")}245 todo = [(u, ln, a.out, a.keep_raw, mode) for u, ln, _ in files if u not in done]246 total_bytes = sum(t[1] for t in todo)247 print(f"[{dt.datetime.now():%H:%M:%S}] manifest {man.get('date')}: {len(files)} files, {len(done)} done, "248 f"{len(todo)} to do ({total_bytes/1e9:.1f} GB), workers={a.workers}", flush=True)249 t0 = time.time(); done_bytes = 0; n_ok = n_err = 0250 with mp.Pool(a.workers, maxtasksperchild=50, initializer=_init_worker, initargs=(a.id_list,)) as pool:251 for r in pool.imap_unordered(process_file, todo):252 now = dt.datetime.now(dt.UTC).isoformat()253 db.execute("UPDATE files SET attempts=attempts+1 WHERE url=?", (r["url"],))254 if r["status"] == "done":255 n_ok += 1; done_bytes += r["expected_len"]256 db.execute("""UPDATE files SET status='done', sha256=?, n_records=?, n_cs=?, n_bad=?, t_download=?,257 t_total=?, finished=?, error=NULL WHERE url=?""",258 (r["sha256"], r["n_records"], r["n_cs"], r["n_bad"], r["t_download"], r["t_total"], now, r["url"]))259 else:260 n_err += 1261 db.execute("UPDATE files SET status='error', error=?, finished=? WHERE url=?", (r["error"], now, r["url"]))262 db.commit()263 el = time.time() - t0264 rate = done_bytes / el if el else 0265 eta = (total_bytes - done_bytes) / rate / 60 if rate else float("nan")266 print(f"[{dt.datetime.now():%H:%M:%S}] {r['status']:5s} {r['name']:40s} rec={r.get('n_records')} "267 f"cs={r.get('n_cs')} | {n_ok} ok {n_err} err | {done_bytes/1e9:.1f}/{total_bytes/1e9:.1f} GB "268 f"{rate/1e6:.1f} MB/s ETA {eta:.0f} min", flush=True)269 print("finished", n_ok, "ok", n_err, "errors", flush=True)270271272if __name__ == "__main__":273 main()274