From 433af1218a38501c17078c22fe80525c969fc5ce Mon Sep 17 00:00:00 2001 From: nuemaan <253263884+nuemaan@users.noreply.github.com> Date: Thu, 13 Aug 2026 22:35:57 +0530 Subject: [PATCH] fix: give each query lexeme its own BM25 candidate budget The FTS prefilter bounded candidates with a single global ts_rank_cd ordering. ts_rank_cd scores term density inside a chunk and ignores how rare a term is across the corpus, while BM25 weights rare terms heavily. A short chunk holding the one rare term in a query therefore sorts near the bottom of that ordering and is truncated first, even though BM25 ranks it top. Measured on 5001 chunks where 5000 densely repeat a common term and one holds a rare term: the rare chunk ranks 5001 of 5001 under ts_rank_cd and 1 of 5001 under rank_rows_by_bm25, so a 2000 candidate limit dropped the best match before BM25 ran. Each lexeme now draws from its own share of the budget through a lateral join, so a lexeme matching few chunks always contributes them. A floor keeps many-lexeme queries from dividing the budget into slivers. The same corpus now yields 1001 candidates including the rare chunk, fewer rows than the old path loaded while keeping the match that matters. Also logs a warning when the pool saturates. The debug line reported candidates == limit whether the corpus held exactly that many or far more, so silent truncation looked identical to a healthy query. The bounded-prefilter test asserted on the literal LIMIT clause as a position marker. Its intent, that scope filters land ahead of any candidate bound, is unchanged and now asserts against the per-lexeme clause. Closes #278 --- .../test_bm25_fts_prefilter_contract.py | 85 +++++++++++++++++++ .../services/retrieval/search/channels.py | 79 ++++++++++++----- .../tests/test_retrieval_search_channels.py | 14 ++- 3 files changed, 153 insertions(+), 25 deletions(-) diff --git a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py index 6b9583298..1b3ee59aa 100644 --- a/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py +++ b/apps/api/tests/contract/test_bm25_fts_prefilter_contract.py @@ -12,6 +12,7 @@ import pytest import pytest_asyncio +from shared.core.config import settings as channel_settings from shared.services.retrieval.search.channels import content_channel, path_channel from shared.testing.contract_runtime import PostgreSQLProcess from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine @@ -53,6 +54,7 @@ """ _NOISE_ROWS = 300 +_COMMON_TERM_ROWS = 1000 @pytest_asyncio.fixture @@ -197,3 +199,86 @@ async def test_exclusions_still_apply_under_the_prefilter( exclude_sections=[], ) assert rows == [] + + +@pytest_asyncio.fixture +async def rare_term_session( + postgresql_proc: PostgreSQLProcess, +) -> AsyncGenerator[AsyncSession, None]: + """A corpus larger than the candidate budget where one chunk holds a rare term. + + ts_rank_cd scores term density inside a chunk and ignores corpus-wide + rarity, so the rare-term chunk sorts last under a single global ordering + even though BM25 ranks it first. + """ + dsn = ( + f"postgresql+asyncpg://{postgresql_proc.user}@" + f"{postgresql_proc.host}:{postgresql_proc.port}/postgres" + ) + engine = create_async_engine(dsn, isolation_level="AUTOCOMMIT") + async with engine.begin() as conn: + await conn.execute(text("DROP SCHEMA IF EXISTS bm25_rare CASCADE")) + await conn.execute(text("CREATE SCHEMA bm25_rare")) + await conn.execute(text("SET search_path TO bm25_rare")) + for statement in filter(None, (s.strip() for s in _SCHEMA.split(";"))): + await conn.execute(text(statement)) + await conn.execute(text("INSERT INTO job_results VALUES (1, 'job1')")) + await conn.execute( + text( + "INSERT INTO documents VALUES " + "('d1', 'u1', 'ns1', 'active', 1, 'sample.pdf')" + ) + ) + await conn.execute(text("INSERT INTO document_sections VALUES ('s1', '/root')")) + await conn.execute( + text( + "INSERT INTO document_chunks " + "(chunk_id, document_id, section_id, chunk_type, content, " + " job_result_id, sort_order, content_search_text, path_search_text) " + "SELECT 'common-' || i, 'd1', 's1', 'text', 'body', 1, i, " + " 'data data data data data filler ' || i, 'p ' || i " + "FROM generate_series(1, :common) AS i" + ), + {"common": _COMMON_TERM_ROWS}, + ) + await conn.execute( + text( + "INSERT INTO document_chunks " + "(chunk_id, document_id, section_id, chunk_type, content, " + " job_result_id, sort_order, content_search_text, path_search_text) " + "VALUES ('rare-zebra', 'd1', 's1', 'text', 'body', 1, 0, " + " 'zebra', 'p rare')" + ) + ) + + session_factory = async_sessionmaker(engine, expire_on_commit=False) + async with session_factory() as session: + await session.execute(text("SET search_path TO bm25_rare")) + yield session + await engine.dispose() + + +@pytest.mark.asyncio +async def test_rare_term_chunk_survives_a_saturated_candidate_budget( + rare_term_session: AsyncSession, + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Budget well under the number of matching chunks, so the pool saturates. + monkeypatch.setattr( + channel_settings, "RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT", 200, raising=False + ) + + rows = await content_channel( + rare_term_session, + user_id="u1", + namespace="ns1", + query="data zebra", + top_k=5, + exclude_document_ids=[], + exclude_sections=[], + ) + + # BM25 weights the rare term far above the common one, so the chunk holding + # it belongs at the top. A single global ts_rank_cd ordering truncates it + # before BM25 ever sees it. + assert [str(row["chunk_id"]) for row in rows][0] == "rare-zebra" diff --git a/packages/shared-python/shared/services/retrieval/search/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py index 1ddea1beb..b21f77b68 100644 --- a/packages/shared-python/shared/services/retrieval/search/channels.py +++ b/packages/shared-python/shared/services/retrieval/search/channels.py @@ -28,6 +28,11 @@ # Guards against pathological queries producing an enormous tsquery. _MAX_FTS_QUERY_TOKENS = 50 +# Floor on each lexeme's share of the candidate budget. A query with many +# lexemes would otherwise divide the budget down to a handful of rows each, +# which throws away candidates the old global ordering would have kept. +_MIN_CANDIDATES_PER_LEXEME = 50 + _TSV_FIELD_BY_SEARCH_FIELD = { "content_search_text": "content_search_tsv", "path_search_text": "path_search_tsv", @@ -339,34 +344,52 @@ async def _bm25_channel( used_fallback = True if fts_tokens: # Postgres lexes the tokens with the same configuration that generated - # the tsvector columns, then ORs the resulting lexemes. Building the - # tsquery server-side keeps the prefilter aligned with the stored - # lexicon and leaves no room for tsquery syntax in user input to - # change the query shape. `fts_query.q` is NULL when no token yields a - # lexeme, which the caller treats as "no usable prefilter". + # the tsvector columns. Building the query server-side keeps the + # prefilter aligned with the stored lexicon and leaves no room for + # tsquery syntax in user input to change the query shape. + # + # Each lexeme draws candidates from its own slice of the budget rather + # than competing in one global ts_rank_cd ordering. ts_rank_cd scores + # term density within a chunk and ignores how rare a term is across the + # corpus, while BM25 weights rare terms heavily. A short chunk holding + # the one rare term in a query therefore sorts near the bottom of a + # global ordering and is truncated first, even though BM25 would rank + # it top. Giving every lexeme its own slice keeps those chunks in the + # pool. See #278. prefilter_sql = ( corpus_cte + f""", - fts_query AS ( - SELECT string_agg(quote_literal(lexeme), ' | ')::tsquery AS q - FROM ( - SELECT DISTINCT - unnest(tsvector_to_array(to_tsvector('{_FTS_CONFIG}', token))) AS lexeme - FROM unnest(CAST(:fts_tokens AS text[])) AS token - ) lexemes + fts_lexemes AS ( + SELECT DISTINCT + unnest(tsvector_to_array(to_tsvector('{_FTS_CONFIG}', token))) AS lexeme + FROM unnest(CAST(:fts_tokens AS text[])) AS token + ), + lexeme_budget AS ( + SELECT + fl.lexeme, + to_tsquery('{_FTS_CONFIG}', quote_literal(fl.lexeme)) AS q, + GREATEST( + :fts_candidate_limit / GREATEST(COUNT(*) OVER (), 1), + :fts_min_per_lexeme + ) AS per_lexeme_limit + FROM fts_lexemes fl ) - SELECT sc.* - FROM scoped_chunks sc, fts_query fq - WHERE COALESCE(sc.{search_field}, '') <> '' - AND fq.q IS NOT NULL - AND sc.{tsv_field} @@ fq.q - ORDER BY ts_rank_cd(sc.{tsv_field}, fq.q) DESC - LIMIT :fts_candidate_limit + SELECT DISTINCT ON (candidates.id) candidates.* + FROM lexeme_budget lb + CROSS JOIN LATERAL ( + SELECT sc.* + FROM scoped_chunks sc + WHERE COALESCE(sc.{search_field}, '') <> '' + AND sc.{tsv_field} @@ lb.q + ORDER BY ts_rank_cd(sc.{tsv_field}, lb.q) DESC + LIMIT lb.per_lexeme_limit + ) candidates """ ) prefilter_params = dict(params) prefilter_params["fts_tokens"] = fts_tokens prefilter_params["fts_candidate_limit"] = candidate_limit + prefilter_params["fts_min_per_lexeme"] = _MIN_CANDIDATES_PER_LEXEME result = await db.execute(text(prefilter_sql), prefilter_params) rows = [_row_to_dict(r) for r in result.all()] used_fallback = not rows @@ -385,15 +408,29 @@ async def _bm25_channel( ranked_rows = rank_rows_by_bm25(rows, query_tokens, search_field=search_field) ranked_rows = ranked_rows[:top_k] + duration_ms = (time.perf_counter() - started_at) * 1000 + saturated = not used_fallback and candidate_count >= candidate_limit logger.debug( - "bm25_channel field={} candidates={} limit={} ranked={} fallback={} duration_ms={:.1f}", + "bm25_channel field={} candidates={} limit={} ranked={} " + "fallback={} saturated={} duration_ms={:.1f}", search_field, candidate_count, candidate_limit, len(ranked_rows), used_fallback, - (time.perf_counter() - started_at) * 1000, + saturated, + duration_ms, ) + if saturated: + # The pool filled the budget, so chunks past it never reached BM25. + # Distinct from the healthy case, which the debug line alone cannot + # convey because both report candidates == limit. + logger.warning( + "bm25_channel candidate budget saturated field={} limit={}; " + "raise RETRIEVAL_POSTGRES_FTS_CANDIDATE_LIMIT if recall looks short", + search_field, + candidate_limit, + ) return ranked_rows diff --git a/packages/shared-python/shared/tests/test_retrieval_search_channels.py b/packages/shared-python/shared/tests/test_retrieval_search_channels.py index e2e488a26..ddc31cd73 100644 --- a/packages/shared-python/shared/tests/test_retrieval_search_channels.py +++ b/packages/shared-python/shared/tests/test_retrieval_search_channels.py @@ -97,13 +97,19 @@ async def test_content_channel_uses_bounded_or_fts_after_scope_filters( assert "sc.content_search_tsv @@" in sql assert "CAST(:fts_tokens AS text[])" in sql assert "ORDER BY ts_rank_cd" in sql - assert "LIMIT :fts_candidate_limit" in sql - assert sql.index("LOWER(dc.chunk_type)") < sql.index("LIMIT :fts_candidate_limit") + # The budget is derived from the configured limit and then spent per + # lexeme, so the bounding clause names the per-lexeme share rather than + # the setting directly. + assert ":fts_candidate_limit" in sql + assert "LIMIT lb.per_lexeme_limit" in sql + # Scope filters still land inside the CTE, ahead of any candidate bound, + # so excluded rows cannot consume the budget. + assert sql.index("LOWER(dc.chunk_type)") < sql.index("LIMIT lb.per_lexeme_limit") assert sql.index("LOWER(COALESCE(ds.section_path") < sql.index( - "LIMIT :fts_candidate_limit" + "LIMIT lb.per_lexeme_limit" ) assert sql.index("POSITION(:_exc_section_path_0") < sql.index( - "LIMIT :fts_candidate_limit" + "LIMIT lb.per_lexeme_limit" ) assert params["fts_tokens"] == ["alpha", "beta"] assert params["fts_candidate_limit"] == 7