feat(build_csr): export the reverse-direction edge list (by_tgt)

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 ``<rel_type>__by_<direction>.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 "
+             "(<rel>__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,

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.

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_<direction> 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()