fix: give each query lexeme its own BM25 candidate budget

[nuemaan]

Summary

• The FTS prefilter from perf: bound classic BM25 candidates with Postgres FTS #252 bounded candidates with one global ts_rank_cd ordering. That ordering scores term density inside a chunk and ignores corpus-wide rarity, while BM25 weights rare terms heavily, so the chunk BM25 ranks first can be the first one truncated.
• Each query lexeme now draws from its own share of the budget through a lateral join, so a lexeme matching few chunks always contributes them. _MIN_CANDIDATES_PER_LEXEME keeps many-lexeme queries from dividing the budget into slivers.
• Adds a warning when the candidate pool saturates. The debug line reported candidates == limit whether the corpus held exactly that many or far more, so silent truncation was indistinguishable from a healthy query.
• Classic route only. Map-nav default path untouched. No API, worker, deployment, or migration impact.
• Resolves Classic BM25 FTS prefilter can truncate the top BM25 match in large namespaces #278, follows up on Optimize BM25 retrieval candidate loading with Postgres FTS prefilter #195 and perf: bound classic BM25 candidates with Postgres FTS #252.

Measured on 5001 chunks where 5000 densely repeat a common term and one holds a rare term, querying data zebra:

ts_rank_cd position of the rare chunk:   5001 of 5001
rank_rows_by_bm25 position:                 1 of 5001

candidate pool, global ordering, limit 2000:  2000 rows, rare chunk absent
candidate pool, per-lexeme budget:            1001 rows, rare chunk present

The per-lexeme pool is smaller and still keeps the match BM25 wants. Cost is one GIN probe per lexeme rather than one overall, and _MAX_FTS_QUERY_TOKENS already caps that at 50.

Verification

• uv run pytest apps/api/tests/unit/test_bm25_channel_tsquery.py apps/api/tests/contract/test_bm25_fts_prefilter_contract.py — 16 passed against real Postgres 16
• uv run pytest packages/shared-python/shared/tests/test_retrieval_search_channels.py — 4 passed
• uvx ruff check packages/shared-python apps/api/tests — clean
• Reverted the channel change and re-ran the new regression test to confirm it fails without the fix: AssertionError: assert 'common-2' == 'rare-zebra'

test_content_channel_uses_bounded_or_fts_after_scope_filters asserted on the literal LIMIT :fts_candidate_limit as a position marker for the candidate bound. That clause is now LIMIT lb.per_lexeme_limit, so the assertions were repointed. The intent from #244, that scope filters land ahead of any candidate bound so exclusions cannot consume the budget, is unchanged and still asserted.

Not tested: production-scale corpora, and lateral-probe cost at the 50 lexeme ceiling on a large namespace. Both need a dataset I do not have locally.

Deployment Notes

• No new environment variables. RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT keeps its meaning as the total budget, now spent per lexeme.
• No database migrations.
• Backwards compatible. The full-scan fallback is untouched, so rolling back is a code-only revert.

Checklist

• Tests were added or updated when behavior changed
• Public docs, examples, or OpenAPI contracts were updated when needed
• Database migrations are idempotent and safe to deploy
• Logs, errors, and validation paths avoid leaking secrets or user data
• The pull request description explains any breaking or user-visible change

[nuemaan]

Heads up that CI here is sitting at action_required rather than a pass or fail, so the required checks never ran. Fork PRs need a maintainer to approve the workflow run, and a fresh push does not clear it. That is what mergeable_state: blocked reflects, not a failing build.

Rebased onto current main just now. No conflicts, channels.py was untouched upstream. Re-ran locally against Postgres 16 after the rebase: 16 passed on the api unit and contract tests, 4 passed on the shared channel tests, ruff clean.

Worth separating the two things regardless of what happens to this PR: the truncation described in #278 is live on main today. If you would rather fix it differently, or fold it into other retrieval work, that is fine by me. The repro in the issue stands on its own.

[suguanYang]

Thanks for the clear repro — I re-ran the two prefilters locally on Postgres 16 (simple FTS + rank_rows_by_bm25) and the rare-term bug is real.

Their case still holds. 2000 dense data chunks + 1 zebra, query data zebra, limit 2000:

ts_rank_cd of rare-zebra:     2001 of 2001
BM25 on full corpus:          1st (11.71 vs 3.45)

global pool:      2000 rows, rare-zebra absent → BM25 top is common-*
per-lexeme pool:  1001 rows, rare-zebra present → BM25 top is rare-zebra

So this is a real fix for that pattern. I still would not merge it as a general recall improvement — the same setup shows a regression on covering / multi-term hits, which global ts_rank_cd keeps and this change drops.

Regression 1 (deterministic). 400 dense-only chunks per term + one chunk that contains all six query terms once. Query one two three four five six, limit 2000 → per_lexeme = 333:

ts_rank_cd of cover-6:        1 of 2401  (0.60 vs 0.50 for a dense single-term chunk)
BM25 on full corpus:          1st (10.30)

global pool:      cover-6 present  → BM25 top is cover-6 (15.08)
per-lexeme pool:  cover-6 absent   → BM25 never sees it; top is one-*

Each term’s slice is filled by denser single-term rows. Combined cover-density still ranks the covering chunk first, which is why current main keeps it.

Regression 2 (common 2-term shape, tie-break dependent). 2000 data filler N + 2000 analysis filler N + 1 data analysis, query data analysis:

ts_rank_cd of cover-both:     1 of 4001  (0.20 vs 0.10)
BM25 on full corpus:          1st (4.33)

global pool:      cover-both present → BM25 top is cover-both
per-lexeme pool:  cover-both absent  if it loses the per-term ORDER BY
                  (same ts_rank_cd 0.10 as the other 2000 hits, limit 1000)

If that chunk happens to be first in the posting list it survives; if it is later, it does not. Production chunks are not inserted first.

I also saw the unsaturated-budget issue: 1500 foo + 400 bar = 1900 matches, under the 2000 cap. Global returned all 1900; per-lexeme returned 1400 (1000+400). In that corpus the BM25 winner did not change, so I am not hanging the review on top-k there — but the old path would not have truncated at all.

Could you add a contract test for the covering chunk (regression 1 is the clean one), and keep a global coverage slice rather than spending the whole budget per lexeme? Rare-term starvation is the actual #278 failure; a hybrid (global ts_rank_cd slice + per-lexeme floor) would fix that without dropping the chunk the combined ranking already puts first. The saturation warning is useful either way.

[Frankie-Xu]

Thanks for the clear write-up and for keeping this on the classic BM25 path.

The rare-term direction is correct. Per-lexeme candidate budget is the right shape for the #278 starvation case, and I independently confirmed it on this branch.

I ran the new tests locally (20 passed on the unit + contract + shared channel files, pytest-postgresql / Postgres 17.10) and also drove the live content_channel / term_channel / path_channel helpers plus merge_channels_rrf on two synthetic corpora. On 2000 dense data chunks + 1 zebra, query data zebra, limit 2000, live content_channel returned rare-zebra at rank 0 with a 1001-row pool. SQL for the old global ts_rank_cd LIMIT 2000 drops that row. Single-lexeme queries are unchanged (GREATEST(limit/1, 50) = limit). Scope filters still sit in scoped_chunks before any LIMIT, and _prepare_fts_tokens still keeps user text out of tsquery syntax. Classic-only is the right blast radius.

I would not treat the new prefilter as a general recall improvement yet — the same LATERAL slice that saves a rare term can still drop a BM25 top match of a different shape.

Where it still drops a BM25 top match

Once each term has more than per_lexeme_limit denser single-term hits, a covering chunk (all query terms, lower per-term density) loses every LATERAL slice. That is the #278 bug inverted: BM25's winner never reaches rank_rows_by_bm25. This covering miss was already sketched earlier on the thread; I independently reproduced it through the live channel helpers, with tokens that survive tokenize_for_retrieval. Function words like one/two/three are stopwords, so content_channel returns [] before SQL and would not make a useful contract test.

Corpus: 400 dense chunks per term × {invoice, payment, contract, vendor, refund, shipment} + one cover-6 holding all six terms once. Limit 2000 → per_lexeme_limit = 333 (channels.py, the lexeme_budget / LATERAL LIMIT lb.per_lexeme_limit block).

path	cover-6
full-corpus BM25	rank 0 (10.30 vs 3.09)
global ts_rank_cd LIMIT 2000	present
this PR's per-lexeme pool	absent (1998 rows)
live content_channel top 20	absent; top is invoice-*
live term_channel	rank 0
default RRF (content=2.0, term=1.5, path=1.0, k=60)	not in fused top 20

term_channel finding the covering chunk does not save classic fusion. A content-channel rank-0 dense hit scores 2/61 ≈ 0.033; a term-only rank-0 hit scores 1.5/61 ≈ 0.025; with path also ranking the dense rows, invoice-1 fused at 0.073. That path is bottom_discovery → merge_channels_rrf in packages/shared-python/shared/services/retrieval/search/scoring.py.

Two related edges on the SQL (channels.py lexeme_budget / LATERAL):

1. ORDER BY ts_rank_cd(...) DESC has no tie-break. On 2000 data filler N + 2000 analysis filler N + 1 data analysis, cover has the same per-term rank as the other 2000 hits (0.1). Whether it survives LIMIT 1000 depends on heap order; inserted last, it was dropped. A secondary id key can save it via one of the two slices, not both, so it is not a complete fix, but the current order is not deterministic either.
2. _MIN_CANDIDATES_PER_LEXEME = 50 means RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT is no longer a cap. At the 50-token ceiling, GREATEST(2000/50, 50) = 50 per lexeme → up to 2500 unique rows. The field description in packages/shared-python/shared/core/config/retrieval.py still says "Maximum rows".

Saturation warning looks at the wrong quantity

saturated = not used_fallback and candidate_count >= candidate_limit

After this change, unique-row count vs the global limit does not mean "a lexeme slice hit its LIMIT":

• Rare-term success path: candidates=1001, limit=2000, saturated=False, while the data slice was truncated at 1000 of 2000 matches.
• Covering-chunk miss: candidates=1998, limit=2000, saturated=False, while every lexeme slice was full and cover-6 was cut.

The warning added here would stay silent on both corpora, including this PR's own rare-term fixture (1000 common + limit 200 → unique ≈ 101). A useful signal is "any LATERAL probe returned per_lexeme_limit rows", not len(rows) >= candidate_limit.

Cost

The extra cost is not "one GIN probe per lexeme". WHERE tsv @@ q ORDER BY ts_rank_cd LIMIT n cannot use GIN for the sort. EXPLAIN ANALYZE on a 20k-row stand-in of this SQL is a nested loop of match scan + ts_rank_cd + top-N, loops = lexeme count. Each common lexeme still materializes its posting list. At _MAX_FTS_QUERY_TOKENS = 50 on a large namespace that is worth measuring before treating it as cheap.

Tests

test_rare_term_chunk_survives_a_saturated_candidate_budget is a good #278 lock (1000 rows vs limit 200 is enough for that shape). It does not catch the covering miss, the warning false-negative, or the floor exceeding the cap. The SQL-shape unit test also does not assert :fts_min_per_lexeme.

Would you be open to:

1. A contract test in the same file as the rare-term one, with a covering chunk vs dense single-term rows, using tokens that survive tokenize_for_retrieval.
2. Keeping a global ts_rank_cd slice and a per-lexeme floor, then DISTINCT-unioning them. Rare-term starvation is the actual #278 bug; combined ranking is what currently keeps covering hits. Spending the whole budget per lexeme is what drops them.
3. Driving the saturation warning off per-lexeme LIMIT hits (and maybe logging per_lexeme_limit / lexeme count).

The rare-term direction is sound; I would just not fold this in as a blanket recall fix until (1) and ideally (2) are in. Happy to re-run the contract file if you push those. Fork CI is still action_required here; the 20 local tests are green.

[nuemaan]

Thanks, that is a much better shape and all three of your cases reproduce here. Switched to the hybrid.

The prefilter now unions a global ts_rank_cd slice with a small per-lexeme floor, _CANDIDATES_PER_LEXEME_FLOOR = 50. The global slice is current main unchanged, so the union is a superset of it and nothing the combined ordering keeps can be dropped. The floor only adds rows.

Re-ran your four corpora on Postgres 16:

                                     main      per-lexeme only   union
#278 rare term, 5000 + 1            missing   present           present
covering chunk, 6 terms, 2401 rows  present   MISSING           present
covering chunk, 2 terms, 4001 rows  present   tie-dependent     present
1900 matches under a 2000 cap       1900      1400              1900

Regression 1 was the sharp one. Your reading was right, each term's slice fills with denser single-term rows and cover-6 never makes the pool. Added the contract test for it, test_chunk_covering_every_term_survives_a_saturated_budget. It asserts the covering chunk is the top BM25 result, and it fails against the previous commit with assert 'alpha-1' == 'cover-all', so it genuinely guards the regression rather than passing either way.

Two notes on the diff.

The bounding clause is LIMIT :fts_candidate_limit again, so the pre-LIMIT scope-filter assertions from #244 hold as written. My earlier commit had repointed them at the per-lexeme clause; that edit is reverted and test_retrieval_search_channels.py is now untouched by this branch.

The saturation warning still fires on candidate_count >= candidate_limit. Since the floor can push the union slightly past the limit, that can over-report when a query has roughly 40 or more lexemes, given the 50-row floor and the 2000 default. Happy to tie it strictly to the global slice if you would rather it be exact, though it costs either an extra column in the CTE or a second query.

Local runs: 17 passed on the api unit and contract tests, 4 passed on the shared channel tests, ruff clean. CI on this fork still needs a maintainer to approve the workflow run, so the checks here will stay unrun until someone does.