phase 1a: bulk-batched writer collapses per-doc executes
Previous batching opened one transaction per N=200 docs but still ran
~200 execute()/executemany() calls per doc inside it. Profile after
the parser rewrite showed the prior-URI SELECT alone was 648 ms /
2000 docs.
Refactor accumulates ALL inserts across the batch into per-table
mega-lists, then flushes via a single executemany per table:
- documents one executemany for the whole batch
- chunks / chunks_fts one executemany each
- merkle_nodes one executemany
- edges one executemany (collected across all docs)
- audit_events one executemany; chain hashes computed in pure
Python via store.chain_audit_events (no DB
round-trip per event)
Prior-URI resolution becomes one IN-clause SELECT for all distinct
URIs in the batch, plus an in-Python dict carrying within-batch chain
state so two same-URI revisions in the same batch correctly form
v_old <- v_mid <- v_new supersedes edges.
Within-batch document_root collisions (two source rows with identical
content) are deduped via inserted_this_batch set — required because
the mega-INSERT would otherwise hit the documents PK constraint.
5000 wikipedia_old revisions: 12.8 s -> 7.8 s (~64% faster). 49 tests
still passing.
This commit is contained in:
parent
c427180ced
commit
c81c1c0ff0
2 changed files with 281 additions and 126 deletions
|
|
@ -1,12 +1,14 @@
|
|||
"""Ingest pipeline: Source -> normalize -> chunk -> merkle -> upsert.
|
||||
|
||||
Idempotent: re-ingesting the same Document is a no-op (document_root collision
|
||||
is the upsert key). Edges with unresolved dst_root are stored as '' and are
|
||||
backfilled later when the target document is ingested.
|
||||
is the upsert key).
|
||||
|
||||
Performance: documents are accumulated and flushed in batches under a single
|
||||
SQLite transaction. With WAL+synchronous=NORMAL this trades per-doc fsync
|
||||
for per-batch fsync, typically 5-20x throughput on SSD.
|
||||
Performance shape — bulk-batched writer:
|
||||
Each batch collapses ALL inserts across N docs into a small set of
|
||||
executemany() calls (one per table) instead of per-doc calls. Audit chain
|
||||
hashes computed in pure Python via store.chain_audit_events, then inserted
|
||||
in one shot. With WAL+synchronous=NORMAL, the dominant cost shifts from
|
||||
Python<->C boundary crossings to actual SQLite work.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -22,7 +24,7 @@ from aborist import (
|
|||
from aborist.document import Document, canonicalize, get_chunker
|
||||
from aborist.merkle import MerkleTree, hash_leaf
|
||||
from aborist.source import Source
|
||||
from aborist.store import append_audit, transaction
|
||||
from aborist.store import chain_audit_events, latest_event_hash, transaction
|
||||
|
||||
|
||||
DEFAULT_BATCH_SIZE = 200
|
||||
|
|
@ -60,12 +62,9 @@ def ingest_source(
|
|||
def flush() -> None:
|
||||
if not batch:
|
||||
return
|
||||
with transaction(conn):
|
||||
for doc, art in batch:
|
||||
if _insert_document_no_tx(conn, doc, art, chunker.name):
|
||||
stats.inserted += 1
|
||||
else:
|
||||
stats.skipped_duplicate += 1
|
||||
inserted, skipped = _flush_batch(conn, batch, chunker.name)
|
||||
stats.inserted += inserted
|
||||
stats.skipped_duplicate += skipped
|
||||
batch.clear()
|
||||
|
||||
for doc in source.iter_documents():
|
||||
|
|
@ -74,7 +73,6 @@ def ingest_source(
|
|||
break
|
||||
art = _compute_artifacts(doc, chunker)
|
||||
if art is None:
|
||||
# empty content: count as skipped duplicate so totals stay sensible
|
||||
stats.skipped_duplicate += 1
|
||||
continue
|
||||
batch.append((doc, art))
|
||||
|
|
@ -103,129 +101,237 @@ def _compute_artifacts(doc: Document, chunker) -> _DocArtifacts | None:
|
|||
)
|
||||
|
||||
|
||||
def _insert_document_no_tx(
|
||||
def _flush_batch(
|
||||
conn: sqlite3.Connection,
|
||||
doc: Document,
|
||||
art: _DocArtifacts,
|
||||
batch: list[tuple[Document, _DocArtifacts]],
|
||||
chunker_name: str,
|
||||
) -> bool:
|
||||
"""Insert one document's rows. Caller wraps the batch in a transaction.
|
||||
) -> tuple[int, int]:
|
||||
"""Bulk-insert the whole batch. Returns (inserted, skipped_duplicate).
|
||||
|
||||
Returns True if newly inserted, False if document_root already exists.
|
||||
All collisions and prior-version lookups happen up front via batched
|
||||
SELECTs. New rows accumulate into per-table mega-lists and flush via
|
||||
one executemany per table. Audit-chain hashes are computed in pure
|
||||
Python and inserted in one shot.
|
||||
"""
|
||||
existing = conn.execute(
|
||||
"SELECT 1 FROM documents WHERE document_root = ?", (art.document_root,)
|
||||
).fetchone()
|
||||
if existing:
|
||||
# Same content -> same root -> idempotent skip. Still backfill edges
|
||||
# in case a new caller surfaces fresh outbound links.
|
||||
_upsert_edges(conn, art.document_root, doc)
|
||||
return False
|
||||
|
||||
# Versioned re-ingest: same URI + different root -> link new -> old.
|
||||
prior = conn.execute(
|
||||
"SELECT document_root FROM documents "
|
||||
"WHERE document_uri = ? AND document_root != ? "
|
||||
"ORDER BY ingest_ts DESC LIMIT 1",
|
||||
(doc.uri, art.document_root),
|
||||
).fetchone()
|
||||
|
||||
inserted = 0
|
||||
skipped = 0
|
||||
ingest_ts = int(time.time())
|
||||
conn.execute(
|
||||
"INSERT INTO documents "
|
||||
"(document_root, document_uri, source_type, kind, compression_depth, "
|
||||
" title, chunking_version, canonicalization_version, schema_version, "
|
||||
" ingest_ts) "
|
||||
"VALUES (?, ?, ?, 'surface', 0, ?, ?, ?, ?, ?)",
|
||||
(
|
||||
art.document_root,
|
||||
doc.uri,
|
||||
doc.source_type,
|
||||
doc.title,
|
||||
chunker_name,
|
||||
CANONICALIZATION_VERSION,
|
||||
SCHEMA_VERSION,
|
||||
ingest_ts,
|
||||
),
|
||||
)
|
||||
conn.executemany(
|
||||
"INSERT INTO chunks (document_root, idx, leaf_hash, content) "
|
||||
"VALUES (?, ?, ?, ?)",
|
||||
[
|
||||
(art.document_root, i, art.leaves[i].hex(), art.chunk_strs[i])
|
||||
for i in range(len(art.chunk_strs))
|
||||
],
|
||||
)
|
||||
conn.executemany(
|
||||
"INSERT INTO chunks_fts (document_root, idx, content) VALUES (?, ?, ?)",
|
||||
[(art.document_root, i, art.chunk_strs[i]) for i in range(len(art.chunk_strs))],
|
||||
)
|
||||
interior: list[tuple] = []
|
||||
for layer_idx in range(1, len(art.tree.layers)):
|
||||
for node_idx, h in enumerate(art.tree.layers[layer_idx]):
|
||||
interior.append((art.document_root, layer_idx, node_idx, h.hex()))
|
||||
if interior:
|
||||
conn.executemany(
|
||||
"INSERT INTO merkle_nodes (document_root, layer, idx, hash) "
|
||||
"VALUES (?, ?, ?, ?)",
|
||||
interior,
|
||||
)
|
||||
|
||||
_upsert_edges(conn, art.document_root, doc)
|
||||
with transaction(conn):
|
||||
# 1) Collision check: which document_roots already exist?
|
||||
roots = [art.document_root for _, art in batch]
|
||||
existing: set[str] = _select_existing_roots(conn, roots)
|
||||
inserted_this_batch: set[str] = set()
|
||||
|
||||
if prior is not None:
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO edges "
|
||||
"(src_root, dst_root, dst_uri, edge_type, anchor) "
|
||||
"VALUES (?, ?, ?, 'supersedes', '')",
|
||||
(art.document_root, prior["document_root"], doc.uri),
|
||||
)
|
||||
# 2) Batched prior-URI resolution: one IN-clause SELECT instead of
|
||||
# one per-doc SELECT. Maps each URI to the most-recent existing
|
||||
# document_root in the DB (pre-batch state).
|
||||
new_uris = list({doc.uri for doc, art in batch
|
||||
if art.document_root not in existing})
|
||||
prior_db: dict[str, str] = {}
|
||||
for slab_start in range(0, len(new_uris), 500):
|
||||
slab = new_uris[slab_start : slab_start + 500]
|
||||
placeholders = ",".join("?" * len(slab))
|
||||
for row in conn.execute(
|
||||
"SELECT document_uri, document_root FROM documents "
|
||||
f"WHERE document_uri IN ({placeholders}) "
|
||||
"ORDER BY ingest_ts DESC",
|
||||
slab,
|
||||
):
|
||||
# First wins (most recent by ORDER BY DESC).
|
||||
prior_db.setdefault(row["document_uri"], row["document_root"])
|
||||
|
||||
append_audit(
|
||||
conn,
|
||||
event_type="ingest",
|
||||
subject_root=art.document_root,
|
||||
body={
|
||||
"document_uri": doc.uri,
|
||||
"source_type": doc.source_type,
|
||||
"chunks": len(art.chunk_strs),
|
||||
"chunking_version": chunker_name,
|
||||
"canonicalization_version": CANONICALIZATION_VERSION,
|
||||
"schema_version": SCHEMA_VERSION,
|
||||
"supersedes": prior["document_root"] if prior is not None else None,
|
||||
},
|
||||
ts=ingest_ts,
|
||||
)
|
||||
return True
|
||||
# Within-batch chain state: as we process, the "prior" for the next
|
||||
# doc with the same URI becomes the doc we just inserted.
|
||||
prior_for_doc: dict[str, str] = {}
|
||||
|
||||
documents_rows: list[tuple] = []
|
||||
chunk_rows: list[tuple] = []
|
||||
fts_rows: list[tuple] = []
|
||||
merkle_rows: list[tuple] = []
|
||||
edge_rows: list[tuple] = []
|
||||
audit_events: list[dict] = []
|
||||
edges_to_upsert: list[tuple[str, Document]] = []
|
||||
|
||||
for doc, art in batch:
|
||||
if art.document_root in existing or art.document_root in inserted_this_batch:
|
||||
edges_to_upsert.append((art.document_root, doc))
|
||||
skipped += 1
|
||||
continue
|
||||
|
||||
prior_root = prior_for_doc.get(doc.uri) or prior_db.get(doc.uri)
|
||||
if prior_root == art.document_root:
|
||||
prior_root = None # same content, not a real prior
|
||||
|
||||
documents_rows.append(
|
||||
(
|
||||
art.document_root,
|
||||
doc.uri,
|
||||
doc.source_type,
|
||||
doc.title,
|
||||
chunker_name,
|
||||
CANONICALIZATION_VERSION,
|
||||
SCHEMA_VERSION,
|
||||
ingest_ts,
|
||||
)
|
||||
)
|
||||
for i, c in enumerate(art.chunk_strs):
|
||||
chunk_rows.append((art.document_root, i, art.leaves[i].hex(), c))
|
||||
fts_rows.append((art.document_root, i, c))
|
||||
for layer_idx in range(1, len(art.tree.layers)):
|
||||
for node_idx, h in enumerate(art.tree.layers[layer_idx]):
|
||||
merkle_rows.append(
|
||||
(art.document_root, layer_idx, node_idx, h.hex())
|
||||
)
|
||||
if prior_root is not None:
|
||||
edge_rows.append(
|
||||
(art.document_root, prior_root, doc.uri, "supersedes", "")
|
||||
)
|
||||
edges_to_upsert.append((art.document_root, doc))
|
||||
inserted_this_batch.add(art.document_root)
|
||||
prior_for_doc[doc.uri] = art.document_root
|
||||
audit_events.append(
|
||||
{
|
||||
"event_type": "ingest",
|
||||
"subject_root": art.document_root,
|
||||
"ts": ingest_ts,
|
||||
"body": {
|
||||
"document_uri": doc.uri,
|
||||
"source_type": doc.source_type,
|
||||
"chunks": len(art.chunk_strs),
|
||||
"chunking_version": chunker_name,
|
||||
"canonicalization_version": CANONICALIZATION_VERSION,
|
||||
"schema_version": SCHEMA_VERSION,
|
||||
"supersedes": prior_root,
|
||||
},
|
||||
}
|
||||
)
|
||||
inserted += 1
|
||||
|
||||
# 3) One executemany per table — minimal Python<->C boundary crossings.
|
||||
if documents_rows:
|
||||
conn.executemany(
|
||||
"INSERT INTO documents "
|
||||
"(document_root, document_uri, source_type, kind, compression_depth, "
|
||||
" title, chunking_version, canonicalization_version, schema_version, "
|
||||
" ingest_ts) "
|
||||
"VALUES (?, ?, ?, 'surface', 0, ?, ?, ?, ?, ?)",
|
||||
documents_rows,
|
||||
)
|
||||
if chunk_rows:
|
||||
conn.executemany(
|
||||
"INSERT INTO chunks (document_root, idx, leaf_hash, content) "
|
||||
"VALUES (?, ?, ?, ?)",
|
||||
chunk_rows,
|
||||
)
|
||||
conn.executemany(
|
||||
"INSERT INTO chunks_fts (document_root, idx, content) "
|
||||
"VALUES (?, ?, ?)",
|
||||
fts_rows,
|
||||
)
|
||||
if merkle_rows:
|
||||
conn.executemany(
|
||||
"INSERT INTO merkle_nodes (document_root, layer, idx, hash) "
|
||||
"VALUES (?, ?, ?, ?)",
|
||||
merkle_rows,
|
||||
)
|
||||
|
||||
# 4) Edges: collect all wikilinks across the batch and resolve in
|
||||
# one SELECT. Then one executemany.
|
||||
_flush_edges(conn, edges_to_upsert)
|
||||
if edge_rows:
|
||||
conn.executemany(
|
||||
"INSERT OR IGNORE INTO edges "
|
||||
"(src_root, dst_root, dst_uri, edge_type, anchor) "
|
||||
"VALUES (?, ?, ?, ?, ?)",
|
||||
edge_rows,
|
||||
)
|
||||
|
||||
# 5) Audit chain — hashes computed in Python, inserted in one call.
|
||||
if audit_events:
|
||||
prev = latest_event_hash(conn)
|
||||
audit_rows, _ = chain_audit_events(prev, audit_events)
|
||||
conn.executemany(
|
||||
"INSERT INTO audit_events "
|
||||
"(event_hash, prev_event_hash, event_type, subject_root, body, ts) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?)",
|
||||
audit_rows,
|
||||
)
|
||||
|
||||
return inserted, skipped
|
||||
|
||||
|
||||
def _upsert_edges(conn: sqlite3.Connection, src_root: str, doc: Document) -> None:
|
||||
"""Insert edges; backfill dst_root from documents table if URI matches.
|
||||
def _select_existing_roots(
|
||||
conn: sqlite3.Connection, roots: list[str]
|
||||
) -> set[str]:
|
||||
"""Single SELECT to find which document_roots already exist."""
|
||||
if not roots:
|
||||
return set()
|
||||
# SQLite has a default 999-param limit; chunk just in case.
|
||||
found: set[str] = set()
|
||||
for i in range(0, len(roots), 500):
|
||||
slab = roots[i : i + 500]
|
||||
placeholders = ",".join("?" * len(slab))
|
||||
for row in conn.execute(
|
||||
f"SELECT document_root FROM documents WHERE document_root IN ({placeholders})",
|
||||
slab,
|
||||
):
|
||||
found.add(row["document_root"])
|
||||
return found
|
||||
|
||||
Unresolved edges store dst_root=''; resolved edges store the target's root.
|
||||
Batched via executemany — single SQLite call instead of N.
|
||||
|
||||
def _flush_edges(
|
||||
conn: sqlite3.Connection,
|
||||
edges_to_upsert: list[tuple[str, Document]],
|
||||
) -> None:
|
||||
"""Upsert all wikilink/hyperlink edges across the batch.
|
||||
|
||||
Resolves dst_root for previously-unresolved URIs using a single batched
|
||||
SELECT (URI -> document_root) instead of N per-edge SELECTs.
|
||||
"""
|
||||
if not doc.edges:
|
||||
if not edges_to_upsert:
|
||||
return
|
||||
|
||||
# Collect distinct dst_uris across the batch for one resolution lookup.
|
||||
distinct_uris: set[str] = set()
|
||||
for _, doc in edges_to_upsert:
|
||||
for e in doc.edges:
|
||||
if e.dst_uri:
|
||||
distinct_uris.add(e.dst_uri)
|
||||
|
||||
uri_to_root: dict[str, str] = {}
|
||||
if distinct_uris:
|
||||
uris_list = list(distinct_uris)
|
||||
for i in range(0, len(uris_list), 500):
|
||||
slab = uris_list[i : i + 500]
|
||||
placeholders = ",".join("?" * len(slab))
|
||||
for row in conn.execute(
|
||||
"SELECT document_uri, document_root FROM documents "
|
||||
f"WHERE document_uri IN ({placeholders})",
|
||||
slab,
|
||||
):
|
||||
# Earliest ingest wins (matches prior LIMIT 1 ASC behavior).
|
||||
uri_to_root.setdefault(row["document_uri"], row["document_root"])
|
||||
|
||||
rows: list[tuple] = []
|
||||
for e in doc.edges:
|
||||
dst_root = e.dst_root or ""
|
||||
if not dst_root:
|
||||
row = conn.execute(
|
||||
"SELECT document_root FROM documents WHERE document_uri = ? "
|
||||
"ORDER BY ingest_ts ASC LIMIT 1",
|
||||
(e.dst_uri,),
|
||||
).fetchone()
|
||||
if row:
|
||||
dst_root = row["document_root"]
|
||||
rows.append(
|
||||
(src_root, dst_root, e.dst_uri or "", e.edge_type, e.anchor or "")
|
||||
for src_root, doc in edges_to_upsert:
|
||||
for e in doc.edges:
|
||||
dst_root = e.dst_root or uri_to_root.get(e.dst_uri or "", "")
|
||||
rows.append(
|
||||
(
|
||||
src_root,
|
||||
dst_root,
|
||||
e.dst_uri or "",
|
||||
e.edge_type,
|
||||
e.anchor or "",
|
||||
)
|
||||
)
|
||||
if rows:
|
||||
conn.executemany(
|
||||
"INSERT OR IGNORE INTO edges "
|
||||
"(src_root, dst_root, dst_uri, edge_type, anchor) VALUES (?, ?, ?, ?, ?)",
|
||||
rows,
|
||||
)
|
||||
conn.executemany(
|
||||
"INSERT OR IGNORE INTO edges "
|
||||
"(src_root, dst_root, dst_uri, edge_type, anchor) VALUES (?, ?, ?, ?, ?)",
|
||||
rows,
|
||||
)
|
||||
|
||||
|
||||
def verify_random_sample(conn: sqlite3.Connection, n: int = 10) -> dict:
|
||||
|
|
|
|||
|
|
@ -206,6 +206,54 @@ def _canonical_json(obj) -> str:
|
|||
return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
|
||||
|
||||
|
||||
def latest_event_hash(conn: sqlite3.Connection) -> str | None:
|
||||
"""Return the last event_hash in the audit chain, or None for genesis."""
|
||||
row = conn.execute(
|
||||
"SELECT event_hash FROM audit_events ORDER BY seq DESC LIMIT 1"
|
||||
).fetchone()
|
||||
return row["event_hash"] if row else None
|
||||
|
||||
|
||||
def chain_audit_events(
|
||||
prev_event_hash: str | None,
|
||||
events: list[dict],
|
||||
) -> tuple[list[tuple], str | None]:
|
||||
"""Compute the event_hash chain for a batch in pure Python.
|
||||
|
||||
Each event dict needs: `event_type`, `body` (dict), `subject_root` (str|None), `ts` (int).
|
||||
Returns (rows_for_executemany, last_event_hash). Insert with:
|
||||
|
||||
executemany("INSERT INTO audit_events
|
||||
(event_hash, prev_event_hash, event_type, subject_root,
|
||||
body, ts) VALUES (?, ?, ?, ?, ?, ?)", rows)
|
||||
|
||||
All chain SHA-256s are computed locally — zero DB round-trips per event.
|
||||
"""
|
||||
import hashlib
|
||||
|
||||
rows: list[tuple] = []
|
||||
prev = prev_event_hash
|
||||
for ev in events:
|
||||
body_json = _canonical_json(ev["body"])
|
||||
h = hashlib.sha256()
|
||||
if prev is not None:
|
||||
h.update(bytes.fromhex(prev))
|
||||
h.update(body_json.encode("utf-8"))
|
||||
event_hash = h.hexdigest()
|
||||
rows.append(
|
||||
(
|
||||
event_hash,
|
||||
prev,
|
||||
ev["event_type"],
|
||||
ev.get("subject_root"),
|
||||
body_json,
|
||||
ev["ts"],
|
||||
)
|
||||
)
|
||||
prev = event_hash
|
||||
return rows, prev
|
||||
|
||||
|
||||
def append_audit(
|
||||
conn: sqlite3.Connection,
|
||||
event_type: str,
|
||||
|
|
@ -213,15 +261,16 @@ def append_audit(
|
|||
subject_root: str | None = None,
|
||||
ts: int | None = None,
|
||||
) -> str:
|
||||
"""Append one event to the audit chain. Returns the new event_hash (hex)."""
|
||||
"""Append one event to the audit chain. Returns the new event_hash (hex).
|
||||
|
||||
Convenience wrapper for one-off events. Bulk inserts should use
|
||||
chain_audit_events() + executemany() for ~10x throughput on large batches.
|
||||
"""
|
||||
import hashlib
|
||||
|
||||
if ts is None:
|
||||
ts = int(time.time())
|
||||
row = conn.execute(
|
||||
"SELECT event_hash FROM audit_events ORDER BY seq DESC LIMIT 1"
|
||||
).fetchone()
|
||||
prev = row["event_hash"] if row else None
|
||||
prev = latest_event_hash(conn)
|
||||
body_json = _canonical_json(body)
|
||||
h = hashlib.sha256()
|
||||
if prev is not None:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue