diff --git a/sync/build_csr.py b/sync/build_csr.py index 22a7fdb..02b18d6 100644 --- a/sync/build_csr.py +++ b/sync/build_csr.py @@ -224,11 +224,16 @@ def _build_csr_duckdb( contiguous range [0, n_unique_nodes) to keep the indptr array small. The remapping is deterministic: sorted by original ID. - DuckDB reads all shards, deduplicates edges, and writes intermediate - results to temporary parquet files. Python reads these in chunks via - pyarrow, keeping peak memory bounded regardless of edge count. - ``work_referenced_works`` (~3B deduplicated edges) fits in 64 GB - with this approach, where ``fetchall()`` needed ~240 GB. + DuckDB does the whole remap. It builds a dense ``id -> idx`` dimension + table (``row_number() OVER (ORDER BY id)``), joins the deduplicated + edges against it twice, and writes the already-dense, already-sorted + ``(src_idx, tgt_idx)`` pairs to a temporary parquet file. Python then + streams that file straight into the preallocated CSR arrays — no + per-batch binary search over the id array, which on the ~3B-edge + work_referenced_works would dominate the build: each such lookup misses + cache on the multi-GB sorted id array. Python holds the id array (for + the id map), ``indptr``, ``indices`` and one batch; the remap join and + its spill stay inside DuckDB, bounded by ``memory_limit``. """ if duckdb is None: raise ImportError("duckdb is required for CSR building") @@ -246,45 +251,61 @@ def _build_csr_duckdb( try: con = duckdb.connect(":memory:") try: + # Give the dim-table window, the joins and the final sort a spill + # target so they can exceed RAM — essential under a tight + # memory_limit, since an in-memory connection won't otherwise + # spill. output_path.parent is the (already-created) CSR dir. + con.execute(f"SET temp_directory='{output_path.parent}'") if memory_limit: con.execute(f"SET memory_limit='{memory_limit}'") - # Step 1: Collect all unique node IDs, sorted deterministically. + # Step 1: Build the dense id -> idx dimension table. The dense + # index is the rank of the ID in sorted order, so the mapping is + # deterministic and id-sorted, exactly as the indptr layout needs. # DuckDB spills to disk if the set exceeds memory_limit. log.info("Collecting unique node IDs...") con.execute(f""" - COPY ( - WITH all_ids AS ( - SELECT DISTINCT CAST("{src_col}" AS UBIGINT) AS id - FROM read_parquet('{glob_pattern}') - WHERE "{src_col}" IS NOT NULL - UNION ALL - SELECT DISTINCT CAST("{tgt_col}" AS UBIGINT) AS id - FROM read_parquet('{glob_pattern}') - WHERE "{tgt_col}" IS NOT NULL - ) - SELECT DISTINCT id - FROM all_ids - ORDER BY id - ) TO '{tmp_ids}' (FORMAT PARQUET) + CREATE TABLE ids AS + SELECT id, (row_number() OVER (ORDER BY id) - 1) AS idx + FROM ( + SELECT CAST("{src_col}" AS UBIGINT) AS id + FROM read_parquet('{glob_pattern}') + WHERE "{src_col}" IS NOT NULL + UNION + SELECT CAST("{tgt_col}" AS UBIGINT) AS id + FROM read_parquet('{glob_pattern}') + WHERE "{tgt_col}" IS NOT NULL + ) """) - n_nodes = con.execute( - f"SELECT COUNT(*) FROM read_parquet('{tmp_ids}')" - ).fetchone()[0] + n_nodes = con.execute("SELECT COUNT(*) FROM ids").fetchone()[0] log.info("Unique nodes: %d", n_nodes) - # Step 2: Deduplicated edges sorted by (src, tgt). - # Still using original IDs — remapping happens in Python. - log.info("Deduplicating and sorting edges...") + # Emit the id map in dense order (idx is monotone in id, so + # ORDER BY idx == ORDER BY id): original_ids[dense_idx] = id. + con.execute( + f"COPY (SELECT id FROM ids ORDER BY idx) " + f"TO '{tmp_ids}' (FORMAT PARQUET)" + ) + + # Step 2: Deduplicate edges, remap both endpoints to dense + # indices via the dimension table, and sort by (src_idx, tgt_idx). + # Because idx is monotone in id, this order equals the (src, tgt) + # id order — the CSR row/column ordering the indptr build assumes. + log.info("Deduplicating, remapping and sorting edges...") 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 + SELECT s.idx AS src_idx, t.idx AS tgt_idx + FROM ( + 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 + ) e + JOIN ids s ON e.src = s.id + JOIN ids t ON e.tgt = t.id + ORDER BY src_idx, tgt_idx ) TO '{tmp_edges}' (FORMAT PARQUET) """) n_edges = con.execute( @@ -298,9 +319,10 @@ def _build_csr_duckdb( if n_edges == 0: return sparse.csr_matrix((0, 0), dtype=np.float64), np.array([], dtype=np.uint64) - # Step 3: Read the sorted unique IDs — this is the dense mapping. - # original_ids[dense_idx] = original OpenAlex ID. - # ~2 GB for 250M nodes (uint64), well within memory. + # Step 3: Read the dense-ordered IDs for the returned id map: + # original_ids[dense_idx] = original OpenAlex ID. The remap itself + # already happened in DuckDB, so this array is only the output + # mapping. ~2 GB for 250M nodes (uint64), well within memory. original_ids = pq.read_table( str(tmp_ids), columns=["id"] ).column("id").to_numpy() @@ -311,40 +333,49 @@ def _build_csr_duckdb( except OSError: pass - # Step 4: Read edges in chunks, remap to dense indices via - # binary search (original_ids is sorted), and accumulate CSR - # components. Peak memory per batch: ~80 MB (5M rows × 16 bytes). + # Step 4: Stream the already-dense, already-sorted edges into the CSR + # components. The remap happened in DuckDB, so per batch this is just + # a copy plus a bincount — no binary search. Peak memory per batch: + # ~40 MB (5M rows × 8 bytes). + # + # The CSR ``indices`` array is the full deduplicated tgt-index set + # (~12 GB for work_referenced_works's ~3B edges), so it is + # preallocated at its final size and filled slice-by-slice as the + # batches stream in. Collecting per-batch arrays in a list and + # ``np.concatenate``-ing at the end would transiently hold both the + # list and the joined array — doubling that 12 GB at the worst + # possible moment. The edge count is known exactly (``n_edges``), + # so the single allocation is safe. indptr = np.zeros(n_nodes + 1, dtype=np.int64) - indices_chunks: list[np.ndarray] = [] + indices = np.empty(n_edges, dtype=np.int32) + offset = 0 batch_count = 0 pf = pq.ParquetFile(str(tmp_edges)) for batch in pf.iter_batches(batch_size=5_000_000): - src = batch.column("src").to_numpy() - tgt = batch.column("tgt").to_numpy() - - # Dense remap via binary search (monotone mapping). - # Since original_ids is sorted and edges are sorted by - # (src, tgt), the resulting (src_idx, tgt_idx) are also - # sorted — no re-sorting needed. - src_idx = np.searchsorted(original_ids, src).astype(np.int32) - tgt_idx = np.searchsorted(original_ids, tgt).astype(np.int32) + src_idx = batch.column("src_idx").to_numpy() + tgt_idx = batch.column("tgt_idx").to_numpy().astype(np.int32) - indices_chunks.append(tgt_idx) + # Append this batch's column indices into the preallocated array. + # Edges arrive in (src_idx, tgt_idx) order, so writing them in + # batch order keeps ``indices`` grouped by source node — exactly + # the CSR layout that ``indptr`` describes. + indices[offset:offset + len(tgt_idx)] = tgt_idx + offset += len(tgt_idx) # Count edges per source node for indptr construction. counts = np.bincount(src_idx, minlength=n_nodes) indptr[1:] += counts batch_count += 1 - log.debug("Processed batch %d (%d rows)", batch_count, len(src)) + log.debug("Processed batch %d (%d rows)", batch_count, len(src_idx)) + + # All deduplicated edges must have been placed exactly once. + assert offset == n_edges, f"placed {offset} indices, expected {n_edges}" # Build CSR indptr (cumulative sum of per-row edge counts). np.cumsum(indptr, out=indptr) - indices = np.concatenate(indices_chunks) - del indices_chunks - data = np.ones(n_edges, dtype=np.float32) csr = sparse.csr_matrix( (data, indices, indptr), diff --git a/tests/test_build_csr.py b/tests/test_build_csr.py new file mode 100644 index 0000000..480593d --- /dev/null +++ b/tests/test_build_csr.py @@ -0,0 +1,198 @@ +"""Tests for sync.build_csr — CSR construction from relationship shards. + +These pin the module's two load-bearing guarantees: + +* **Correctness** — the dense-remapped CSR matches an independently computed + reference, with nulls dropped, duplicate edges collapsed, and the sparse + OpenAlex IDs mapped to a contiguous ``[0, n_nodes)`` range sorted by ID. +* **Determinism** — the same shards always produce a byte-identical ``.npz`` + (a documented invariant: explicit ``ORDER BY`` plus a sorted shard list). + +build_csr depends on numpy, scipy, and duckdb, which are optional relative to +the core sync pipeline; the tests skip where those are unavailable. +""" + +from __future__ import annotations + +import hashlib + +import pytest + +# build_csr's deps (numpy/scipy/duckdb) are optional relative to the core sync +# pipeline; importorskip yields the modules and skips the file where any is +# absent. The first-party import then necessarily follows the guards (E402). +np = pytest.importorskip("numpy") +pa = pytest.importorskip("pyarrow") +pq = pytest.importorskip("pyarrow.parquet") +sparse = pytest.importorskip("scipy.sparse") +pytest.importorskip("duckdb") + +from sync import build_csr # noqa: E402 + +REL = "work_referenced_works" +SRC_COL, TGT_COL = "work_id", "referenced_work_id" + + +def _sha256(path): + return hashlib.sha256(path.read_bytes()).hexdigest() + + +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 _random_graph(seed, *, n_nodes=2_000, n_edges=8_000): + """A reproducible sparse graph over OpenAlex-scale (uint64) IDs.""" + rng = np.random.default_rng(seed) + ids = rng.choice(np.arange(1, 7_000_000_000), size=n_nodes, replace=False) + return [ + (int(ids[i]), int(ids[j])) + for i, j in zip( + rng.integers(0, n_nodes, n_edges), rng.integers(0, n_nodes, n_edges) + ) + ] + + +def _reference_csr(edges): + """Build the expected CSR the long way, independent of build_csr's path. + + Mirrors the module contract: node set is every non-null src plus every + non-null tgt; edges are the distinct rows with both endpoints present; + dense index is the rank of the ID in sorted order. + """ + nodes = {s for s, _ in edges if s is not None} + nodes |= {t for _, t in edges if t is not None} + sorted_ids = sorted(nodes) + index = {nid: i for i, nid in enumerate(sorted_ids)} + valid = sorted({(s, t) for s, t in edges if s is not None and t is not None}) + rows = [index[s] for s, _ in valid] + cols = [index[t] for _, t in valid] + n = len(sorted_ids) + ref = sparse.csr_matrix( + (np.ones(len(valid), dtype=np.float32), (rows, cols)), + shape=(n, n), + ) + ref.sort_indices() + return ref, np.array(sorted_ids, dtype=np.uint64) + + +def _build(tmp_path, edges=None, *, shards=None, force=True): + """Build CSR for ``edges`` (one shard) or ``shards`` (one shard per list).""" + 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_csr( + REL, parquet_dir=parquet_dir, output_dir=out_dir, force=force + ) + return result, out_dir / f"{REL}.npz" + + +def test_csr_matches_independent_reference(tmp_path): + # Sparse, OpenAlex-scale IDs; duplicate edge (10, 30); a null on each side. + edges = [ + (10, 30), (10, 20), (40, 10), (10, 30), # dup + (30, 30), (20, 40), (40, 20), + (50, None), (None, 60), # endpoints seed nodes 50 & 60 but no edge + (7_000_000_001, 10), (10, 7_000_000_001), + ] + _, npz = _build(tmp_path, edges) + got = sparse.load_npz(npz) + id_map = np.load(npz.with_suffix(".id_map.npy")) + + ref, ref_ids = _reference_csr(edges) + assert got.shape == ref.shape + assert got.nnz == ref.nnz + assert got.indices.dtype == np.int32 + np.testing.assert_array_equal(got.indptr, ref.indptr) + np.testing.assert_array_equal(got.indices, ref.indices) + np.testing.assert_array_equal(got.data, ref.data) + np.testing.assert_array_equal(id_map, ref_ids) + # {10,20,30,40,50,60,7000000001}: 50 and 60 are seeded by null-partnered + # rows, so they are nodes with no incident edge. + assert got.shape[0] == 7 + + +def test_node_set_includes_null_partnered_endpoints(tmp_path): + edges = [(10, 20), (30, None), (None, 40)] + _, npz = _build(tmp_path, edges) + got = sparse.load_npz(npz) + id_map = np.load(npz.with_suffix(".id_map.npy")) + # 10,20 from the real edge; 30 and 40 seeded by their null-partnered rows + np.testing.assert_array_equal(id_map, np.array([10, 20, 30, 40], dtype=np.uint64)) + assert got.shape == (4, 4) + assert got.nnz == 1 # only (10,20) is a real edge + + +def test_duplicate_edges_collapse(tmp_path): + edges = [(1, 2)] * 5 + [(2, 1)] + _, npz = _build(tmp_path, edges) + got = sparse.load_npz(npz) + assert got.nnz == 2 # (1,2) and (2,1), the four duplicates dropped + + +def test_empty_relationship_yields_empty_matrix(tmp_path): + result, npz = _build(tmp_path, []) + got = sparse.load_npz(npz) + assert got.shape == (0, 0) + assert got.nnz == 0 + assert result["n_edges"] == 0 + + +def test_duplicate_edges_collapse_across_shards(tmp_path): + # Production always reads many shards; an edge repeated in *different* shards + # must dedup globally, not just within one shard. + shards = [[(1, 2), (3, 4)], [(1, 2), (5, 6)]] + _, npz = _build(tmp_path, shards=shards) + got = sparse.load_npz(npz) + ref, ref_ids = _reference_csr([e for shard in shards for e in shard]) + assert got.nnz == 3 # (1, 2) once, plus (3, 4) and (5, 6) + np.testing.assert_array_equal(got.indptr, ref.indptr) + np.testing.assert_array_equal(got.indices, ref.indices) + np.testing.assert_array_equal(np.load(npz.with_suffix(".id_map.npy")), ref_ids) + + +def test_larger_graph_matches_reference(tmp_path): + # Validate the remap against the independent reference at a scale where + # dense-index ordering bugs would surface, not just on the ~7-node fixture. + edges = _random_graph(20260620) + _, npz = _build(tmp_path, edges) + got = sparse.load_npz(npz) + ref, ref_ids = _reference_csr(edges) + assert got.shape == ref.shape + assert got.nnz == ref.nnz + np.testing.assert_array_equal(got.indptr, ref.indptr) + np.testing.assert_array_equal(got.indices, ref.indices) + np.testing.assert_array_equal(got.data, ref.data) + np.testing.assert_array_equal(np.load(npz.with_suffix(".id_map.npy")), ref_ids) + + +def test_unchanged_inputs_skip_and_preserve_output(tmp_path): + # Idempotency: a second build with force=False detects unchanged inputs via + # the provenance fingerprint, skips, and leaves the .npz byte-identical. + edges = [(10, 20), (20, 30), (10, 30)] + result1, npz = _build(tmp_path, edges) + assert result1["status"] == "built" + before = _sha256(npz) + result2, _ = _build(tmp_path, edges, force=False) + assert result2["status"] == "skipped" + assert _sha256(npz) == before + + +def test_output_is_byte_identical_across_runs(tmp_path): + edges = _random_graph(20260620) + _, npz_a = _build(tmp_path / "a", edges) + _, npz_b = _build(tmp_path / "b", edges) + assert _sha256(npz_a) == _sha256(npz_b) + assert _sha256(npz_a.with_suffix(".id_map.npy")) == _sha256( + npz_b.with_suffix(".id_map.npy") + )