diff --git a/sync/build_csr.py b/sync/build_csr.py index 22a7fdb..ef96ab7 100644 --- a/sync/build_csr.py +++ b/sync/build_csr.py @@ -526,6 +526,201 @@ def build_all_csr( return results +# ── Edge-list export ───────────────────────────────────────────────────── + + +def build_edge_list( + rel_type: str, + *, + parquet_dir: Path | None = None, + output_dir: Path | None = None, + force: bool = False, + memory_limit: str | None = None, + direction: str = "src", +) -> dict: + """Export a relationship as a sorted, deduplicated edge-list Parquet. + + Writes ``__by_.parquet`` with columns ``(src, tgt)`` + in the original OpenAlex IDs — so it joins the ``main`` and relationship + tables directly, no dense-index/id-map indirection. The edges are + deduplicated and sorted with bounded row groups so DuckDB (or any Parquet + reader) prunes lookups to a few row groups instead of loading the graph: + + * ``direction="src"`` sorts by ``(src, tgt)`` — prunes ``WHERE src = X`` + ("what X cites"). + * ``direction="tgt"`` sorts by ``(tgt, src)`` — prunes ``WHERE tgt = X`` + ("who cites X"). + + A complement to the CSR ``.npz``: the matrix is built for in-RAM linear + algebra, this is for out-of-core SQL queries over the same edges. + + Idempotent (provenance fingerprint of the input shards) and atomic + (temp file + ``os.replace``), matching :func:`build_csr`. + """ + if duckdb is None: + raise ImportError("duckdb is required for edge-list export") + if direction not in ("src", "tgt"): + raise ValueError(f"direction must be 'src' or 'tgt', got {direction!r}") + + _parquet_dir = parquet_dir or _resolve_parquet_dir() + _output_dir = output_dir or _resolve_csr_dir() + + if rel_type not in _CSR_RELATIONSHIP_TYPES: + return {"error": f"Unknown relationship type: {rel_type}"} + + src_col, tgt_col = _CSR_RELATIONSHIP_TYPES[rel_type] + parquet_files = _find_parquet_shards(_parquet_dir, rel_type) + if not parquet_files: + return {"error": f"No parquet shards found for {rel_type}"} + + # Sort by the queried endpoint first so its row-group zonemaps prune. + order_by = "src, tgt" if direction == "src" else "tgt, src" + output_path = _output_dir / f"{rel_type}__by_{direction}.parquet" + fingerprint = _input_fingerprint(parquet_files) + + # Idempotency check + if not force and output_path.exists(): + existing = _load_provenance(output_path) + if existing and existing.get("input_fingerprint") == fingerprint: + log.info( + "%s by_%s: up-to-date (%d shards, fingerprint %s), skipping", + rel_type, direction, len(parquet_files), fingerprint, + ) + return { + "rel_type": rel_type, + "direction": direction, + "status": "skipped", + "n_edges": existing.get("n_edges", 0), + "shard_count": len(parquet_files), + } + + log.info( + "%s by_%s: exporting from %d shards (fingerprint %s)", + rel_type, direction, len(parquet_files), fingerprint, + ) + t0 = time.time() + + _output_dir.mkdir(parents=True, exist_ok=True) + glob_pattern = str(parquet_files[0].parent / "*.parquet") + + # Atomic: write to a temp file, then os.replace into place. + fd, tmp_path = tempfile.mkstemp( + dir=_output_dir, prefix=".tmp-edges-", suffix=".parquet" + ) + os.close(fd) + try: + con = duckdb.connect(":memory:") + try: + # Give the DISTINCT and ORDER BY a spill target so they can + # exceed RAM (an in-memory connection won't otherwise spill). + con.execute(f"SET temp_directory='{_output_dir}'") + if memory_limit: + con.execute(f"SET memory_limit='{memory_limit}'") + con.execute(f""" + COPY ( + SELECT DISTINCT + CAST("{src_col}" AS UBIGINT) AS src, + CAST("{tgt_col}" AS UBIGINT) AS tgt + FROM read_parquet('{glob_pattern}') + WHERE "{src_col}" IS NOT NULL + AND "{tgt_col}" IS NOT NULL + ORDER BY {order_by} + ) TO '{tmp_path}' ( + FORMAT PARQUET, + COMPRESSION zstd, + COMPRESSION_LEVEL 3, + ROW_GROUP_SIZE 1000000 + ) + """) + n_edges = con.execute( + f"SELECT COUNT(*) FROM read_parquet('{tmp_path}')" + ).fetchone()[0] + finally: + con.close() + os.replace(tmp_path, output_path) + except BaseException: + try: + os.unlink(tmp_path) + except OSError: + pass + raise + + _write_provenance( + output_path, + rel_type=rel_type, + n_edges=n_edges, + n_nodes=0, # not meaningful for an edge list + shard_count=len(parquet_files), + input_fingerprint=fingerprint, + ) + + elapsed = time.time() - t0 + output_size = output_path.stat().st_size + log.info( + "%s by_%s: %d edges, %.2f GB, %.1fs", + rel_type, direction, n_edges, output_size / 1e9, elapsed, + ) + + return { + "rel_type": rel_type, + "direction": direction, + "status": "built", + "n_edges": n_edges, + "shard_count": len(parquet_files), + "output_path": str(output_path), + "output_size_bytes": output_size, + "elapsed_seconds": elapsed, + } + + +def build_all_edge_lists( + *, + parquet_dir: Path | None = None, + output_dir: Path | None = None, + force: bool = False, + memory_limit: str | None = None, + rel_types: list[str] | None = None, + directions: tuple[str, ...] = ("src", "tgt"), +) -> list[dict]: + """Export edge-list Parquets for all (or specified) relationship types. + + Emits one file per ``(rel_type, direction)``: ``by_src`` for + ``WHERE src = X`` queries and ``by_tgt`` for ``WHERE tgt = X``. Pass a + single-element ``directions`` to export only one orientation. + """ + _output_dir = output_dir or _resolve_csr_dir() + _output_dir.mkdir(parents=True, exist_ok=True) + + _clean_orphan_temps(_output_dir) + + types = rel_types or list(_CSR_RELATIONSHIP_TYPES.keys()) + results: list[dict] = [] + + for rel_type in types: + for direction in directions: + result = build_edge_list( + rel_type, + parquet_dir=parquet_dir, + output_dir=_output_dir, + force=force, + memory_limit=memory_limit, + direction=direction, + ) + results.append(result) + if "error" in result: + log.warning("%s by_%s: %s", rel_type, direction, result["error"]) + + built = sum(1 for r in results if r.get("status") == "built") + skipped = sum(1 for r in results if r.get("status") == "skipped") + failed = sum(1 for r in results if "error" in r) + log.info( + "Edge-list export complete: %d built, %d skipped, %d failed", + built, skipped, failed, + ) + + return results + + # ── CLI ───────────────────────────────────────────────────────────────── @@ -570,6 +765,14 @@ def main(argv: list[str] | None = None) -> int: help="DuckDB memory limit (e.g. '32GB') " f"(default: env {_MEMORY_LIMIT_ENV} or none)", ) + parser.add_argument( + "--edge-list", + action="store_true", + help="Export sorted, deduplicated edge-list Parquet " + "(__by_src.parquet and __by_tgt.parquet) instead of CSR " + "matrices — queryable directly in DuckDB without loading the " + "whole graph", + ) args = parser.parse_args(argv) @@ -584,7 +787,8 @@ def main(argv: list[str] | None = None) -> int: memory = args.memory_limit or os.environ.get(_MEMORY_LIMIT_ENV) - results = build_all_csr( + builder = build_all_edge_lists if args.edge_list else build_all_csr + results = builder( parquet_dir=args.parquet_dir, output_dir=args.output_dir, force=args.force, diff --git a/sync/readme.py b/sync/readme.py index 270ab17..dc21659 100644 --- a/sync/readme.py +++ b/sync/readme.py @@ -230,6 +230,25 @@ def _yaml_dump_configs(configs: list[dict]) -> str: python3 -m sync.build_csr --rel-type work_referenced_works ``` +### Edge-list Parquet + +For querying the graph in DuckDB/Arrow without loading a whole matrix into memory, `--edge-list` exports each relationship as a sorted, deduplicated edge list in the original OpenAlex IDs (so it joins the `main` and relationship tables directly), in both directions: + +``` +csr/ + work_referenced_works__by_src.parquet # (src, tgt), sorted by src + work_referenced_works__by_tgt.parquet # (src, tgt), sorted by tgt + ... +``` + +Each file's bounded, sorted row groups let zonemaps prune that direction's lookups to a few row groups — millisecond range scans against the full graph. `by_src` answers "what X cites" (`WHERE src = X`); `by_tgt` answers "who cites X" (`WHERE tgt = X`). + +```bash +python3 -m sync.build_csr --all --edge-list +duckdb -c "SELECT tgt FROM 'csr/work_referenced_works__by_src.parquet' WHERE src = 2741809807" # what it cites +duckdb -c "SELECT src FROM 'csr/work_referenced_works__by_tgt.parquet' WHERE tgt = 2741809807" # who cites it +``` + ### Example: Work record fields `id`, `doi`, `title`, `display_name`, `publication_year`, `type`, `language`, `authorships`, `concepts`, `topics`, `keywords`, `cited_by_count`, `referenced_works`, `related_works`, `locations`, `open_access`, `funders`, `awards`, `mesh`, `sustainable_development_goals`, `counts_by_year`, `updated_date`, and more. diff --git a/tests/test_build_edges.py b/tests/test_build_edges.py new file mode 100644 index 0000000..db7f8c8 --- /dev/null +++ b/tests/test_build_edges.py @@ -0,0 +1,144 @@ +"""Tests for sync.build_csr.build_edge_list — the edge-list Parquet export. + +Pins the artifact's contract: a deduplicated, sorted edge list in original +OpenAlex IDs, queryable directly in DuckDB. build_csr's deps (duckdb/pyarrow) +are optional relative to the core pipeline; the file skips where absent. +""" + +from __future__ import annotations + +import pytest + +pa = pytest.importorskip("pyarrow") +pq = pytest.importorskip("pyarrow.parquet") +duckdb = pytest.importorskip("duckdb") + +from sync import build_csr # noqa: E402 + +REL = "work_referenced_works" +SRC_COL, TGT_COL = "work_id", "referenced_work_id" + + +def _write_shard(shard_dir, edges, index=0): + """Write ``edges`` (list of (src, tgt), None allowed) as one parquet shard.""" + shard_dir.mkdir(parents=True, exist_ok=True) + src = pa.array([e[0] for e in edges], type=pa.uint64()) + tgt = pa.array([e[1] for e in edges], type=pa.uint64()) + pq.write_table( + pa.table({SRC_COL: src, TGT_COL: tgt}), + shard_dir / f"part_{index:04d}.parquet", + ) + + +def _build(tmp_path, edges=None, *, shards=None, force=True, direction="src"): + """Export the by_ edge list for ``edges`` (one shard) or ``shards``.""" + groups = shards if shards is not None else [edges] + parquet_dir = tmp_path / "data" + shard_dir = build_csr.rt_dir(parquet_dir, REL) + for i, grp in enumerate(groups): + _write_shard(shard_dir, grp, i) + out_dir = tmp_path / "csr" + result = build_csr.build_edge_list( + REL, parquet_dir=parquet_dir, output_dir=out_dir, force=force, + direction=direction, + ) + return result, out_dir / f"{REL}__by_{direction}.parquet" + + +def _read(path): + """Return the edge list as a list of (src, tgt) tuples, file order.""" + tbl = pq.read_table(path) + return list(zip(tbl.column("src").to_pylist(), tbl.column("tgt").to_pylist())) + + +def test_dedups_sorts_and_drops_nulls(tmp_path): + edges = [ + (10, 30), (10, 20), (40, 10), (10, 30), # (10, 30) duplicated + (30, 30), (20, 40), + (50, None), (None, 60), # null endpoints — dropped entirely + ] + result, path = _build(tmp_path, edges) + got = _read(path) + expected = sorted({(s, t) for s, t in edges if s is not None and t is not None}) + assert got == expected # deduplicated, nulls dropped + assert got == sorted(got) # sorted by (src, tgt) for zonemap pruning + assert result["n_edges"] == len(expected) + + +def test_uses_original_ids_not_dense_indices(tmp_path): + # Sparse, OpenAlex-scale IDs must survive verbatim — no dense remapping. + edges = [(7_000_000_001, 12), (12, 7_000_000_001)] + _, path = _build(tmp_path, edges) + got = _read(path) + assert got == [(12, 7_000_000_001), (7_000_000_001, 12)] + + +def test_dedups_across_shards(tmp_path): + # The same edge in different shards must collapse globally. + shards = [[(1, 2), (3, 4)], [(1, 2), (5, 6)]] + _, path = _build(tmp_path, shards=shards) + assert _read(path) == [(1, 2), (3, 4), (5, 6)] + + +def test_is_queryable_by_src_in_duckdb(tmp_path): + edges = [(10, 20), (10, 30), (10, 40), (50, 60)] + _, path = _build(tmp_path, edges) + con = duckdb.connect() + rows = con.execute( + f"SELECT tgt FROM read_parquet('{path}') WHERE src = 10 ORDER BY tgt" + ).fetchall() + assert [r[0] for r in rows] == [20, 30, 40] + + +def test_skips_unchanged_inputs(tmp_path): + edges = [(10, 20), (20, 30)] + result1, path = _build(tmp_path, edges) + assert result1["status"] == "built" + mtime = path.stat().st_mtime_ns + result2, _ = _build(tmp_path, edges, force=False) + assert result2["status"] == "skipped" + assert path.stat().st_mtime_ns == mtime # not rewritten + + +def test_empty_relationship_yields_empty_edge_list(tmp_path): + result, path = _build(tmp_path, []) + assert result["n_edges"] == 0 + assert _read(path) == [] + + +def test_by_tgt_is_sorted_by_tgt(tmp_path): + edges = [(10, 30), (40, 10), (20, 30), (10, 20), (40, 10)] # dup (40, 10) + _, path = _build(tmp_path, edges, direction="tgt") + got = _read(path) + expected = sorted({(s, t) for s, t in edges}, key=lambda e: (e[1], e[0])) + assert got == expected # deduped, sorted by (tgt, src) + assert got == sorted(got, key=lambda e: (e[1], e[0])) + + +def test_by_tgt_is_queryable_by_tgt_in_duckdb(tmp_path): + edges = [(10, 99), (20, 99), (30, 99), (40, 50)] + _, path = _build(tmp_path, edges, direction="tgt") + con = duckdb.connect() + rows = con.execute( + f"SELECT src FROM read_parquet('{path}') WHERE tgt = 99 ORDER BY src" + ).fetchall() + assert [r[0] for r in rows] == [10, 20, 30] + + +def test_both_directions_hold_the_same_edge_set(tmp_path): + edges = [(10, 30), (40, 10), (20, 30), (10, 20)] + _, src_path = _build(tmp_path / "s", edges, direction="src") + _, tgt_path = _build(tmp_path / "t", edges, direction="tgt") + assert set(_read(src_path)) == set(_read(tgt_path)) + + +def test_build_all_emits_both_directions(tmp_path): + parquet_dir = tmp_path / "data" + _write_shard(build_csr.rt_dir(parquet_dir, REL), [(1, 2), (2, 3)]) + out_dir = tmp_path / "csr" + results = build_csr.build_all_edge_lists( + parquet_dir=parquet_dir, output_dir=out_dir, rel_types=[REL] + ) + assert {r["direction"] for r in results} == {"src", "tgt"} + assert (out_dir / f"{REL}__by_src.parquet").exists() + assert (out_dir / f"{REL}__by_tgt.parquet").exists()