engineering/arborist
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

#000061: bound memory in edges fan-in dump (index + batched groupby)

Browse source 
Live v2 corpus run showed ~5 GB RSS per worker — RssAnon dominant, so
process heap, not mmap. Traced to two unfixed memory pits in the edges
fan-in dump path:

1. SQLite ORDER BY on edges (22M rows, no covering index for the v2
   sort order dst_uri+edge_type+anchor) allocates a multi-GB in-memory
   sort area before spilling. Adding:

       CREATE INDEX IF NOT EXISTS idx_edges_dst_uri_type_anchor
       ON edges(dst_uri, edge_type, anchor, dst_root, src_root)

   means the ORDER BY walks the index in order — no in-memory sort.
   First create takes ~30-60 s on a 22M-row shard; idempotent on
   subsequent dumps. Disk cost ~1 GB per shard (4 shards × 1 GB ≈
   2-3 % corpus footprint increase). Worth it.

2. Python groupby accumulator: src_roots = [row[4] for row in group]
   materializes the entire src_root list per destination. For
   en.wikipedia.org/wiki/* destinations with millions of inbound links,
   this list is itself ~GB-sized. Switch to bounded batches:

       _FAN_IN_BATCH = 10_000     # max src_roots per fan-in JSON row

   A destination with N inbound links splits into ceil(N / batch) rows.
   Restore path (INSERT OR IGNORE) handles multi-row destinations
   correctly because PK includes src_root — accidental duplicates
   collapse cleanly.

New regression test test_edges_fan_in_batches_huge_destinations builds
an edges table with FAN_IN_BATCH+137 rows pointing at one dst_uri,
verifies the dump produces the expected number of split rows and the
restore reconstructs all N edges with no loss or duplication.

25 cold-object + evict tests pass.
This commit is contained in:
russell@unturf.com 2026-05-26 06:18:33 -04:00
parent ea743565de
commit bc7efe4434
No known key found for this signature in database
2 changed files with 119 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

59
arborist/cold_pack_metadata.py
View file

@ -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

69
tests/test_cold_object.py
View file

@ -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