qa/query: fan shard search out across a thread pool
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.
This commit is contained in:
parent
0c4fbb53f8
commit
2b9d1f0b72
1 changed files with 196 additions and 159 deletions
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue