perf(retrieval): optimize map-nav snapshot and scoring

apps/api/scripts/backfill_map_unit_indexes.py

@@ -67,7 +67,7 @@ def _bootstrap_python_path() -> None:
     lock_namespace_generation,
 )
 from shared.services.retrieval.serving_manifest import (
-    decode_serving_manifest,
+    decode_namespace_map_snapshot,
     persist_revision_serving_state,
 )
 
@@ -160,7 +160,7 @@ def check_fallback_readiness(*, document_id: str = "") -> list[NamespaceFallback
                 snapshot_status = "missing"
             else:
                 try:
-                    payload = decode_serving_manifest(
+                    payload = decode_namespace_map_snapshot(
                         bytes(snapshot.payload_zlib),
                         checksum=str(snapshot.checksum),
                         format_version=int(snapshot.format_version),

apps/api/tests/contract/test_retrieval_map_score_parity_contract.py

@@ -0,0 +1,88 @@
+"""Contract tests for map score pooling semantics."""
+
+from __future__ import annotations
+
+from typing import Final
+
+import pytest
+
+from shared.services.retrieval.nav.nav_map_scores import _pool_unit_scores_to_tree
+
+
+_CASES: Final[tuple[tuple[dict[str, list[str]], set[str], dict[str, float]], ...]] = (
+    (
+        {"root-a": ["section-a", "section-b"], "section-a": [], "section-b": []},
+        {"section-a", "section-b"},
+        {"section-a": 0.4, "section-b": 0.8, "root-a__self": 0.2},
+    ),
+    (
+        {
+            "root-a": ["parent-a"],
+            "parent-a": ["leaf-a", "leaf-b"],
+            "leaf-a": [],
+            "leaf-b": [],
+            "root-b": ["leaf-c"],
+            "leaf-c": [],
+        },
+        {"leaf-a", "leaf-b", "leaf-c"},
+        {
+            "leaf-a": 0.9,
+            "leaf-b": 0.3,
+            "leaf-c": 0.7,
+            "parent-a__self": 0.95,
+        },
+    ),
+)
+
+
+def _legacy_pool(
+    children_map: dict[str, list[str]],
+    leaves: set[str],
+    unit_scores: dict[str, float],
+) -> dict[str, float]:
+    map_scores = {
+        leaf_id: float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in leaves
+    }
+
+    def score_node(section_id: str) -> float:
+        if section_id in map_scores:
+            return map_scores[section_id]
+        children = children_map.get(section_id) or []
+        if not children:
+            score = float(unit_scores.get(section_id, 0.0) or 0.0)
+            map_scores[section_id] = score
+            return score
+        descendants: list[str] = []
+
+        def collect(section: str) -> None:
+            nested = children_map.get(section) or []
+            if not nested:
+                if section in leaves:
+                    descendants.append(section)
+                return
+            for child in nested:
+                collect(child)
+
+        collect(section_id)
+        parts = [float(unit_scores.get(leaf_id, 0.0) or 0.0) for leaf_id in descendants]
+        self_key = f"{section_id}__self"
+        if self_key in unit_scores:
+            parts.append(float(unit_scores[self_key]))
+        score = float(max(parts)) if parts else 0.0
+        map_scores[section_id] = score
+        return score
+
+    for section_id in children_map:
+        score_node(section_id)
+    return map_scores
+
+
+@pytest.mark.parametrize("children_map, leaves, unit_scores", _CASES)
+def test_map_score_pooling_preserves_legacy_semantics(
+    children_map: dict[str, list[str]],
+    leaves: set[str],
+    unit_scores: dict[str, float],
+) -> None:
+    assert _pool_unit_scores_to_tree(children_map, leaves, unit_scores) == _legacy_pool(
+        children_map, leaves, unit_scores
+    )

apps/api/tests/contract/test_retrieval_map_unit_index_contract.py

@@ -6,7 +6,7 @@
 from uuid import uuid4
 
 from httpx import AsyncClient
-from sqlalchemy import delete, select, text
+from sqlalchemy import Engine, delete, event, select, text
 
 from shared.models.database.document import (
     DocumentMapUnit,
@@ -21,6 +21,7 @@
     compute_corpus_map_and_unit_scores,
     select_map_highlights,
 )
+from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows
 from shared.services.retrieval.nav.nav_knowhere import (
     KnowhereProvider,
     LazyKnowhereProvider,
@@ -105,11 +106,15 @@ def close(self) -> None:
     ]
     assert len(frequency_executions) == 1
     statement, parameters = frequency_executions[0]
-    assert "map_unit_id = ANY" in statement
+    assert "scoped_units AS MATERIALIZED" in statement
+    assert "JOIN scoped_units" in statement
+    assert "channel = ANY" in statement
     assert "token_hash = ANY" in statement
+    assert "map_unit_id = ANY" not in statement
     assert isinstance(parameters, list)
     assert parameters[0] == ["unit-frequency"]
-    assert parameters[1] == [
+    assert parameters[1] == ["path", "content"]
+    assert parameters[2] == [
         "6e51d6a3d90b6a3243d38e6da6b3f31f49867c1360beba83da8ca9630f9672c7"
     ]
 
@@ -493,6 +498,100 @@ def record_reference_load(
         snapshot.close()
 
 
+async def test_connected_hydration_does_not_load_legacy_job_chunks(
+    developer_api_client_factory: Callable[
+        [], AbstractAsyncContextManager[AsyncClient]
+    ],
+) -> None:
+    identifier = uuid4().hex[:8]
+    namespace = f"connected-job-{identifier}"
+    document_id = f"doc_connected_{identifier}"
+    job_id = f"job_connected_{identifier}"
+    job_result_id = f"result_connected_{identifier}"
+    statements: list[str] = []
+
+    def capture_job_chunk_query(
+        _connection: Any,
+        _cursor: Any,
+        statement: str,
+        _parameters: Any,
+        _context: Any,
+        _executemany: bool,
+    ) -> None:
+        if "job_chunks" in statement.lower():
+            statements.append(statement)
+
+    async with developer_api_client_factory():
+        await _seed_revision(
+            namespace=namespace,
+            document_id=document_id,
+            job_id=job_id,
+            job_result_id=job_result_id,
+        )
+        scope = DocumentPublicationScope(
+            user_id=_USER_ID,
+            namespace=namespace,
+            document_id=document_id,
+            job_result_id=job_result_id,
+            source_file_name="connected.pdf",
+        )
+        chunks = [
+            {
+                "chunk_id": "body-connected",
+                "type": "text",
+                "content": "body connected evidence",
+                "path": "connected.pdf/Root/Section/body",
+                "order": 1,
+                "metadata": {"connect_to": [{"target": "asset-connected"}]},
+            },
+            {
+                "chunk_id": "asset-connected",
+                "type": "image",
+                "content": "asset connected summary",
+                "path": "images/asset-connected.png",
+                "order": 2,
+                "file_path": "images/asset-connected.png",
+                "metadata": {},
+            },
+        ]
+        async with contract_db_session() as db:
+            await db.run_sync(
+                lambda sync_db: _publish_revision_with_generation_lock(
+                    sync_db,
+                    scope=scope,
+                    chunks=chunks,
+                )
+            )
+            await db.commit()
+
+        event.listen(Engine, "before_cursor_execute", capture_job_chunk_query)
+        try:
+            async with contract_db_session() as db:
+                hydrated = await hydrate_connected_target_rows(
+                    db=db,
+                    rows=[
+                        {
+                            "document_id": document_id,
+                            "job_result_id": job_result_id,
+                            "chunk_id": "body-connected",
+                            "chunk_type": "text",
+                            "chunk_metadata": {
+                                "connect_to": [{"target": "asset-connected"}]
+                            },
+                        }
+                    ],
+                    exclude_document_ids=[],
+                    exclude_sections=[],
+                    revision_pins={document_id: job_result_id},
+                )
+        finally:
+            event.remove(Engine, "before_cursor_execute", capture_job_chunk_query)
+
+    assert [row["chunk_id"] for row in hydrated] == ["asset-connected"]
+    assert hydrated[0]["job_id"] == job_id
+    assert statements == []
+
+
 def test_incomplete_index_returns_empty_scores() -> None:
     first_sections = [
         SectionRow("root-a", None, "Root A", "Root A", 0, "", 0),

apps/api/tests/contract/test_retrieval_serving_manifest_contract.py

@@ -5,8 +5,11 @@
 import pytest
 
 from shared.services.retrieval.serving_manifest import (
+    NAMESPACE_MAP_SNAPSHOT_FORMAT_VERSION,
     SERVING_MANIFEST_FORMAT_VERSION,
+    decode_namespace_map_snapshot,
     decode_serving_manifest,
+    encode_namespace_map_snapshot,
     encode_serving_manifest,
 )
 
@@ -49,3 +52,55 @@ def test_serving_manifest_rejects_unknown_version() -> None:
             checksum=checksum,
             format_version=SERVING_MANIFEST_FORMAT_VERSION + 1,
         )
+
+
+def test_namespace_snapshot_uses_routing_only_v2_and_reads_legacy_v1() -> None:
+    payload = {
+        "documents": {
+            "doc_1": {
+                "job_result_id": "result_1",
+                "job_id": "job_1",
+                "source_file_name": "private.pdf",
+                "sections": [
+                    {
+                        "section_id": "sec_1",
+                        "section_path": "Root",
+                        "section_title": "Root",
+                        "section_level": 0,
+                        "summary": "summary",
+                        "sort_order": 0,
+                        "unused": "drop",
+                    }
+                ],
+                "chunks": [
+                    {
+                        "chunk_id": "chunk_1",
+                        "section_id": "sec_1",
+                        "chunk_type": "text",
+                        "sort_order": 0,
+                        "connect_to": [],
+                        "content": "drop",
+                    }
+                ],
+            }
+        }
+    }
+    compressed, checksum, version = encode_namespace_map_snapshot(payload)
+
+    assert version == NAMESPACE_MAP_SNAPSHOT_FORMAT_VERSION
+    decoded = decode_namespace_map_snapshot(
+        compressed, checksum=checksum, format_version=version
+    )
+    document = decoded["documents"]["doc_1"]
+    assert "source_file_name" not in document
+    assert "unused" not in document["sections"][0]
+    assert "content" not in document["chunks"][0]
+
+    legacy_compressed, legacy_checksum, legacy_version = encode_serving_manifest(
+        payload
+    )
+    assert decode_namespace_map_snapshot(
+        legacy_compressed,
+        checksum=legacy_checksum,
+        format_version=legacy_version,
+    ) == payload

apps/api/tests/contract/test_retrieval_snapshot_consistency_contract.py

@@ -8,18 +8,46 @@
 from httpx import AsyncClient
 import pytest
 from sqlalchemy import Executable, Result
+from sqlalchemy.exc import SQLAlchemyError
 
 from shared.services.retrieval.execution.reference_resolver import (
     resolve_workflow_references,
 )
 from shared.services.retrieval.nav_snapshot import SnapshotSession, load_nav_snapshot
+from shared.services.retrieval.nav_snapshot import _resolve_namespace_snapshot_entries
 from tests.support.retrieval_snapshot_support import contract_db_session
 from tests.support.contract_database import ContractDatabase
 
 
 _USER_ID = "local-dev-user"
 
 
+class _GenerationUnavailableSession:
+    def __init__(self) -> None:
+        self.rollback_count = 0
+
+    async def execute(self, _statement: Executable) -> Result[tuple[object, ...]]:
+        raise SQLAlchemyError("generation table unavailable")
+
+    async def rollback(self) -> None:
+        self.rollback_count += 1
+
+
+@pytest.mark.asyncio
+async def test_snapshot_loader_falls_back_when_generation_cannot_be_verified() -> None:
+    session = _GenerationUnavailableSession()
+
+    result = await _resolve_namespace_snapshot_entries(
+        session,
+        user_id=_USER_ID,
+        namespace="default",
+        document_revisions=[("doc-a", "result-a")],
+    )
+
+    assert result is None
+    assert session.rollback_count == 1
+
+
 class _PublishingSession:
     def __init__(
         self,