perf(core): count admitted vector chunks from the manifest alone

src/basic_memory/repository/pgvector_index.py

@@ -58,6 +58,25 @@ def hnsw_ef_search_for(limit: int) -> int:
     return min(max(limit, HNSW_EF_SEARCH_DEFAULT), HNSW_EF_SEARCH_MAX)
 
 
+def admission_count_sql(chunks_in_scope: str, admission: VectorAdmission) -> str:
+    """Count a filter's admitted chunks, up to ``:exact_cap``, from the manifest alone.
+
+    The count only chooses the plan, so it reads just what the admission names: scope and
+    the chunk columns in ``search_vector_chunks AS c``. Joining the embeddings to apply
+    the full manifest filter made Postgres hash-join both tables whole, which LIMIT cannot
+    stop; on a 480k-chunk project that read 2.5 GB per filtered search (#1791). Without
+    the join it is an index-only scan of the chunk location index that stops at the cap.
+    Pending or stale chunks are counted too; that can only tip a search near the cap
+    toward the HNSW scan, which is correct for any count.
+    """
+    return (
+        "SELECT count(*) FROM ("
+        "SELECT 1 FROM search_vector_chunks c "
+        f"WHERE {chunks_in_scope} AND {admission.predicate} "
+        "LIMIT :exact_cap) admitted"
+    )
+
+
 # Initializations of one database can run at once: a cloud API pre-warm and a worker
 # touching the same tenant, or several workers starting together. CREATE ... IF NOT
 # EXISTS is not safe under concurrency (the loser fails on a catalog unique index), so
@@ -455,13 +474,7 @@ async def search(
             exact = False
             if admission is not None:
                 counted = await session.execute(
-                    text(
-                        "SELECT count(*) FROM ("
-                        "SELECT 1 FROM search_vector_chunks c "
-                        "JOIN search_vector_embeddings e ON e.chunk_id = c.id "
-                        f"WHERE {manifest}{admitted} "
-                        "LIMIT :exact_cap) admitted"
-                    ),
+                    text(admission_count_sql(chunks_in_scope, admission)),
                     {**params, "exact_cap": EXACT_ADMISSION_CHUNK_LIMIT + 1},
                 )
                 exact = int(counted.scalar_one()) <= EXACT_ADMISSION_CHUNK_LIMIT

tests/repository/test_postgres_search_repository.py

@@ -31,6 +31,8 @@
     _strip_nul_from_row,
 )
 from basic_memory.repository.search_index_row import SearchIndexRow
+from basic_memory.repository.search_scope import ProjectScope
+from basic_memory.repository.semantic_vector_index import VectorAdmission
 from basic_memory.runtime.vector_sync import VectorSyncBatchResult
 from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode
 from typing import Any, override
@@ -1890,3 +1892,80 @@ async def spy(query, *, limit, projects, admission=None):
     [call] = calls
     assert call["admission"] is not None
     assert call["entities"] == notes
+
+
+def _plan_relations(plan: dict[str, Any]) -> set[str]:
+    """Every table an EXPLAIN (FORMAT JSON) plan node reads."""
+    relations = {plan["Relation Name"]} if "Relation Name" in plan else set()
+    for child in plan.get("Plans", []):
+        relations |= _plan_relations(child)
+    return relations
+
+
+@pytest.mark.asyncio
+async def test_postgres_admission_count_reads_only_the_chunk_manifest(session_maker, test_project):
+    """The filtered-search admission count must not touch the embeddings (#1791).
+
+    It only chooses between exact ranking and the HNSW scan. Joining the embeddings made
+    Postgres hash-join both tables whole, which LIMIT cannot stop, so every filtered
+    search on a 480k-chunk project read 2.5 GB. The count still counts the admitted
+    chunks, and only them.
+    """
+    await _skip_if_pgvector_unavailable(session_maker)
+    app_config = BasicMemoryConfig(
+        env="test",
+        projects={"test-project": "/tmp/basic-memory-test"},
+        default_project="test-project",
+        database_backend=DatabaseBackend.POSTGRES,
+    )
+    repo = PostgresSearchRepository(
+        session_maker,
+        project_id=test_project.id,
+        app_config=app_config,
+        embedding_provider=StubEmbeddingProvider(),
+    )
+    await repo.init_search_index()
+    now = datetime.now(timezone.utc)
+    await repo.bulk_index_items(
+        [
+            SearchIndexRow(
+                project_id=test_project.id,
+                id=row_id,
+                title=f"Admission {row_id}",
+                content_stems="auth token archive",
+                content_snippet="auth token archive",
+                permalink=f"admission/entry-{row_id}",
+                file_path=f"admission/entry-{row_id}.md",
+                type=SearchItemType.ENTITY.value,
+                entity_id=row_id,
+                metadata={"note_type": "note"},
+                created_at=now,
+                updated_at=now,
+            )
+            for row_id in (601, 602, 603)
+        ]
+    )
+    for row_id in (601, 602, 603):
+        await repo.sync_entity_vectors(row_id)
+
+    params: dict[str, object] = {
+        "admit_type_0": SearchItemType.ENTITY.value,
+        "exact_cap": 10,
+    }
+    sql = pgvector_index_module.admission_count_sql(
+        ProjectScope.single(test_project.id).predicate("c.project_id", params),
+        VectorAdmission(predicate="c.source_type IN (:admit_type_0)", params={}),
+    )
+    async with db.scoped_session(session_maker) as session:
+        explained = await session.execute(text(f"EXPLAIN (FORMAT JSON) {sql}"), params)
+        plan = explained.scalar_one()[0]["Plan"]
+        counted = (await session.execute(text(sql), params)).scalar_one()
+        manifest_rows = (
+            await session.execute(
+                text("SELECT count(*) FROM search_vector_chunks WHERE project_id = :project_id"),
+                {"project_id": test_project.id},
+            )
+        ).scalar_one()
+
+    assert _plan_relations(plan) == {"search_vector_chunks"}
+    assert counted == manifest_rows > 0