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

search/fts5: progressive-AND fallback + DF filter at OR-pool

Browse source 
The fan-out commit (2b9d1f0) exposed a 13.5s shard-002 fts5_body call
on the Gundremmingen query and hypothesised the synonym OR-pool was
blowing the FTS5 candidate set. Profiling falsified that hypothesis:
synonym_expand returned no synonyms, the OR pool was just the three
query tokens, and the bottleneck was a single high-DF QUERY token
("located": 286,160 matches on a 1.5M-chunk wiki shard) carried into
OR-mode after AND-mode found zero co-occurrences. BM25 ranked all
~290k matches just to pick the top-32.

Two layered fixes:

A. Progressive-AND fallback. When AND returns zero, drop the shortest
   token (input order breaks ties) and retry AND. Repeat until hits or
   one token left. Only after every chain returns zero do we fall to
   OR-mode. On the Gundremmingen case, dropping "located" leaves
   "Gundremmingen AND Bavaria" which intersects to 3 docs in 11ms
   instead of the 290k-match OR-mode wall.

B. Document-frequency filter at OR-fallback time. ``COUNT(MATCH "tok")``
   per OR-pool token; drop any whose corpus DF exceeds
   ``_OR_FALLBACK_MAX_TOKEN_DF`` (default 50,000). ~15ms warm per
   probe. Only fires on the rare path where every progressive-AND
   chain still returned zero. Backstops A for queries where the
   answer genuinely requires OR (synonym-anchored retrieval, queries
   for content that uses different vocabulary than the question) but
   one of the OR clauses is a high-DF stopword-adjacent verb.

Both are deletion-first per the five-step algorithm: A deletes the
"jump straight to OR" path, B deletes high-DF tokens that contribute
~zero IDF anyway. No magic constants for A; B has one tunable knob
(threshold).

Bench (cold-cache, n=3, serial workers=1, query "where is
Gundremmingen located? where is Bavaria?"):

  metric                    BEFORE         AFTER (A+B)    delta
  total search wall         57.20s ± 0.22  1.67s ± 0.13   -97% / 34x
  shard 002 fts5_body cold  28.07s         0.05s          ~560x
  shard 002 fts5_body hits  32             3              -29 (the
                                                          dropped
                                                          were
                                                          "located"-
                                                          only noise)

Top-K=8 chosen sources unchanged before/after — the dropped fts5_body
candidates were filtered by the title-relevance step downstream
anyway.

Tests (tests/test_search_fts5.py, 11 cases):
- Helper: 5 cases on _progressive_and_token_chains (single-token,
  empty, shortest-first, strict length sort, always-keeps-one).
- Search behaviour: 4 cases (progressive-AND drops high-DF token;
  full-AND succeeds without progression; OR fallback when no chain
  hits; empty result when corpus has neither token nor synonym).
- DF filter: 2 cases (drops high-DF token, keeps input when all
  candidates would otherwise be dropped).

Verification:
- make test → 1605 passed, 28 skipped (was 1597 pre-change)
- make chain-check-shards → 0 breaks across all 7 shards
- arborist query "where is Gundremmingen located? where is Bavaria?"
  returns the same top-8 sources before/after
This commit is contained in:
russell@unturf.com 2026-05-09 18:31:24 -04:00
parent 2487b1c05c
commit 416f956734
No known key found for this signature in database
2 changed files with 394 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

168
arborist/search/fts5.py
View file

@ -42,6 +42,54 @@ _FTS5_STOPWORDS = frozenset(
# specificity — "neurotechnology" matters; "soon" doesn't).
_OR_FALLBACK_MAX_TOKENS = 5
# Per-token document-frequency cap at OR-fallback time. A token whose
# corpus DF exceeds this threshold (e.g. "located" matching 286k chunks
# on a 1.5M-chunk wiki shard) contributes ~zero IDF to BM25 yet
# dominates the candidate set the engine has to score. Drop it from the
# OR pool before MATCH. Cost: one ``COUNT(MATCH "tok")`` per candidate
# (~15ms warm). Backstops the progressive-AND fallback in case AND
# returns zero on every chain — that path lands in OR-mode where this
# filter prevents a single high-DF token from blowing the wall.
_OR_FALLBACK_MAX_TOKEN_DF = 50_000
def _progressive_and_token_chains(tokens: list[str]) -> list[list[str]]:
"""Yield successively-narrower AND token sets for progressive fallback.
First chain is the full token list. Each subsequent chain drops one
additional token, picked SHORTEST-FIRST with input order as the
tie-break. Stops when only one token remains.
Why shortest-first: short tokens are usually high-DF verbs/connectors
that collapse the AND set to zero ("located" matches 286k chunks on
a 1.5M-chunk wiki shard but contributes ~zero topical signal). The
rare topical token ("Gundremmingen") is what keeps AND tight; we
drop the cheap-signal tokens first so AND stays narrow and BM25
isn't forced to rank a full-corpus union via OR-fallback.
Concrete example, query tokens ["Gundremmingen", "located", "Bavaria"]:
chain 0: Gundremmingen AND located AND Bavaria (zero hits)
chain 1: Gundremmingen AND Bavaria ("located" dropped)
chain 2: Gundremmingen ("Bavaria" dropped)
The first chain that returns at least one row wins; OR-fallback only
fires if every progressive-AND chain returns zero.
"""
if len(tokens) <= 1:
return [list(tokens)]
# Decorate-sort-undecorate: sort by (length-asc, original-position-asc),
# then drop one at a time from the front of the sort.
indexed = sorted(enumerate(tokens), key=lambda iv: (len(iv[1]), iv[0]))
drop_order = [i for i, _ in indexed]
chains: list[list[str]] = []
dropped: set[int] = set()
chains.append(list(tokens))
for drop_idx in drop_order[:-1]: # always keep at least one token
dropped.add(drop_idx)
chains.append([t for i, t in enumerate(tokens) if i not in dropped])
return chains
def _query_tokens(query: str) -> list[str]:
raw = _FTS5_TOKEN_RE.findall(query)
@ -52,6 +100,46 @@ def _quote(t: str) -> str:
return '"' + t.replace('"', '""') + '"'
def _filter_or_pool_by_df(
conn,
tokens: list[str],
*,
threshold: int | None = None,
) -> list[str]:
"""Keep only tokens whose corpus document-frequency is <= ``threshold``.
Probe via ``COUNT(*) FROM chunks_fts WHERE chunks_fts MATCH '"tok"'``
— same FTS5 path the real search would take, so the cost mirrors a
real lookup (~15ms warm, ~50ms cold per token on a 10GB shard).
``threshold=None`` resolves to the module global at call time so a
test can monkeypatch ``_OR_FALLBACK_MAX_TOKEN_DF`` and have the
helper see the new value (function-default args bind at def time,
which would freeze the constant at import).
If the filter would empty the pool, fall back to returning the input
unchanged: OR-fallback is the last-resort retrieval path, and "some
hits, slow" is still strictly better than "no hits". Quietly swallow
malformed-MATCH SQL errors (rare; tokenizer-strange chars) — those
would have been dropped by the search anyway.
"""
if threshold is None:
threshold = _OR_FALLBACK_MAX_TOKEN_DF
keep: list[str] = []
for t in tokens:
try:
row = conn.execute(
"SELECT COUNT(*) FROM chunks_fts WHERE chunks_fts MATCH ?",
(_quote(t),),
).fetchone()
except Exception:
continue
n = row[0] if row else 0
if 0 < n <= threshold:
keep.append(t)
return keep or list(tokens)
def _escape_fts5(
query: str,
*,
@ -196,13 +284,22 @@ class FTS5Backend(SearchBackend):
"""
if not query.strip():
return []
# Try strict AND first; fall back to OR if it returns nothing.
for mode in ("and", "or"):
fts_query = _escape_fts5(
query,
mode=mode,
extra_or_tokens=extra_or_tokens if mode == "or" else None,
)
# Progressive-AND fallback before OR. AND-mode is signal-dense
# but brittle: a single high-DF token in the query (e.g.
# "located") can collapse the intersection to zero even though
# the rest of the tokens uniquely identify the topic. Old
# behavior fell straight to OR-mode, which on a 1.5M-chunk
# shard pulls 286k matches when "located" is one of the OR
# clauses and forces BM25 to rank them all (~27s cold I/O).
# Progressive-AND drops shortest-first and retries before
# surrendering to OR.
rows: list = []
and_tokens = _query_tokens(query)
for chain in _progressive_and_token_chains(and_tokens):
if not chain:
continue
fts_query = " AND ".join(_quote(t) for t in chain)
try:
rows = self.conn.execute(
"""
@ -226,6 +323,63 @@ class FTS5Backend(SearchBackend):
rows = []
if rows:
break
# OR-mode fallback. Folds in synonym OR-pool, then drops any
# token whose corpus DF exceeds ``_OR_FALLBACK_MAX_TOKEN_DF`` —
# backstops progressive-AND when every chain still returned
# zero. Without this filter a single high-DF query token
# (e.g. "located": 286k matches on a 1.5M-chunk shard) would
# force BM25 to score ~all of those matches just to pick the
# top-K.
if not rows:
or_pool = _query_tokens(query)
if not or_pool:
# All-stopword query: keep raw single-char-survivor
# tokens (current _escape_fts5 behavior).
raw = _FTS5_TOKEN_RE.findall(query)
or_pool = [t for t in raw if len(t) > 1] or ['""']
else:
# Synonym merge — same dedup as the legacy _escape_fts5
# OR branch.
if extra_or_tokens:
seen_lower = {t.lower() for t in or_pool}
for t in extra_or_tokens:
tl = t.lower()
if tl and tl not in seen_lower and tl not in _FTS5_STOPWORDS:
or_pool.append(tl)
seen_lower.add(tl)
# B: DF filter. ~15ms warm per token; only fires on the
# rare path where every progressive-AND chain returned
# zero, so the cost is amortized over the entire query.
or_pool = _filter_or_pool_by_df(self.conn, or_pool)
# Top-N-longest cap (proxy for rarity / topical specificity).
if len(or_pool) > _OR_FALLBACK_MAX_TOKENS:
or_pool = sorted(or_pool, key=len, reverse=True)[
:_OR_FALLBACK_MAX_TOKENS
]
fts_query = " OR ".join(_quote(t) for t in or_pool)
try:
rows = self.conn.execute(
"""
SELECT
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 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 ?
""",
(fts_query, limit),
).fetchall()
except Exception:
rows = []
return [
Hit(
document_root=r["document_root"],

233
tests/test_search_fts5.py Normal file
View file

@ -0,0 +1,233 @@
"""Unit tests for progressive-AND fallback in arborist.search.fts5."""
from __future__ import annotations
from typing import Iterator
from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.search import FTS5Backend
from arborist.search.fts5 import _progressive_and_token_chains
from arborist.source import Source
from arborist.store import connect
class _FakeSource(Source):
source_type = "fake"
def __init__(self, docs: list[Document]):
self.docs = docs
def iter_documents(self) -> Iterator[Document]:
yield from self.docs
def _doc(uri: str, content: str, title: str | None = None) -> Document:
return Document(
uri=uri,
content=content,
source_type="fake",
title=title if title is not None else uri.rsplit("/", 1)[-1],
edges=[],
)
# ----------------------------------------------------------- helper unit tests
def test_progressive_chains_single_token_passthrough():
assert _progressive_and_token_chains(["alpha"]) == [["alpha"]]
def test_progressive_chains_empty_input():
assert _progressive_and_token_chains([]) == [[]]
def test_progressive_chains_drops_shortest_first():
# Lengths: Gundremmingen=13, located=7, Bavaria=7. Tie between
# "located" and "Bavaria" — input order ("located" first) breaks it.
chains = _progressive_and_token_chains(["Gundremmingen", "located", "Bavaria"])
assert chains == [
["Gundremmingen", "located", "Bavaria"],
["Gundremmingen", "Bavaria"],
["Gundremmingen"],
]
def test_progressive_chains_strict_length_ordering():
# Pure length sort, no ties.
chains = _progressive_and_token_chains(["aaa", "bb", "ccccc", "dddd"])
assert chains == [
["aaa", "bb", "ccccc", "dddd"],
["aaa", "ccccc", "dddd"], # "bb" (shortest) dropped first
["ccccc", "dddd"], # "aaa" dropped second
["ccccc"], # "dddd" dropped third
]
def test_progressive_chains_always_keeps_one_token():
# With N tokens, we yield N chains: full, full-1, full-2, ..., 1.
chains = _progressive_and_token_chains(["x", "yy", "zzz"])
assert len(chains) == 3
assert all(c for c in chains) # no empty chain
assert len(chains[-1]) == 1
# ----------------------------------------------------------- end-to-end search
def _ingest(tmp_path, docs: list[Document]):
db_path = tmp_path / "test.db"
conn = connect(db_path)
ingest_source(conn, _FakeSource(docs))
return conn
def test_search_progressive_and_drops_high_df_token(tmp_path):
"""The "located" failure mode: full AND zero-hits; drop "located"
and the topical AND ("Gundremmingen AND Bavaria") finds the answer.
The old straight-to-OR fallback would have included every doc with
"located", "Bavaria", or "Gundremmingen"; progressive-AND keeps the
intersection tight.
"""
docs = [
_doc("test://gundremmingen-bavaria",
"Gundremmingen is a town in Bavaria.",
title="Gundremmingen"),
_doc("test://bavaria",
"Bavaria is a state in Germany.",
title="Bavaria"),
# Many decoy docs that contain "located" but are off-topic.
# Old OR-mode would surface these alongside the real answer.
*[_doc(f"test://decoy-{i}",
f"The widget {i} is located on the third shelf.",
title=f"Widget {i}")
for i in range(20)],
]
conn = _ingest(tmp_path, docs)
try:
backend = FTS5Backend(conn)
hits = backend.search("where is Gundremmingen located in Bavaria?", limit=8)
assert hits, "expected progressive-AND to surface the topical doc"
# Top hit must be the topical doc, not a decoy.
assert hits[0].title == "Gundremmingen"
# Decoys ("located" but not "Gundremmingen") must not appear at all,
# because progressive-AND dropping "located" still requires
# "Gundremmingen AND Bavaria".
assert all("Widget" not in h.title for h in hits)
finally:
conn.close()
def test_search_full_and_succeeds_no_progressive_drop(tmp_path):
"""When the full AND already returns hits, progressive-AND
is a no-op — the topical chain wins on chain 0.
"""
docs = [
_doc("test://a", "alpha beta gamma delta", title="A"),
_doc("test://b", "alpha beta gamma", title="B"),
_doc("test://c", "alpha", title="C"),
]
conn = _ingest(tmp_path, docs)
try:
backend = FTS5Backend(conn)
hits = backend.search("alpha beta gamma delta")
# Only doc A contains all four; full AND should return exactly it.
assert len(hits) == 1
assert hits[0].title == "A"
finally:
conn.close()
def test_search_falls_back_to_or_when_no_chain_hits(tmp_path):
"""If even single-token AND returns zero, OR-mode fallback fires.
OR can match any one of the synonyms / tokens.
"""
docs = [
_doc("test://a", "alpha is one fact", title="Alpha"),
_doc("test://b", "beta is another fact", title="Beta"),
]
conn = _ingest(tmp_path, docs)
try:
backend = FTS5Backend(conn)
# No doc contains "zeta" — every AND chain (including single-token
# "zeta") returns zero. OR-mode with synonym "alpha" fires.
hits = backend.search("zeta", extra_or_tokens={"alpha"})
assert hits, "OR-mode synonym fallback should still find alpha"
assert hits[0].title == "Alpha"
finally:
conn.close()
def test_search_or_fallback_when_synonym_not_in_corpus(tmp_path):
"""OR-mode with no synonyms and no matching tokens returns empty —
not an exception.
"""
docs = [_doc("test://a", "completely unrelated content", title="A")]
conn = _ingest(tmp_path, docs)
try:
backend = FTS5Backend(conn)
hits = backend.search("Gundremmingen")
assert hits == []
finally:
conn.close()
# ----------------------------------------------------------- B: DF filter at OR
def test_or_pool_df_filter_drops_high_df_token(tmp_path, monkeypatch):
"""When OR-mode fires and one of the tokens is high-DF, the DF
filter drops it before MATCH so it doesn't dominate the candidate
set. Synonym path: no AND-able tokens exist; OR-fallback fires.
"""
from arborist.search import fts5 as _fts5_mod
monkeypatch.setattr(_fts5_mod, "_OR_FALLBACK_MAX_TOKEN_DF", 5)
# Build a corpus where "common" appears in many docs (above
# threshold of 5) and "rare" appears in few.
docs: list[Document] = []
for i in range(15):
docs.append(_doc(f"test://common-{i}",
f"this document contains common token number {i}",
title=f"common-doc-{i}"))
docs.append(_doc("test://rare",
"rare needle is buried somewhere",
title="needle-doc"))
conn = _ingest(tmp_path, docs)
try:
backend = FTS5Backend(conn)
# Query has no AND-success path because "zeta" doesn't exist —
# progressive-AND will hit zero, OR fires. extra_or_tokens
# carries both "common" (DF=15, >threshold) and "rare" (DF=1).
hits = backend.search(
"zeta",
extra_or_tokens={"common", "rare"},
)
assert hits, "OR-mode should have surfaced the rare needle"
# The needle doc must be top — "common" got DF-filtered out so
# only "rare" actually matched.
assert hits[0].title == "needle-doc"
# No common-doc-N should appear: filter dropped "common".
assert all(not h.title.startswith("common-doc-") for h in hits)
finally:
conn.close()
def test_or_pool_df_filter_keeps_input_when_all_above_threshold(tmp_path, monkeypatch):
"""If every candidate is above threshold the filter must NOT empty
the pool — we'd rather have a slow OR than zero hits. The keep-all
fallback fires.
"""
from arborist.search import fts5 as _fts5_mod
monkeypatch.setattr(_fts5_mod, "_OR_FALLBACK_MAX_TOKEN_DF", 0) # nothing passes
docs = [_doc(f"test://d{i}", "alpha bravo charlie", title=f"d{i}")
for i in range(3)]
conn = _ingest(tmp_path, docs)
try:
backend = FTS5Backend(conn)
# extra_or_tokens carries content tokens; under threshold=0 all
# candidates would be dropped, but the fallback keeps them so
# OR-mode still finds the docs.
hits = backend.search("zeta", extra_or_tokens={"alpha"})
assert hits, "all-above-threshold case must keep tokens, not empty pool"
finally:
conn.close()