From 2b9d1f0b723e1130dd5d4e231908ce807e959e20 Mon Sep 17 00:00:00 2001 From: "russell@unturf.com" Date: Sat, 9 May 2026 17:12:34 -0400 Subject: [PATCH] qa/query: fan shard search out across a thread pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-shard search work in _search_corpus is independent (separate SQLite connection, separate accumulators, no shared mutation until merge), so the old serial loop was leaving wall on the table — the slowest fts5_body call on a wide-pool query took 30s on its own shard while the others sat idle. Extract the four routes (fts5_body, title, phrase, core_keyword) into _search_one_shard and fan over shards via ThreadPoolExecutor. Shard arrival order now interleaves; the downstream raw.sort + dedup makes order irrelevant. Worker count caps at min(8, len(paths)) or via ARBORIST_SHARD_WORKERS for cgroup-bound runners. Measured against ~/.arborist/shards (4 wiki shards + 3 small): query: "where is Gundremmingen located? where is Bavaria?" before: 66.32s total (search 62.22s, serial) after: 15.54s total (search 14.60s, parallel) -76% wall, 4.3x search speedup query: "is Gundremmingen planned to close?" before: 60.78s total (search 56.63s, serial) after: ~12s total (extrapolated from same shape) Slowest single shard now caps the wall (shard 002 at ~13.5s on the Gundremmingen query). The second-order question — why does fts5_body take 13.5s for 32 hits — is now isolated and worth a follow-up round (likely synonym OR-pool blowing the FTS5 candidate set before BM25 truncates). Hoisted accept_stems out of the per-shard loop body too — it's shard-invariant, no reason to recompute per shard. Tests: 1,560 passing under make test-ci (no regressions). Full suite (1,684) also green. --- arborist/qa/query.py | 355 ++++++++++++++++++++++++------------------- 1 file changed, 196 insertions(+), 159 deletions(-) diff --git a/arborist/qa/query.py b/arborist/qa/query.py index f0a6548..513cad1 100644 --- a/arborist/qa/query.py +++ b/arborist/qa/query.py @@ -34,8 +34,10 @@ chunk's proof on demand. from __future__ import annotations import json +import os import re import time +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path @@ -962,6 +964,159 @@ def _body_density_passes( return distinct_present >= breadth_threshold and total_mentions >= min_mentions +def _search_one_shard( + p: Path, + *, + question: str, + over_fetch: int, + or_synonym_pool: list[str], + accept_tokens: set[str], + accept_stems: set[str], + progress: Progress, +) -> tuple[ + list[tuple], + dict[str, str], + set[str], + set[str], +]: + """Run all four search routes against a single shard. + + Returns (raw_rows, root_to_shard, core_match_roots, + phrase_match_roots) — shard-local buffers that the caller merges. + Each shard worker opens its own connection so this is safe to + invoke from a thread pool. + """ + shard_name = p.name + raw: list[tuple] = [] + root_to_shard: dict[str, str] = {} + core_match_roots: set[str] = set() + phrase_match_roots: set[str] = set() + sp_str = str(p.resolve()) + + conn = connect(p) + try: + backend = FTS5Backend(conn) + _route_hits = 0 + for h in backend.search( + question, limit=over_fetch, extra_or_tokens=or_synonym_pool + ): + raw.append( + ( + h.score, + h.document_root, + h.document_uri, + h.title, + h.chunk_idx, + sp_str, + ) + ) + root_to_shard[h.document_root] = sp_str + _route_hits += 1 + progress.emit( + f"search.shard.{shard_name}.fts5_body", hits=_route_hits + ) + # Parallel title search using synonym-expanded tokens. + # The score starts deliberately low so this signal can't drown + # FTS5 BM25 + body relevance. _rerank_by_title later adds + # `overlap*10` to every hit (FTS5 and title-search alike) that + # has title-token overlap, so a doc whose title genuinely IS + # the topic ends up rewarded twice (once here, once in rerank). + # Title search now uses documents_fts FTS5 index (~0.05s/shard + # regardless of token count) — the prior >5-token bypass + # existed because LIKE '%tok%' was O(corpus × |tokens|). + # FTS5 MATCH makes this an O(K) hash lookup, so synonym- + # expanded title search is affordable at any query length. + title_search_rows = _search_titles( + conn, list(accept_tokens), over_fetch + ) + _title_hits = 0 + for r in title_search_rows: + title_norm = (r["title"] or "").replace("_", " ") + title_stems = { + _stem_token_for_match(t) + for t in _title_query_tokens(title_norm) + } + overlap = len(accept_stems & title_stems) + if overlap == 0: + continue + title_score = overlap * 10.0 + raw.append( + ( + title_score, + r["document_root"], + r["document_uri"], + r["title"], + 0, + sp_str, + ) + ) + root_to_shard[r["document_root"]] = sp_str + _title_hits += 1 + progress.emit( + f"search.shard.{shard_name}.title", hits=_title_hits + ) + # Phrase-pattern search — verbatim multi-token sequences from + # the question. Catches allusions / idioms / fictional-world + # references whose diagnostic signal is the exact sequence + # including function words. Two passes for layered specificity: + # - n=6 (highest specificity): "oceania always been at war + # with" is essentially unique to Orwell. Score 100. + # - n=5 (high specificity): "oceania always been at war" + # still strongly Orwell-anchored. Score 90. + # 4-grams were tried (2026-05-01) and dropped: too noisy. + _phrase_hits_total = 0 + for n in (6, 5): + phrase_score = 100.0 if n == 6 else 90.0 + for r in _search_phrases( + conn, _question_phrases(question, n=n), over_fetch + ): + raw.append( + ( + phrase_score, + r["document_root"], + r["document_uri"], + r["title"], + r["idx"], + sp_str, + ) + ) + root_to_shard[r["document_root"]] = sp_str + phrase_match_roots.add(r["document_root"]) + _phrase_hits_total += 1 + progress.emit( + f"search.shard.{shard_name}.phrase", hits=_phrase_hits_total + ) + # Core-keyword search: docs whose TF-IDF core keywords match. + # The query token doesn't need to be in title or even in body — + # being a TF-IDF keyword of the doc's core is enough signal. + _kw_hits = 0 + for r in _docs_with_core_keyword_match( + conn, list(accept_tokens), over_fetch + ): + core_match_roots.add(r["document_root"]) + match_count = r["match_count"] or 1 + kw_score = 40.0 + 25.0 * match_count + raw.append( + ( + kw_score, + r["document_root"], + r["document_uri"], + r["title"], + 0, + sp_str, + ) + ) + root_to_shard[r["document_root"]] = sp_str + _kw_hits += 1 + progress.emit( + f"search.shard.{shard_name}.core_keyword", hits=_kw_hits + ) + finally: + conn.close() + + return raw, root_to_shard, core_match_roots, phrase_match_roots + + def _search_corpus( shards_dir: Path | None, single_db: Path | None, @@ -1015,166 +1170,48 @@ def _search_corpus( # in accept-path 3. Lets us reach back to the source shard cheaply. root_to_shard: dict[str, str] = {} - for p in paths: - shard_name = p.name - conn = connect(p) - try: - backend = FTS5Backend(conn) - _route_hits = 0 - for h in backend.search( - question, limit=over_fetch, extra_or_tokens=or_synonym_pool - ): - raw.append( - ( - h.score, - h.document_root, - h.document_uri, - h.title, - h.chunk_idx, - str(p.resolve()), - ) - ) - root_to_shard[h.document_root] = str(p.resolve()) - _route_hits += 1 - progress.emit( - f"search.shard.{shard_name}.fts5_body", hits=_route_hits + # Hoisted out of the per-shard loop — accept_stems is shard-invariant. + accept_stems = {_stem_token_for_match(t) for t in accept_tokens} + + # Per-shard work is independent: each shard opens its own SQLite + # connection (no shared mutable state until merge) and runs four + # routes whose results we accumulate in shard-local buffers. Fan + # the shards out across a thread pool so the slowest shard caps + # the search wall instead of the sum-of-shards capping it. The + # downstream `raw.sort` + dedup step makes per-shard arrival + # order irrelevant. Worker count is capped via env override + # (`ARBORIST_SHARD_WORKERS`) for cgroup-bound runners; default + # is `min(8, len(paths))` which leaves headroom on commodity + # boxes while delivering ~4× on a 4-shard cluster. + if not paths: + max_workers = 1 + else: + env_cap = os.environ.get("ARBORIST_SHARD_WORKERS", "").strip() + if env_cap.isdigit() and int(env_cap) > 0: + max_workers = min(int(env_cap), len(paths)) + else: + max_workers = min(8, len(paths)) + + with ThreadPoolExecutor(max_workers=max_workers) as ex: + futures = [ + ex.submit( + _search_one_shard, + p, + question=question, + over_fetch=over_fetch, + or_synonym_pool=or_synonym_pool, + accept_tokens=accept_tokens, + accept_stems=accept_stems, + progress=progress, ) - # Parallel title search using synonym-expanded tokens. - # The score starts deliberately low so this signal can't drown - # FTS5 BM25 + body relevance. _rerank_by_title later adds - # `overlap*10` to every hit (FTS5 and title-search alike) that - # has title-token overlap, so a doc whose title genuinely IS - # the topic ends up rewarded twice (once here, once in rerank). - # Single-token title matches against generic terms ("intel", - # "cpu") used to score 60+ standalone, dominating top-K with - # legacy 80486-era articles for queries like "fastest intel - # CPU?". Now that contribution is the same scale as FTS5 body - # BM25, which lets Pentium_4 (high body relevance, no title - # overlap) win on its actual topical fit. - # Word-boundary overlap check (was substring-match, - # which falsely passed `"out" in "south"`, - # `"date" in "candidate"`, `"come" in "outcome"`, etc.). - # 2026-05-02 case fox surfaced: "what date did back to - # the future come out?" — the substring check admitted - # "Aberdeen, South Dakota" (south contains "out"), - # "Backplane" (back), "Outline of biology" (out), and - # consumed the over_fetch budget so the legitimate - # `Back to the Future` film article never made the - # rerank cut. Tokenizing the title via - # `_title_query_tokens` + stem-aware comparison keeps - # the title-LIKE pass returning only docs whose title - # has an actual matching word. - accept_stems = { - _stem_token_for_match(t) for t in accept_tokens - } - # Title search now uses documents_fts FTS5 index (~0.05s/shard - # regardless of token count) — the prior >5-token bypass - # existed because LIKE '%tok%' was O(corpus × |tokens|). - # FTS5 MATCH makes this an O(K) hash lookup, so synonym- - # expanded title search is affordable at any query length. - title_search_rows = _search_titles( - conn, list(accept_tokens), over_fetch - ) - _title_hits = 0 - for r in title_search_rows: - title_norm = (r["title"] or "").replace("_", " ") - title_stems = { - _stem_token_for_match(t) - for t in _title_query_tokens(title_norm) - } - overlap = len(accept_stems & title_stems) - if overlap == 0: - continue - title_score = overlap * 10.0 - raw.append( - ( - title_score, - r["document_root"], - r["document_uri"], - r["title"], - 0, - str(p.resolve()), - ) - ) - root_to_shard[r["document_root"]] = str(p.resolve()) - _title_hits += 1 - progress.emit( - f"search.shard.{shard_name}.title", hits=_title_hits - ) - # Phrase-pattern search — verbatim multi-token sequences - # from the question. Catches allusions / idioms / fictional- - # world references whose diagnostic signal is the exact - # sequence including function words. Two passes for layered - # specificity: - # - n=6 (highest specificity): "oceania always been at - # war with" is essentially unique to Orwell. Score 100. - # - n=5 (high specificity): "oceania always been at war" - # still strongly Orwell-anchored. Score 90. - # 4-grams were tried (2026-05-01) and dropped: "always been - # at war" matches generic war-history articles too often, - # creating retrieval noise that the rerank pipeline can't - # cleanly separate from the actual allusion. Empirically - # 5+ grams trade recall for precision — most allusions - # ("may the force be with you", "to be or not to be", - # "winter is coming") survive at length 5 or 3-with-light- - # tokens, but the 4-gram floor is where diagnostic-ness - # collapses. - _phrase_hits_total = 0 - for n in (6, 5): - phrase_score = 100.0 if n == 6 else 90.0 - for r in _search_phrases( - conn, _question_phrases(question, n=n), over_fetch - ): - raw.append( - ( - phrase_score, - r["document_root"], - r["document_uri"], - r["title"], - r["idx"], - str(p.resolve()), - ) - ) - root_to_shard[r["document_root"]] = str(p.resolve()) - phrase_match_roots.add(r["document_root"]) - _phrase_hits_total += 1 - progress.emit( - f"search.shard.{shard_name}.phrase", hits=_phrase_hits_total - ) - # Core-keyword search: docs whose TF-IDF core keywords match. - # The query token doesn't need to be in title or even in body — - # being a TF-IDF keyword of the doc's core is enough signal. - _kw_hits = 0 - for r in _docs_with_core_keyword_match( - conn, list(accept_tokens), over_fetch - ): - core_match_roots.add(r["document_root"]) - # 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( - ( - kw_score, - r["document_root"], - r["document_uri"], - r["title"], - 0, - str(p.resolve()), - ) - ) - root_to_shard[r["document_root"]] = str(p.resolve()) - _kw_hits += 1 - progress.emit( - f"search.shard.{shard_name}.core_keyword", hits=_kw_hits - ) - finally: - conn.close() + for p in paths + ] + for fut in futures: + shard_raw, shard_r2s, shard_core, shard_phrase = fut.result() + raw.extend(shard_raw) + root_to_shard.update(shard_r2s) + core_match_roots.update(shard_core) + phrase_match_roots.update(shard_phrase) raw.sort(key=lambda r: -r[0]) seen: set[str] = set()