diff --git a/arborist/search/fts5.py b/arborist/search/fts5.py index 16423a9..4e0204a 100644 --- a/arborist/search/fts5.py +++ b/arborist/search/fts5.py @@ -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"], diff --git a/tests/test_search_fts5.py b/tests/test_search_fts5.py new file mode 100644 index 0000000..0fb14e6 --- /dev/null +++ b/tests/test_search_fts5.py @@ -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()