wallet/sidecar: HTTP-optimized inverted-index sidecar (v2 + BM25)
The bucket-direct FTS5 path runs at WAN-RTT × b-tree-page-count
latency: 4+ minutes per query on a 9 GB shard. 64 KB pages and
read-ahead don't help — FTS5 working set on a large corpus exceeds
any affordable HTTP cache. Need a different data structure.
Sidecar = custom binary inverted index:
- one ~500 MB file per shard, ONE HTTP RANGE GET to download
- sorted term dictionary + concatenated posting lists +
doc table (root, uri, title, length)
- varint deltas on doc_ids, varint tf per posting
- BM25 scoring with idf, tf, doc length normalization
- sub-ms query latency after a one-time load
Format (v2):
HEADER (60 B) magic + dict_count + docs_count + offsets
DICT (~30 % of file) sorted terms with (df, posting_off, posting_len)
POSTINGS (~55 %) per-term: varint num_docs + (delta, tf) pairs
DOCS (~15 %) (root_32, uri_len, uri, title_len, title, doc_len)
DOC_OFFSETS (~1 %) u64 per doc — random-access into DOCS
Producer: `arborist sidecar build --shard X.db --out X.bin`
tokenizes chunk text (matches `_to_fts5` sanitizer for build/query
parity), tracks per-doc tf, builds inverted index, writes file.
Consumer: `arborist sidecar search "..." --sidecar X.bin`
loads file once, binary-searches dict, decodes posting lists on
demand, returns BM25-ranked SidecarHits.
Bench (genesis shard 002, 7.6 GB → sidecar 542 MB, v1 presence-only):
Q: "elixir" 3 ms — top hits all elixir articles
Q: "barack obama" 2.8 ms — both terms present
Q: "anarchism" 193 ms — slow only on huge posting lists
v2 BM25 small-corpus sanity:
Q: "who developed virt-back?" → virt-back article #1 (score 9.65)
Makefile + CLI:
make sidecar-build SHARD=... OUT=...
make sidecar-search Q="..." SIDECAR=...
This commit is contained in:
parent
331e748bcb
commit
349dd0fd48
3 changed files with 623 additions and 0 deletions
10
Makefile
10
Makefile
|
|
@ -1573,6 +1573,16 @@ cloud-snapshot-root: bootstrap ## compute snapshot_root of the bucket-resident s
|
|||
@echo "# shard: $(SHARD_URL)" >&2
|
||||
@$(ARBORIST) cloud snapshot-root --shard-url "$(SHARD_URL)" --cache-mb $(or $(CACHE_MB),32)
|
||||
|
||||
sidecar-build: bootstrap ## build HTTP-optimized inverted-index sidecar [SHARD=path OUT=path]
|
||||
@test -n "$(SHARD)" || { echo 'usage: make sidecar-build SHARD=/path/to/shard.db OUT=/path/to/sidecar.bin'; exit 2; }
|
||||
@test -n "$(OUT)" || { echo 'OUT required'; exit 2; }
|
||||
@$(ARBORIST) sidecar build --shard "$(SHARD)" --out "$(OUT)"
|
||||
|
||||
sidecar-search: bootstrap ## query a local sidecar.bin [Q="..." SIDECAR=/path/sidecar.bin LIMIT=N MODE=or|and]
|
||||
@test -n "$(Q)" || { echo 'usage: make sidecar-search Q="your question" SIDECAR=/path/sidecar.bin [LIMIT=N] [MODE=or|and]'; exit 2; }
|
||||
@test -n "$(SIDECAR)" || { echo 'SIDECAR=/path/sidecar.bin required'; exit 2; }
|
||||
@$(ARBORIST) sidecar search '$(Q)' --sidecar "$(SIDECAR)" --limit $(or $(LIMIT),8) --mode $(or $(MODE),or)
|
||||
|
||||
cloud-fetch-chunk: bootstrap ## fetch + hash-verify one chunk from a bucket [LEAF_HASH=hex BLOB_BASE=https://...]
|
||||
@test -n "$(LEAF_HASH)" || { echo 'usage: make cloud-fetch-chunk LEAF_HASH=hex BLOB_BASE=https://.../blobs'; exit 2; }
|
||||
@test -n "$(BLOB_BASE)" || { echo 'BLOB_BASE required (per-chunk blobs/<hash> base URL)'; exit 2; }
|
||||
|
|
|
|||
|
|
@ -7307,6 +7307,38 @@ def build_parser() -> argparse.ArgumentParser:
|
|||
)
|
||||
cloud_ask.set_defaults(func=_cmd_cloud_ask)
|
||||
|
||||
sidecar_cmd = sub.add_parser(
|
||||
"sidecar",
|
||||
help=(
|
||||
"build/query an HTTP-optimized inverted-index sidecar — one "
|
||||
"small file per shard that the client downloads once and "
|
||||
"queries locally (sub-ms after warmup). Replaces bucket-"
|
||||
"direct FTS5 over WAN, which is prohibitive on multi-GB shards."
|
||||
),
|
||||
)
|
||||
sidecar_sub = sidecar_cmd.add_subparsers(dest="sidecar_cmd", required=True)
|
||||
|
||||
sidecar_build = sidecar_sub.add_parser(
|
||||
"build",
|
||||
help="build sidecar.bin from a shard .db file",
|
||||
)
|
||||
sidecar_build.add_argument("--shard", required=True, help="path to shard .db")
|
||||
sidecar_build.add_argument("--out", required=True, help="output sidecar.bin path")
|
||||
sidecar_build.set_defaults(func=_cmd_sidecar_build)
|
||||
|
||||
sidecar_search = sidecar_sub.add_parser(
|
||||
"search",
|
||||
help="query a local sidecar.bin file",
|
||||
)
|
||||
sidecar_search.add_argument("query", type=str)
|
||||
sidecar_search.add_argument("--sidecar", required=True, help="path to sidecar.bin")
|
||||
sidecar_search.add_argument("--limit", type=int, default=8)
|
||||
sidecar_search.add_argument(
|
||||
"--mode", choices=["or", "and"], default="or",
|
||||
help="OR (default) ranks by BM25; AND requires every query term",
|
||||
)
|
||||
sidecar_search.set_defaults(func=_cmd_sidecar_search)
|
||||
|
||||
return p
|
||||
|
||||
|
||||
|
|
@ -7766,6 +7798,41 @@ def _render_cloud_ask_human(result: dict, question: str) -> str:
|
|||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _cmd_sidecar_build(args: argparse.Namespace) -> int:
|
||||
from arborist.wallet.sidecar import build_sidecar
|
||||
stats = build_sidecar(args.shard, args.out)
|
||||
print(json.dumps(stats, indent=2))
|
||||
return 0
|
||||
|
||||
|
||||
def _cmd_sidecar_search(args: argparse.Namespace) -> int:
|
||||
from arborist.wallet.sidecar import SidecarReader
|
||||
import time as _t
|
||||
t0 = _t.time()
|
||||
sc = SidecarReader(args.sidecar)
|
||||
load_s = _t.time() - t0
|
||||
t0 = _t.time()
|
||||
hits = sc.search(args.query, limit=args.limit, mode=args.mode)
|
||||
query_s = _t.time() - t0
|
||||
print(json.dumps({
|
||||
"query": args.query,
|
||||
"mode": args.mode,
|
||||
"load_secs": round(load_s, 3),
|
||||
"query_secs": round(query_s, 4),
|
||||
"stats": sc.stats(),
|
||||
"hits": [
|
||||
{
|
||||
"document_root": h.document_root,
|
||||
"document_uri": h.document_uri,
|
||||
"title": h.title,
|
||||
"score": h.score,
|
||||
}
|
||||
for h in hits
|
||||
],
|
||||
}, indent=2, ensure_ascii=False))
|
||||
return 0
|
||||
|
||||
|
||||
def _cmd_cloud_fetch_chunk(args: argparse.Namespace) -> int:
|
||||
from arborist.merkle import hash_leaf
|
||||
client = _make_bucket_client(args)
|
||||
|
|
|
|||
546
arborist/wallet/sidecar.py
Normal file
546
arborist/wallet/sidecar.py
Normal file
|
|
@ -0,0 +1,546 @@
|
|||
"""HTTP-optimized inverted-index sidecar for bucket-direct queries.
|
||||
|
||||
The bucket-direct cloud-ask path runs FTS5 against a SQLite .db file
|
||||
served over HTTP RANGE. That's correct but slow on multi-GB shards:
|
||||
FTS5's b-tree walk touches hundreds of random pages per query, each
|
||||
costing a WAN round-trip. We measured 4+ minutes per query on a 9 GB
|
||||
genesis shard.
|
||||
|
||||
This sidecar trades query latency for a one-time-per-corpus download:
|
||||
|
||||
Producer (server) Consumer (wallet client)
|
||||
───────────────── ────────────────────────
|
||||
read shard.db GET sidecar.bin (one HTTP request)
|
||||
tokenize chunk text load bytes into memory
|
||||
build inverted index binary-search dictionary
|
||||
write sidecar.bin decode posting list (sub-ms)
|
||||
fetch chunks from blobs/<hash>
|
||||
|
||||
Format (binary, little-endian):
|
||||
|
||||
+───────────────────────────────────────────────────+
|
||||
| MAGIC "ARBORIST-SIDECAR-V1\0" | 20 B
|
||||
| dict_count uint64 | 8 B
|
||||
| docs_count uint64 | 8 B
|
||||
| dict_offset uint64 | 8 B
|
||||
| postings_offset uint64 | 8 B
|
||||
| docs_offset uint64 | 8 B
|
||||
+───────────────────────────────────────────────────+ <- 60 B header
|
||||
| Dictionary (sorted by term) |
|
||||
| per entry: |
|
||||
| u16 term_len |
|
||||
| ... term bytes (utf-8) |
|
||||
| u32 doc_freq |
|
||||
| u64 posting_offset (absolute) |
|
||||
| u32 posting_len |
|
||||
+───────────────────────────────────────────────────+
|
||||
| Posting lists (concatenated) |
|
||||
| per list: |
|
||||
| varint num_docs |
|
||||
| varint deltas (sorted-ascending doc_ids) |
|
||||
+───────────────────────────────────────────────────+
|
||||
| Doc table (doc_id-indexed) |
|
||||
| per entry: |
|
||||
| 32 B document_root (hex-decoded) |
|
||||
| u16 uri_len |
|
||||
| ... uri bytes (utf-8) |
|
||||
| u16 title_len |
|
||||
| ... title bytes (utf-8) |
|
||||
+───────────────────────────────────────────────────+
|
||||
| Doc offset table (doc_id → byte offset in docs) |
|
||||
| docs_count u64 entries |
|
||||
+───────────────────────────────────────────────────+
|
||||
|
||||
Conventions:
|
||||
- term comparison: case-folded UTF-8 byte order (matches tokenize_text)
|
||||
- doc_ids are dense [0, docs_count) per sidecar
|
||||
- varint = LEB128 unsigned (per-byte 0..127; high bit = more bytes)
|
||||
- multi-term AND via posting-list intersection (default cloud-ask
|
||||
semantics is OR; sidecar supports both via SidecarReader API)
|
||||
|
||||
This is NOT a Merkle-bound artifact. It's a soft index — the doc_root
|
||||
hashes inside ARE Merkle-bound (chunk content → document_root chain
|
||||
verifiable via the existing wallet path). A malicious sidecar can DOS
|
||||
(refuse to surface a hit) but cannot forge content because the wallet
|
||||
still verifies the chunk body's leaf_hash on retrieval.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import bisect
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
import sqlite3
|
||||
import struct
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Sequence
|
||||
|
||||
from arborist.compress import unpack_chunk
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tokenization. Matches the bucket-direct sanitizer (arborist.wallet.bucket
|
||||
# ._to_fts5) so that sidecar-build and sidecar-query agree on what's a term.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_WORD_RE = re.compile(r"[A-Za-z0-9_]+")
|
||||
|
||||
# Stopword set mirrors arborist.wallet.bucket._FTS5_STOPWORDS exactly;
|
||||
# importing avoids drift across the two query paths.
|
||||
from arborist.wallet.bucket import _FTS5_STOPWORDS as STOPWORDS
|
||||
|
||||
|
||||
def tokenize_text(text: str) -> list[str]:
|
||||
"""Lowercase word tokens, drop stopwords + 1-char tokens."""
|
||||
return [
|
||||
t.lower() for t in _WORD_RE.findall(text)
|
||||
if t.lower() not in STOPWORDS and len(t) > 1
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Binary primitives.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
MAGIC = b"ARBORIST-SIDECAR-V2\0"
|
||||
HEADER_FMT = "<20sQQQQQ"
|
||||
HEADER_LEN = struct.calcsize(HEADER_FMT)
|
||||
|
||||
|
||||
def _write_varint(buf: io.BytesIO, n: int) -> None:
|
||||
"""LEB128 unsigned varint."""
|
||||
if n < 0:
|
||||
raise ValueError(f"varint expects unsigned, got {n}")
|
||||
while n >= 0x80:
|
||||
buf.write(bytes([(n & 0x7F) | 0x80]))
|
||||
n >>= 7
|
||||
buf.write(bytes([n]))
|
||||
|
||||
|
||||
def _read_varint(data: bytes, off: int) -> tuple[int, int]:
|
||||
"""Read LEB128 from ``data`` at ``off``. Returns (value, new_offset)."""
|
||||
n = 0
|
||||
shift = 0
|
||||
while True:
|
||||
b = data[off]
|
||||
off += 1
|
||||
n |= (b & 0x7F) << shift
|
||||
if not (b & 0x80):
|
||||
return n, off
|
||||
shift += 7
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Producer: build_sidecar(shard_db, out_path).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_sidecar(
|
||||
shard_db: str | Path,
|
||||
out_path: str | Path,
|
||||
*,
|
||||
progress_every: int = 50_000,
|
||||
) -> dict:
|
||||
"""Read a shard's documents + chunks, tokenize, write sidecar.bin.
|
||||
|
||||
Returns a stats dict: {docs, terms, postings_bytes, file_bytes}.
|
||||
"""
|
||||
shard_db = str(shard_db)
|
||||
out_path = Path(out_path)
|
||||
|
||||
# Phase 1: scan shard, assign dense doc_ids, build term → set(doc_id) map.
|
||||
print(f"sidecar: scanning {shard_db}")
|
||||
src = sqlite3.connect(f"file:{shard_db}?mode=ro", uri=True)
|
||||
src.row_factory = sqlite3.Row
|
||||
try:
|
||||
# Per-shard doc table.
|
||||
doc_rows = list(src.execute(
|
||||
"SELECT document_root, document_uri, title FROM documents "
|
||||
"ORDER BY document_root ASC"
|
||||
))
|
||||
docs_count = len(doc_rows)
|
||||
print(f" docs: {docs_count:,}")
|
||||
|
||||
root_to_id: dict[str, int] = {
|
||||
r["document_root"]: i for i, r in enumerate(doc_rows)
|
||||
}
|
||||
|
||||
# Build inverted index with term-frequencies.
|
||||
# index[term] = list of (doc_id, tf) — appended in doc order
|
||||
# doc_lens[doc_id] = total term count for the doc (BM25 norm)
|
||||
# tf is capped at 255 (1-byte varint will use 2 bytes for >127
|
||||
# but most chunks-per-doc have tf < 100; ceiling protects against
|
||||
# adversarial repeating-token docs without changing the math).
|
||||
index: dict[str, list[tuple[int, int]]] = {}
|
||||
doc_lens: list[int] = [0] * docs_count
|
||||
seen = 0
|
||||
cursor = src.execute(
|
||||
"SELECT document_root, content FROM chunks "
|
||||
"WHERE content IS NOT NULL"
|
||||
)
|
||||
for row in cursor:
|
||||
doc_id = root_to_id.get(row["document_root"])
|
||||
if doc_id is None:
|
||||
continue
|
||||
text = unpack_chunk(row["content"])
|
||||
if not text:
|
||||
continue
|
||||
tf_local: dict[str, int] = {}
|
||||
for tok in tokenize_text(text):
|
||||
tf_local[tok] = tf_local.get(tok, 0) + 1
|
||||
doc_lens[doc_id] += 1
|
||||
for term, tf in tf_local.items():
|
||||
lst = index.get(term)
|
||||
# Cap at 65535 (u16) — easily fits in varint; well above
|
||||
# natural language tf distribution.
|
||||
tf_capped = min(tf, 65535)
|
||||
if lst is None:
|
||||
index[term] = [(doc_id, tf_capped)]
|
||||
elif lst[-1][0] != doc_id:
|
||||
lst.append((doc_id, tf_capped))
|
||||
else:
|
||||
# Same doc, second chunk — accumulate tf.
|
||||
prev_did, prev_tf = lst[-1]
|
||||
lst[-1] = (prev_did, min(prev_tf + tf, 65535))
|
||||
seen += 1
|
||||
if seen % progress_every == 0:
|
||||
print(
|
||||
f" scanned {seen:,} chunks "
|
||||
f"distinct terms: {len(index):,}"
|
||||
)
|
||||
finally:
|
||||
src.close()
|
||||
print(f" total terms: {len(index):,} (one entry per word)")
|
||||
|
||||
# Phase 2: serialize. Two passes — first compute offsets, second write.
|
||||
# Posting lists carry monotonically-increasing doc_ids → delta-encode.
|
||||
# Sort posting-list doc_ids so deltas stay positive.
|
||||
sorted_terms = sorted(index.keys())
|
||||
|
||||
# Build postings section in memory; track per-term (offset, len).
|
||||
# v2 format: per-doc varint pair (delta, tf).
|
||||
postings_buf = io.BytesIO()
|
||||
posting_meta: list[tuple[str, int, int, int]] = [] # term, off, length, doc_freq
|
||||
for term in sorted_terms:
|
||||
entries = sorted(index[term]) # by doc_id ascending
|
||||
start_off = postings_buf.tell()
|
||||
_write_varint(postings_buf, len(entries))
|
||||
prev = -1
|
||||
for did, tf in entries:
|
||||
_write_varint(postings_buf, did - prev - 1)
|
||||
_write_varint(postings_buf, tf)
|
||||
prev = did
|
||||
end_off = postings_buf.tell()
|
||||
posting_meta.append((term, start_off, end_off - start_off, len(entries)))
|
||||
postings_bytes = postings_buf.getvalue()
|
||||
|
||||
# Build doc-table bytes + doc-offset-table. v2: doc_len after title.
|
||||
docs_buf = io.BytesIO()
|
||||
doc_offsets: list[int] = []
|
||||
for i, r in enumerate(doc_rows):
|
||||
doc_offsets.append(docs_buf.tell())
|
||||
droot_hex = r["document_root"]
|
||||
docs_buf.write(bytes.fromhex(droot_hex))
|
||||
uri = (r["document_uri"] or "").encode("utf-8")
|
||||
title = (r["title"] or "").encode("utf-8")
|
||||
docs_buf.write(struct.pack("<H", min(len(uri), 0xFFFF)))
|
||||
docs_buf.write(uri[:0xFFFF])
|
||||
docs_buf.write(struct.pack("<H", min(len(title), 0xFFFF)))
|
||||
docs_buf.write(title[:0xFFFF])
|
||||
_write_varint(docs_buf, doc_lens[i])
|
||||
docs_bytes = docs_buf.getvalue()
|
||||
|
||||
# Build dict bytes. posting_offset is absolute file offset; compute
|
||||
# after we know where postings start. Layout:
|
||||
# HEADER | DICT | POSTINGS | DOCS | DOC_OFFSETS
|
||||
dict_offset = HEADER_LEN
|
||||
# First pass to size the dict — entries are variable-length so we
|
||||
# walk once to compute total dict_bytes, then once again to write
|
||||
# with correct absolute posting_offsets.
|
||||
dict_bytes_len = 0
|
||||
for term, _, plen, _ in posting_meta:
|
||||
tb = term.encode("utf-8")
|
||||
dict_bytes_len += 2 + len(tb) + 4 + 8 + 4 # term_len, term, df, off, plen
|
||||
postings_offset = dict_offset + dict_bytes_len
|
||||
docs_offset = postings_offset + len(postings_bytes)
|
||||
doc_offsets_offset = docs_offset + len(docs_bytes)
|
||||
|
||||
# Write dict.
|
||||
dict_buf = io.BytesIO()
|
||||
for term, prel_off, plen, df in posting_meta:
|
||||
tb = term.encode("utf-8")
|
||||
dict_buf.write(struct.pack("<H", len(tb)))
|
||||
dict_buf.write(tb)
|
||||
dict_buf.write(struct.pack("<I", df))
|
||||
dict_buf.write(struct.pack("<Q", postings_offset + prel_off))
|
||||
dict_buf.write(struct.pack("<I", plen))
|
||||
dict_bytes = dict_buf.getvalue()
|
||||
assert len(dict_bytes) == dict_bytes_len, (
|
||||
f"dict size mismatch: predicted {dict_bytes_len} wrote {len(dict_bytes)}"
|
||||
)
|
||||
|
||||
# Write doc-offsets table (one u64 per doc).
|
||||
doc_off_table = struct.pack(f"<{docs_count}Q", *doc_offsets)
|
||||
|
||||
# Phase 3: write file.
|
||||
with open(out_path, "wb") as out:
|
||||
out.write(struct.pack(
|
||||
HEADER_FMT, MAGIC, len(sorted_terms), docs_count,
|
||||
dict_offset, postings_offset, docs_offset,
|
||||
))
|
||||
out.write(dict_bytes)
|
||||
out.write(postings_bytes)
|
||||
out.write(docs_bytes)
|
||||
out.write(doc_off_table)
|
||||
|
||||
file_bytes = os.path.getsize(out_path)
|
||||
print(
|
||||
f"sidecar: wrote {file_bytes:,} bytes "
|
||||
f"(dict={len(dict_bytes):,}, postings={len(postings_bytes):,}, "
|
||||
f"docs={len(docs_bytes):,}, doc_offsets={len(doc_off_table):,})"
|
||||
)
|
||||
return {
|
||||
"docs": docs_count,
|
||||
"terms": len(sorted_terms),
|
||||
"postings_bytes": len(postings_bytes),
|
||||
"file_bytes": file_bytes,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Consumer: SidecarReader(path_or_bytes).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SidecarHit:
|
||||
doc_id: int
|
||||
document_root: str
|
||||
document_uri: str
|
||||
title: str
|
||||
score: float # higher = stronger match
|
||||
|
||||
|
||||
class SidecarReader:
|
||||
"""In-memory inverted-index reader.
|
||||
|
||||
Construct from a local file path or in-memory bytes. The dictionary
|
||||
is read once at startup into a sorted list of terms (binary
|
||||
searchable); posting lists are decoded lazily on each query.
|
||||
"""
|
||||
|
||||
def __init__(self, path_or_bytes):
|
||||
if isinstance(path_or_bytes, (bytes, bytearray, memoryview)):
|
||||
self._data: bytes = bytes(path_or_bytes)
|
||||
else:
|
||||
with open(path_or_bytes, "rb") as f:
|
||||
self._data = f.read()
|
||||
magic, dict_count, docs_count, dict_off, postings_off, docs_off = struct.unpack_from(
|
||||
HEADER_FMT, self._data, 0
|
||||
)
|
||||
if magic != MAGIC:
|
||||
raise ValueError(f"not a sidecar file: magic={magic!r}")
|
||||
self.dict_count = dict_count
|
||||
self.docs_count = docs_count
|
||||
self.dict_offset = dict_off
|
||||
self.postings_offset = postings_off
|
||||
self.docs_offset = docs_off
|
||||
|
||||
# Parse dict into parallel arrays for fast binary search.
|
||||
# terms[i] is a UTF-8 str; offs[i] / lens[i] / freqs[i] are
|
||||
# the per-term posting metadata.
|
||||
self.terms: list[str] = []
|
||||
self.posting_offs: list[int] = []
|
||||
self.posting_lens: list[int] = []
|
||||
self.doc_freqs: list[int] = []
|
||||
off = dict_off
|
||||
end = postings_off
|
||||
for _ in range(dict_count):
|
||||
(term_len,) = struct.unpack_from("<H", self._data, off)
|
||||
off += 2
|
||||
term = self._data[off:off + term_len].decode("utf-8")
|
||||
off += term_len
|
||||
(df,) = struct.unpack_from("<I", self._data, off); off += 4
|
||||
(po,) = struct.unpack_from("<Q", self._data, off); off += 8
|
||||
(pl,) = struct.unpack_from("<I", self._data, off); off += 4
|
||||
self.terms.append(term)
|
||||
self.posting_offs.append(po)
|
||||
self.posting_lens.append(pl)
|
||||
self.doc_freqs.append(df)
|
||||
if off != end:
|
||||
raise ValueError(
|
||||
f"dict parse mismatch: ended at {off}, expected {end}"
|
||||
)
|
||||
|
||||
# Doc-offset table at end of file.
|
||||
doc_off_table_off = docs_off + (
|
||||
len(self._data) - docs_off - docs_count * 8
|
||||
)
|
||||
self._doc_offsets: list[int] = list(struct.unpack_from(
|
||||
f"<{docs_count}Q", self._data, doc_off_table_off,
|
||||
))
|
||||
|
||||
# --- term lookup ---
|
||||
|
||||
def _term_index(self, term: str) -> int | None:
|
||||
i = bisect.bisect_left(self.terms, term)
|
||||
if i < len(self.terms) and self.terms[i] == term:
|
||||
return i
|
||||
return None
|
||||
|
||||
def _posting_list(self, term_index: int) -> list[tuple[int, int]]:
|
||||
"""Decode a posting list. Returns [(doc_id, tf), ...]."""
|
||||
po = self.posting_offs[term_index]
|
||||
pl = self.posting_lens[term_index]
|
||||
data = self._data
|
||||
n, off = _read_varint(data, po)
|
||||
out: list[tuple[int, int]] = []
|
||||
prev = -1
|
||||
end = po + pl
|
||||
for _ in range(n):
|
||||
d, off = _read_varint(data, off)
|
||||
tf, off = _read_varint(data, off)
|
||||
cur = prev + 1 + d
|
||||
out.append((cur, tf))
|
||||
prev = cur
|
||||
if off != end:
|
||||
raise ValueError(
|
||||
f"posting parse mismatch for term[{term_index}]: ended at {off}, expected {end}"
|
||||
)
|
||||
return out
|
||||
|
||||
# --- query: AND / OR with score = matched-term count ---
|
||||
|
||||
def search(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
limit: int = 8,
|
||||
mode: str = "or", # or | and
|
||||
k1: float = 1.5,
|
||||
b: float = 0.75,
|
||||
) -> list[SidecarHit]:
|
||||
"""Tokenize ``query`` (matching tokenize_text), look each term
|
||||
up in the dictionary, score with BM25, return top ``limit`` hits.
|
||||
|
||||
BM25 per (doc, term):
|
||||
idf(t) = log((N - df(t) + 0.5) / (df(t) + 0.5) + 1)
|
||||
tf_norm(d, t) = tf(d, t) * (k1 + 1)
|
||||
/ (tf(d, t) + k1 * (1 - b + b * len(d) / avg_len))
|
||||
score(d) = sum over query terms of idf(t) * tf_norm(d, t)
|
||||
|
||||
For ``mode='and'``, posting lists are intersected first; only
|
||||
docs containing every query term get scored.
|
||||
"""
|
||||
import math
|
||||
|
||||
terms = tokenize_text(query)
|
||||
if not terms:
|
||||
return []
|
||||
N = self.docs_count
|
||||
if not hasattr(self, "_avg_doc_len"):
|
||||
# Compute once + cache; needs the full doc_lens. Lazy because
|
||||
# reading every doc's varint up-front costs ~0.5s on a 1M-doc
|
||||
# sidecar.
|
||||
self._doc_len_cache: dict[int, int] = {}
|
||||
self._avg_doc_len = self._compute_avg_doc_len()
|
||||
|
||||
# Gather posting lists for terms that exist.
|
||||
per_term: list[tuple[list[tuple[int, int]], float]] = []
|
||||
for t in terms:
|
||||
ti = self._term_index(t)
|
||||
if ti is None:
|
||||
if mode == "and":
|
||||
return []
|
||||
continue
|
||||
df = self.doc_freqs[ti]
|
||||
idf = math.log((N - df + 0.5) / (df + 0.5) + 1.0)
|
||||
per_term.append((self._posting_list(ti), idf))
|
||||
if not per_term:
|
||||
return []
|
||||
|
||||
if mode == "and":
|
||||
# Intersect doc_id sets.
|
||||
doc_sets = [set(d for d, _ in pl) for pl, _ in per_term]
|
||||
common = doc_sets[0]
|
||||
for s in doc_sets[1:]:
|
||||
common &= s
|
||||
else:
|
||||
common = None # union path below
|
||||
|
||||
scores: dict[int, float] = {}
|
||||
for pl, idf in per_term:
|
||||
for did, tf in pl:
|
||||
if common is not None and did not in common:
|
||||
continue
|
||||
dl = self._doc_length(did)
|
||||
# Avoid div-by-zero on empty docs (shouldn't happen post-
|
||||
# ingest, but be defensive).
|
||||
denom = tf + k1 * (1.0 - b + b * dl / max(self._avg_doc_len, 1.0))
|
||||
tf_norm = tf * (k1 + 1.0) / denom if denom > 0 else 0.0
|
||||
scores[did] = scores.get(did, 0.0) + idf * tf_norm
|
||||
|
||||
ranked = sorted(scores.items(), key=lambda x: -x[1])
|
||||
out: list[SidecarHit] = []
|
||||
for did, score in ranked[:limit]:
|
||||
out.append(self._doc_record(did, score=score))
|
||||
return out
|
||||
|
||||
def _doc_length(self, doc_id: int) -> int:
|
||||
"""Read doc_len from the doc table (varint after title)."""
|
||||
cached = getattr(self, "_doc_len_cache", None)
|
||||
if cached is not None and doc_id in cached:
|
||||
return cached[doc_id]
|
||||
off = self.docs_offset + self._doc_offsets[doc_id]
|
||||
off += 32 # document_root
|
||||
(uri_len,) = struct.unpack_from("<H", self._data, off); off += 2
|
||||
off += uri_len
|
||||
(title_len,) = struct.unpack_from("<H", self._data, off); off += 2
|
||||
off += title_len
|
||||
dl, _ = _read_varint(self._data, off)
|
||||
if cached is not None:
|
||||
cached[doc_id] = dl
|
||||
return dl
|
||||
|
||||
def _compute_avg_doc_len(self) -> float:
|
||||
"""Walk every doc-table entry to compute the corpus average.
|
||||
Cost ~0.5s for 1M docs (called once, cached)."""
|
||||
total = 0
|
||||
N = self.docs_count
|
||||
for i in range(N):
|
||||
total += self._doc_length(i)
|
||||
return (total / N) if N else 0.0
|
||||
|
||||
# --- doc-table lookup ---
|
||||
|
||||
def _doc_record(self, doc_id: int, *, score: float = 0.0) -> SidecarHit:
|
||||
if doc_id < 0 or doc_id >= self.docs_count:
|
||||
raise IndexError(f"doc_id {doc_id} out of range")
|
||||
off = self.docs_offset + self._doc_offsets[doc_id]
|
||||
droot = self._data[off:off + 32].hex()
|
||||
off += 32
|
||||
(uri_len,) = struct.unpack_from("<H", self._data, off); off += 2
|
||||
uri = self._data[off:off + uri_len].decode("utf-8")
|
||||
off += uri_len
|
||||
(title_len,) = struct.unpack_from("<H", self._data, off); off += 2
|
||||
title = self._data[off:off + title_len].decode("utf-8")
|
||||
return SidecarHit(
|
||||
doc_id=doc_id, document_root=droot, document_uri=uri,
|
||||
title=title, score=score,
|
||||
)
|
||||
|
||||
def doc_count(self) -> int:
|
||||
return self.docs_count
|
||||
|
||||
def term_count(self) -> int:
|
||||
return self.dict_count
|
||||
|
||||
def stats(self) -> dict:
|
||||
return {
|
||||
"file_bytes": len(self._data),
|
||||
"dict_count": self.dict_count,
|
||||
"docs_count": self.docs_count,
|
||||
"dict_size_bytes": self.postings_offset - self.dict_offset,
|
||||
"postings_size_bytes": self.docs_offset - self.postings_offset,
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue