1#!/usr/bin/env python32"""WP44 — normalise the filtered OpenAlex CS records (cs/*.jsonl.zst) into Parquet tables.34Tables (one Parquet file per input part under <out>/tables/<table>/<part>.parquet, resumable):5 works, authorships, work_topics, work_keywords, references, abstracts, locations, grants, counts_by_year6Run on the compute node after (or while) ingestion/openalex_snapshot.py runs; parts whose outputs exist are skipped.7"""8import argparse, glob, io, multiprocessing as mp, os, re, time, datetime as dt9import orjson, zstandard as zstd10import pyarrow as pa, pyarrow.parquet as pq1112ARXIV_RE = re.compile(r"arxiv\.org/(?:abs|pdf)/([0-9]{4}\.[0-9]{4,5}(?:v\d+)?|[a-z\-]+(?:\.[A-Z]{2})?/[0-9]{7}(?:v\d+)?)", re.I)13ARXIV_DOI_RE = re.compile(r"^10\.48550/arxiv\.(.+)$", re.I)1415S = {16 "works": pa.schema([("id", pa.int64()), ("doi", pa.string()), ("title", pa.string()), ("publication_year", pa.int32()),17 ("publication_date", pa.string()), ("created_date", pa.string()), ("updated_date", pa.string()), ("type", pa.string()),18 ("type_crossref", pa.string()), ("language", pa.string()), ("cited_by_count", pa.int32()), ("fwci", pa.float32()),19 ("citation_percentile", pa.float32()), ("is_top1", pa.bool_()), ("is_top10", pa.bool_()), ("is_oa", pa.bool_()), ("oa_status", pa.string()),20 ("license", pa.string()), ("is_retracted", pa.bool_()), ("is_paratext", pa.bool_()), ("is_xpac", pa.bool_()), ("n_refs", pa.int32()),21 ("n_authors", pa.int32()), ("n_institutions", pa.int32()), ("n_countries", pa.int32()), ("countries", pa.list_(pa.string())),22 ("institution_types", pa.list_(pa.string())), ("primary_topic_id", pa.int32()), ("primary_subfield_id", pa.int32()),23 ("primary_field_id", pa.int32()), ("primary_domain_id", pa.int32()), ("cs_primary", pa.bool_()), ("cs_any", pa.bool_()),24 ("source_id", pa.int64()), ("source_name", pa.string()), ("source_type", pa.string()), ("host_organization", pa.string()),25 ("is_in_doaj", pa.bool_()), ("arxiv_id", pa.string()), ("pmid", pa.string()), ("mag", pa.string()), ("has_abstract", pa.bool_()),26 ("has_fulltext", pa.bool_()), ("n_locations", pa.int32()), ("indexed_in", pa.list_(pa.string())), ("referenced_works_count", pa.int32()),27 ("n_grants", pa.int32()), ("apc_paid_usd", pa.int32()), ("is_authors_truncated", pa.bool_()), ("part", pa.string())]),28 "authorships": pa.schema([("work_id", pa.int64()), ("position", pa.int32()), ("author_position", pa.string()), ("author_id", pa.int64()),29 ("author_name", pa.string()), ("orcid", pa.string()), ("is_corresponding", pa.bool_()), ("institution_ids", pa.list_(pa.int64())),30 ("institution_names", pa.list_(pa.string())), ("institution_types", pa.list_(pa.string())), ("countries", pa.list_(pa.string())),31 ("n_raw_affiliations", pa.int32())]),32 "work_topics": pa.schema([("work_id", pa.int64()), ("rank", pa.int16()), ("topic_id", pa.int32()), ("score", pa.float32()),33 ("subfield_id", pa.int32()), ("field_id", pa.int32()), ("domain_id", pa.int32())]),34 "work_keywords": pa.schema([("work_id", pa.int64()), ("keyword", pa.string()), ("score", pa.float32())]),35 "references": pa.schema([("work_id", pa.int64()), ("ref_id", pa.int64())]),36 "abstracts": pa.schema([("work_id", pa.int64()), ("abstract", pa.string()), ("n_words", pa.int32())]),37 "locations": pa.schema([("work_id", pa.int64()), ("rank", pa.int16()), ("is_primary", pa.bool_()), ("is_best_oa", pa.bool_()), ("source_id", pa.int64()),38 ("source_name", pa.string()), ("source_type", pa.string()), ("landing_page_url", pa.string()), ("pdf_url", pa.string()), ("is_oa", pa.bool_()),39 ("license", pa.string()), ("version", pa.string()), ("is_accepted", pa.bool_()), ("is_published", pa.bool_())]),40 "grants": pa.schema([("work_id", pa.int64()), ("funder_id", pa.int64()), ("funder_name", pa.string()), ("award_id", pa.string())]),41 "counts_by_year": pa.schema([("work_id", pa.int64()), ("year", pa.int16()), ("cited_by_count", pa.int32())]),42}434445def oid(u):46 if not u:47 return None48 t = u.rsplit("/", 1)[-1]49 i = 050 while i < len(t) and not t[i].isdigit():51 i += 152 try:53 return int(t[i:])54 except ValueError:55 return None565758def abstract_text(inv):59 if not inv:60 return None61 pos = []62 for w, ps in inv.items():63 for p in ps:64 pos.append((p, w))65 pos.sort()66 return " ".join(w for _, w in pos)676869def arxiv_from(w):70 ids = w.get("ids") or {}71 if ids.get("arxiv"):72 return str(ids["arxiv"]).rsplit("/", 1)[-1]73 doi = w.get("doi") or ""74 m = ARXIV_DOI_RE.match(doi.replace("https://doi.org/", ""))75 if m:76 return m.group(1)77 for loc in (w.get("locations") or []):78 for k in ("landing_page_url", "pdf_url"):79 u = loc.get(k) or ""80 m = ARXIV_RE.search(u)81 if m:82 return re.sub(r"v\d+$", "", m.group(1))83 return None848586def process(args):87 path, out, cs_field = args88 part = os.path.basename(path).replace(".jsonl.zst", "")89 outs = {t: os.path.join(out, "tables", t, part + ".parquet") for t in S}90 if all(os.path.exists(p) for p in outs.values()):91 return part, "skip", 092 for p in outs.values():93 os.makedirs(os.path.dirname(p), exist_ok=True)94 rows = {t: {c: [] for c in S[t].names} for t in S}95 n = 096 with open(path, "rb") as fh, zstd.ZstdDecompressor().stream_reader(fh) as r:97 for line in io.TextIOWrapper(r, encoding="utf-8"):98 if not line.strip():99 continue100 w = orjson.loads(line); n += 1101 wid = oid(w["id"])102 doi = (w.get("doi") or "").replace("https://doi.org/", "").lower() or None103 pt = w.get("primary_topic") or {}104 topics = w.get("topics") or []105 cs_any = any(((t.get("field") or {}).get("id") == cs_field) for t in topics) or ((pt.get("field") or {}).get("id") == cs_field)106 loc = w.get("primary_location") or {}107 src = loc.get("source") or {}108 best = w.get("best_oa_location") or {}109 oa = w.get("open_access") or {}110 auths = w.get("authorships") or []111 inst_ids = set(); countries = set(); itypes = []112 for i, a in enumerate(auths):113 ins = a.get("institutions") or []114 ii = [oid(x.get("id")) for x in ins]115 inst_ids.update(x for x in ii if x)116 cc = [c for c in (a.get("countries") or []) if c]117 countries.update(cc)118 itypes.extend(x.get("type") for x in ins if x.get("type"))119 au = a.get("author") or {}120 R = rows["authorships"]121 R["work_id"].append(wid); R["position"].append(i); R["author_position"].append(a.get("author_position"))122 R["author_id"].append(oid(au.get("id"))); R["author_name"].append(au.get("display_name")); R["orcid"].append(au.get("orcid"))123 R["is_corresponding"].append(bool(a.get("is_corresponding"))); R["institution_ids"].append(ii)124 R["institution_names"].append([x.get("display_name") for x in ins]); R["institution_types"].append([x.get("type") for x in ins])125 R["countries"].append(cc); R["n_raw_affiliations"].append(len(a.get("raw_affiliation_strings") or []))126 for k, t in enumerate(topics):127 R = rows["work_topics"]128 R["work_id"].append(wid); R["rank"].append(k); R["topic_id"].append(oid(t.get("id"))); R["score"].append(t.get("score"))129 R["subfield_id"].append(oid((t.get("subfield") or {}).get("id"))); R["field_id"].append(oid((t.get("field") or {}).get("id")))130 R["domain_id"].append(oid((t.get("domain") or {}).get("id")))131 for kw in (w.get("keywords") or []):132 R = rows["work_keywords"]; R["work_id"].append(wid); R["keyword"].append(kw.get("display_name")); R["score"].append(kw.get("score"))133 refs = w.get("referenced_works") or []134 for rr in refs:135 rows["references"]["work_id"].append(wid); rows["references"]["ref_id"].append(oid(rr))136 ab = abstract_text(w.get("abstract_inverted_index"))137 if ab:138 R = rows["abstracts"]; R["work_id"].append(wid); R["abstract"].append(ab); R["n_words"].append(len(ab.split()))139 locs = w.get("locations") or []140 best_url = best.get("landing_page_url") if best else None141 for k, L in enumerate(locs):142 s_ = L.get("source") or {}143 R = rows["locations"]144 R["work_id"].append(wid); R["rank"].append(k); R["is_primary"].append(L.get("landing_page_url") == loc.get("landing_page_url") and L.get("pdf_url") == loc.get("pdf_url"))145 R["is_best_oa"].append(bool(best_url) and L.get("landing_page_url") == best_url)146 R["source_id"].append(oid(s_.get("id"))); R["source_name"].append(s_.get("display_name")); R["source_type"].append(s_.get("type"))147 R["landing_page_url"].append(L.get("landing_page_url")); R["pdf_url"].append(L.get("pdf_url")); R["is_oa"].append(bool(L.get("is_oa")))148 R["license"].append(L.get("license")); R["version"].append(L.get("version")); R["is_accepted"].append(bool(L.get("is_accepted"))); R["is_published"].append(bool(L.get("is_published")))149 for g in (w.get("grants") or []):150 R = rows["grants"]; R["work_id"].append(wid); R["funder_id"].append(oid(g.get("funder"))); R["funder_name"].append(g.get("funder_display_name")); R["award_id"].append(g.get("award_id"))151 for c in (w.get("counts_by_year") or []):152 R = rows["counts_by_year"]; R["work_id"].append(wid); R["year"].append(c.get("year")); R["cited_by_count"].append(c.get("cited_by_count"))153 cnp = w.get("citation_normalized_percentile") or {}154 ids = w.get("ids") or {}155 apc = (w.get("apc_paid") or {}).get("value_usd") if isinstance(w.get("apc_paid"), dict) else None156 R = rows["works"]157 vals = dict(id=wid, doi=doi, title=w.get("title") or w.get("display_name"), publication_year=w.get("publication_year"), publication_date=w.get("publication_date"),158 created_date=w.get("created_date"), updated_date=(w.get("updated_date") or "")[:10] or None, type=w.get("type"), type_crossref=w.get("type_crossref"),159 language=w.get("language"), cited_by_count=w.get("cited_by_count"), fwci=w.get("fwci"), citation_percentile=cnp.get("value"),160 is_top1=bool(cnp.get("is_in_top_1_percent")), is_top10=bool(cnp.get("is_in_top_10_percent")), is_oa=bool(oa.get("is_oa")), oa_status=oa.get("oa_status"),161 license=(best.get("license") if best else None) or loc.get("license"), is_retracted=bool(w.get("is_retracted")), is_paratext=bool(w.get("is_paratext")),162 is_xpac=bool(w.get("is_xpac")), n_refs=len(refs), n_authors=w.get("authors_count") if w.get("authors_count") is not None else len(auths),163 n_institutions=len(inst_ids), n_countries=len(countries), countries=sorted(countries), institution_types=sorted(set(itypes)),164 primary_topic_id=oid(pt.get("id")), primary_subfield_id=oid((pt.get("subfield") or {}).get("id")), primary_field_id=oid((pt.get("field") or {}).get("id")),165 primary_domain_id=oid((pt.get("domain") or {}).get("id")), cs_primary=(pt.get("field") or {}).get("id") == cs_field, cs_any=bool(cs_any),166 source_id=oid(src.get("id")), source_name=src.get("display_name"), source_type=src.get("type"), host_organization=src.get("host_organization_name"),167 is_in_doaj=bool(src.get("is_in_doaj")), arxiv_id=arxiv_from(w), pmid=(ids.get("pmid") or "").rsplit("/", 1)[-1] or None, mag=str(ids.get("mag")) if ids.get("mag") else None,168 has_abstract=ab is not None, has_fulltext=bool(w.get("has_fulltext") or w.get("has_content")), n_locations=len(locs), indexed_in=w.get("indexed_in") or [],169 referenced_works_count=w.get("referenced_works_count") if w.get("referenced_works_count") is not None else len(refs), n_grants=len(w.get("grants") or []),170 apc_paid_usd=apc, is_authors_truncated=bool(w.get("is_authors_truncated")), part=part)171 for k, v in vals.items():172 R[k].append(v)173 for t, sch in S.items():174 tbl = pa.Table.from_pydict(rows[t], schema=sch)175 tmp = outs[t] + ".part"176 pq.write_table(tbl, tmp, compression="zstd"); os.replace(tmp, outs[t])177 return part, "done", n178179180def main():181 ap = argparse.ArgumentParser(); ap.add_argument("--cs-dir", required=True); ap.add_argument("--out", required=True)182 ap.add_argument("--workers", type=int, default=8); ap.add_argument("--limit", type=int, default=0)183 a = ap.parse_args()184 files = sorted(glob.glob(os.path.join(a.cs_dir, "*.jsonl.zst")))185 if a.limit:186 files = files[: a.limit]187 t0 = time.time(); tot = 0; done = 0188 with mp.Pool(a.workers, maxtasksperchild=40) as pool:189 for part, st, n in pool.imap_unordered(process, [(f, a.out, "https://openalex.org/fields/17") for f in files]):190 tot += n; done += 1191 if st == "done" or done % 100 == 0:192 print(f"[{dt.datetime.now():%H:%M:%S}] {st} {part} n={n} | {done}/{len(files)} files, {tot} works, {time.time()-t0:.0f}s", flush=True)193 print("finished", done, "files", tot, "works")194195196if __name__ == "__main__":197 main()198