perf(core): read only the projects in scope from the vector indexes

CHANGELOG.md

@@ -23,6 +23,19 @@
   database is skipped with a warning and an inexact total, and a retryable outage
   or a server too old to attribute its results fails the whole page.
 
+- **#1558**: Vector retrieval reads only the projects in scope and fills the window it
+  asks for. The sqlite-vec table gains a `project_id` partition key, so a scoped
+  nearest-neighbour query ranks each project's own vectors instead of taking the k
+  nearest across the whole database and discarding the out-of-scope ones, which
+  could leave a small project with an empty page for a query its notes answered.
+  Existing local storage is carried into the partitioned table without re-embedding.
+  On Postgres the nearest-neighbour statement now runs on the HNSW index (its
+  tie-break sort keys had kept the planner on an exact scan of every vector), with
+  `hnsw.ef_search` sized to the candidate window and an iterative scan that keeps
+  going until the scope and manifest filters have admitted enough rows. That scan
+  needs pgvector 0.8 or later; an older extension is reported as a dependency error
+  instead of quietly returning short windows.
+
 - **#1512**: Word, PowerPoint, and CSV files get the same sidecar Markdown note a
   PDF gets. `bm import document <path>` indexes the project, extracts the file,
   and writes `<file>.<ext>.md` next to it plus a run note under

src/basic_memory/models/search.py

@@ -223,11 +223,18 @@
 
 
 def create_sqlite_search_vector_embeddings(dimensions: int) -> DDL:
-    """Build sqlite-vec virtual table DDL for the configured embedding dimension."""
+    """Build sqlite-vec virtual table DDL for the configured embedding dimension.
+
+    ``project_id`` is a vec0 partition key: a scoped nearest-neighbour query reads
+    only the partitions in scope, instead of ranking every project's vectors in the
+    database and discarding the out-of-scope ones afterwards, which left a small
+    project with an under-filled window whenever a larger neighbour sat closer.
+    """
     return DDL(
         f"""
 CREATE VIRTUAL TABLE IF NOT EXISTS search_vector_embeddings
 USING vec0(
+    project_id integer partition key,
     embedding float[{dimensions}],
     +source_hash text
 )

src/basic_memory/repository/pgvector_index.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import re
 from collections.abc import Sequence
 
 from loguru import logger
@@ -23,6 +24,32 @@
 )
 
 
+# pgvector's HNSW scan hands back at most ``hnsw.ef_search`` rows, 40 by default,
+# whatever the LIMIT asks for; the server caps the setting at 1000.
+HNSW_EF_SEARCH_DEFAULT = 40
+HNSW_EF_SEARCH_MAX = 1000
+
+_PGVECTOR_VERSION = re.compile(r"^(\d+)\.(\d+)")
+
+
+def pgvector_supports_iterative_scan(extversion: str) -> bool:
+    """Whether this pgvector keeps scanning an HNSW index until the LIMIT is filled.
+
+    Iterative index scans arrived in pgvector 0.8.0. Before that, a scan stops at
+    ``hnsw.ef_search`` candidates, and a filter applied afterwards (the manifest
+    join, a scope narrower than the table) leaves the window under-filled, so the
+    adapter cannot promise the nearest ``limit`` rows in scope. A version string
+    the pattern cannot read counts as older.
+    """
+    match = _PGVECTOR_VERSION.match(extversion)
+    return match is not None and (int(match.group(1)), int(match.group(2))) >= (0, 8)
+
+
+def hnsw_ef_search_for(limit: int) -> int:
+    """The candidate-list size that lets one HNSW scan return ``limit`` rows."""
+    return min(max(limit, HNSW_EF_SEARCH_DEFAULT), HNSW_EF_SEARCH_MAX)
+
+
 class PgVectorIndex:
     """Persist and query semantic vectors in PostgreSQL with pgvector."""
 
@@ -56,6 +83,19 @@ async def initialize(self) -> None:
                     raise SemanticDependenciesMissingError(
                         "pgvector extension is unavailable for this Postgres database."
                     ) from exc
+                version = await session.execute(
+                    text("SELECT extversion FROM pg_extension WHERE extname = 'vector'")
+                )
+                extversion = str(version.scalar_one())
+                # Trigger: the installed pgvector predates iterative index scans.
+                # Why: without them a scoped query silently returns fewer rows than
+                #   the window asked for; a deployment gap should read as one.
+                # Outcome: a typed dependency error the API reports as a bad request.
+                if not pgvector_supports_iterative_scan(extversion):
+                    raise SemanticDependenciesMissingError(
+                        f"pgvector {extversion} predates iterative index scans; semantic "
+                        "search needs pgvector 0.8 or later (ALTER EXTENSION vector UPDATE)."
+                    )
 
                 existing_dimensions = await self._existing_dimensions(session)
                 storage_missing = existing_dimensions is None
@@ -322,10 +362,25 @@ async def search(
         embeddings_in_scope = projects.predicate("e.project_id", params)
         chunks_in_scope = projects.predicate("c.project_id", params)
         async with db.scoped_session(self._session_maker) as session:
+            # Both settings are transaction-local, so they last exactly as long as
+            # this scoped session. The scan is sized to the window it must fill
+            # and continues past that until the scope and manifest filters have
+            # admitted enough rows, instead of stopping at the first ef_search
+            # candidates and returning whatever of them survived.
+            await session.execute(
+                text("SELECT set_config('hnsw.ef_search', :ef_search, true)"),
+                {"ef_search": str(hnsw_ef_search_for(limit))},
+            )
+            await session.execute(
+                text("SELECT set_config('hnsw.iterative_scan', 'relaxed_order', true)")
+            )
+            # A relaxed iterative scan may hand rows back slightly out of distance
+            # order, so the window is taken by distance alone and sorted once more.
             result = await session.execute(
                 text(
+                    "WITH nearest AS MATERIALIZED ("
                     "SELECT c.entity_id, c.chunk_key, "
-                    "1 - (e.embedding <=> CAST(:query AS vector)) AS similarity "
+                    "e.embedding <=> CAST(:query AS vector) AS distance "
                     "FROM search_vector_embeddings e "
                     "JOIN search_vector_chunks c ON c.id = e.chunk_id "
                     f"WHERE {embeddings_in_scope} "
@@ -335,9 +390,12 @@ async def search(
                     "AND c.embedding_status = 'ready' "
                     "AND c.embedding_model = :embedding_identity "
                     "AND e.source_hash = c.source_hash "
-                    "ORDER BY e.embedding <=> CAST(:query AS vector), "
-                    "c.entity_id ASC, c.chunk_key ASC "
+                    "ORDER BY e.embedding <=> CAST(:query AS vector) "
                     "LIMIT :limit"
+                    ") "
+                    "SELECT entity_id, chunk_key, 1 - distance AS similarity "
+                    "FROM nearest "
+                    "ORDER BY distance ASC, entity_id ASC, chunk_key ASC"
                 ),
                 params,
             )

src/basic_memory/repository/sqlite_vec_index.py

@@ -115,6 +115,7 @@ async def initialize(self) -> None:
                 expected_dimensions = f"float[{self.scope.dimensions}]"
                 dimensions_changed = bool(vector_sql and expected_dimensions not in vector_sql)
                 source_hash_missing = bool(vector_sql and "+source_hash text" not in vector_sql)
+                partitions_missing = bool(vector_sql and "partition key" not in vector_sql)
                 if dimensions_changed or source_hash_missing:
                     logger.warning(
                         "SQLite vector storage schema mismatch "
@@ -124,6 +125,19 @@ async def initialize(self) -> None:
                         source_hash_missing=source_hash_missing,
                     )
                     await session.execute(text("DROP TABLE IF EXISTS search_vector_embeddings"))
+                elif partitions_missing:
+                    # Trigger: storage predates the project_id partition key, and its
+                    #   vectors are otherwise current.
+                    # Why: vec0 cannot add a column in place, and re-embedding a whole
+                    #   vault only to change how rows are partitioned would cost every
+                    #   local user a full embedding pass for nothing new.
+                    # Outcome: the vectors are carried into partitioned storage, each
+                    #   keyed by the project its manifest row names; manifests stay ready.
+                    logger.info(
+                        "SQLite vector storage predates project partitions; "
+                        "carrying vectors into partitioned storage"
+                    )
+                    await self._partition_existing_storage(session)
 
                 await session.execute(create_sqlite_search_vector_embeddings(self.scope.dimensions))
                 # Missing or dimension-rebuilt vec storage has no vectors, so ready
@@ -138,6 +152,35 @@ async def initialize(self) -> None:
                 await session.commit()
             self._initialized = True
 
+    async def _partition_existing_storage(self, session: AsyncSession) -> None:
+        """Rebuild vec storage with the partition key, keeping every current vector.
+
+        SQLite DDL is transactional, so the copy out, drop, recreate, and copy back
+        either all land or none do. A vector whose manifest row is gone has no
+        project to file under and is left behind, which is what the orphan sweep
+        would have done to it anyway.
+        """
+        await session.execute(
+            text(
+                "CREATE TEMP TABLE search_vector_embeddings_carry AS "
+                "SELECT e.rowid AS id, c.project_id AS project_id, "
+                "e.embedding AS embedding, e.source_hash AS source_hash "
+                "FROM search_vector_embeddings e "
+                "JOIN search_vector_chunks c ON c.id = e.rowid"
+            )
+        )
+        await session.execute(text("DROP TABLE search_vector_embeddings"))
+        await session.execute(create_sqlite_search_vector_embeddings(self.scope.dimensions))
+        await session.execute(
+            text(
+                "INSERT INTO search_vector_embeddings "
+                "(rowid, project_id, embedding, source_hash) "
+                "SELECT id, project_id, embedding, source_hash "
+                "FROM search_vector_embeddings_carry"
+            )
+        )
+        await session.execute(text("DROP TABLE search_vector_embeddings_carry"))
+
     async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None:
         if not records:
             return
@@ -192,12 +235,14 @@ async def upsert(self, project_id: int, records: Sequence[VectorRecord]) -> None
             )
             await session.execute(
                 text(
-                    "INSERT INTO search_vector_embeddings (rowid, embedding, source_hash) "
-                    "VALUES (:rowid, :embedding, :source_hash)"
+                    "INSERT INTO search_vector_embeddings "
+                    "(rowid, project_id, embedding, source_hash) "
+                    "VALUES (:rowid, :project_id, :embedding, :source_hash)"
                 ),
                 [
                     {
                         "rowid": rowids_by_key[record.key],
+                        "project_id": project_id,
                         "embedding": json.dumps(record.values),
                         "source_hash": record.source_hash,
                     }
@@ -323,21 +368,24 @@ async def search(
             "embedding_identity": self.scope.embedding_identity,
             "limit": limit,
         }
-        chunks_in_scope = projects.predicate("c.project_id", params)
+        # vec0 ranks the k nearest within each partition in scope, so a small
+        # project is never crowded out of its own window by a larger neighbour that
+        # shares the database; the outer ORDER BY merges the partitions.
+        partitions_in_scope = projects.predicate("project_id", params)
         async with db.scoped_session(self._session_maker) as session:
             await self._ensure_loaded(session)
             result = await session.execute(
                 text(
                     "WITH vector_matches AS MATERIALIZED ("
                     " SELECT rowid, distance, source_hash FROM search_vector_embeddings "
-                    " WHERE embedding MATCH :query AND k = :vector_k"
+                    f" WHERE {partitions_in_scope} "
+                    " AND embedding MATCH :query AND k = :vector_k"
                     ") "
                     "SELECT c.entity_id, c.chunk_key, vector_matches.distance "
                     "FROM vector_matches "
                     "JOIN search_vector_chunks c ON c.id = vector_matches.rowid "
                     "AND c.source_hash = vector_matches.source_hash "
-                    f"WHERE {chunks_in_scope} "
-                    "AND c.vector_index = 'sqlite-vec' "
+                    "WHERE c.vector_index = 'sqlite-vec' "
                     "AND c.embedding_status = 'ready' "
                     "AND c.embedding_model = :embedding_identity "
                     "ORDER BY vector_matches.distance ASC, "

tests/repository/test_pgvector_index.py

@@ -9,7 +9,11 @@
 import pytest
 
 from basic_memory.repository import pgvector_index as pgvector_index_module
-from basic_memory.repository.pgvector_index import PgVectorIndex
+from basic_memory.repository.pgvector_index import (
+    PgVectorIndex,
+    hnsw_ef_search_for,
+    pgvector_supports_iterative_scan,
+)
 from basic_memory.repository.search_scope import ProjectScope
 from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError
 from basic_memory.repository.semantic_vector_index import (
@@ -40,6 +44,10 @@ def fetchone(self) -> object | None:
     def scalar_one_or_none(self) -> object | None:
         return self._scalar
 
+    def scalar_one(self) -> object:
+        assert self._scalar is not None
+        return self._scalar
+
     def mappings(self) -> FakeResult:
         return self
 
@@ -59,13 +67,15 @@ def __init__(
         chunk_rows: list[dict[str, object]] | None = None,
         search_rows: list[dict[str, object]] | None = None,
         fail_extension: bool = False,
+        pgvector_version: str = "0.8.0",
     ) -> None:
         self.table_exists = table_exists
         self.dimensions = dimensions
         self.has_source_hash = has_source_hash
         self.chunk_rows = chunk_rows or []
         self.search_rows = search_rows or []
         self.fail_extension = fail_extension
+        self.pgvector_version = pgvector_version
         self.calls: list[tuple[str, dict[str, object] | None]] = []
         self.commit_count = 0
 
@@ -78,6 +88,8 @@ async def execute(
         self.calls.append((sql, params))
         if "CREATE EXTENSION" in sql and self.fail_extension:
             raise RuntimeError("extension unavailable")
+        if "SELECT extversion" in sql:
+            return FakeResult(scalar=self.pgvector_version)
         if "information_schema.tables" in sql:
             return FakeResult(fetchone=(1,) if self.table_exists else None)
         if "attname = 'source_hash'" in sql:
@@ -330,7 +342,7 @@ async def test_search_returns_normalized_stable_matches(monkeypatch) -> None:
         (15, 0.0),
     ]
     search_call = next(call for call in session.calls if "AS similarity" in call[0])
-    assert "c.entity_id ASC, c.chunk_key ASC" in search_call[0]
+    assert "ORDER BY distance ASC, entity_id ASC, chunk_key ASC" in search_call[0]
     assert search_call[1] == {
         "query": "[1,0,0,0]",
         "scope_0": 7,
@@ -356,3 +368,73 @@ async def test_search_binds_every_project_in_scope(monkeypatch) -> None:
     assert "AND c.project_id IN (:scope_0, :scope_1)" in sql
     assert params["scope_0"] == 7
     assert params["scope_1"] == 9
+
+
+def _settings(session: FakeSession) -> list[tuple[str, dict[str, object] | None]]:
+    return [call for call in session.calls if "set_config" in call[0]]
+
+
+@pytest.mark.parametrize(
+    ("extversion", "expected"),
+    [("0.7.4", False), ("0.8.0", True), ("0.8.1", True), ("1.0.0", True), ("garbage", False)],
+)
+def test_iterative_scan_arrived_in_pgvector_0_8(extversion: str, expected: bool) -> None:
+    assert pgvector_supports_iterative_scan(extversion) is expected
+
+
+def test_ef_search_is_sized_to_the_window_within_the_server_bounds() -> None:
+    assert hnsw_ef_search_for(5) == 40
+    assert hnsw_ef_search_for(40) == 40
+    assert hnsw_ef_search_for(250) == 250
+    assert hnsw_ef_search_for(5000) == 1000
+
+
+@pytest.mark.asyncio
+async def test_initialize_requires_a_pgvector_that_can_keep_scanning(monkeypatch) -> None:
+    """An extension too old to fill a filtered window is a deployment error, not a quiet gap."""
+    older = FakeSession(pgvector_version="0.7.4")
+    _install_session(monkeypatch, older)
+    index = PgVectorIndex(MagicMock(), _scope())
+
+    with pytest.raises(SemanticDependenciesMissingError, match="pgvector 0.7.4 predates"):
+        await index.initialize()
+
+    assert not any("CREATE TABLE" in sql for sql in _sql_calls(older))
+    assert index._initialized is False
+
+
+@pytest.mark.asyncio
+async def test_search_sizes_the_scan_to_the_window_it_must_fill(monkeypatch) -> None:
+    """A 250-row candidate window asks HNSW for 250 candidates, not the default 40."""
+    session = FakeSession(search_rows=[])
+    _install_session(monkeypatch, session)
+    index = PgVectorIndex(MagicMock(), _scope())
+    index._initialized = True
+
+    await index.search([1.0, 0.0, 0.0, 0.0], limit=250, projects=PROJECTS)
+
+    ef_search = next(call for call in _settings(session) if "hnsw.ef_search" in call[0])
+    assert "set_config('hnsw.ef_search', :ef_search, true)" in ef_search[0]
+    assert ef_search[1] == {"ef_search": "250"}
+    search_sql = next(sql for sql, _params in session.calls if "AS similarity" in sql)
+    # The window is taken by distance inside the CTE, then re-sorted with tie-breaks.
+    assert "ORDER BY e.embedding <=> CAST(:query AS vector) LIMIT :limit" in search_sql
+
+
+@pytest.mark.asyncio
+async def test_search_keeps_scanning_until_the_window_fills(monkeypatch) -> None:
+    session = FakeSession(search_rows=[])
+    _install_session(monkeypatch, session)
+    index = PgVectorIndex(MagicMock(), _scope())
+    index._initialized = True
+
+    await index.search([1.0, 0.0, 0.0, 0.0], limit=5, projects=PROJECTS)
+
+    settings = [sql for sql, _params in _settings(session)]
+    assert any("'hnsw.ef_search'" in sql for sql in settings)
+    assert any("'hnsw.iterative_scan', 'relaxed_order', true" in sql for sql in settings)
+    # Settings precede the scan they configure, inside the same session.
+    ordered = [sql for sql, _params in session.calls]
+    assert max(ordered.index(sql) for sql in settings) < next(
+        position for position, sql in enumerate(ordered) if "AS similarity" in sql
+    )

tests/repository/test_postgres_search_repository_unit.py

@@ -225,10 +225,13 @@ async def fake_scoped_session(session_maker):
         )
         missing_table = MagicMock()
         missing_table.fetchone.return_value = None
+        pgvector_version = MagicMock()
+        pgvector_version.scalar_one.return_value = "0.8.0"
         session.execute.side_effect = [
             MagicMock(),
             MagicMock(),
             MagicMock(),
+            pgvector_version,
             missing_table,
             MagicMock(),
             MagicMock(),