From cc67a90f88afb58a084ca2df95329b2dad52cd91 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 09:55:09 +0800 Subject: [PATCH 1/2] perf: bound retrieval snapshot payload reads --- .../services/retrieval/execution/routes.py | 31 +++------- .../shared/services/retrieval/nav_snapshot.py | 61 +++++++++++-------- .../retrieval/search/scoped_corpus.py | 13 +++- 3 files changed, 56 insertions(+), 49 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index d2b6d211..0998a8df 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -23,7 +23,6 @@ ) from shared.services.retrieval.search.ranking import rank_retrieval_candidates from shared.services.retrieval.search.scoped_corpus import ( - count_manifest_chunks, count_scoped_chunks, load_all_scoped_chunks, ) @@ -58,27 +57,15 @@ async def _try_run_small_corpus_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome | None: total_chunk_count: int | None = None - if ( - context.use_agentic is not False - and context.revision_pins is not None - and not context.exclude_document_ids - and not context.exclude_sections - and context.allowed_chunk_types is None - and not context.signal_paths - ): - total_chunk_count = await count_manifest_chunks( - context.db, - revision_pins=context.revision_pins, - ) - if total_chunk_count is None: - total_chunk_count = await count_scoped_chunks( - context.db, - user_id=context.user_id, - namespace=context.namespace, - exclude_document_ids=context.exclude_document_ids, - allowed_chunk_types=context.allowed_chunk_types, - revision_pins=context.revision_pins, - ) + total_chunk_count = await count_scoped_chunks( + context.db, + user_id=context.user_id, + namespace=context.namespace, + exclude_document_ids=context.exclude_document_ids, + allowed_chunk_types=context.allowed_chunk_types, + revision_pins=context.revision_pins, + max_count=context.top_k + 1, + ) logger.info(f"\n Total chunks in scope: {total_chunk_count}") if total_chunk_count > context.top_k: diff --git a/packages/shared-python/shared/services/retrieval/nav_snapshot.py b/packages/shared-python/shared/services/retrieval/nav_snapshot.py index c68d6d8c..5d34e451 100644 --- a/packages/shared-python/shared/services/retrieval/nav_snapshot.py +++ b/packages/shared-python/shared/services/retrieval/nav_snapshot.py @@ -55,6 +55,12 @@ # Keep revision predicates bounded while reducing round trips for large # namespaces. Keyset paging still caps each payload query at 10,000 rows. _REVISION_GROUP_SIZE = 64 +# Compressed manifests are efficient for small namespaces, but transferring a +# large set of binary payloads can exceed the asyncpg statement timeout before +# PostgreSQL has done meaningful work. For larger snapshots, the normalized +# section/chunk index loaders below transfer only the fields map-nav needs and +# page them by revision. +_MANIFEST_MAX_REVISION_COUNT = 32 _logger = logging.getLogger(__name__) @@ -206,12 +212,14 @@ async def load_nav_snapshot( if job_result_id and job_id } - manifest_sections = await _load_manifest_sections( - db, - document_revisions=document_revisions, - exclude_sections=excluded_secs, - job_id_by_result_id=job_id_by_result_id, - ) + manifest_sections = None + if len(document_revisions) <= _MANIFEST_MAX_REVISION_COUNT: + manifest_sections = await _load_manifest_sections( + db, + document_revisions=document_revisions, + exclude_sections=excluded_secs, + job_id_by_result_id=job_id_by_result_id, + ) if manifest_sections is None: sections_by_doc, section_path_by_id = await _load_sections( db, @@ -358,26 +366,31 @@ async def _load_manifest_sections( if len(manifest_entries) != len(document_revisions): return None else: - statement = select( - RetrievalServingRevisionManifest.document_id, - RetrievalServingRevisionManifest.job_result_id, - RetrievalServingRevisionManifest.payload_zlib, - RetrievalServingRevisionManifest.checksum, - RetrievalServingRevisionManifest.format_version, - ).where( - tuple_( + manifest_entries = [] + for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE): + revision_group = document_revisions[ + group_start : group_start + _REVISION_GROUP_SIZE + ] + statement = select( RetrievalServingRevisionManifest.document_id, RetrievalServingRevisionManifest.job_result_id, - ).in_(document_revisions) - ) - try: - rows = (await db.execute(statement)).all() - except Exception: - await db.rollback() - return None - if len(rows) != len(document_revisions): - return None - manifest_entries = [tuple(row) for row in rows] + RetrievalServingRevisionManifest.payload_zlib, + RetrievalServingRevisionManifest.checksum, + RetrievalServingRevisionManifest.format_version, + ).where( + tuple_( + RetrievalServingRevisionManifest.document_id, + RetrievalServingRevisionManifest.job_result_id, + ).in_(revision_group) + ) + try: + rows = (await db.execute(statement)).all() + except Exception: + await db.rollback() + return None + if len(rows) != len(revision_group): + return None + manifest_entries.extend(tuple(row) for row in rows) by_doc: dict[str, list[SectionRow]] = {} path_by_id: dict[str, str] = {} ids_by_doc: dict[str, list[str]] = {} diff --git a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py index ed6f2719..0e087549 100644 --- a/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py +++ b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py @@ -77,10 +77,11 @@ async def count_scoped_chunks( exclude_document_ids: list[str], allowed_chunk_types: set[str] | None, revision_pins: Mapping[str, str] | None = None, + max_count: int | None = None, ) -> int: if revision_pins is None: stmt = ( - select(func.count(DocumentChunk.id)) + select(DocumentChunk.id) .join( Document, (Document.document_id == DocumentChunk.document_id) @@ -92,7 +93,7 @@ async def count_scoped_chunks( ) else: stmt = ( - select(func.count(DocumentChunk.id)) + select(DocumentChunk.id) .join(Document, Document.document_id == DocumentChunk.document_id) .where(Document.user_id == user_id) .where(Document.namespace == namespace) @@ -106,7 +107,13 @@ async def count_scoped_chunks( stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids))) if allowed_chunk_types is not None: stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types))) - result = await db.execute(stmt) + + if max_count is not None: + stmt = stmt.limit(max_count) + + result = await db.execute( + select(func.count()).select_from(stmt.order_by(None).subquery()) + ) return result.scalar() or 0 From b683dede3166970d873534f281bd0b6d04b8e512 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 10:03:44 +0800 Subject: [PATCH 2/2] test: lock large snapshot manifest bypass --- .../test_retrieval_snapshot_large_corpus_contract.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py index a414e44b..0b45a703 100644 --- a/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py +++ b/apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py @@ -6,6 +6,7 @@ from uuid import uuid4 from httpx import AsyncClient +import pytest from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job_result import JobResult @@ -15,6 +16,7 @@ _REVISION_GROUP_SIZE, load_nav_snapshot, ) +import shared.services.retrieval.nav_snapshot as nav_snapshot_module from sqlalchemy import Executable, Result, select from sqlalchemy.engine import Row from sqlalchemy.sql.selectable import Select @@ -252,10 +254,19 @@ async def test_large_snapshot_keeps_all_retrieval_inputs_after_bounded_sql_load( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], + monkeypatch: pytest.MonkeyPatch, ) -> None: namespace = f"large-corpus-{uuid4().hex[:8]}" async with developer_api_client_factory(): await _seed_large_retrieval_corpus(namespace) + async def unexpected_manifest_load(*_args: object, **_kwargs: object) -> None: + raise AssertionError("large snapshots must use normalized retrieval rows") + + monkeypatch.setattr( + nav_snapshot_module, + "_load_manifest_sections", + unexpected_manifest_load, + ) legacy_rows = await _load_legacy_rows(namespace) async with contract_db_session() as db: counting_db = _CountingSession(db)