diff --git a/arborist/qa/corpus.py b/arborist/qa/corpus.py index 4114436..422bf21 100644 --- a/arborist/qa/corpus.py +++ b/arborist/qa/corpus.py @@ -344,15 +344,23 @@ class SqliteShardCorpus: def fts_title(self, query: str, *, limit: int = 8) -> list[Hit]: """Title-only bm25() via the documents_fts virtual table. - Per-token fan-out (one MATCH per non-stopword token, dedupe - by doc_root keeping best bm25). Single-MATCH ``tok1 OR tok2`` - is biased toward the rarer token — for "what is the capitol - of paris", "capitol" (DF ~1k titles) outranks "paris" (DF ~5k - titles) so a unioned MATCH returns 12 Capitol-* titles and - zero Paris-* hits, squeezing the actual "Paris" article out - of the top-K entirely. Per-token MATCH gives each token its - own candidate space; the merge then takes the best bm25 each - doc earned in any token's run. + Two-pass union (AND + per-token), merged by best bm25: + + - **AND pass**: MATCH on all tokens joined with AND ("spider + AND man") — rewards multi-token titles like "Spider-Man" + that contain BOTH tokens. Without this, a single-token MATCH + on "spider" alone surfaces "Spider", "Wolf spider", "Spiders + (album)" etc. ahead of "Spider-Man" because single-token + titles have shorter doc length → higher bm25. + + - **Per-token pass**: one MATCH per non-stopword token — + guarantees each query token gets independent surface area. + Necessary because single-MATCH ``tok1 OR tok2`` is biased + toward the rarer token (e.g. "capitol" DF~1k outranks + "paris" DF~5k titles in OR mode and squeezes Paris out + entirely). + + The merge takes the best bm25 each doc earned in any pass. """ from arborist.wallet.bucket import ( _to_fts5, @@ -360,9 +368,6 @@ class SqliteShardCorpus: _FTS5_WORD_RE, ) import unicodedata - # Single-token MATCH list — preserve FTS5-safe tokenization - # via the same sanitizer fts_body uses (NFKD-fold, drop - # stopwords, ≥2 char filter). if not _to_fts5(query).strip(): return [] folded = "".join( @@ -384,15 +389,35 @@ class SqliteShardCorpus: "ORDER BY score LIMIT ?" ) best: dict[str, tuple] = {} + # AND pass first when there are ≥2 tokens. FTS5 returns nothing + # if any token has zero matches across documents_fts; that's + # fine — the per-token pass below picks up the slack. + if len(tokens) >= 2: + and_match = " AND ".join(tokens) + try: + for r in self._conn.execute(sql, (and_match, limit)): + droot = r["document_root"] + cur = best.get(droot) + if cur is None or r["score"] < cur[0]: + best[droot] = ( + r["score"], r["document_uri"] or "", + r["title"] or "", + ) + except Exception: + pass + # Per-token pass. for tok in tokens: - for r in self._conn.execute(sql, (tok, limit)): - droot = r["document_root"] - cur = best.get(droot) - if cur is None or r["score"] < cur[0]: - best[droot] = ( - r["score"], r["document_uri"] or "", - r["title"] or "", - ) + try: + for r in self._conn.execute(sql, (tok, limit)): + droot = r["document_root"] + cur = best.get(droot) + if cur is None or r["score"] < cur[0]: + best[droot] = ( + r["score"], r["document_uri"] or "", + r["title"] or "", + ) + except Exception: + pass ranked = sorted(best.items(), key=lambda kv: kv[1][0])[:limit] return [ Hit( @@ -653,18 +678,17 @@ class MultiShardSqliteCorpus: return merged[:limit] def fts_title(self, query: str, *, limit: int = 8) -> list[Hit]: - """Per-token title fan-out across all shards. + """Two-pass title fan-out: AND-mode plus per-token, merged. - Single-MATCH ``tok1 OR tok2`` over documents_fts is biased - toward whichever token has lower corpus DF — rarer tokens get - higher bm25 per match, so a single-MATCH return is dominated - by their hits. For "what is the capitol of paris" that meant - Capitol-* titles filled top-16 and the "Paris" article never - made it into the merge pool. Per-shard per-token MATCH gets - balanced per-token at the shard level, BUT MultiShard's - generic _fanout then re-sorts globally by bm25 — undoing the - balance. Fix: run per-token here too, giving each token its - own slot in the final return. + - **AND pass**: each shard runs a progressive-AND MATCH + ("tok1 AND tok2 AND …") so multi-token titles like + "Spider-Man" rank where their multi-token content earns — + single-token title floods don't crowd them out. + + - **Per-token pass**: each token gets its own shard fan-out + bucket; round-robin merge so every query token contributes + even when one token's bm25 dominates (the "Paris" case — + "capitol" alone would otherwise squeeze "Paris" out). """ from arborist.wallet.bucket import ( _to_fts5, _FTS5_STOPWORDS, _FTS5_WORD_RE, @@ -685,11 +709,6 @@ class MultiShardSqliteCorpus: if not tokens: return [] - # Per-token: fan out across shards in parallel; each token - # gets its OWN top-(limit*2) bucket. The final return is the - # union of all tokens' buckets — guarantees every token has - # surface area even when one token (e.g. "capitol", low-DF) - # outranks another ("paris", higher-DF) on raw bm25. def _shard_token(item, token): sh_path, sc = item try: @@ -698,6 +717,48 @@ class MultiShardSqliteCorpus: except Exception: return sh_path, [] + # AND pass: progressive-AND single MATCH per shard. The shard's + # own fts_title runs its own AND when len(tokens)>=2 + per-token, + # so passing the original query in gets the full union — but we + # explicitly want JUST the AND-mode hits in this bucket, so call + # the shard with the joined AND expression directly (treating it + # as a single MATCH token from the shard's perspective). + and_bucket: list[Hit] = [] + if len(tokens) >= 2: + and_match = " AND ".join(tokens) + def _shard_and(item): + sh_path, sc = item + try: + # Direct SQL query to bypass shard's own re-tokenizer. + sql = ( + "SELECT d.document_root, d.document_uri, d.title, " + " bm25(documents_fts) AS score " + "FROM documents_fts " + "JOIN documents d ON d.rowid = documents_fts.rowid " + "WHERE documents_fts MATCH ? " + "ORDER BY score LIMIT ?" + ) + rows = list(sc._conn.execute(sql, (and_match, limit * 2))) + return sh_path, rows + except Exception: + return sh_path, [] + with ThreadPoolExecutor( + max_workers=max(1, len(self._shards)) + ) as ex: + futures = [ex.submit(_shard_and, item) for item in self._shards] + for f in futures: + sh_path, rows = f.result() + for r in rows: + and_bucket.append(Hit( + document_root=r["document_root"], + document_uri=r["document_uri"] or "", + title=r["title"] or "", + score=r["score"], + shard_id=sh_path, + )) + and_bucket.sort(key=lambda h: h.score) + and_bucket = and_bucket[: limit * 2] + per_token_buckets: list[list[Hit]] = [] with ThreadPoolExecutor( max_workers=max(1, len(self._shards)) @@ -721,22 +782,35 @@ class MultiShardSqliteCorpus: )) bucket.sort(key=lambda h: h.score) per_token_buckets.append(bucket[: limit * 2]) - # Union: each token gets equal allocation. Dedupe by - # document_root keeping best bm25 across tokens. The final - # sort is per-bucket-aware: walk all buckets, take the - # top-N from each round-robin, until we hit `limit`. + # Merge: dedupe by document_root keeping best bm25 across + # all buckets (AND + per-token). AND-mode usually wins + # — its bm25 reflects matching MULTIPLE tokens which gives + # a stronger score for the multi-token title. best: dict[str, Hit] = {} + for h in and_bucket: + cur = best.get(h.document_root) + if cur is None or h.score < cur.score: + best[h.document_root] = h for bucket in per_token_buckets: for h in bucket: cur = best.get(h.document_root) if cur is None or h.score < cur.score: best[h.document_root] = h - # Round-robin across token buckets so each token contributes - # representatives even when one token's bm25 dominates the - # global sort. Within a bucket, walk by bm25-ascending order. - bucket_iters = [iter(b) for b in per_token_buckets] - seen_roots: set[str] = set() + # Output: AND hits come first (best multi-token matches), then + # round-robin across per-token buckets so no token's surface + # area starves. The AND pass takes priority because its hits + # are the strongest title-relevance signal we have. ranked: list[Hit] = [] + seen_roots: set[str] = set() + for h in and_bucket: + if len(ranked) >= limit: + break + if h.document_root in seen_roots: + continue + seen_roots.add(h.document_root) + ranked.append(best[h.document_root]) + # Round-robin per-token fill. + bucket_iters = [iter(b) for b in per_token_buckets] active = True while active and len(ranked) < limit: active = False @@ -748,9 +822,6 @@ class MultiShardSqliteCorpus: if h.document_root in seen_roots: continue seen_roots.add(h.document_root) - # Use the merged best-score (in case the same - # doc surfaced from multiple tokens with better - # bm25 elsewhere). ranked.append(best[h.document_root]) break return ranked