diff --git a/arborist/evict.py b/arborist/evict.py index 9a0702e..32db79e 100644 --- a/arborist/evict.py +++ b/arborist/evict.py @@ -861,48 +861,79 @@ def _pull_pack_inner_routed( finally: shutil.rmtree(tables_dir, ignore_errors=True) - # Phase 2: fill chunk content. Each chunk's metadata row landed on - # ONE target (the one its document_root hashes to). Iterate targets; - # only the owning shard finds a row for any given leaf_hash. The - # idx_chunks_leaf index makes the per-target lookup ~O(log N). + # Phase 2: fill chunk content. Bench-max'd 2026-05-26 after the + # initial naive impl took ~2h wall on a 6M-chunk corpus (same speed + # as a fresh ingest — the whole point of a pack restore is to be + # MUCH faster than that). Three improvements: + # + # (1) Build leaf_hash → (target_idx, chunk_id) map ONCE up front + # from each target's chunks table — replaces M=4 SELECTs per + # chunk_body (24M lookups → 6M scan-once). + # (2) Batched executemany UPDATE per target — replaces per-row + # BEGIN/COMMIT (was 1-row-per-txn = 6M transactions). + # (3) Skip chunks_fts inserts entirely during the fill loop; + # defer to a single bulk rebuild after all chunks are in + # place. Each per-row chunks_fts insert costs an + # inverted-index update; bulk rebuild streams chunks → FTS + # in one pass per target, ~10× faster. + # + # Plus PRAGMA synchronous=OFF on each target — genesis is + # rebuildable end-to-end (a partial crash leaves us with empty + # shards we can re-pull), so durability of intermediate WAL pages + # is not required. The PRAGMA wal_checkpoint(TRUNCATE) at the end + # of each phase keeps WAL bounded. + for target in targets: + target.execute("PRAGMA synchronous = OFF") + + # Phase 2a — build the leaf_hash → owning-target map. + leaf_to_target: dict[str, tuple[int, int]] = {} + for idx, target in enumerate(targets): + cur = target.execute( + "SELECT leaf_hash, chunk_id FROM chunks WHERE leaf_hash IS NOT NULL" + ) + for leaf_hash, chunk_id in cur: + leaf_to_target[leaf_hash] = (idx, chunk_id) + + # Phase 2b — batched UPDATE per target, no chunks_fts touch. + BATCH = 5000 + target_batches: list[list[tuple]] = [[] for _ in range(M)] restored = 0 skipped_unknown = 0 - skipped_already_hot = 0 + skipped_already_hot = 0 # kept for backward-compat with caller log + UPDATE_SQL = "UPDATE chunks SET content = ?, tier = 'hot' WHERE chunk_id = ?" for leaf_hash, chunk_body in chunks_to_restore: - landed = False - for target in targets: - rows = target.execute( - "SELECT chunk_id, tier FROM chunks WHERE leaf_hash = ?", - (leaf_hash,), - ).fetchall() - if not rows: - continue - landed = True - text = chunk_body.decode("utf-8") - for row in rows: - chunk_id = row["chunk_id"] if isinstance(row, sqlite3.Row) else row[0] - tier = row["tier"] if isinstance(row, sqlite3.Row) else row[1] - if tier == "hot" and chunk_id is not None: - existing = target.execute( - "SELECT content IS NOT NULL FROM chunks WHERE chunk_id = ?", - (chunk_id,), - ).fetchone() - if existing and existing[0]: - skipped_already_hot += 1 - continue - with transaction(target): - target.execute( - "UPDATE chunks SET content = ?, tier = 'hot' WHERE chunk_id = ?", - (pack_chunk(text), chunk_id), - ) - target.execute( - "INSERT OR REPLACE INTO chunks_fts (rowid, content) " - "VALUES (?, ?)", - (chunk_id, text), - ) - restored += 1 - if not landed: + owner = leaf_to_target.get(leaf_hash) + if owner is None: skipped_unknown += 1 + continue + idx, chunk_id = owner + text = chunk_body.decode("utf-8") + target_batches[idx].append((pack_chunk(text), chunk_id)) + restored += 1 + if len(target_batches[idx]) >= BATCH: + with transaction(targets[idx]): + targets[idx].executemany(UPDATE_SQL, target_batches[idx]) + target_batches[idx].clear() + for idx, batch in enumerate(target_batches): + if batch: + with transaction(targets[idx]): + targets[idx].executemany(UPDATE_SQL, batch) + + # Phase 2c — bulk FTS5 rebuild per target. Same primitive as the + # reshard executor (arborist.migrate._rebuild_fts_on_target): + # stream chunks → unpack_chunk → batched INSERT INTO chunks_fts. + # Each target shard's FTS only sees rows from the chunks ITS pack + # owned, so this rebuilds incrementally without touching other + # shards' FTS state. + from arborist.migrate import _rebuild_fts_on_target + for target in targets: + _rebuild_fts_on_target(target) + + # Hygiene: leave WAL files bounded after each phase. The migrate + # module already does this implicitly via _rebuild_fts_on_target + # but we call it explicitly here for the post-Phase-2b state. + for target in targets: + target.execute("PRAGMA wal_checkpoint(TRUNCATE)") chunk_pack_hashes_referenced: tuple[str, ...] = () if manifest_bytes is not None: