From 73032cb5ba66ba0eda22c1250ebb9e4b575e0800 Mon Sep 17 00:00:00 2001 From: ak2k <19240940+ak2k@users.noreply.github.com> Date: Sat, 20 Jun 2026 16:01:16 -0400 Subject: [PATCH 1/2] feat(build_csr): export a queryable edge-list Parquet (by_src) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CSR .npz is built for in-RAM linear algebra: scipy.sparse.load_npz pulls the whole matrix into memory and is Python-specific. For querying the graph in DuckDB/Arrow — the rest of the dataset's format — there was no out-of-core option. Add build_edge_list / --edge-list, which exports each relationship as a deduplicated edge list sorted by (src, tgt), in the original OpenAlex IDs (so it joins the main and relationship tables directly, no dense-index/id-map hop). Bounded row groups + the sort let DuckDB zonemaps prune WHERE src = X ('what X cites') to a few row groups. On a 60M-edge sample: 288 MB (zstd-3), and an out-neighbours lookup is ~2 ms straight off the Parquet with no full load. Idempotent (provenance fingerprint) and atomic (temp + os.replace), matching build_csr; temp_directory is set so the DISTINCT/ORDER BY can spill under a tight memory_limit. The reverse direction (who-cites-X, WHERE tgt = X) wants a tgt-sorted copy; that is a separate change. Refs #4. --- sync/build_csr.py | 186 +++++++++++++++++++++++++++++++++++++- sync/readme.py | 17 ++++ tests/test_build_edges.py | 105 +++++++++++++++++++++ 3 files changed, 307 insertions(+), 1 deletion(-) create mode 100644 tests/test_build_edges.py diff --git a/sync/build_csr.py b/sync/build_csr.py index 22a7fdb..e4e0476 100644 --- a/sync/build_csr.py +++ b/sync/build_csr.py @@ -526,6 +526,182 @@ 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, +) -> dict: + """Export a relationship as a sorted, deduplicated edge-list Parquet. + + Writes ``__by_src.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 by ``(src, tgt)`` with bounded row groups, so DuckDB (or any + Parquet reader) prunes ``WHERE src = X`` lookups ("what X cites") to a + few row groups instead of loading the whole graph. + + 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") + + _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}"} + + output_path = _output_dir / f"{rel_type}__by_src.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 edge-list: up-to-date (%d shards, fingerprint %s), skipping", + rel_type, len(parquet_files), fingerprint, + ) + return { + "rel_type": rel_type, + "status": "skipped", + "n_edges": existing.get("n_edges", 0), + "shard_count": len(parquet_files), + } + + log.info( + "%s edge-list: exporting from %d shards (fingerprint %s)", + rel_type, 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 src, tgt + ) 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 edge-list: %d edges, %.2f GB, %.1fs", + rel_type, n_edges, output_size / 1e9, elapsed, + ) + + return { + "rel_type": rel_type, + "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, +) -> list[dict]: + """Export edge-list Parquets for all (or specified) relationship types.""" + _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: + result = build_edge_list( + rel_type, + parquet_dir=parquet_dir, + output_dir=_output_dir, + force=force, + memory_limit=memory_limit, + ) + results.append(result) + if "error" in result: + log.warning("%s: %s", rel_type, 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 +746,13 @@ 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) instead of CSR matrices — queryable " + "directly in DuckDB without loading the whole graph", + ) args = parser.parse_args(argv) @@ -584,7 +767,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..5feedcb 100644 --- a/sync/readme.py +++ b/sync/readme.py @@ -230,6 +230,23 @@ 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): + +``` +csr/ + work_referenced_works__by_src.parquet # (src, tgt), sorted by src + ... +``` + +Sorted with bounded row groups, so `WHERE src = X` ("what X cites") prunes to a few row groups — millisecond range scans against the full graph. + +```bash +python3 -m sync.build_csr --all --edge-list +duckdb -c "SELECT tgt FROM 'csr/work_referenced_works__by_src.parquet' WHERE src = 2741809807" +``` + ### 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..e9ba284 --- /dev/null +++ b/tests/test_build_edges.py @@ -0,0 +1,105 @@ +"""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): + """Export the by_src 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 + ) + return result, out_dir / f"{REL}__by_src.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) == [] From 1c7261523efc64b4ea55fe96980359a443e79460 Mon Sep 17 00:00:00 2001 From: ak2k <19240940+ak2k@users.noreply.github.com> Date: Sat, 20 Jun 2026 16:04:41 -0400 Subject: [PATCH 2/2] feat(build_csr): export the reverse-direction edge list (by_tgt) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A by_src edge list (sorted by src) prunes 'what X cites' (WHERE src = X) but not 'who cites X' (WHERE tgt = X) — a tgt predicate scans every row group of a src-sorted file. who-cites-X is the more common citation query (forward citations, supersession, discovery), so add a tgt-sorted copy. Generalize build_edge_list with a direction ('src'|'tgt'): by_src sorts (src, tgt), by_tgt sorts (tgt, src), each emitted as __by_.parquet. build_all_edge_lists / --edge-list now emit both; pass directions=('src',) to keep only the forward copy. Same edge set, just two sort orders, so both directions stay O(ms): on a 60M-edge sample, a who-cites-X lookup is ~2 ms straight off by_tgt. The cost is storage — a second sorted copy roughly doubles the graph's footprint (~14 GB -> ~28 GB at full work_referenced_works scale), and it is a separate dedup+sort pass. directions=('src',) opts out if that trade-off isn't wanted. Closes #4. --- sync/build_csr.py | 74 +++++++++++++++++++++++++-------------- sync/readme.py | 8 +++-- tests/test_build_edges.py | 47 ++++++++++++++++++++++--- 3 files changed, 95 insertions(+), 34 deletions(-) diff --git a/sync/build_csr.py b/sync/build_csr.py index e4e0476..ef96ab7 100644 --- a/sync/build_csr.py +++ b/sync/build_csr.py @@ -536,15 +536,20 @@ def build_edge_list( 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_src.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 by ``(src, tgt)`` with bounded row groups, so DuckDB (or any - Parquet reader) prunes ``WHERE src = X`` lookups ("what X cites") to a - few row groups instead of loading the whole graph. + 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. @@ -554,6 +559,8 @@ def build_edge_list( """ 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() @@ -566,7 +573,9 @@ def build_edge_list( if not parquet_files: return {"error": f"No parquet shards found for {rel_type}"} - output_path = _output_dir / f"{rel_type}__by_src.parquet" + # 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 @@ -574,19 +583,20 @@ def build_edge_list( existing = _load_provenance(output_path) if existing and existing.get("input_fingerprint") == fingerprint: log.info( - "%s edge-list: up-to-date (%d shards, fingerprint %s), skipping", - rel_type, len(parquet_files), fingerprint, + "%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 edge-list: exporting from %d shards (fingerprint %s)", - rel_type, len(parquet_files), fingerprint, + "%s by_%s: exporting from %d shards (fingerprint %s)", + rel_type, direction, len(parquet_files), fingerprint, ) t0 = time.time() @@ -614,7 +624,7 @@ def build_edge_list( FROM read_parquet('{glob_pattern}') WHERE "{src_col}" IS NOT NULL AND "{tgt_col}" IS NOT NULL - ORDER BY src, tgt + ORDER BY {order_by} ) TO '{tmp_path}' ( FORMAT PARQUET, COMPRESSION zstd, @@ -647,12 +657,13 @@ def build_edge_list( elapsed = time.time() - t0 output_size = output_path.stat().st_size log.info( - "%s edge-list: %d edges, %.2f GB, %.1fs", - rel_type, n_edges, output_size / 1e9, elapsed, + "%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), @@ -669,8 +680,14 @@ def build_all_edge_lists( 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.""" + """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) @@ -680,16 +697,18 @@ def build_all_edge_lists( results: list[dict] = [] for rel_type in types: - result = build_edge_list( - rel_type, - parquet_dir=parquet_dir, - output_dir=_output_dir, - force=force, - memory_limit=memory_limit, - ) - results.append(result) - if "error" in result: - log.warning("%s: %s", rel_type, result["error"]) + 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") @@ -750,8 +769,9 @@ def main(argv: list[str] | None = None) -> int: "--edge-list", action="store_true", help="Export sorted, deduplicated edge-list Parquet " - "(__by_src.parquet) instead of CSR matrices — queryable " - "directly in DuckDB without loading the whole graph", + "(__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) diff --git a/sync/readme.py b/sync/readme.py index 5feedcb..dc21659 100644 --- a/sync/readme.py +++ b/sync/readme.py @@ -232,19 +232,21 @@ def _yaml_dump_configs(configs: list[dict]) -> str: ### 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): +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 ... ``` -Sorted with bounded row groups, so `WHERE src = X` ("what X cites") prunes to a few row groups — millisecond range scans against the full graph. +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" +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 diff --git a/tests/test_build_edges.py b/tests/test_build_edges.py index e9ba284..db7f8c8 100644 --- a/tests/test_build_edges.py +++ b/tests/test_build_edges.py @@ -30,8 +30,8 @@ def _write_shard(shard_dir, edges, index=0): ) -def _build(tmp_path, edges=None, *, shards=None, force=True): - """Export the by_src edge list for ``edges`` (one shard) or ``shards``.""" +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) @@ -39,9 +39,10 @@ def _build(tmp_path, edges=None, *, shards=None, force=True): _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 + REL, parquet_dir=parquet_dir, output_dir=out_dir, force=force, + direction=direction, ) - return result, out_dir / f"{REL}__by_src.parquet" + return result, out_dir / f"{REL}__by_{direction}.parquet" def _read(path): @@ -103,3 +104,41 @@ 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()