diff --git a/arborist/cold_pack_metadata.py b/arborist/cold_pack_metadata.py index 5334720..de31480 100644 --- a/arborist/cold_pack_metadata.py +++ b/arborist/cold_pack_metadata.py @@ -189,33 +189,74 @@ def _dump_generic_table( # because many src documents link to the same dst URI. _EDGES_FAN_IN_COLUMNS = ["dst_uri", "edge_type", "anchor", "dst_root", "src_roots"] +# Cap src_roots per fan-in row so a popular dst_uri (e.g. en.wikipedia +# .org/wiki/USA with millions of inbound links) doesn't materialize a +# multi-GB Python list in memory before being flushed to disk. A single +# destination with N inbound links is split into ceil(N / FAN_IN_BATCH) +# rows; restore handles multi-row destinations correctly because the live +# edges PK includes src_root, so INSERT OR IGNORE collapses any +# accidental duplicates from chunk boundaries. +_FAN_IN_BATCH = 10_000 + def _dump_edges_fan_in(conn: sqlite3.Connection, out_path: Path) -> int: """Group edges by destination, store src_roots as an array per row. - Streams rows from SQLite in (dst_uri, edge_type, anchor, dst_root) - sorted order, then groups in Python with itertools.groupby. Both - sides are streaming, so memory stays bounded regardless of edge count. + Two memory bounds: + + 1. SQLite-side sort: we create an index on (dst_uri, edge_type, + anchor) if missing so ORDER BY walks the index in order. Without + this index, SQLite allocates a multi-GB in-memory sort area on + a 22M-row table before spilling to disk — observed live as ~5 GB + per worker on the v2 run. Index creation is one-time per shard; + subsequent dumps reuse it. Disk cost ~1.1 GB per 22M-row shard. + + 2. Python-side groupby: src_roots are accumulated in batches of + _FAN_IN_BATCH and flushed to disk; a single popular destination + becomes multiple JSON rows instead of one giant in-memory list. Returns the number of fan-in rows written. """ - # ORDER BY the group key first so itertools.groupby works correctly. + # One-time index creation. CREATE INDEX IF NOT EXISTS is idempotent + # and atomic; takes ~30-60 s on a 22M-row table the first time, a + # no-op every time after. Costs ~1 GB disk per shard but lets every + # subsequent dump skip the in-memory sort entirely. + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_edges_dst_uri_type_anchor " + "ON edges(dst_uri, edge_type, anchor, dst_root, src_root)" + ) + conn.commit() + cursor = conn.execute( "SELECT dst_uri, edge_type, anchor, dst_root, src_root " "FROM edges " "ORDER BY dst_uri, edge_type, anchor, dst_root, src_root" ) + + def _write_row(f, key, batch: list[str]) -> None: + dst_uri, edge_type, anchor, dst_root = key + row_arr = [dst_uri, edge_type, anchor, dst_root, batch] + f.write(json.dumps(row_arr, ensure_ascii=False, separators=(",", ":")) + "\n") + n = 0 with open(out_path, "w", encoding="utf-8") as f: f.write(json.dumps({"_columns": _EDGES_FAN_IN_COLUMNS}, sort_keys=True) + "\n") for key, group in itertools.groupby( cursor, key=lambda r: (r[0], r[1], r[2], r[3]) ): - dst_uri, edge_type, anchor, dst_root = key - src_roots = [row[4] for row in group] - row_arr = [dst_uri, edge_type, anchor, dst_root, src_roots] - f.write(json.dumps(row_arr, ensure_ascii=False, separators=(",", ":")) + "\n") - n += 1 + # Stream src_roots into bounded batches so we never hold more + # than _FAN_IN_BATCH refs in Python heap at any time, no matter + # how popular the destination. + batch: list[str] = [] + for row in group: + batch.append(row[4]) + if len(batch) >= _FAN_IN_BATCH: + _write_row(f, key, batch) + n += 1 + batch = [] + if batch: + _write_row(f, key, batch) + n += 1 return n diff --git a/tests/test_cold_object.py b/tests/test_cold_object.py index 13521f7..97dfcc6 100644 --- a/tests/test_cold_object.py +++ b/tests/test_cold_object.py @@ -299,6 +299,75 @@ def test_push_pack_local_dir_writes_files(tmp_path): conn.close() +def test_edges_fan_in_batches_huge_destinations(tmp_path): + """Regression test for the memory-bounded edges dump path. + + A destination with > _FAN_IN_BATCH inbound links should be split + into multiple JSON rows on dump, and restored to the live schema + with no missing or duplicated edges. Without this batching the + groupby accumulator held the whole src_roots list in Python heap + (~5 GB per worker observed on the v2 corpus run). + """ + import sqlite3 + from arborist.cold_pack_metadata import ( + _dump_edges_fan_in, + _restore_edges_fan_out, + _FAN_IN_BATCH, + ) + from arborist.store import connect as store_connect + + # Build a minimal schema-bearing DB and stuff edges directly. + src_db = tmp_path / "edges.db" + conn = store_connect(src_db) + try: + # Insert N_LINKS edges all pointing at the same dst_uri so the + # groupby for this destination spans multiple batches. + N_LINKS = _FAN_IN_BATCH + 137 # one full batch + a partial + dst_uri = "https://example.com/popular" + edge_type = "wikilink" + anchor = "" + dst_root = "" + with conn: + conn.executemany( + "INSERT OR IGNORE INTO edges " + "(src_root, dst_root, dst_uri, edge_type, anchor) " + "VALUES (?, ?, ?, ?, ?)", + [ + (f"src{i:08d}" + "0" * 24, dst_root, dst_uri, edge_type, anchor) + for i in range(N_LINKS) + ], + ) + + # Dump + dump_path = tmp_path / "edges.jsonl" + n_rows_written = _dump_edges_fan_in(conn, dump_path) + # One destination split into ceil(N_LINKS / FAN_IN_BATCH) rows. + expected_rows = (N_LINKS + _FAN_IN_BATCH - 1) // _FAN_IN_BATCH + assert n_rows_written == expected_rows, ( + f"expected {expected_rows} fan-in rows for {N_LINKS} edges, " + f"got {n_rows_written}" + ) + + # Wipe edges and restore. INSERT OR IGNORE handles any duplicate + # PK rows from chunk boundaries cleanly. + with conn: + conn.execute("DELETE FROM edges") + n_restored = _restore_edges_fan_out(conn, dump_path) + conn.commit() + assert n_restored == N_LINKS, ( + f"expected {N_LINKS} edges restored, got {n_restored}" + ) + + # Verify the live table has exactly N_LINKS rows for that destination. + final_count = conn.execute( + "SELECT COUNT(*) FROM edges WHERE dst_uri = ?", + (dst_uri,), + ).fetchone()[0] + assert final_count == N_LINKS + finally: + conn.close() + + def test_push_pack_v2_hydrates_fresh_empty_db(tmp_path): """Pack format v2 self-sufficiency test: build a pack from a populated DB, then unpack it into a FRESH empty DB and verify every table is