diff --git a/CONTEXT.md b/CONTEXT.md index 8609579f..029d2c64 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -228,7 +228,10 @@ generation and retries or falls back if publication changes it during capture. The compatibility requirement that a serving-index retrieval returns the same selected chunk IDs, ordering, rounded scores, citations, and asset references -as the legacy retrieval path for the same request. +as the legacy retrieval path for the same request. Every retrieval optimization +must preserve this quality contract; a latency improvement without validated +semantic parity is not shippable. Validation also compares source sections, +evidence content, router, and stop reason for each pinned request. ### Retrieval Revision Pin diff --git a/apps/api/alembic/versions/a0b1c2d3e4f5_add_token_leading_map_unit_covering_index.py b/apps/api/alembic/versions/a0b1c2d3e4f5_add_token_leading_map_unit_covering_index.py new file mode 100644 index 00000000..b051b8f0 --- /dev/null +++ b/apps/api/alembic/versions/a0b1c2d3e4f5_add_token_leading_map_unit_covering_index.py @@ -0,0 +1,44 @@ +"""Add the token-leading covering index for map-unit lookups.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op + + +revision: str = "a0b1c2d3e4f5" +down_revision: str | None = "9f0a1b2c3d4e" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + +_INDEX_NAME = "idx_document_map_unit_tokens_token_lookup" + + +def upgrade() -> None: + """Create the additive index without taking a table-wide write lock.""" + uses_external_transaction: bool = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + statement: str = ( + f"CREATE INDEX {{concurrently}}IF NOT EXISTS {_INDEX_NAME} " + "ON document_map_unit_tokens (channel, token_hash, map_unit_id) " + "INCLUDE (token, frequency)" + ) + if uses_external_transaction: + op.execute(statement.format(concurrently="")) + return + with op.get_context().autocommit_block(): + op.execute(statement.format(concurrently="CONCURRENTLY ")) + + +def downgrade() -> None: + """Remove only the index introduced by this migration.""" + uses_external_transaction: bool = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + if uses_external_transaction: + op.execute(f"DROP INDEX IF EXISTS {_INDEX_NAME}") + return + with op.get_context().autocommit_block(): + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_INDEX_NAME}") diff --git a/apps/api/alembic/versions/b1c2d3e4f5a6_add_channel_bm25_statistics.py b/apps/api/alembic/versions/b1c2d3e4f5a6_add_channel_bm25_statistics.py new file mode 100644 index 00000000..bb826334 --- /dev/null +++ b/apps/api/alembic/versions/b1c2d3e4f5a6_add_channel_bm25_statistics.py @@ -0,0 +1,42 @@ +"""Add persisted per-channel BM25 corpus statistics.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op + + +revision: str = "b1c2d3e4f5a6" +down_revision: str | None = "a0b1c2d3e4f5" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + + +def upgrade() -> None: + """Add nullable statistics so incomplete legacy rows keep the fallback.""" + column_names: tuple[str, ...] = ( + "path_document_count", + "path_total_length", + "content_document_count", + "content_total_length", + ) + for column_name in column_names: + op.execute( + f"ALTER TABLE document_map_unit_indexes " + f"ADD COLUMN IF NOT EXISTS {column_name} INTEGER" + ) + + +def downgrade() -> None: + """Remove the additive statistics columns.""" + column_names: tuple[str, ...] = ( + "content_total_length", + "content_document_count", + "path_total_length", + "path_document_count", + ) + for column_name in column_names: + op.execute( + f"ALTER TABLE document_map_unit_indexes DROP COLUMN IF EXISTS {column_name}" + ) diff --git a/apps/api/alembic/versions/c2d3e4f5a6b7_repair_token_leading_map_unit_covering_index.py b/apps/api/alembic/versions/c2d3e4f5a6b7_repair_token_leading_map_unit_covering_index.py new file mode 100644 index 00000000..241f1193 --- /dev/null +++ b/apps/api/alembic/versions/c2d3e4f5a6b7_repair_token_leading_map_unit_covering_index.py @@ -0,0 +1,62 @@ +"""Repair a missing or invalid token-leading map-unit covering index.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op +from sqlalchemy import text + + +revision: str = "c2d3e4f5a6b7" +down_revision: str | None = "b1c2d3e4f5a6" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + +_INDEX_NAME = "idx_document_map_unit_tokens_token_lookup" +_INDEX_COLUMNS = "(channel, token_hash, map_unit_id) INCLUDE (token, frequency)" + + +def _is_index_ready() -> bool: + """Return whether the current schema contains a usable covering index.""" + is_ready = op.get_bind().execute( + text( + "SELECT indexes.indisvalid AND indexes.indisready " + "FROM pg_index AS indexes " + "JOIN pg_class AS classes ON classes.oid = indexes.indexrelid " + "JOIN pg_namespace AS namespaces " + "ON namespaces.oid = classes.relnamespace " + "WHERE namespaces.nspname = current_schema() " + "AND classes.relname = :index_name" + ), + {"index_name": _INDEX_NAME}, + ).scalar_one_or_none() + return bool(is_ready) + + +def _repair_index(*, concurrently: bool) -> None: + """Replace a missing or unusable index using the allowed DDL mode.""" + if _is_index_ready(): + return + concurrent_clause: str = "CONCURRENTLY " if concurrently else "" + op.execute(f"DROP INDEX {concurrent_clause}IF EXISTS {_INDEX_NAME}") + op.execute( + f"CREATE INDEX {concurrent_clause}{_INDEX_NAME} " + f"ON document_map_unit_tokens {_INDEX_COLUMNS}" + ) + + +def upgrade() -> None: + """Ensure the additive covering index exists and is usable.""" + uses_external_transaction: bool = bool( + op.get_context().opts.get("knowhere_external_transaction", False) + ) + if uses_external_transaction: + _repair_index(concurrently=False) + return + with op.get_context().autocommit_block(): + _repair_index(concurrently=True) + + +def downgrade() -> None: + """Keep the index owned by the preceding additive migration.""" diff --git a/apps/api/scripts/backfill_map_unit_indexes.py b/apps/api/scripts/backfill_map_unit_indexes.py index f37afb66..48d9bef1 100644 --- a/apps/api/scripts/backfill_map_unit_indexes.py +++ b/apps/api/scripts/backfill_map_unit_indexes.py @@ -155,7 +155,11 @@ class NamespaceFallbackReport: @property def ready(self) -> bool: - return not self.would_hit_snapshot_fallback and not self.scoring_incomplete + return ( + not self.would_hit_snapshot_fallback + and not self.scoring_incomplete + and self.missing_revision_manifest == 0 + ) def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallbackReport]: @@ -224,10 +228,14 @@ def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallback select( DocumentMapUnitIndex.document_id, DocumentMapUnitIndex.job_result_id, + DocumentMapUnitIndex.format_version, DocumentMapUnitIndex.unit_count, DocumentMapUnitIndex.average_idf_path, DocumentMapUnitIndex.average_idf_content, - DocumentMapUnitIndex.format_version, + DocumentMapUnitIndex.path_document_count, + DocumentMapUnitIndex.path_total_length, + DocumentMapUnitIndex.content_document_count, + DocumentMapUnitIndex.content_total_length, ).where( DocumentMapUnitIndex.document_id.in_( [document_id_value for document_id_value, _ in revisions] @@ -237,18 +245,26 @@ def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallback ) index_by_revision = { (str(document_id_value), str(job_result_id)): ( + int(format_version or 0), int(unit_count or 0), float(average_idf_path or 0.0), float(average_idf_content or 0.0), - int(format_version or 0), + path_document_count, + path_total_length, + content_document_count, + content_total_length, ) for ( document_id_value, job_result_id, + format_version, unit_count, average_idf_path, average_idf_content, - format_version, + path_document_count, + path_total_length, + content_document_count, + content_total_length, ) in index_rows } missing_map_index = 0 @@ -258,8 +274,23 @@ def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallback if stats is None: missing_map_index += 1 continue - unit_count, average_idf_path, average_idf_content, format_version = stats - if format_version != MAP_UNIT_INDEX_FORMAT_VERSION: + ( + format_version, + unit_count, + average_idf_path, + average_idf_content, + path_document_count, + path_total_length, + content_document_count, + content_total_length, + ) = stats + if ( + format_version != MAP_UNIT_INDEX_FORMAT_VERSION + or path_document_count is None + or path_total_length is None + or content_document_count is None + or content_total_length is None + ): missing_map_index += 1 continue if ( @@ -294,7 +325,11 @@ def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallback would_hit_snapshot_fallback = ( snapshot_status != "ok" or missing_from_snapshot > 0 ) - scoring_incomplete = missing_map_index > 0 or suspicious_zero_idf > 0 + # A zero average IDF is mathematically valid, notably for a + # two-unit corpus where every token appears in exactly one unit. + # Keep the count as diagnostic output, but readiness is determined + # by the format marker and required persisted statistics above. + scoring_incomplete = missing_map_index > 0 reports.append( NamespaceFallbackReport( user_id=user_id, diff --git a/apps/api/scripts/backfill_map_unit_statistics.py b/apps/api/scripts/backfill_map_unit_statistics.py new file mode 100644 index 00000000..a4ec4d28 --- /dev/null +++ b/apps/api/scripts/backfill_map_unit_statistics.py @@ -0,0 +1,221 @@ +# ruff: noqa: E402 + +"""Backfill persisted per-channel BM25 statistics without rebuilding tokens. + +This maintenance command aggregates existing ``document_map_units`` rows and +updates the four nullable statistics columns on the current active revision. +It never rewrites map-unit tokens, serving manifests, or namespace snapshots. +Each revision is committed independently so interruption is safe. +""" + +from __future__ import annotations + +import argparse +import os +import sys +from dataclasses import dataclass +from pathlib import Path + + +def _bootstrap_python_path() -> None: + api_root = Path(__file__).resolve().parents[1] + candidate_roots = ( + api_root / "packages" / "shared-python", + api_root.parents[1] / "packages" / "shared-python", + ) + shared_root = next( + (path for path in candidate_roots if path.is_dir()), None + ) + if shared_root is None: + raise RuntimeError("Could not locate shared-python package") + for path in (api_root, shared_root): + value = os.fspath(path) + if value not in sys.path: + sys.path.insert(0, value) + + +_bootstrap_python_path() + +from sqlalchemy import func, select, update +from sqlalchemy.orm import Session + +from shared.core.database_sync import get_sync_session_factory +from shared.models.database.document import ( + Document, + DocumentMapUnit, + DocumentMapUnitIndex, +) +from shared.services.retrieval.nav.knowhere_hybrid import MAP_UNIT_INDEX_FORMAT_VERSION + + +@dataclass(frozen=True) +class RevisionStatistics: + path_document_count: int + path_total_length: int + content_document_count: int + content_total_length: int + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--apply", action="store_true", help="Write statistics.") + parser.add_argument( + "--check", action="store_true", help="Report missing or mismatched statistics." + ) + parser.add_argument("--document-id", default="") + parser.add_argument("--user-id", default="") + parser.add_argument("--namespace", default="") + parser.add_argument("--batch-size", type=int, default=100) + return parser + + +def _load_documents( + *, document_id: str, user_id: str, namespace: str +) -> list[Document]: + session_factory = get_sync_session_factory() + with session_factory() as db: + statement = ( + select(Document) + .where(Document.status == "active") + .where(Document.current_job_result_id.is_not(None)) + .order_by(Document.document_id) + ) + if document_id: + statement = statement.where(Document.document_id == document_id) + if user_id: + statement = statement.where(Document.user_id == user_id) + if namespace: + statement = statement.where(Document.namespace == namespace) + return list(db.scalars(statement).all()) + + +def _aggregate_statistics( + db: Session, *, document_id: str, job_result_id: str +) -> RevisionStatistics: + statement = ( + select( + func.count() + .filter(DocumentMapUnit.path_token_count > 0) + .label("path_document_count"), + func.coalesce(func.sum(DocumentMapUnit.path_token_count), 0).label( + "path_total_length" + ), + func.count() + .filter(DocumentMapUnit.content_token_count > 0) + .label("content_document_count"), + func.coalesce(func.sum(DocumentMapUnit.content_token_count), 0).label( + "content_total_length" + ), + ) + .where(DocumentMapUnit.document_id == document_id) + .where(DocumentMapUnit.job_result_id == job_result_id) + ) + row = db.execute(statement).one() + return RevisionStatistics( + path_document_count=int(row.path_document_count or 0), + path_total_length=int(row.path_total_length or 0), + content_document_count=int(row.content_document_count or 0), + content_total_length=int(row.content_total_length or 0), + ) + + +def _is_complete(index: DocumentMapUnitIndex | None, stats: RevisionStatistics) -> bool: + return bool( + index + and index.format_version == MAP_UNIT_INDEX_FORMAT_VERSION + and index.path_document_count == stats.path_document_count + and index.path_total_length == stats.path_total_length + and index.content_document_count == stats.content_document_count + and index.content_total_length == stats.content_total_length + ) + + +def _is_check_ready( + *, would_update: int, complete: int, skipped: int, documents: int +) -> bool: + """Return whether a read-only inventory proves every document is ready.""" + return would_update == 0 and skipped == 0 and complete == documents + + +def _process_batch( + documents: list[Document], *, apply_changes: bool +) -> tuple[int, int, int]: + session_factory = get_sync_session_factory() + updated = 0 + complete = 0 + skipped = 0 + with session_factory() as db: + for document in documents: + job_result_id = str(document.current_job_result_id or "") + stats = _aggregate_statistics( + db, document_id=document.document_id, job_result_id=job_result_id + ) + index_statement = ( + select(DocumentMapUnitIndex) + .where(DocumentMapUnitIndex.document_id == document.document_id) + .where(DocumentMapUnitIndex.job_result_id == job_result_id) + ) + if apply_changes: + index_statement = index_statement.with_for_update() + index = db.scalar(index_statement) + if index is None or index.format_version != MAP_UNIT_INDEX_FORMAT_VERSION: + skipped += 1 + print(f"skip document={document.document_id} reason=missing_or_legacy_index") + db.rollback() + continue + if _is_complete(index, stats): + complete += 1 + db.rollback() + continue + if not apply_changes: + updated += 1 + db.rollback() + continue + db.execute( + update(DocumentMapUnitIndex) + .where(DocumentMapUnitIndex.id == index.id) + .values( + path_document_count=stats.path_document_count, + path_total_length=stats.path_total_length, + content_document_count=stats.content_document_count, + content_total_length=stats.content_total_length, + ) + ) + db.commit() + updated += 1 + return updated, complete, skipped + + +def main() -> None: + args = _build_parser().parse_args() + if args.apply == args.check: + raise SystemExit("choose exactly one of --apply or --check") + if args.batch_size <= 0: + raise SystemExit("--batch-size must be positive") + documents = _load_documents( + document_id=args.document_id.strip(), + user_id=args.user_id.strip(), + namespace=args.namespace.strip(), + ) + totals = [0, 0, 0] + for offset in range(0, len(documents), args.batch_size): + batch_totals = _process_batch( + documents[offset : offset + args.batch_size], apply_changes=args.apply + ) + totals = [left + right for left, right in zip(totals, batch_totals)] + action = "applied" if args.apply else "would_update" + print( + f"{action}={totals[0]} complete={totals[1]} skipped={totals[2]} " + f"documents={len(documents)}" + ) + if args.check and not _is_check_ready( + would_update=totals[0], + complete=totals[1], + skipped=totals[2], + documents=len(documents), + ): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/apps/api/tests/contract/test_backfill_map_unit_indexes_contract.py b/apps/api/tests/contract/test_backfill_map_unit_indexes_contract.py index 72c61f91..3f61f463 100644 --- a/apps/api/tests/contract/test_backfill_map_unit_indexes_contract.py +++ b/apps/api/tests/contract/test_backfill_map_unit_indexes_contract.py @@ -1,6 +1,10 @@ from __future__ import annotations from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +from sqlalchemy.dialects import postgresql def test_backfill_script_resolves_shared_package_from_runtime_image_layout( @@ -26,3 +30,126 @@ def test_backfill_script_resolves_shared_package_from_source_checkout_layout( shared_root.mkdir(parents=True) assert _resolve_shared_root(api_root) == shared_root + + +def test_statistics_backfill_aggregates_positive_lengths_per_channel() -> None: + from scripts.backfill_map_unit_statistics import ( + RevisionStatistics, + _aggregate_statistics, + ) + + class AggregateResult: + def one(self) -> SimpleNamespace: + return SimpleNamespace( + path_document_count=2, + path_total_length=9, + content_document_count=1, + content_total_length=7, + ) + + class AggregateSession: + statement_sql: str = "" + + def execute(self, statement: Any) -> AggregateResult: + self.statement_sql = str( + statement.compile( + dialect=postgresql.dialect(), + compile_kwargs={"literal_binds": True}, + ) + ) + return AggregateResult() + + session = AggregateSession() + statistics = _aggregate_statistics( + cast(Any, session), document_id="doc_1", job_result_id="result_1" + ) + + assert statistics == RevisionStatistics( + path_document_count=2, + path_total_length=9, + content_document_count=1, + content_total_length=7, + ) + assert "path_token_count > 0" in session.statement_sql + assert "content_token_count > 0" in session.statement_sql + + +def test_statistics_backfill_readiness_requires_every_document_complete() -> None: + from scripts.backfill_map_unit_statistics import _is_check_ready + + assert _is_check_ready( + would_update=0, complete=4, skipped=0, documents=4 + ) + assert not _is_check_ready( + would_update=1, complete=3, skipped=0, documents=4 + ) + assert not _is_check_ready( + would_update=0, complete=3, skipped=1, documents=4 + ) + + +def test_statistics_backfill_completion_rejects_missing_or_legacy_indexes() -> None: + from scripts.backfill_map_unit_statistics import RevisionStatistics, _is_complete + + statistics = RevisionStatistics( + path_document_count=1, + path_total_length=2, + content_document_count=1, + content_total_length=3, + ) + legacy_index = SimpleNamespace( + format_version=1, + path_document_count=1, + path_total_length=2, + content_document_count=1, + content_total_length=3, + ) + complete_index = SimpleNamespace( + format_version=2, + path_document_count=1, + path_total_length=2, + content_document_count=1, + content_total_length=3, + ) + + assert not _is_complete(None, statistics) + assert not _is_complete(cast(Any, legacy_index), statistics) + assert _is_complete(cast(Any, complete_index), statistics) + + +def test_full_backfill_readiness_requires_revision_manifests() -> None: + from scripts.backfill_map_unit_indexes import NamespaceFallbackReport + + report = NamespaceFallbackReport( + user_id="user_1", + namespace="default", + active_docs=1, + snapshot_status="ok", + missing_from_snapshot=0, + missing_map_index=0, + missing_revision_manifest=1, + suspicious_zero_idf=0, + would_hit_snapshot_fallback=False, + scoring_incomplete=False, + ) + + assert not report.ready + + +def test_full_backfill_readiness_allows_mathematically_valid_zero_idf() -> None: + from scripts.backfill_map_unit_indexes import NamespaceFallbackReport + + report = NamespaceFallbackReport( + user_id="user_1", + namespace="default", + active_docs=1, + snapshot_status="ok", + missing_from_snapshot=0, + missing_map_index=0, + missing_revision_manifest=0, + suspicious_zero_idf=1, + would_hit_snapshot_fallback=False, + scoring_incomplete=False, + ) + + assert report.ready diff --git a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py index 63941a1f..fe5b732c 100644 --- a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -5,7 +5,8 @@ from typing import Any, cast from uuid import uuid4 -from httpx import AsyncClient +import pytest +from httpx import AsyncClient, Response from sqlalchemy import Engine, event, select from shared.models.database.document import DocumentMapUnit @@ -13,6 +14,10 @@ replace_document_revision_content, ) from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.search.map_unit_discovery import ( + DiscoveryResult, + map_unit_discovery, +) from shared.services.retrieval.serving_generation import lock_namespace_generation from tests.support.contract_database import ContractDatabase from tests.support.retrieval_snapshot_support import contract_db_session @@ -173,6 +178,561 @@ def capture_frequency_query( assert "FROM matching_tokens" in statements[-1] +async def test_classic_route_only_uses_token_selective_projection_for_unfiltered_scope( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-projection-{identifier}" + projection_statements: list[str] = [] + + def capture_unit_projection( + _connection: Any, + _cursor: Any, + statement: str, + _parameters: Any, + _context: Any, + _executemany: bool, + ) -> None: + if "SELECT DISTINCT scoped_units.*" in statement: + projection_statements.append("token_selective") + elif "SELECT * FROM scoped_units" in statement: + projection_statements.append("legacy") + + event.listen(Engine, "before_cursor_execute", capture_unit_projection) + try: + async with developer_api_client_factory() as api_client: + await _publish_document( + namespace=namespace, + source_file_name="projection.pdf", + chunks=[ + { + "chunk_id": f"projection-hit-{identifier}", + "type": "text", + "content": "projection token marker", + "path": "projection.pdf/Root/Section/hit", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"projection-filler-a-{identifier}", + "type": "text", + "content": "unrelated filler a", + "path": "projection.pdf/Root/Section/a", + "order": 2, + "metadata": {}, + }, + { + "chunk_id": f"projection-filler-b-{identifier}", + "type": "text", + "content": "unrelated filler b", + "path": "projection.pdf/Root/Section/b", + "order": 3, + "metadata": {}, + }, + { + "chunk_id": f"projection-filler-c-{identifier}", + "type": "text", + "content": "unrelated filler c", + "path": "projection.pdf/Root/Section/c", + "order": 4, + "metadata": {}, + }, + ], + ) + unfiltered_response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "projection token marker", + "top_k": 1, + "use_agentic": False, + }, + ) + filtered_response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "projection token marker", + "top_k": 1, + "use_agentic": False, + "signal_paths": ["Section"], + "filter_mode": "keep", + }, + ) + finally: + event.remove(Engine, "before_cursor_execute", capture_unit_projection) + + assert unfiltered_response.status_code == 200 + assert filtered_response.status_code == 200 + assert projection_statements[:2] == ["token_selective", "legacy"] + + +async def test_unfiltered_revision_pins_reuse_index_metadata_without_scoped_revision_cte( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-pinned-index-{identifier}" + index_statements: list[str] = [] + + def capture_index_query( + _connection: Any, + _cursor: Any, + statement: str, + _parameters: Any, + _context: Any, + _executemany: bool, + ) -> None: + if ( + "document_map_unit_indexes" in statement + and "average_idf_path" in statement + ): + index_statements.append(statement) + + async with developer_api_client_factory(): + document = await _publish_document( + namespace=namespace, + source_file_name="pinned-index.pdf", + chunks=[ + { + "chunk_id": f"pinned-hit-{identifier}", + "type": "text", + "content": "pinned revision semantic marker", + "path": "pinned-index.pdf/Root/Section/hit", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"pinned-filler-{identifier}", + "type": "text", + "content": "unrelated filler", + "path": "pinned-index.pdf/Root/Section/filler", + "order": 2, + "metadata": {}, + }, + ], + ) + + def result_signature(result: Any) -> list[tuple[Any, ...]]: + rows = list(result.payload.get("fused_rows") or []) + return [ + ( + row.get("chunk_id"), + row.get("document_id"), + row.get("job_result_id"), + row.get("section_path"), + row.get("score"), + row.get("discovery_score"), + row.get("content"), + ) + for row in rows + ] + + event.listen(Engine, "before_cursor_execute", capture_index_query) + try: + async with contract_db_session() as db: + unpinned = await map_unit_discovery( + db, + user_id=_USER_ID, + namespace=namespace, + query="pinned revision semantic marker", + top_k=10, + exclude_document_ids=[], + exclude_sections=[], + revision_pins=None, + ) + + index_statements.clear() + + async with contract_db_session() as db: + pinned = await map_unit_discovery( + db, + user_id=_USER_ID, + namespace=namespace, + query="pinned revision semantic marker", + top_k=10, + exclude_document_ids=[], + exclude_sections=[], + revision_pins={ + document["document_id"]: document["job_result_id"] + }, + ) + finally: + event.remove(Engine, "before_cursor_execute", capture_index_query) + + assert result_signature(pinned) == result_signature(unpinned) + assert index_statements + assert all("JOIN (VALUES" in statement for statement in index_statements) + assert all("scoped_units AS" not in statement for statement in index_statements) + assert all( + "SELECT DISTINCT document_id, job_result_id" not in statement + for statement in index_statements + ) + + +async def test_classic_discovery_returns_empty_for_an_empty_revision_pin( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + namespace: str = f"classic-empty-pins-{uuid4().hex[:8]}" + async with developer_api_client_factory(): + async with contract_db_session() as db: + result: DiscoveryResult = await map_unit_discovery( + db, + user_id=_USER_ID, + namespace=namespace, + query="empty revision marker", + top_k=1, + exclude_document_ids=[], + exclude_sections=[], + revision_pins={}, + ) + + assert result.payload["fused_rows"] == [] + + +async def test_classic_route_falls_back_for_v1_index_with_excluded_document( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-v1-fallback-{identifier}" + legacy_queries: list[str] = [] + + def capture_legacy_query( + _connection: Any, + _cursor: Any, + statement: str, + _parameters: Any, + _context: Any, + _executemany: bool, + ) -> None: + if "plainto_tsquery('simple'" in statement: + legacy_queries.append(statement) + + event.listen(Engine, "before_cursor_execute", capture_legacy_query) + try: + async with developer_api_client_factory() as api_client: + first = await _publish_document( + namespace=namespace, + source_file_name="legacy-fallback.pdf", + chunks=[ + { + "chunk_id": f"legacy-hit-{identifier}", + "type": "text", + "content": "legacy fallback marker", + "path": "legacy-fallback.pdf/Root/Section/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"legacy-filler-a-{identifier}", + "type": "text", + "content": "unrelated legacy filler a", + "path": "legacy-fallback.pdf/Root/Section/a", + "order": 2, + "metadata": {}, + }, + { + "chunk_id": f"legacy-filler-b-{identifier}", + "type": "text", + "content": "unrelated legacy filler b", + "path": "legacy-fallback.pdf/Root/Section/b", + "order": 3, + "metadata": {}, + }, + ], + ) + excluded = await _publish_document( + namespace=namespace, + source_file_name="excluded.pdf", + chunks=[ + { + "chunk_id": f"excluded-{identifier}", + "type": "text", + "content": "unrelated filler", + "path": "excluded.pdf/Root/Section/body", + "order": 1, + "metadata": {}, + } + ], + ) + await ContractDatabase.execute( + """ + UPDATE document_map_unit_indexes + SET format_version = 1 + WHERE document_id = :document_id + """, + {"document_id": first["document_id"]}, + ) + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "legacy fallback marker", + "top_k": 1, + "use_agentic": False, + "exclude_document_ids": [excluded["document_id"]], + }, + ) + finally: + event.remove(Engine, "before_cursor_execute", capture_legacy_query) + + assert response.status_code == 200 + body = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], body["results"]) + assert len(results) == 1 + assert results[0]["chunk_id"] == f"legacy-hit-{identifier}" + assert legacy_queries + + +async def test_classic_discovery_preserves_results_before_statistics_backfill( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-statistics-parity-{identifier}" + legacy_queries: list[str] = [] + + def capture_legacy_query( + _connection: Any, + _cursor: Any, + statement: str, + _parameters: Any, + _context: Any, + _executemany: bool, + ) -> None: + if "plainto_tsquery('simple'" in statement: + legacy_queries.append(statement) + + def result_signature(result: DiscoveryResult) -> list[tuple[Any, ...]]: + return [ + ( + row.get("chunk_id"), + row.get("document_id"), + row.get("job_result_id"), + row.get("section_path"), + row.get("source_file_name"), + row.get("score"), + row.get("discovery_score"), + row.get("content"), + row.get("chunk_metadata"), + ) + for row in list(result.payload.get("fused_rows") or []) + ] + + event.listen(Engine, "before_cursor_execute", capture_legacy_query) + try: + async with developer_api_client_factory(): + document = await _publish_document( + namespace=namespace, + source_file_name="statistics-parity.pdf", + chunks=[ + { + "chunk_id": f"statistics-hit-{identifier}", + "type": "text", + "content": "statistics parity retrieval marker", + "path": "statistics-parity.pdf/Root/Section/hit", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"statistics-filler-a-{identifier}", + "type": "text", + "content": "unrelated statistics filler a", + "path": "statistics-parity.pdf/Root/Section/a", + "order": 2, + "metadata": {}, + }, + { + "chunk_id": f"statistics-filler-b-{identifier}", + "type": "text", + "content": "unrelated statistics filler b", + "path": "statistics-parity.pdf/Root/Section/b", + "order": 3, + "metadata": {}, + }, + ], + ) + query = "statistics parity retrieval marker" + async with contract_db_session() as db: + post_backfill = await map_unit_discovery( + db, + user_id=_USER_ID, + namespace=namespace, + query=query, + top_k=3, + exclude_document_ids=[], + exclude_sections=[], + ) + + await ContractDatabase.execute( + """ + UPDATE document_map_unit_indexes + SET path_document_count = NULL, + path_total_length = NULL, + content_document_count = NULL, + content_total_length = NULL + WHERE document_id = :document_id + """, + {"document_id": document["document_id"]}, + ) + + async with contract_db_session() as db: + pre_backfill = await map_unit_discovery( + db, + user_id=_USER_ID, + namespace=namespace, + query=query, + top_k=3, + exclude_document_ids=[], + exclude_sections=[], + ) + finally: + event.remove(Engine, "before_cursor_execute", capture_legacy_query) + + assert result_signature(pre_backfill) == result_signature(post_backfill) + assert legacy_queries == [] + + +@pytest.mark.parametrize( + "incomplete_index_kind", ["legacy_format", "missing_index", "missing_tokens"] +) +async def test_unfiltered_classic_route_falls_back_when_selective_rows_are_unavailable( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + incomplete_index_kind: str, +) -> None: + identifier: str = uuid4().hex[:8] + namespace: str = f"classic-token-fallback-{incomplete_index_kind}-{identifier}" + legacy_queries: list[str] = [] + + def capture_legacy_query( + _connection: object, + _cursor: object, + statement: str, + _parameters: object, + _context: object, + _executemany: bool, + ) -> None: + if "plainto_tsquery('simple'" in statement: + legacy_queries.append(statement) + + event.listen(Engine, "before_cursor_execute", capture_legacy_query) + try: + async with developer_api_client_factory() as api_client: + document: dict[str, str] = await _publish_document( + namespace=namespace, + source_file_name="legacy-token-hash.pdf", + chunks=[ + { + "chunk_id": f"legacy-token-hit-{identifier}", + "type": "text", + "content": "legacy token fallback marker", + "path": "legacy-token-hash.pdf/Root/Section/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"legacy-token-filler-a-{identifier}", + "type": "text", + "content": "unrelated legacy filler a", + "path": "legacy-token-hash.pdf/Root/Section/a", + "order": 2, + "metadata": {}, + }, + { + "chunk_id": f"legacy-token-filler-b-{identifier}", + "type": "text", + "content": "unrelated legacy filler b", + "path": "legacy-token-hash.pdf/Root/Section/b", + "order": 3, + "metadata": {}, + }, + ], + ) + if incomplete_index_kind == "legacy_format": + await ContractDatabase.execute( + """ + UPDATE document_map_unit_indexes + SET format_version = 1 + WHERE document_id = :document_id + """, + {"document_id": document["document_id"]}, + ) + await ContractDatabase.execute( + """ + UPDATE document_map_unit_tokens + SET token_hash = :legacy_token_hash + WHERE map_unit_id IN ( + SELECT id + FROM document_map_units + WHERE document_id = :document_id + ) + """, + { + "document_id": document["document_id"], + "legacy_token_hash": "legacy-token-hash", + }, + ) + elif incomplete_index_kind == "missing_tokens": + await ContractDatabase.execute( + """ + DELETE FROM document_map_unit_tokens + WHERE map_unit_id IN ( + SELECT id + FROM document_map_units + WHERE document_id = :document_id + ) + """, + {"document_id": document["document_id"]}, + ) + else: + await ContractDatabase.execute( + """ + DELETE FROM document_map_unit_indexes + WHERE document_id = :document_id + """, + {"document_id": document["document_id"]}, + ) + await ContractDatabase.execute( + """ + DELETE FROM document_map_unit_tokens + WHERE map_unit_id IN ( + SELECT id + FROM document_map_units + WHERE document_id = :document_id + ) + """, + {"document_id": document["document_id"]}, + ) + response: Response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "legacy token fallback marker", + "top_k": 1, + "use_agentic": False, + }, + ) + finally: + event.remove(Engine, "before_cursor_execute", capture_legacy_query) + + assert response.status_code == 200 + body = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], body["results"]) + assert len(results) == 1 + assert results[0]["chunk_id"] == f"legacy-token-hit-{identifier}" + assert legacy_queries + + async def test_classic_route_image_filter_scores_only_units_with_images( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] diff --git a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py index 287a6b4c..5128382f 100644 --- a/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py +++ b/apps/api/tests/contract/test_retrieval_map_unit_index_contract.py @@ -57,7 +57,11 @@ def __init__(self) -> None: def execute(self, statement: str, parameters: object = None) -> None: executions.append((statement, parameters)) if "document_map_unit_indexes" in statement: - self.rows = [(document_id, job_result_id, 2, 1, 0.0, 0.0)] + self.rows = [ + (document_id, job_result_id, 2, 1, 0.0, 0.0, 1, 1, 1, 1) + ] + elif "FROM document_sections" in statement: + self.rows = [(document_id, job_result_id, 1)] elif "FROM document_map_units AS units" in statement: self.rows = [("unit-frequency", document_id, "chunk-frequency", "section-frequency", 1, 1)] elif "FROM document_map_unit_tokens" in statement: @@ -99,10 +103,18 @@ def close(self) -> None: ) assert corpus is not None + selective_executions = [ + statement + for statement, _parameters in executions + if "SELECT DISTINCT map_unit_id" in statement + ] + assert len(selective_executions) == 1 + assert "JOIN (VALUES" in selective_executions[0] frequency_executions = [ (statement, parameters) for statement, parameters in executions if "FROM document_map_unit_tokens" in statement + and "SELECT DISTINCT map_unit_id" not in statement ] assert len(frequency_executions) == 1 statement, parameters = frequency_executions[0] @@ -286,6 +298,18 @@ async def test_published_map_units_preserve_scores_without_chunk_payload_reads( ) assert index.unit_count == len(expected_units) + assert index.path_document_count == sum( + unit.path_token_count > 0 for unit in persisted_units + ) + assert index.path_total_length == sum( + unit.path_token_count for unit in persisted_units + ) + assert index.content_document_count == sum( + unit.content_token_count > 0 for unit in persisted_units + ) + assert index.content_total_length == sum( + unit.content_token_count for unit in persisted_units + ) assert [unit.unit_id for unit in persisted_units] == [ str(unit["chunk_id"]) for unit in expected_units ] diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index 3ead1217..8b4dc85d 100644 --- a/apps/api/tests/migrations/test_schema_contract.py +++ b/apps/api/tests/migrations/test_schema_contract.py @@ -49,6 +49,11 @@ def _upgrade_to_snapshot_parents(*, engine: Engine) -> None: command.upgrade(config, "fbe1c2d3e4f5") +def _upgrade_to_channel_statistics(*, engine: Engine) -> None: + config = _build_alembic_command_config(engine=engine) + command.upgrade(config, "b1c2d3e4f5a6") + + def _insert_job( connection: Connection, *, @@ -255,6 +260,161 @@ def test_should_index_document_chunks_in_lazy_section_order( ) +def test_should_create_token_leading_map_unit_covering_index( + migrated_head_engine: Engine, +) -> None: + with migrated_head_engine.begin() as connection: + index_row = connection.execute( + text( + """ + SELECT pg_get_indexdef(indexes.indexrelid), + indexes.indisvalid, + indexes.indisready + FROM pg_index AS indexes + JOIN pg_class AS classes ON classes.oid = indexes.indexrelid + JOIN pg_namespace AS namespaces + ON namespaces.oid = classes.relnamespace + WHERE namespaces.nspname = current_schema() + AND classes.relname = 'idx_document_map_unit_tokens_token_lookup' + """ + ) + ).one() + + definition = str(index_row[0]) + assert "(channel, token_hash, map_unit_id)" in definition + assert "INCLUDE (token, frequency)" in definition + assert index_row[1] is True + assert index_row[2] is True + + +def test_should_repair_a_missing_token_leading_map_unit_covering_index( + alembic_engine: Engine, +) -> None: + _upgrade_to_channel_statistics(engine=alembic_engine) + with alembic_engine.begin() as connection: + connection.execute( + text("DROP INDEX idx_document_map_unit_tokens_token_lookup") + ) + + _upgrade_to_heads(engine=alembic_engine) + + with alembic_engine.begin() as connection: + index_state = connection.execute( + text( + """ + SELECT indexes.indisvalid, indexes.indisready + FROM pg_index AS indexes + JOIN pg_class AS classes ON classes.oid = indexes.indexrelid + JOIN pg_namespace AS namespaces + ON namespaces.oid = classes.relnamespace + WHERE namespaces.nspname = current_schema() + AND classes.relname = 'idx_document_map_unit_tokens_token_lookup' + """ + ) + ).one() + + assert index_state[0] is True + assert index_state[1] is True + + +def test_should_repair_a_missing_covering_index_with_a_caller_owned_connection( + alembic_engine: Engine, +) -> None: + _upgrade_to_channel_statistics(engine=alembic_engine) + with alembic_engine.begin() as connection: + connection.execute( + text("DROP INDEX idx_document_map_unit_tokens_token_lookup") + ) + + _upgrade_to_heads_with_external_connection(engine=alembic_engine) + + with alembic_engine.begin() as connection: + index_state = connection.execute( + text( + """ + SELECT indexes.indisvalid, indexes.indisready + FROM pg_index AS indexes + JOIN pg_class AS classes ON classes.oid = indexes.indexrelid + JOIN pg_namespace AS namespaces + ON namespaces.oid = classes.relnamespace + WHERE namespaces.nspname = current_schema() + AND classes.relname = 'idx_document_map_unit_tokens_token_lookup' + """ + ) + ).one() + + assert index_state[0] is True + assert index_state[1] is True + + +def test_should_repair_an_invalid_token_leading_map_unit_covering_index( + alembic_engine: Engine, +) -> None: + _upgrade_to_channel_statistics(engine=alembic_engine) + with alembic_engine.begin() as connection: + connection.execute( + text( + """ + UPDATE pg_index + SET indisvalid = FALSE, + indisready = FALSE + WHERE indexrelid = + 'idx_document_map_unit_tokens_token_lookup'::regclass + """ + ) + ) + + _upgrade_to_heads(engine=alembic_engine) + + with alembic_engine.begin() as connection: + index_state = connection.execute( + text( + """ + SELECT indexes.indisvalid, indexes.indisready + FROM pg_index AS indexes + JOIN pg_class AS classes ON classes.oid = indexes.indexrelid + JOIN pg_namespace AS namespaces + ON namespaces.oid = classes.relnamespace + WHERE namespaces.nspname = current_schema() + AND classes.relname = 'idx_document_map_unit_tokens_token_lookup' + """ + ) + ).one() + + assert index_state[0] is True + assert index_state[1] is True + + +def test_should_add_per_channel_map_unit_bm25_statistics( + migrated_head_engine: Engine, +) -> None: + with migrated_head_engine.begin() as connection: + columns = { + str(row[0]): str(row[1]) + for row in connection.execute( + text( + """ + SELECT column_name, is_nullable + FROM information_schema.columns + WHERE table_schema = current_schema() + AND table_name = 'document_map_unit_indexes' + AND column_name IN ( + 'path_document_count', 'path_total_length', + 'content_document_count', 'content_total_length' + ) + """ + ) + ).all() + } + + assert columns == { + "path_document_count": "YES", + "path_total_length": "YES", + "content_document_count": "YES", + "content_total_length": "YES", + } + + def test_should_upgrade_with_a_caller_owned_connection( alembic_engine: Engine, ) -> None: diff --git a/deploy/ecs/README.md b/deploy/ecs/README.md index 7b7d6170..37a8c572 100644 --- a/deploy/ecs/README.md +++ b/deploy/ecs/README.md @@ -77,31 +77,44 @@ services. It does not create or delete AWS resources. ## Required post-deploy backfill -The map-nav lexical-index migration creates the derived index tables, but it does -not rebuild indexes for revisions that already exist. Until those revisions are -backfilled, retrieval remains quality-preserving but uses the legacy scoring -path. Every release containing the map-nav index change must include the -following DevOps action in its release notification. +Follow the complete +[`retrieval-serving-index-rollout-runbook.md`](../../docs/design/retrieval-serving-index-rollout-runbook.md) +for schema verification, statistics maintenance, readiness gates, parity, +monitoring, pause/resume, and rollback. + +The additive migration does not populate the four per-channel statistics for +existing revisions. Until those revisions are ready, retrieval remains +quality-preserving but uses the legacy scoring path. Every release containing +this change must include the following DevOps action in its release +notification. Run the commands as a one-off container using the newly deployed API image and the production database secret. Do not run them inside the long-lived API task. ```bash -# Read-only inventory -python /app/scripts/backfill_map_unit_indexes.py +# Read-only statistics inventory, after migration +python /app/scripts/backfill_map_unit_statistics.py --check --batch-size 100 # Optional canary: apply one affected document first -python /app/scripts/backfill_map_unit_indexes.py \ +python /app/scripts/backfill_map_unit_statistics.py \ --document-id \ + --batch-size 100 \ --apply -# Apply to all current document revisions -python /app/scripts/backfill_map_unit_indexes.py --apply +# Apply statistics to complete format-v2 indexes +python /app/scripts/backfill_map_unit_statistics.py --apply --batch-size 100 + +# Final statistics and full serving-readiness checks +python /app/scripts/backfill_map_unit_statistics.py --check --batch-size 100 +python /app/scripts/backfill_map_unit_indexes.py --check ``` -The script commits each document revision independently and is safe to rerun. -Verify the canary retrieval before starting the full apply. New or republished -documents build their index automatically during publication. +The statistics script commits each document revision independently and is safe +to rerun. Missing or legacy indexes reported as `skipped` require the existing +full `backfill_map_unit_indexes.py --apply --document-id ` path. +Run only one maintenance process at a time. Verify the canary retrieval before +starting the full apply. New or republished documents build their index +automatically during publication. ## Manual staging availability diff --git a/docs/adr/0009-use-token-leading-covering-index-for-map-unit-lookup.md b/docs/adr/0009-use-token-leading-covering-index-for-map-unit-lookup.md new file mode 100644 index 00000000..da5d9c99 --- /dev/null +++ b/docs/adr/0009-use-token-leading-covering-index-for-map-unit-lookup.md @@ -0,0 +1,11 @@ +# Use a token-leading covering index for map-unit lookup + +**Status: accepted.** Token-selective retrieval will use an additive, +idempotent PostgreSQL covering index led by `(channel, token_hash, map_unit_id)` +and including `(token, frequency)`. This preserves the token-index-driven query +shape and allows index-only plans where PostgreSQL visibility permits them. The +migration is additive and keeps the existing lookup index until a separate +production-plan and load review proves it redundant. Requests fall back to the +exact legacy reader when the index or serving data is missing or incomplete. +The additional storage and write-maintenance cost is accepted in exchange for +lower first-request retrieval latency and row transfer. diff --git a/docs/adr/0010-backfill-serving-index-statistics-in-place.md b/docs/adr/0010-backfill-serving-index-statistics-in-place.md new file mode 100644 index 00000000..75e2ac4c --- /dev/null +++ b/docs/adr/0010-backfill-serving-index-statistics-in-place.md @@ -0,0 +1,18 @@ +# Backfill serving-index statistics in place + +**Status: accepted.** Existing retrieval-visible data will be maintained by +backfilling only the current revision of each active Document. Complete serving +indexes receive the new per-channel BM25 statistics through the statistics-only +`apps/api/scripts/backfill_map_unit_statistics.py --apply` command. It aggregates +existing map units and does not regenerate token rows, manifests, or namespace +snapshots. Missing or legacy indexes use the existing full +`apps/api/scripts/backfill_map_unit_indexes.py --apply --document-id ...` +rebuild. v1 and v2 are internal read paths under the same Retrieval contract, +so each revision commits independently. While the four statistics are NULL, +the reader keeps the full scope-first map-unit projection and derives the +existing per-channel denominators from those rows; it does not switch to a +semantically different FTS query. Once statistics are complete, the +token-selective projection may be used. Missing, legacy, or storage- +inconsistent index data still uses the exact legacy reader. An interrupted or +failed backfill therefore preserves retrieval semantics and avoids unnecessary +regeneration of correct data. diff --git a/docs/adr/README.md b/docs/adr/README.md index 00d7b9da..4d535f15 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -23,3 +23,5 @@ Use this shape: | [0006](0006-atomically-publish-retrieval-serving-index.md) | Atomically publish the retrieval-serving index | | [0007](0007-use-coherent-retrieval-serving-generations.md) | Use coherent retrieval-serving generations | | [0008](0008-use-a-maintenance-window-for-serving-index-rollout.md) | Roll out the serving index online | +| [0009](0009-use-token-leading-covering-index-for-map-unit-lookup.md) | Use a token-leading covering index for map-unit lookup | +| [0010](0010-backfill-serving-index-statistics-in-place.md) | Backfill serving-index statistics in place | diff --git a/docs/design/retrieval-heart-vessel-trace-optimization-plan.md b/docs/design/retrieval-heart-vessel-trace-optimization-plan.md new file mode 100644 index 00000000..ed62a57b --- /dev/null +++ b/docs/design/retrieval-heart-vessel-trace-optimization-plan.md @@ -0,0 +1,617 @@ +# Retrieval optimization plan from the fresh `心血管` trace + +## Scope and evidence + +The detailed stage breakdown below is based on a newly issued production +request, not a historical request: + +```json +{ + "namespace": "default", + "query": "心血管", + "top_k": 3, + "use_agentic": null +} +``` + +- Trace: `01a05e530f7787333c0e32ca16633b85` +- Route: `/api/v1/retrieval/query` +- HTTP status: `200` +- Router: `mapnav` +- Stop reason: `completed` +- Server span: `42.365 s` +- Client wall time: `44.690 s` +- Result: `29` referenced chunks, approximately `11,998` evidence characters + +The current production baseline is restricted to the last-night and today +window in Beijing time (`2026-09-01 20:00` through `2026-09-02 11:00`, or +`2026-09-01 12:00Z` through `2026-09-02 03:00Z`). Logfire recorded 203 v1 +retrieval roots and one v2 root in that window; all successful roots were from +`deployment.environment=production`. The v1 traffic includes both the direct +`/v1/retrieval/query` path and the externally prefixed `/api/v1/retrieval/query` +path. The FastAPI route is `/v1/retrieval/query` in both cases. + +Logfire does not currently record a git commit in `service_version` or another +resource attribute. These traces are therefore the latest observed production +behavior, but they are not proof that production matches a particular local +HEAD. The runtime reports Python `3.12.14`, built on `2026-09-01`, which is +useful deployment context but not a source revision identifier. + +v1 and v2 share the same retrieval execution plan. v2 only adds the optional +`llm_config` request field and passes it into the shared plan; when it is absent, +the retrieval SQL, route selection, scoring, hydration, and public projection +are the same. The one v2 request in this baseline had `llm_config=null`. + +The timings below are parent/child timings. Parent stages include their child +stages and must not be summed with those children. + +## Non-negotiable quality invariant + +Every optimization in this plan must preserve retrieval semantic parity for the +same pinned request: selected chunk IDs, ordering, rounded scores, source +sections, citations, evidence content, and asset references must remain +unchanged. Score comparisons use a maximum absolute tolerance of `1e-4` to +avoid treating harmless floating-point accumulation-order differences as a +retrieval change; IDs, ordering, source sections, citations, evidence content, +and asset references have no tolerance. A latency improvement that has not +passed the parity checks is not shippable, and the exact legacy retrieval path +remains the fallback. + +## Current optimization boundary + +This phase covers non-LLM work only: SQL plans, indexes, serving-index row +transfer, cache behavior, snapshot decoding, map scoring, and instrumentation. +Planner, Harvest, and Control prompts, models, thinking settings, and model-call +orchestration are deferred until this phase reaches semantic-parity and latency +gates. + +The optimized reader is an internal implementation shared by v1 and v2. Public +API versioning is not a performance variant: every non-LLM change must preserve +the same Retrieval contract for v1 and v2. v2's optional `llm_config` may change +model settings, but it must not change lexical retrieval semantics. + +Implement and validate slices sequentially on the immutable local production +copy. Instrumentation is measurement-only and must first prove that the +baseline result is unchanged. Each subsequent SQL, index, cache, or local +scoring change is evaluated independently against the temporary parity set +before the next change begins; several individually validated slices may be +deployed together afterward. + +## Stage breakdown + +| Stage | Observed time | Evidence | +| --- | ---: | --- | +| Scope threshold probe | 1.838 s | `count_scoped_chunks` applies `LIMIT top_k + 1` before the outer `count(*)`; this is a bounded small-corpus gate, not an exact count | +| Snapshot load | 3.064 s | 644 documents, 115,892 references | +| Planner | 19.996 s | One subgoal; Planner LLM 19.995 s | +| Tree build | 0.966 s | 50,112 sections | +| Map-index load | 8.523 s | 644 index rows; 42,794 unit rows; 3 query tokens | +| Unit scoring | 0.190 s | 16,662 scored units | +| Map scoring aggregate | 9.802 s | Parent aggregate around scoring/pooling work | +| Harvest | 5.260 s | Harvest LLM 2.762 s; one wave/subgoal | +| Control | 1.256 s | Control LLM 1.248 s | +| Orchestration | 6.518 s | Parent around Harvest and Control | +| Evidence pack | 0.262 s | 23 chunks | +| Final hydration | 0.444 s | 23 results | +| Episode | 36.633 s | Parent from episode start through evidence pack | + +The map-index sub-stages were: + +- index metadata: `0.028 s`, 644 rows; +- map-unit rows: `4.288 s`, 42,794 rows, `cache_hit=false`; +- frequency rows: `3.866 s`, 42,794 units and 3 tokens. + +The `sections=50112` value in the current trace is now accounted for at the +source boundary: both the `tree_build` and `map_pooling` log sites computed it +as `sum(len(tree_by_doc[doc_id][0]) ...)`, where the first tuple member is the +reachable section-node map. It therefore means **section nodes**, not +section-child edges, leaf sections, or chunks. The local instrumentation now +also records `section_edges` and `leaf_sections` for the map scorer, and +`section_rows`, `section_paths`, `root_sections`, and `leaf_sections` for the +snapshot loader. A previous isolated fixture contained the target namespace's +644 documents, 50,112 sections, 60,187 chunks, and 42,794 map units. That +fixture has been discarded after CSV transfer caused excessive local write +amplification; the replacement DevOps dump is pending and must be treated as +the only current end-to-end parity corpus. + +For planner work, the discarded fixture contained 41,105 real rows for the +three-token probe plus 1,390,687 non-matching filler rows. Its 1,431,792 total +rows preserved the observed production channel ratio (content 1,275,676; path +156,116) and kept the real target-token distribution. Production had +13,948,605 rows (content 12,524,037; path 1,423,950), so those measurements were +from a scaled production-shaped copy rather than an exact row-count clone. +They remain historical investigation evidence only; repeat them on the +replacement dump before using them for an implementation decision. + +On that discarded copy, with all 42,794 map units in scope, the existing +scope-first frequency query returned 36,478 rows in every run and measured +`103–121 ms` across five warm repetitions. The materialized token-first +candidate also returned 36,478 rows (symmetric difference `0`) but measured +`323–378 ms`; PostgreSQL used the existing non-covering lookup index by +default. A token-selective unit projection driven by a materialized +matching-token CTE returned 13,573 units and measured `205–223 ms`; PostgreSQL +still chose a unit-leading lookup for the join after an index-only token scan. +The scaled copy has a much higher target-hash fraction than production, so +this does not predict the full-scale join order. These results validate row +parity but do not yet demonstrate an end-to-end latency benefit; the reader +must not change until rows/bytes transferred and retrieval quality are measured +through the complete caller path. + +The historical SQL-output proxy confirms the potential transfer reduction: the +existing +unit projection emitted 42,794 rows / 4.05 MB, while the token-selective +projection emitted 13,573 rows / 1.29 MB. This is a database-client output +measurement, not an application latency claim; Python decoding, filtering, +scoring, and cache behavior still need to be measured together. + +Using the existing persisted BM25 scorer with the same full-corpus channel +denominators, the complete and token-selective projections both produced +13,573 positive-scoring units for `心血管`; the maximum score delta was `0.0` +and their top-100 ordering was identical. This confirms that dropping +zero-frequency units is score-safe for this probe, but it is not yet a public +API retrieval-quality gate. + +The actual `ReadOnlyChunkStore.load_persisted_score_corpus` caller was also run +against 508 revisions in that discarded fixture whose serving index had format +version `1`. +It loaded 44,894 section rows, returned the same 13,573 score units, and took +`1.089 s` for loading plus `0.168 s` for scoring. Revisions with incomplete +serving indexes correctly returned the existing fallback (`None`); they were +not silently included in the optimized sample. + +A temporary scorer baseline was previously frozen at +`/tmp/knowhere-retrieval-parity-baseline.json.gz`. That artifact was generated +before rebasing onto the current main tokenizer and is now historical evidence +only; it must not be used as the acceptance baseline for this branch. Generate +a new baseline from the current `origin/main` behavior after the local copy has +been rebuilt with v2 tokens. + +The classic-route SQL shape was then measured against the same historical +namespace and revision scope. After one cold run, five repetitions of the legacy projection +measured `716–763 ms`; the token-selective projection measured `781–844 ms` +(one outlier at `1.73 s`). Output fell from 42,794 rows / 9.53 MB to 13,573 +rows / 2.93 MB. This is a substantial transfer reduction with no stable SQL +latency win yet, so the next gate is application-level p95 rather than a claim +that the query itself is faster. + +The earlier discovery-level comparison and three-query digests were also +captured before the rebase, against the old tokenizer/data state. They are +retained only as historical investigation notes and do not establish current +latency or retrieval parity. Re-run both the baseline and optimized discovery +paths after the v2 local backfill, using the current main code as the baseline. + +The replacement local dump is now available on PostgreSQL port `55433`. +After applying the additive schema migration, populating the four channel +statistics from the existing map units, and running `VACUUM (ANALYZE)`, the +token-leading covering index uses an index-only scan (`Heap Fetches=0`). For +the `心血管` probe, the isolated SQL projection measured approximately +`150 ms` for scope-first versus `77 ms` for token-selective, and the frequency +lookup measured `2.7 ms`. Ten repeated classic application calls across +`心血管`, `心脏`, and `肺` preserved chunk IDs, ordering, sources, and evidence; +the maximum score delta was `1.7e-5`. Warm p50 improved by about `24 ms` for +`心血管`, was effectively unchanged for `心脏`, and regressed by less than +`1 ms` for the empty `肺` probe. Treat this as a SQL/transfer improvement with +no yet-established broad end-to-end latency win; production rollout still +requires the same migration, backfill, and rollback checks below. + +The current classic reader now reuses the immutable revision pins captured at +request start when validating an unfiltered scope. This removes a redundant +`SELECT DISTINCT` over scoped map units; requests without pins retain the +original query. On the restored dump, the removed validation query measured +approximately `0.35–0.6 s` before the change. Stage instrumentation shows the +remaining warm classic work is dominated by database reads and hydration, not +BM25 statistics or Python scoring. Result IDs, ordering, sources, evidence, +and score deltas remain within the existing `1e-4` parity tolerance. + +For the local restored PostgreSQL (which does not enable SSL), agentic smoke +must set `DB_SSL_MODE=disable` for the global async engine and provide the +plain libpq form through `KNOWHERE_DATABASE_URL` for the synchronous map-nav +reader. Without these local-only settings, final reference hydration attempts +an SSL upgrade and fails even though the database is healthy. + +The reference trace above is not representative of request mix. In the current +baseline window, 196 successful v1 roots used `use_agentic=false` (classic), +three used the default map-nav route (`use_agentic=null`), and one explicitly +used `use_agentic=true`. The classic roots had P50 `10.99 s`, P90 `25.46 s`, +and a maximum of `36.65 s`. There were also three unauthenticated v1 attempts +(`401`) and one successful v2 root. The explicit `use_agentic=true` v1 request +took `41.31 s`, and the single v2 request took `109.36 s`; neither is combined +with the classic latency distribution because they exercise different route +and/or model settings. +Across their SQL children, `WITH knowhere` spans had P50 `1.15 s`, P90 +`8.24 s`, P95 `10.46 s`, and a maximum of `18.90 s`. In a representative +27.91-second classic request, map-unit discovery took `4.02 s` and the +token-frequency query took `18.90 s`. These current-window classic measurements +are part of the baseline for the shared v1/v2 reader. + +The resource log still renders literal format placeholders +(`cpu_seconds=%.3f process_max_rss_kb=%d`), so this trace cannot prove CPU +saturation or request-level memory usage. + +## Findings + +### Confirmed bottlenecks + +1. In the reference map-nav trace, the Planner LLM consumes approximately 20 + seconds. Outbound HTTP spans are only a few hundred milliseconds, so most of + the elapsed time is provider generation/thinking or uninstrumented client + wait, not network transfer. +2. The first request transfers all 42,794 map-unit rows before applying the + three query tokens. This costs 4.288 seconds and creates unnecessary Python + allocations. +3. The frequency lookup still costs 3.866 seconds even for only three tokens. + The current reader already filters `document_map_unit_tokens` by + `token_hash` before joining the pinned scope; the remaining question is + whether a token-leading covering index or a different join shape helps on a + full-scale corpus. The local production-shaped benchmark above does not + show a benefit yet. +4. The map scoring aggregate is 9.802 seconds, while the explicitly measured + unit scoring and pooling work is only 0.190 and 0.094 seconds. The gap is an + instrumentation boundary and/or additional map-scoring work that must be + measured before CPU tuning. +5. The bounded scope threshold probe costs 1.838 seconds. The default map-nav + route does not need an exact chunk count for ranking, but the probe still + distinguishes corpora at or below `top_k` from larger corpora before route + selection. The trace alone does not establish that an `EXISTS` rewrite or + snapshot metadata would be faster. +6. Classic discovery previously re-scanned scoped map units solely to derive + revision keys after the request had already captured revision pins. Reusing + those pins removes that duplicate read without changing the completeness + check; the no-pin path remains unchanged. + The same pins now drive the unfiltered index-metadata lookup directly, + avoiding another scoped-unit CTE. On the restored dump this reduced the + warm index stage from roughly `0.16–0.8 s` to `0.02–0.04 s` for the sampled + namespace. The classic result parity gate still passes for the temporary + query set. +7. The current production request mix is primarily classic retrieval. The + token-selective reader must therefore be exercised through both the classic + and map-nav callers; a map-nav-only benchmark would not represent the + dominant workload. + +### Already healthy or lower priority + +- Snapshot load is material but not the largest stage at 3.064 seconds. +- Evidence pack and final hydration together are below one second in this + request. They are not the current optimization priority. +- The previous connected-hydration change is active: no full `JobResult` + chunk graph load appears in the trace. +- No prompt or retrieval-quality behavior should be changed as part of the + first implementation slices. + +## Prioritized optimization plan + +### Validation gate: bounded scope threshold probe (not an implementation P0) + +The current implementation already uses `LIMIT top_k + 1` inside +`count_scoped_chunks`, then counts that bounded subquery. Do not describe this +as an exact count or assume a 1.5–1.8 second saving from an `EXISTS` rewrite. +Benchmark the current probe against any equivalent early-exit query or trusted, +generation-pinned snapshot metadata. Retain this slice only if +`EXPLAIN (ANALYZE, BUFFERS)` and production-shaped repetitions show a meaningful +improvement; otherwise leave the current implementation unchanged and keep +this out of the optimization queue. + +Acceptance criteria: + +- route selection is unchanged for corpora below, equal to, and above + `top_k`; +- no stale snapshot can incorrectly classify a small corpus; +- the baseline and candidate plans make the `top_k + 1` bound and early-exit + behavior explicit; +- any candidate replacement is retained only when a latency improvement is + demonstrated; otherwise this remains a validation note, not an optimization + slice; +- selected chunk IDs, scores, ordering, and evidence remain unchanged; +- benchmark results record the route family (`classic`, `mapnav`, or + `small_corpus`) and separate cold, warm, and response-cache-hit requests; + cache-hit timings are not mixed into cold-request latency claims. + +### P0: Make map-unit projection token-selective + +The current map-nav reader already makes its frequency lookup token-selective, +but it still loads every revision-scoped map unit before applying the query +tokens. Change only the unit projection: start from +`document_map_unit_tokens` filtered by `channel` and `token_hash`, then join the +pinned revision/unit scope, and return only units matching at least one query +token. Keep the existing frequency SQL shape until a production-shaped plan +proves that a different join order is better. + +Roll this out in two stages. The first stage may enable the token-selective +reader only when the request has no section exclusions, where revision-scoped +channel statistics are sufficient. Requests with section filters continue to +use the existing scope-first map-unit reader until section-scoped denominator +statistics are implemented and pass semantic-parity checks. A missing, +incompatible, or storage-inconsistent serving index uses the exact legacy FTS +fallback. A complete index with NULL channel statistics keeps the full scope +map-unit rows and derives the two channel denominators from those rows until +the statistics backfill completes. + +BM25 denominators must remain corpus-wide. Obtain exact corpus statistics using +persisted serving-index statistics while preserving the current per-channel +semantics. Extend `DocumentMapUnitIndex` with +`path_document_count`, `path_total_length`, `content_document_count`, and +`content_total_length`, calculated atomically when publishing each revision. +Units with zero path length must not contribute to path `document_count` or +`total_length`, and the same rule applies independently to content. Do not +reuse `DocumentMapUnitIndex.unit_count` for both channels, and never calculate +average length from only the matching subset. + +Hypothesis to validate: this should reduce map-index latency, rows transferred, +and temporary memory for this workload; no fixed seconds-saving claim is made +before a full-scale or statistically equivalent local benchmark. The scaled +copy above is a negative/shape-control result, not an approval to change the +reader. + +The current implementation slice only adds the per-channel statistics fields, +publication-time population, and the DevOps backfill/readiness contract. It +enables token-selective projection when the scope is unfiltered and serving +statistics are present. Before statistics backfill, an unfiltered revision +uses the full scope map-unit rows with row-derived per-channel denominators, so +the lexical result remains unchanged; section-filtered requests retain their +existing scope-first map-unit reader. Only a missing, incompatible, or +storage-inconsistent index uses the exact legacy FTS fallback. Both v1 and v2 +call the same shared readers; no route-specific lexical algorithm was +introduced. + +Acceptance criteria: + +- a new idempotent migration adds the token-leading covering index required by + this query, for example `(channel, token_hash, map_unit_id) INCLUDE + (token, frequency)`. Keep the existing lookup index during this rollout; + evaluate removal separately only after `EXPLAIN (ANALYZE, BUFFERS)` and + production load confirm the new index is safe (or document the visibility + conditions that prevent an index-only scan); +- serving-index statistics include positive-length `document_count` and + `total_length` separately for path and content, aggregated over the exact + pinned revision and section scope; +- frequency maps and scores are identical to the current implementation on a + fixed production-data copy; +- duplicate token matches do not duplicate units; +- rows and bytes transferred are measured before and after; +- a safe fallback remains available if the persisted index is incomplete; +- v1 and v2 requests with equivalent retrieval fields produce identical public + retrieval results; v2's optional `llm_config` is recorded as a model-setting + difference, not as a separate lexical retrieval algorithm; +- the benchmark uses current-window classic requests as the primary sample and + includes at least one map-nav and one v2 request as shared-plan regression + samples; +- revision-scope cardinality and the size of the `(document_id, + job_result_id)` parameter set are recorded, so large `IN` lists are not hidden + inside token-query timing. + +### P1: Optimize the frequency SQL plan + +Remove or avoid a materialized full `scoped_units` side when the planner picks +the wrong join order. Drive from the token index, then apply channel, +token-hash, pinned revision, active-document, and namespace predicates. + +Hypothesis to validate: the revised join order may reduce frequency-stage +latency; claim no fixed saving until production-shaped `EXPLAIN (ANALYZE, +BUFFERS)` and repeated local runs confirm it. + +Acceptance criteria: + +- returned `(map_unit_id, channel, token, frequency)` rows are identical; +- no archived or wrong-generation revision can enter the result; +- statement timeout is not approached on the production read-only database; +- the query remains correct for one token, many tokens, and no matching token; +- repeated measurements cover the observed frequency-latency range (warm and + cold cache, one and many query tokens, and small and large revision scopes) + before any saving claim is made. + +### P1: Measure and then reduce the unaccounted map-scoring time + +Add structured timers around map-pooling, section aggregation, ranking, and +projection so the `9.802 s` parent stage can be reconciled with its children. +Only after that measurement should we change data structures or algorithms. + +Likely follow-up: precompute section owners, children, and postorder arrays +while decoding the snapshot, then reuse them during scoring instead of +rebuilding dictionaries and sets. + +Hypothesis to validate: precomputation should reduce local map-scoring time; +the magnitude is intentionally left to measurement. + +### Deferred: Planner and other LLM latency + +The Planner is the largest single stage, but this phase makes no changes to +Planner, Harvest, or Control model calls. Keep their existing behavior and +fallbacks intact. Revisit model timing, thinking, prompts, and call +orchestration only after the non-LLM slices below have passed their semantic- +parity and latency gates. + +### Deferred: overlap planner and query-independent map scoring (experimental) + +The current route waits for the Planner before starting map scoring. For a +planner result whose `retrieval_query` is exactly the user query and whose +scope/filters are unchanged, map scoring can be started concurrently on a +separate read-only connection. The planner result is then used only to decide +whether the already-computed candidate set is sufficient. + +This could hide part of the approximately `9.8 s` map-scoring aggregate behind +the approximately `20 s` Planner wait, but it changes LLM-boundary +orchestration. It is deferred until the Planner phase is explicitly in scope; +it must not run when the planner rewrites the query, adds a node filter, or +changes the revision/scope. + +Potential saving: up to the overlapping map-scoring time for eligible simple +queries. This is an experimental concurrency change with risks around +connection-pool pressure, duplicate work, cancellation, and revision +coherence. + +Acceptance criteria: + +- the candidate scores are byte-for-byte equal to the serial path; +- the concurrent read-only store is isolated from mutable navigation state; +- planner rewrites and filters correctly disable the overlap; +- cancellation and database connection cleanup are covered by contract tests; +- concurrency load tests show no increase in statement timeouts or RSS. + +### P2: Reduce snapshot parse allocations + +Redis stores compressed snapshot bytes, but every request still decompresses +and JSON-decodes 115,892 references. The current parsing path creates both +canonical and owner-qualified reference keys and then copies the mapping. + +Use one canonical representation with owner-aware fallback lookup and avoid +the second full mapping copy. Consider a versioned binary format only after a +benchmark; it is a larger migration and is not required for the first slices. + +Hypothesis to validate: the representation change should reduce decode time and +temporary RSS; the magnitude is intentionally left to measurement. + +Acceptance criteria: + +- both existing lookup forms remain valid; +- revision pinning and selected references are unchanged; +- checksum/version validation and generation invalidation remain intact; +- cold Redis-miss and Redis-hit memory are measured separately. + +### P2: Redis map-unit projection cache + +The current unit cache is episode/process-local; the snapshot Redis cache does +not eliminate the map-unit SQL transfer. Add a generation- and revision-scoped +Redis cache for a compact map-unit projection and index statistics, with a +bounded TTL and size limit. Bind each key to `user_id`, `namespace`, serving +generation, `document_id`, `job_result_id`, and index format version. A +generation change naturally invalidates old entries; Redis misses, version +mismatches, and decode failures fall back to PostgreSQL. + +This is not a guaranteed cold-first-request optimization. It helps only when a +previous publisher or request has populated the generation key. PostgreSQL +remains the source of truth and is the fallback on misses or Redis failures. + +## Existing-data maintenance + +Backfill only the currently retrieval-visible revision for each active +Document. Add the channel-stat columns with an additive migration and introduce +serving-index format version 2. v1 and v2 remain two internal read paths under +the same Retrieval contract and must produce identical results. For a complete +existing index, a resumable per-revision statistics backfill aggregates the existing +`document_map_units` rows, writes the four path/content values, and marks the +revision v2 in the same transaction. It must not regenerate token rows or +snapshots unnecessarily. Revisions with a missing or legacy index use the +existing full `backfill_map_unit_indexes --apply` path. + +Each revision is committed independently under the existing generation and +active-document locks. The job is idempotent and safe to interrupt. While the +four statistics are NULL, retrieval keeps the full scope-first map-unit reader +and derives the existing per-channel denominators from its rows. Once the +backfill check reports complete, coherent active revisions, the unfiltered +reader can use token-selective projection and persisted denominators. Missing, +incompatible, or storage-inconsistent index data still uses the exact legacy +FTS fallback. + +### DevOps handoff + +DevOps runs the maintenance explicitly; API startup and user requests never +trigger it. The executable rollout procedure is documented in +[`retrieval-serving-index-rollout-runbook.md`](retrieval-serving-index-rollout-runbook.md). +The runbook: + +1. deploy the additive migration and verify the reader still falls back safely; + Build the new index without dropping the existing lookup index, and monitor + lock waits, build duration, and disk headroom; +2. deploy the application code that writes v2 for new publications and selects + the token-selective reader only for complete v2 revisions; revisions with + NULL statistics continue on the full scope-first map-unit reader, while v1 + or otherwise unusable index data continues on the legacy FTS reader; +3. run the read-only inventory/check command and record active/current revision + counts; +4. run the resumable statistics backfill in bounded batches (with optional + document selection), monitoring database load and generation changes. The + command is: + + ```bash + python /app/scripts/backfill_map_unit_statistics.py \ + --apply --batch-size 100 + ``` + + Use `--user-id`, `--namespace`, or `--document-id` to narrow a rehearsal. + `--check` is read-only and should return `would_update=0` before rollout. + This command only aggregates existing `document_map_units`; it does not + regenerate token rows or snapshots. Revisions with a missing or legacy + index remain on the existing full backfill path. +5. rerun the check until every retrieval-visible revision has a coherent v2 + marker and no serving fallback is reported; +6. monitor semantic-parity probes, latency, errors, and timeouts. The reader's + checks retain the full scope-first map-unit path while only statistics are + incomplete, and use legacy FTS when index data is unusable. + +The handoff must document the exact command, batch/concurrency limits, +pause/resume procedure, application-version rollback procedure, and the final +check output. A partial +backfill is an expected intermediate state, not a failed deployment, provided +that incomplete statistics do not select token-only projection and unusable +indexes remain on the legacy reader. + +### Deferred: LLM orchestration improvements + +- Do not add a deterministic shortcut around Control. `plan_control` remains + the sole checklist-reconciliation authority; non-empty or seemingly related + evidence is not sufficient to replace its accept/widen/drop/replan decision. +- Do not change subgoal parallelism or Harvest/Control materialization in this + phase. These remain candidates only after the LLM boundary is explicitly + revisited. + +These changes must not alter prompt text or citation selection semantics. + +## Validation protocol + +Every implementation slice must be validated against a local copy of the +production namespace before deployment. Production rollout then relies on the +automatic legacy fallback for any incomplete or incompatible revision. + +Validation has two layers. First, maintain a small temporary retrieval-parity +set (two or three representative queries, including the current `心血管` trace) +containing current legacy-reader outputs, and run deterministic non-LLM parity +checks on an immutable production-data copy with fixed revision pins, section +scope, and queries. Compare loader rows, frequency maps, channel statistics, +per-unit scores, and ordering against that set. This set is a quality guard for +the optimization work, not a new product behavior contract; do not refresh it +just to hide a regression. Generate it once from the current legacy reader and +freeze it for the whole optimization series; refresh it only after an explicit +decision to accept a retrieval behavior change. Second, run end-to-end +Retrieval requests to measure latency, rows/bytes, and memory. Do not use +nondeterministic Planner/Harvest timing as evidence for this non-LLM change. + +1. Capture a baseline with the exact request, a cold process, and an empty or + generation-scoped Redis key. Run at least 10 cold repetitions; measure p50, + p95, stage timings, peak RSS, rows, and bytes transferred. +2. Run a separate warm-cache series. Never present warm-cache results as the + first-request improvement. +3. For every SQL change, compare `EXPLAIN (ANALYZE, BUFFERS)` and exact returned + rows against the current query using the read-only production database or + an immutable local copy. +4. Compare retrieval quality per request, requiring zero public-result + differences: selected chunk IDs, order, rounded scores, source sections, + evidence content/hash, asset references, router, stop reason, and citation + validity. Aggregate quality metrics cannot replace this exact parity check. +5. Exercise edge cases: empty result, one token, multiple tokens, small + corpus, incomplete index, archived revision, namespace generation change, + and Redis miss/failure. +6. Fix the malformed CPU/RSS log fields before making any CPU or memory claim. + Use per-stage CPU deltas and cgroup metrics rather than process high-water + RSS alone. +7. Roll out one optimization at a time, with automatic legacy fallback and + production p50/p95/error/timeout monitoring. If rollback is required, + redeploy the previous application version; leave additive schema/data in + place for a later retry. + +For each slice, retain it only when repeated local runs show a stable +improvement in the targeted stage and no regression in total non-LLM p95. The +temporary parity set is a hard gate regardless of the measured speedup. + +## Recommended implementation order + +1. Add missing structured timing and CPU/RSS instrumentation. +2. Benchmark the existing bounded `top_k + 1` scope probe; keep a replacement + only if it demonstrates a measured improvement. +3. Implement one token-selective serving-index slice: add the covering index, + persist per-channel BM25 statistics, and make unfiltered map-unit discovery + token-selective with the exact legacy fallback for filtered/incomplete + scopes. +4. Optimize the frequency join order and verify the production query plan. +5. Reconcile and optimize the unaccounted map-scoring work. +6. Consider snapshot allocation and Redis map-unit projection work as follow-up + improvements. diff --git a/docs/design/retrieval-serving-index-rollout-runbook.md b/docs/design/retrieval-serving-index-rollout-runbook.md new file mode 100644 index 00000000..56d2c3cc --- /dev/null +++ b/docs/design/retrieval-serving-index-rollout-runbook.md @@ -0,0 +1,404 @@ +# Retrieval Serving Index Rollout Runbook + +## Scope + +This runbook rolls out the retrieval serving-index performance changes without +changing retrieval semantics. It covers the additive schema migration, the +existing-data statistics backfill, readiness checks, retrieval parity probes, +monitoring, pause/resume, and application rollback. + +This rollout does not use a runtime feature flag. The reader selects the +optimized path only when every revision required by a request has a coherent +format-v2 index and all four channel statistics. Otherwise it uses the existing +legacy reader automatically. + +Planner, Harvest, Control, prompts, BM25 formulas, and citation-selection +semantics are outside this rollout. + +## Preconditions + +- Deploy only a build produced from the reviewed retrieval optimization branch. +- Confirm the target is the intended production environment before every + command. +- Use the API image for Alembic and maintenance commands. It contains: + - `/app/alembic` + - `/app/scripts/backfill_map_unit_statistics.py` + - `/app/scripts/backfill_map_unit_indexes.py` +- Confirm database backups and the normal application-version rollback path are + available. +- Record current retrieval p50/p95, errors, timeouts, and the output of the + temporary two-or-three-query quality set before deployment. +- Check free database disk space. The token-leading index is additive and the + existing indexes must remain in place. + +## Local acceptance evidence + +The production-shaped local restore used for final verification contained: + +- 2,996 documents; +- 125,835 sections; +- 170,280 chunks; +- 107,032 map units; +- 12,423,317 map-unit token rows; +- 644 active documents in the benchmark namespace. + +Final readiness on that namespace: + +```text +alembic head: c2d3e4f5a6b7 +covering index: idx_document_map_unit_tokens_token_lookup +statistics check: would_update=0 complete=644 skipped=0 documents=644 +coherent current indexes: 644/644 +``` + +Classic parity queries preserved chunk IDs, ordering, sources, and evidence. +The maximum observed score delta was below the accepted `1e-4` tolerance. +The v1 and v2 classic entry points returned the same chunk IDs and evidence +hash. A complete map-nav smoke returned `stop_reason=completed`, 23 result rows, +and 27 referenced chunks. + +These figures describe the local restored copy. They are not substitutes for +the production checks below. + +## Phase 1: Preflight inventory on the existing schema + +Before migration, record active/current revision counts using only the existing +schema. Do not run the statistics checker yet because it reads columns added by +revision `b1c2d3e4f5a6`. + +```sql +SELECT count(*) AS active_current_documents +FROM documents +WHERE status = 'active' + AND current_job_result_id IS NOT NULL; + +SELECT count(*) AS current_map_indexes +FROM documents +JOIN document_map_unit_indexes AS indexes + ON indexes.document_id = documents.document_id + AND indexes.job_result_id = documents.current_job_result_id +WHERE documents.status = 'active' + AND documents.current_job_result_id IS NOT NULL; +``` + +Record both counts and investigate any pre-existing difference. This inventory +does not determine format-v2 readiness. + +## Phase 2: Apply the additive migrations + +The production release workflow runs migrations before updating ECS services. +Record the workflow job URL and output. For a manual rehearsal, run Alembic from +the API image: + +```bash +cd /app +python -m alembic upgrade heads +python -m alembic current +``` + +The expected head for this rollout is: + +```text +c2d3e4f5a6b7 +``` + +The three relevant additive migrations are: + +1. `a0b1c2d3e4f5`: creates + `idx_document_map_unit_tokens_token_lookup` on + `(channel, token_hash, map_unit_id) INCLUDE (token, frequency)`; +2. `b1c2d3e4f5a6`: adds nullable path/content document-count and total-length + columns to `document_map_unit_indexes`; +3. `c2d3e4f5a6b7`: repairs the covering index if it is missing or PostgreSQL + reports `indisvalid=false` or `indisready=false` after an interrupted build. + +The index migration uses `CREATE INDEX CONCURRENTLY` in the normal Alembic +execution path. Monitor lock waits, database CPU, I/O, replication lag, and +free disk space while it runs. Do not drop either pre-existing map-unit-token +index during this rollout. + +Verify the schema: + +```sql +SELECT version_num FROM alembic_version; + +SELECT + classes.relname AS index_name, + indexes.indisvalid, + indexes.indisready, + pg_get_indexdef(indexes.indexrelid) AS index_definition +FROM pg_index AS indexes +JOIN pg_class AS classes ON classes.oid = indexes.indexrelid +WHERE classes.relname = 'idx_document_map_unit_tokens_token_lookup'; + +SELECT column_name +FROM information_schema.columns +WHERE table_name = 'document_map_unit_indexes' + AND column_name IN ( + 'path_document_count', + 'path_total_length', + 'content_document_count', + 'content_total_length' + ) +ORDER BY column_name; +``` + +The covering index must exist with `indisvalid=true`, `indisready=true`, and +the expected key/include columns. A failed concurrent build can leave an +invalid index. If either flag is false, stop the rollout, drop only that invalid +index with `DROP INDEX CONCURRENTLY`, and rerun the migration. + +## Phase 3: Deploy the application build + +Deploy the application after the additive migrations finish. New publications +will write coherent format-v2 statistics. Existing revisions with NULL channel +statistics remain on the full scope-first map-unit reader until maintenance +completes; missing, legacy, or unusable indexes remain on the legacy reader. + +Immediately verify: + +- API health checks pass; +- no migration or model-loading error appears in API logs; +- classic and map-nav requests still complete; +- incomplete-index warnings distinguish statistics-incomplete map-unit serving + from `fallback=legacy_fts`; neither case may return partial or empty results; +- no increase appears in retrieval errors or timeouts. + +## Phase 4: Backfill existing format-v2 indexes + +After the migration, run the statistics inventory: + +```bash +python /app/scripts/backfill_map_unit_statistics.py \ + --check \ + --batch-size 100 +``` + +Record the final summary line: + +```text +would_update= complete= skipped= documents= +``` + +Interpretation: + +- `complete`: already coherent format-v2 revisions; +- `would_update`: existing format-v2 indexes whose four statistics need an + in-place update; +- `skipped`: missing or legacy indexes that require the full index backfill. + +Run exactly one maintenance process against the database. `--batch-size` is a +serial session grouping, not a concurrency setting. + +First rehearse one namespace or document: + +```bash +python /app/scripts/backfill_map_unit_statistics.py \ + --apply \ + --batch-size 100 \ + --user-id \ + --namespace +``` + +Then run the full statistics backfill: + +```bash +python /app/scripts/backfill_map_unit_statistics.py \ + --apply \ + --batch-size 100 +``` + +The command: + +- reads existing `document_map_units`; +- computes positive-length document counts and total token lengths separately + for the path and content channels; +- commits each revision independently; +- does not regenerate token rows, manifests, or namespace snapshots; +- is idempotent and safe to restart. + +If `skipped` is non-zero, list and process those documents separately with the +existing full index command: + +```bash +python /app/scripts/backfill_map_unit_indexes.py \ + --apply \ + --document-id +``` + +Do not use `--tokens-only` unless investigation proves that only map-unit token +data is missing and the serving manifest and namespace snapshot are already +coherent. + +## Phase 5: Final readiness gate + +Repeat the read-only check until it exits successfully and reports: + +```text +would_update=0 skipped=0 complete= documents= +``` + +```bash +python /app/scripts/backfill_map_unit_statistics.py \ + --check \ + --batch-size 100 +``` + +Also verify that all retrieval-visible active revisions are coherent: + +```sql +SELECT + count(*) FILTER ( + WHERE indexes.format_version = 2 + AND indexes.path_document_count IS NOT NULL + AND indexes.path_total_length IS NOT NULL + AND indexes.content_document_count IS NOT NULL + AND indexes.content_total_length IS NOT NULL + ) AS ready_revisions, + count(*) AS current_revisions +FROM documents +LEFT JOIN document_map_unit_indexes AS indexes + ON indexes.document_id = documents.document_id + AND indexes.job_result_id = documents.current_job_result_id +WHERE documents.status = 'active' + AND documents.current_job_result_id IS NOT NULL; +``` + +With the `LEFT JOIN`, `current_revisions` includes active documents with no +index. `ready_revisions` must equal `current_revisions`. Also require: + +```sql +SELECT count(*) AS missing_index_rows +FROM documents +LEFT JOIN document_map_unit_indexes AS indexes + ON indexes.document_id = documents.document_id + AND indexes.job_result_id = documents.current_job_result_id +WHERE documents.status = 'active' + AND documents.current_job_result_id IS NOT NULL + AND indexes.id IS NULL; +``` + +`missing_index_rows` must be zero. + +Run the complete serving-index readiness checker: + +```bash +python /app/scripts/backfill_map_unit_indexes.py --check +``` + +For every namespace, require `status=READY` and inspect the report rather than +only its exit code. Require: + +```text +missing_from_snapshot=0 +missing_map_index=0 +missing_revision_manifest=0 +``` + +`suspicious_zero_idf` is diagnostic only. Zero average IDF is valid for some +small corpora, including a two-unit corpus where each token occurs in exactly +one unit; it must not independently block readiness. + +## Phase 6: Retrieval-quality gate + +Run the frozen temporary quality set. Use the same two or three queries, +namespace, top-k, filters, and revision generation captured before deployment. + +For classic retrieval, require no change in: + +- router; +- selected chunk IDs and order; +- source document and section; +- evidence content/hash; +- asset references; +- rounded scores, with an absolute tolerance of `1e-4`. + +Exercise at least: + +1. v1 `use_agentic=false`; +2. v2 `use_agentic=false` with equivalent retrieval fields; +3. one `use_agentic=true` map-nav smoke; +4. one request with `use_agentic` omitted, confirming it routes to map-nav; +5. one filtered request, confirming filtered-scope semantics and the safe + fallback where required. + +Map-nav LLM output is nondeterministic. For production smoke, require successful +completion, valid citations, expected namespace isolation, and relevant +evidence. Do not require byte-identical ordering between independent Planner +runs. Deterministic map-score parity remains covered by the contract suite. + +Stop the rollout if classic result parity fails. Do not refresh the baseline to +hide a difference. + +## Phase 7: Performance and reliability observation + +Monitor at least one normal traffic window after the backfill. Classic public +result parity must have zero differences, total non-LLM p95 must not exceed the +recorded baseline, and retrieval error/timeout rates must not regress. + +- retrieval request p50/p95 and maximum latency, separated by `router_used`; +- classic `search.map_unit_discovery` stages: units, frequencies, indexes, + statistics, scoring, and hydration; +- map-nav snapshot, episode, and hydration stages; +- PostgreSQL statement timeouts, lock waits, CPU, I/O, and connection usage; +- Redis errors and namespace snapshot cache misses; +- retrieval errors, incomplete-index fallbacks, and response timeouts; +- process CPU and maximum RSS from the corrected map-nav resource log. + +Do not mix classic and map-nav latency distributions. Do not treat Redis-warm +snapshot measurements as cold-request performance. + +## Pause and resume + +To pause, stop launching new maintenance command processes. Wait for the +current revision to finish if database health permits, then send `SIGINT` or +`SIGTERM` to the one-off task. Each completed revision has already committed; +the interrupted transaction rolls back and remains on the safe reader path. + +To resume, rerun the same `--apply` command. Completed revisions are detected +and skipped. Follow it with `--check` and retain both summary outputs in the +deployment record. + +## Rollback + +If application errors, timeouts, or quality regressions occur: + +1. stop the backfill process; +2. redeploy the previous application version; +3. verify classic and map-nav requests using the frozen quality set; +4. retain the additive columns, index, and already-computed statistics unless + database health specifically requires their removal. + +Application rollback is sufficient because the previous application ignores +the additive schema. Avoid running Alembic downgrade during an incident: a +concurrent index drop or table alteration adds operational risk and is not +required to restore the previous behavior. + +## Completion record + +Attach the following to the deployment ticket: + +- deployed application image digest and Git commit; +- pre-migration active/current inventory; +- post-migration statistics `--check` output; +- Alembic `current` output; +- covering-index and column verification output; +- rehearsal and full `--apply` summaries; +- final `--check` output; +- ready/current revision counts; +- frozen-query parity results; +- classic and map-nav latency summaries; +- observed fallback, error, and timeout counts; +- rollback decision or explicit confirmation that rollback was not required. + +## ECS one-off task requirements + +Run maintenance as a one-off task created from the newly registered production +API task definition. Reuse its task role, execution role, VPC subnets, security +groups, Secrets Manager injection, and CloudWatch log configuration. Override +only the container command with one of the commands in this runbook. + +Do not execute maintenance inside a long-lived API task. Do not copy, export, or +place the production database URL in the ECS command, shell history, workflow +input, deployment ticket, or logs. The one-off task must receive it through the +same production secret as the API task. diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index 3d630b91..6a46fdd2 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -328,6 +328,13 @@ class DocumentMapUnitToken(Base): "token_hash", "map_unit_id", ), + Index( + "idx_document_map_unit_tokens_token_lookup", + "channel", + "token_hash", + "map_unit_id", + postgresql_include=["token", "frequency"], + ), Index( "idx_document_map_unit_tokens_unit_lookup", "map_unit_id", @@ -340,7 +347,7 @@ class DocumentMapUnitToken(Base): class DocumentMapUnitIndex(Base): - """Completeness marker for a revision's materialized map-unit index.""" + """Completeness marker and corpus statistics for a materialized index.""" __tablename__ = "document_map_unit_indexes" @@ -362,6 +369,16 @@ class DocumentMapUnitIndex(Base): average_idf_content: Mapped[float] = mapped_column( Float, nullable=False, default=0.0 ) + path_document_count: Mapped[Optional[int]] = mapped_column( + Integer, nullable=True + ) + path_total_length: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) + content_document_count: Mapped[Optional[int]] = mapped_column( + Integer, nullable=True + ) + content_total_length: Mapped[Optional[int]] = mapped_column( + Integer, nullable=True + ) created_at: Mapped[datetime] = mapped_column( DateTime, default=utc_now_naive, nullable=False ) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 6e0bee4a..6ed99293 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -364,8 +364,8 @@ async def _run_mapnav_route( completion_detail = f"chunks | evidence={len(evidence_text)} chars | router=mapnav" process_finished = resource.getrusage(resource.RUSAGE_SELF) logger.info( - "retrieval mapnav stage=process_resources cpu_seconds=%.3f " - "process_max_rss_kb=%d", + "retrieval mapnav stage=process_resources cpu_seconds={:.3f} " + "process_max_rss_kb={}", (process_finished.ru_utime + process_finished.ru_stime) - (process_started.ru_utime + process_started.ru_stime), int(process_finished.ru_maxrss), diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py index ee18434d..ab7a0523 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -88,6 +88,10 @@ def replace_document_map_units( token_count = 0 path_unit_df: Counter[str] = Counter() content_unit_df: Counter[str] = Counter() + path_document_count: int = 0 + path_total_length: int = 0 + content_document_count: int = 0 + content_total_length: int = 0 for sort_order, unit in enumerate(score_units): unit_id = str(unit.get("chunk_id") or "").strip() section_id = str(unit.get("section_id") or "").strip() @@ -96,6 +100,12 @@ def replace_document_map_units( map_unit_id = f"dmu_{uuid4().hex}" path_tokens = str(unit.get("path_search_text") or "").split() content_tokens = str(unit.get("content_search_text") or "").split() + if path_tokens: + path_document_count += 1 + path_total_length += len(path_tokens) + if content_tokens: + content_document_count += 1 + content_total_length += len(content_tokens) path_unit_df.update(set(path_tokens)) content_unit_df.update(set(content_tokens)) # ``provider.self_units`` already reflects root-asset remount (assets @@ -152,6 +162,10 @@ def replace_document_map_units( unit_count=persisted_count, token_document_frequency=content_unit_df, ), + path_document_count=path_document_count, + path_total_length=path_total_length, + content_document_count=content_document_count, + content_total_length=content_total_length, ) ) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py index 81179b25..cd569631 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_knowhere.py @@ -335,7 +335,9 @@ def load_persisted_score_corpus( cur.execute( "SELECT indexes.document_id, indexes.job_result_id, " "indexes.format_version, indexes.unit_count, " - "indexes.average_idf_path, indexes.average_idf_content " + "indexes.average_idf_path, indexes.average_idf_content, " + "indexes.path_document_count, indexes.path_total_length, " + "indexes.content_document_count, indexes.content_total_length " "FROM document_map_unit_indexes AS indexes " f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " "ON indexes.document_id = revisions.document_id " @@ -353,7 +355,8 @@ def load_persisted_score_corpus( len(index_rows), ) if len(index_rows) != len(revisions) or any( - len(row) < 6 or int(row[2]) != MAP_UNIT_INDEX_FORMAT_VERSION + len(row) < 10 or int(row[2]) != MAP_UNIT_INDEX_FORMAT_VERSION + or any(value is None for value in row[6:10]) for row in index_rows ): return None @@ -386,8 +389,67 @@ def load_persisted_score_corpus( for document_id, section_ids in allowed_by_document.items() for section_id in section_ids } + cur.execute( + "SELECT sections.document_id, sections.job_result_id, count(*) " + "FROM document_sections AS sections " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON sections.document_id = revisions.document_id " + "AND sections.job_result_id = revisions.job_result_id " + "GROUP BY sections.document_id, sections.job_result_id", + revision_params, + ) + section_counts = { + (str(document_id), str(job_result_id)): int(count) + for document_id, job_result_id, count in cur.fetchall() + } + has_complete_section_scope = all( + len(allowed_by_document.get(document_id, set())) + == section_counts.get((document_id, job_result_id), 0) + for document_id, job_result_id in revisions + ) all_unit_rows = self._score_unit_rows_cache.get(revision_key) - if all_unit_rows is None: + if has_complete_section_scope and query_token_hashes: + stage_started = time.perf_counter() + cur.execute( + "WITH matching_tokens AS MATERIALIZED (" + "SELECT DISTINCT map_unit_id FROM document_map_unit_tokens " + "WHERE channel = ANY(%s) AND token_hash = ANY(%s)" + "), scoped_units AS MATERIALIZED (" + f"SELECT units.id, units.document_id, units.unit_id, units.section_id, " + "units.path_token_count, units.content_token_count " + "FROM document_map_units AS units " + f"JOIN (VALUES {values_sql}) AS revisions(document_id, job_result_id) " + "ON units.document_id = revisions.document_id " + "AND units.job_result_id = revisions.job_result_id" + ") SELECT scoped_units.id, scoped_units.document_id, " + "scoped_units.unit_id, scoped_units.section_id, " + "scoped_units.path_token_count, scoped_units.content_token_count " + "FROM matching_tokens JOIN scoped_units " + "ON scoped_units.id = matching_tokens.map_unit_id", + [ + list(_MAP_SCORE_CHANNELS), + list(query_token_hashes), + *revision_params, + ], + ) + unit_rows = [ + { + "map_unit_id": str(row[0]), + "document_id": str(row[1]), + "unit_id": str(row[2]), + "section_id": str(row[3] or ""), + "path_token_count": int(row[4] or 0), + "content_token_count": int(row[5] or 0), + } + for row in cur.fetchall() + ] + _logger.info( + "retrieval map-index load stage=units-selective seconds=%.3f rows=%d", + time.perf_counter() - stage_started, + len(unit_rows), + ) + all_unit_rows = None + elif all_unit_rows is None: stage_started = time.perf_counter() cur.execute( "SELECT units.id, units.document_id, units.unit_id, " @@ -427,11 +489,13 @@ def load_persisted_score_corpus( True, ) - unit_rows = [ - row - for row in all_unit_rows - if (str(row["document_id"]), str(row["section_id"])) in allowed_pairs_set - ] + if not has_complete_section_scope or not query_token_hashes: + unit_rows = [ + row + for row in all_unit_rows or [] + if (str(row["document_id"]), str(row["section_id"])) + in allowed_pairs_set + ] frequencies: Dict[Tuple[str, str], Dict[str, int]] = {} if unit_rows and query_tokens: stage_started = time.perf_counter() @@ -482,6 +546,16 @@ def load_persisted_score_corpus( query_tokens=query_tokens, frequencies=frequencies, average_idf=average_idf_path, + document_count_override=( + sum(int(row[6] or 0) for row in index_rows) + if has_complete_section_scope and query_token_hashes + else None + ), + total_length_override=( + sum(int(row[7] or 0) for row in index_rows) + if has_complete_section_scope and query_token_hashes + else None + ), ) content_stats = build_channel_bm25_stats( unit_rows=unit_rows, @@ -491,6 +565,16 @@ def load_persisted_score_corpus( query_tokens=query_tokens, frequencies=frequencies, average_idf=average_idf_content, + document_count_override=( + sum(int(row[8] or 0) for row in index_rows) + if has_complete_section_scope and query_token_hashes + else None + ), + total_length_override=( + sum(int(row[9] or 0) for row in index_rows) + if has_complete_section_scope and query_token_hashes + else None + ), ) return PersistedScoreCorpus( units=[ diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py index c3637464..eb207ce3 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_map_scores.py @@ -20,6 +20,22 @@ _logger = logging.getLogger(__name__) +def _count_tree_shape( + tree_by_doc: Dict[ + str, + Tuple[Dict[str, List[str]], Set[str], Dict[str, str]], + ], +) -> Tuple[int, int, int]: + """Return reachable section nodes, parent-child edges, and leaves.""" + section_nodes: int = sum(len(value[0]) for value in tree_by_doc.values()) + section_edges: int = sum( + sum(len(children) for children in value[0].values()) + for value in tree_by_doc.values() + ) + leaf_sections: int = sum(len(value[1]) for value in tree_by_doc.values()) + return section_nodes, section_edges, leaf_sections + + def _build_legacy_score_corpus(ts: Any, doc_ids: Sequence[str]) -> PersistedScoreCorpus: """Build the retired in-memory scorer input when persisted indexes are absent.""" raw_units: List[dict] = [] @@ -414,11 +430,15 @@ def compute_corpus_map_and_unit_scores_many( cached = _walk_tree(ts, doc_id, root_ids) tree_cache[doc_id] = cached tree_by_doc[doc_id] = cached + section_nodes, section_edges, leaf_sections = _count_tree_shape(tree_by_doc) _logger.info( - "retrieval mapnav phase=tree_build seconds=%.3f documents=%d sections=%d", + "retrieval mapnav phase=tree_build seconds=%.3f documents=%d " + "section_nodes=%d section_edges=%d leaf_sections=%d", time.perf_counter() - tree_started, len(valid_doc_ids), - sum(len(value[0]) for value in tree_by_doc.values()), + section_nodes, + section_edges, + leaf_sections, ) persisted_loader = getattr(ts, "load_persisted_score_corpus", None) @@ -471,10 +491,13 @@ def compute_corpus_map_and_unit_scores_many( map_scores[doc_id] = doc_max results[query] = (map_scores, unit_scores) _logger.info( - "retrieval mapnav phase=map_pooling seconds=%.3f documents=%d sections=%d", + "retrieval mapnav phase=map_pooling seconds=%.3f documents=%d " + "section_nodes=%d section_edges=%d leaf_sections=%d", time.perf_counter() - pooling_started, len(valid_doc_ids), - sum(len(value[0]) for value in tree_by_doc.values()), + section_nodes, + section_edges, + leaf_sections, ) return results diff --git a/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py b/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py index 9ce10c2e..712f23b5 100644 --- a/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py +++ b/packages/shared-python/shared/services/retrieval/nav/persisted_score_load.py @@ -47,12 +47,19 @@ def build_channel_bm25_stats( query_tokens: Sequence[str], frequencies: Mapping[tuple[str, str], Mapping[str, int]], average_idf: float, + document_count_override: int | None = None, + total_length_override: int | None = None, ) -> PersistedBm25Stats: """Build channel stats from already-fetched unit rows and query-token freqs.""" lengths = [ int(row[length_field]) for row in unit_rows if int(row[length_field]) > 0 ] - document_count = len(lengths) + document_count = ( + len(lengths) if document_count_override is None else document_count_override + ) + total_length = ( + sum(lengths) if total_length_override is None else total_length_override + ) document_frequency = { token: sum( 1 @@ -66,7 +73,7 @@ def build_channel_bm25_stats( } return PersistedBm25Stats( document_count=document_count, - total_length=sum(lengths), + total_length=total_length, document_frequency=document_frequency, average_idf=float(average_idf), ) diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index b15425e3..9f92617c 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -86,6 +86,26 @@ def _is_snapshot_timing_enabled() -> bool: } +def _count_section_shape( + sections_by_doc: Mapping[str, list[SectionRow]], +) -> tuple[int, int, int, int]: + """Return section rows, distinct paths, roots, and leaves for diagnostics.""" + rows: list[SectionRow] = [ + row for document_rows in sections_by_doc.values() for row in document_rows + ] + parent_ids: set[str] = { + str(row.parent_section_id).strip() + for row in rows + if str(row.parent_section_id or "").strip() + } + return ( + len(rows), + len({str(row.section_path) for row in rows}), + sum(1 for row in rows if not str(row.parent_section_id or "").strip()), + sum(1 for row in rows if str(row.section_id) not in parent_ids), + ) + + class SnapshotSession(Protocol): """Minimal database interface required by the snapshot loader.""" @@ -369,17 +389,25 @@ async def load_nav_snapshot( document_revisions=dict(revisions), ) if _is_snapshot_timing_enabled(): + section_rows, section_paths, root_sections, leaf_sections = ( + _count_section_shape(sections_by_doc) + ) _logger.info( "retrieval snapshot timing total_seconds=%.3f parse_seconds=%.3f " "ref_index_copy_seconds=%.3f doc_query_seconds=%.3f " - "job_query_seconds=%.3f documents=%d sections=%d refs=%d mode=lazy", + "job_query_seconds=%.3f documents=%d section_rows=%d " + "section_paths=%d root_sections=%d leaf_sections=%d refs=%d " + "mode=lazy", time.perf_counter() - snapshot_started, parse_seconds, ref_index_seconds, doc_query_seconds, job_query_seconds, len(snapshot.document_ids), - sum(len(rows) for rows in sections_by_doc.values()), + section_rows, + section_paths, + root_sections, + leaf_sections, len(snapshot.chunk_ref_index), ) return snapshot @@ -414,16 +442,23 @@ async def load_nav_snapshot( }, ) if _is_snapshot_timing_enabled(): + section_rows, section_paths, root_sections, leaf_sections = ( + _count_section_shape(sections_by_doc) + ) _logger.info( "retrieval snapshot timing total_seconds=%.3f parse_seconds=%.3f " "doc_query_seconds=%.3f job_query_seconds=%.3f documents=%d " - "sections=%d refs=%d mode=eager", + "section_rows=%d section_paths=%d root_sections=%d " + "leaf_sections=%d refs=%d mode=eager", time.perf_counter() - snapshot_started, parse_seconds, doc_query_seconds, job_query_seconds, len(snapshot.document_ids), - sum(len(rows) for rows in sections_by_doc.values()), + section_rows, + section_paths, + root_sections, + leaf_sections, len(snapshot.chunk_ref_index), ) return snapshot diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index 4e7d1a4a..68e27402 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -20,7 +20,7 @@ from collections.abc import Mapping from dataclasses import dataclass, field from hashlib import sha256 -from typing import Any +from typing import Any, cast from loguru import logger from sqlalchemy import text @@ -193,6 +193,8 @@ async def map_unit_discovery( query_tokens = tokenize_query_for_ranker(query) if not query_tokens: return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + if revision_pins is not None and not revision_pins: + return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) query_token_hashes = [ sha256(token.encode("utf-8")).hexdigest() for token in query_tokens ] @@ -211,6 +213,15 @@ async def map_unit_discovery( params.update(type_params) params.update(signal_params) + is_unfiltered_scope: bool = not any( + ( + chunk_types, + signal_paths, + exclude_sections, + exclude_document_ids, + ) + ) + cte = _SCOPED_UNITS_CTE.format( revision_join=revision_join, revision_clause=revision_clause, @@ -218,7 +229,25 @@ async def map_unit_discovery( type_clause=type_clause, signal_clause=signal_clause, ) - unit_result = await db.execute(text(cte + "SELECT * FROM scoped_units"), params) + unit_statement = cte + "SELECT * FROM scoped_units" + if is_unfiltered_scope: + unit_statement = ( + "WITH matching_tokens AS MATERIALIZED (" + "SELECT DISTINCT map_unit_id FROM document_map_unit_tokens " + "WHERE channel = ANY(:channels) AND token_hash = ANY(:token_hashes)" + "), " + + cte.lstrip().removeprefix("WITH ") + + " SELECT DISTINCT scoped_units.* " + "FROM matching_tokens JOIN scoped_units " + "ON scoped_units.map_unit_id = matching_tokens.map_unit_id" + ) + params = { + **params, + "channels": list(_MAP_SCORE_CHANNELS), + "token_hashes": query_token_hashes, + } + stage_started = time.monotonic() + unit_result = await db.execute(text(unit_statement), params) unit_rows = [dict(row._mapping) for row in unit_result.all()] unit_rows = [ row @@ -229,9 +258,11 @@ async def map_unit_discovery( exclude_sections=exclude_sections, ) ] - - if not unit_rows: - return DiscoveryResult(status="discovery_done", payload={"fused_rows": []}) + logger.info( + "retrieval map-unit stage=units seconds={:.3f} rows={}", + time.monotonic() - stage_started, + len(unit_rows), + ) frequency_scope_cte = cte if signal_paths else _SCOPED_UNIT_IDS_CTE.format( revision_join=revision_join, @@ -255,6 +286,7 @@ async def map_unit_discovery( ON scoped_units.map_unit_id = matching_tokens.map_unit_id """ ) + stage_started = time.monotonic() frequency_result = await db.execute( frequency_query, { @@ -268,9 +300,37 @@ async def map_unit_discovery( frequencies.setdefault((str(map_unit_id), str(channel)), {})[str(token)] = ( int(frequency) ) + logger.info( + "retrieval map-unit stage=frequencies seconds={:.3f} rows={} units={}", + time.monotonic() - stage_started, + sum(len(values) for values in frequencies.values()), + len(frequencies), + ) - index_result = await db.execute( - text( + if is_unfiltered_scope and revision_pins is not None: + pinned_pairs = [ + (str(document_id).strip(), str(job_result_id).strip()) + for document_id, job_result_id in revision_pins.items() + if str(document_id).strip() and str(job_result_id).strip() + ] + pinned_values_sql = ", ".join( + f"(:_pin_document_{index}, :_pin_revision_{index})" + for index, _pair in enumerate(pinned_pairs) + ) + index_statement = f""" + SELECT indexes.average_idf_path, indexes.average_idf_content, + indexes.unit_count, indexes.format_version, + indexes.path_document_count, indexes.path_total_length, + indexes.content_document_count, indexes.content_total_length, + indexes.token_count + FROM document_map_unit_indexes AS indexes + JOIN (VALUES {pinned_values_sql}) + AS scoped_revisions(document_id, job_result_id) + ON indexes.document_id = scoped_revisions.document_id + AND indexes.job_result_id = scoped_revisions.job_result_id + """ + else: + index_statement = ( ( cte if signal_paths @@ -283,7 +343,10 @@ async def map_unit_discovery( ) + """ SELECT indexes.average_idf_path, indexes.average_idf_content, - indexes.unit_count, indexes.format_version + indexes.unit_count, indexes.format_version, + indexes.path_document_count, indexes.path_total_length, + indexes.content_document_count, indexes.content_total_length, + indexes.token_count FROM document_map_unit_indexes AS indexes JOIN ( SELECT DISTINCT document_id, job_result_id FROM scoped_units @@ -291,32 +354,150 @@ async def map_unit_discovery( ON indexes.document_id = scoped_revisions.document_id AND indexes.job_result_id = scoped_revisions.job_result_id """ - ), - params, - ) - # Only count indexes written with the current tokenizer (v2 = word-level). - # Stale char-level (v1) rows look complete by unit_count but cannot match - # word-level query hashes — treat them as missing so readiness fails closed. + ) + stage_started = time.monotonic() + index_result = await db.execute(text(index_statement), params) index_parts = [ - (float(path_idf or 0.0), float(content_idf or 0.0), int(unit_count or 0)) - for path_idf, content_idf, unit_count, format_version in index_result.all() - if int(format_version or 0) == MAP_UNIT_INDEX_FORMAT_VERSION - ] - expected_revisions = { - (str(row["document_id"]), str(row["job_result_id"])) for row in unit_rows - } - unfiltered_scope = not any( ( - chunk_types, - signal_paths, - exclude_sections, - exclude_document_ids, + float(path_idf or 0.0), + float(content_idf or 0.0), + int(unit_count or 0), + int(format_version or 0), + path_document_count, + path_total_length, + content_document_count, + content_total_length, + int(token_count or 0), ) + for ( + path_idf, + content_idf, + unit_count, + format_version, + path_document_count, + path_total_length, + content_document_count, + content_total_length, + token_count, + ) in index_result.all() + ] + logger.info( + "retrieval map-unit stage=indexes seconds={:.3f} rows={}", + time.monotonic() - stage_started, + len(index_parts), + ) + if is_unfiltered_scope: + if revision_pins is not None: + # The route captures the active revision set before discovery. Reuse + # that immutable set instead of scanning scoped map units a second + # time just to derive the same revision keys. Missing index rows + # still fail the existing completeness check below. + expected_revisions = { + (str(document_id), str(job_result_id)) + for document_id, job_result_id in revision_pins.items() + } + revision_check_seconds = 0.0 + else: + stage_started = time.monotonic() + revision_result = await db.execute( + text(cte + "SELECT DISTINCT document_id, job_result_id FROM scoped_units"), + params, + ) + expected_revisions = { + (str(row[0]), str(row[1])) for row in revision_result.all() + } + revision_check_seconds = time.monotonic() - stage_started + logger.info( + "retrieval map-unit stage=revision-check seconds={:.3f} rows={}", + revision_check_seconds, + len(expected_revisions), + ) + else: + expected_revisions = { + (str(row["document_id"]), str(row["job_result_id"])) + for row in unit_rows + } + indexed_unit_count = sum( + unit_count for _path_idf, _content_idf, unit_count, *_rest in index_parts + ) + has_index_storage_mismatch: bool = False + if is_unfiltered_scope and not unit_rows and expected_revisions: + storage_counts: tuple[int | None, int | None] = cast( + tuple[int | None, int | None], + ( + await db.execute( + text( + _SCOPED_UNIT_IDS_CTE.format( + revision_join=revision_join, + revision_clause=revision_clause, + exclude_clause=exclude_clause, + type_clause=type_clause, + ) + + """ + SELECT + (SELECT COUNT(*) FROM scoped_units) AS unit_count, + ( + SELECT COUNT(*) + FROM document_map_unit_tokens AS tokens + JOIN scoped_units + ON scoped_units.map_unit_id = tokens.map_unit_id + ) AS token_count + """ + ), + params, + ) + ).one(), + ) + actual_unit_count: int | None = storage_counts[0] + actual_token_count: int | None = storage_counts[1] + indexed_token_count: int = sum(int(part[8]) for part in index_parts) + has_index_storage_mismatch = indexed_unit_count != int( + actual_unit_count or 0 + ) or indexed_token_count != int(actual_token_count or 0) + has_index_unit_count_mismatch = ( + indexed_unit_count < len(unit_rows) + if is_unfiltered_scope + else indexed_unit_count != len(unit_rows) + ) + is_index_format_incompatible = any( + format_version != MAP_UNIT_INDEX_FORMAT_VERSION + for ( + _path_idf, + _content_idf, + _unit_count, + format_version, + _path_document_count, + _path_total_length, + _content_document_count, + _content_total_length, + _token_count, + ) in index_parts + ) + has_incomplete_index_statistics = is_unfiltered_scope and any( + format_version != MAP_UNIT_INDEX_FORMAT_VERSION + or path_document_count is None + or path_total_length is None + or content_document_count is None + or content_total_length is None + for ( + _path_idf, + _content_idf, + _unit_count, + format_version, + path_document_count, + path_total_length, + content_document_count, + content_total_length, + _token_count, + ) in index_parts + ) + has_unusable_index = ( + len(index_parts) != len(expected_revisions) + or has_index_unit_count_mismatch + or has_index_storage_mismatch + or is_index_format_incompatible ) - index_unit_count_mismatch = unfiltered_scope and sum( - unit_count for _path_idf, _content_idf, unit_count in index_parts - ) != len(unit_rows) - if len(index_parts) != len(expected_revisions) or index_unit_count_mismatch: + if has_unusable_index: try: await record_retrieval_index_readiness( user_id=user_id, @@ -348,11 +529,46 @@ async def map_unit_discovery( filter_mode=filter_mode, revision_pins=revision_pins, ) + if has_incomplete_index_statistics: + logger.warning( + "retrieval map index statistics incomplete user_id=%s namespace=%s " + "using row-derived BM25 denominators", + user_id, + namespace, + ) + # Token-selective projection is only sufficient when persisted channel + # denominators are available. Before the statistics backfill, restore + # the old full-scope unit projection so row-derived BM25 statistics use + # the same corpus as the legacy map-unit reader. + stage_started = time.monotonic() + full_unit_params = { + key: value + for key, value in params.items() + if key not in {"channels", "token_hashes"} + } + full_unit_result = await db.execute( + text(cte + "SELECT * FROM scoped_units"), full_unit_params + ) + unit_rows = [dict(row._mapping) for row in full_unit_result.all()] + unit_rows = [ + row + for row in unit_rows + if not is_excluded_section( + document_id=row.get("document_id"), + section_path=row.get("section_path"), + exclude_sections=exclude_sections, + ) + ] + logger.info( + "retrieval map-unit stage=full-units-for-stats seconds={:.3f} rows={}", + time.monotonic() - stage_started, + len(unit_rows), + ) try: await record_retrieval_index_readiness( user_id=user_id, namespace=namespace, - ready=True, + ready=not has_incomplete_index_statistics, expected_revisions=len(expected_revisions), indexed_revisions=len(index_parts), ) @@ -361,16 +577,17 @@ async def map_unit_discovery( average_idf_path = combine_average_idf( [ (path_idf, unit_count) - for path_idf, _content_idf, unit_count in index_parts + for path_idf, _content_idf, unit_count, *_rest in index_parts ] ) average_idf_content = combine_average_idf( [ (content_idf, unit_count) - for _path_idf, content_idf, unit_count in index_parts + for _path_idf, content_idf, unit_count, *_rest in index_parts ] ) + stage_started = time.monotonic() path_stats = build_channel_bm25_stats( unit_rows=unit_rows, map_unit_id_field="map_unit_id", @@ -379,6 +596,16 @@ async def map_unit_discovery( query_tokens=query_tokens, frequencies=frequencies, average_idf=average_idf_path, + document_count_override=( + sum(int(part[4] or 0) for part in index_parts) + if is_unfiltered_scope and not has_incomplete_index_statistics + else None + ), + total_length_override=( + sum(int(part[5] or 0) for part in index_parts) + if is_unfiltered_scope and not has_incomplete_index_statistics + else None + ), ) content_stats = build_channel_bm25_stats( unit_rows=unit_rows, @@ -388,6 +615,21 @@ async def map_unit_discovery( query_tokens=query_tokens, frequencies=frequencies, average_idf=average_idf_content, + document_count_override=( + sum(int(part[6] or 0) for part in index_parts) + if is_unfiltered_scope and not has_incomplete_index_statistics + else None + ), + total_length_override=( + sum(int(part[7] or 0) for part in index_parts) + if is_unfiltered_scope and not has_incomplete_index_statistics + else None + ), + ) + logger.info( + "retrieval map-unit stage=stats seconds={:.3f} units={}", + time.monotonic() - stage_started, + len(unit_rows), ) corpus = PersistedScoreCorpus( @@ -406,7 +648,13 @@ async def map_unit_discovery( path_stats=path_stats, content_stats=content_stats, ) + stage_started = time.monotonic() scores_by_unit = score_persisted_corpus_many(corpus, [query]).get(query, {}) + logger.info( + "retrieval map-unit stage=scoring seconds={:.3f} units={}", + time.monotonic() - stage_started, + len(scores_by_unit), + ) rows_by_unit_id = {row["map_unit_id"]: row for row in unit_rows} ranked_unit_ids = sorted( @@ -415,6 +663,7 @@ async def map_unit_discovery( reverse=True, )[:top_k] + stage_started = time.monotonic() fused_rows = await _hydrate_winning_units( db, ranked_unit_ids=ranked_unit_ids, @@ -425,6 +674,11 @@ async def map_unit_discovery( exclude_sections=exclude_sections, revision_pins=revision_pins, ) + logger.info( + "retrieval map-unit stage=hydration seconds={:.3f} rows={}", + time.monotonic() - stage_started, + len(fused_rows), + ) if fused_rows: normalize_row_scores( fused_rows,