From 0c94f63077372000dba6512d15720b896e3a1931 Mon Sep 17 00:00:00 2001 From: Thiago Henrique Nogueira Date: Sat, 19 Sep 2026 12:57:33 -0300 Subject: [PATCH] perf(core): bound sqlite-vec reconciliation work Signed-off-by: Thiago Henrique Nogueira --- .../repository/sqlite_vec_index.py | 44 ++++---- ...t_sqlite_vec_reconciliation_performance.py | 105 ++++++++++++++++++ 2 files changed, 129 insertions(+), 20 deletions(-) create mode 100644 tests/repository/test_sqlite_vec_reconciliation_performance.py diff --git a/src/basic_memory/repository/sqlite_vec_index.py b/src/basic_memory/repository/sqlite_vec_index.py index 9f81baf06..66f6d0262 100644 --- a/src/basic_memory/repository/sqlite_vec_index.py +++ b/src/basic_memory/repository/sqlite_vec_index.py @@ -27,6 +27,7 @@ SQLITE_VEC_MAX_K = 4096 +SQLITE_DELETE_BATCH_SIZE = 500 class SQLiteVecIndex: @@ -321,33 +322,36 @@ async def delete_orphans(self, project_id: int, _live_keys: Sequence[VectorKey]) ) ) orphan_rowids = [int(rowid) for rowid in orphan_result.scalars().all()] - if orphan_rowids: - params = { - f"orphan_rowid_{index}": rowid for index, rowid in enumerate(orphan_rowids) - } - placeholders = ", ".join( - f":orphan_rowid_{index}" for index in range(len(orphan_rowids)) - ) - await session.execute( - text(f"DELETE FROM search_vector_embeddings WHERE rowid IN ({placeholders})"), - params, - ) - await session.execute( + stale_result = await session.execute( text( - "DELETE FROM search_vector_embeddings WHERE rowid IN (" - "SELECT id FROM search_vector_chunks " - "WHERE project_id = :project_id AND NOT (" - "vector_index = 'sqlite-vec' " - "AND embedding_model = :embedding_identity " - "AND search_vector_embeddings.source_hash = " - "search_vector_chunks.source_hash " - "AND embedding_status = 'ready'))" + "SELECT chunks.id FROM search_vector_chunks AS chunks " + "JOIN search_vector_embeddings AS embeddings " + "ON embeddings.rowid = chunks.id " + "WHERE chunks.project_id = :project_id AND NOT (" + "chunks.vector_index = 'sqlite-vec' " + "AND chunks.embedding_model = :embedding_identity " + "AND embeddings.source_hash = chunks.source_hash " + "AND chunks.embedding_status = 'ready')" ), { "project_id": project_id, "embedding_identity": self.scope.embedding_identity, }, ) + stale_rowids = [int(rowid) for rowid in stale_result.scalars().all()] + + # sqlite-vec evaluates a correlated DELETE by repeatedly scanning the + # virtual table. Resolve rowids from the indexed manifest first, then + # delete by primary rowid in bounded batches. + rowids = orphan_rowids + stale_rowids + for offset in range(0, len(rowids), SQLITE_DELETE_BATCH_SIZE): + batch = rowids[offset : offset + SQLITE_DELETE_BATCH_SIZE] + params = {f"rowid_{index}": rowid for index, rowid in enumerate(batch)} + placeholders = ", ".join(f":rowid_{index}" for index in range(len(batch))) + await session.execute( + text(f"DELETE FROM search_vector_embeddings WHERE rowid IN ({placeholders})"), + params, + ) await session.commit() async def search( diff --git a/tests/repository/test_sqlite_vec_reconciliation_performance.py b/tests/repository/test_sqlite_vec_reconciliation_performance.py new file mode 100644 index 000000000..42e13ff95 --- /dev/null +++ b/tests/repository/test_sqlite_vec_reconciliation_performance.py @@ -0,0 +1,105 @@ +"""Regression coverage for bounded sqlite-vec reconciliation work.""" + +from collections.abc import Callable +from typing import cast + +import pytest +from sqlalchemy import text +from sqlalchemy.exc import OperationalError as SAOperationalError +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker + +from basic_memory import db +from basic_memory.models.project import Project +from basic_memory.repository.semantic_vector_index import VectorIndexScope +from basic_memory.repository.sqlite_vec_index import SQLiteVecIndex + + +@pytest.mark.asyncio +async def test_sqlite_vec_reconciliation_is_bounded_by_project_manifest( + engine_factory: tuple[AsyncEngine, async_sessionmaker[AsyncSession]], + test_project: Project, +) -> None: + """A project reconciliation must not correlate every vector with every manifest row.""" + engine, session_maker = engine_factory + index = SQLiteVecIndex( + session_maker, + VectorIndexScope( + namespace="test", + embedding_identity="stub:model", + dimensions=4, + ), + ) + await index.initialize() + + manifest_rows = [ + { + "id": row_id, + "entity_id": row_id, + "project_id": test_project.id, + "chunk_key": f"entity:{row_id}:0", + "embedding_status": "ready" if row_id <= 2_000 else "pending", + } + for row_id in range(1, 3_002) + ] + vector_rows = [ + { + "rowid": row_id, + "project_id": test_project.id, + "embedding": "[1,0,0,0]", + } + for row_id in range(1, 3_003) + ] + async with db.scoped_session(session_maker) as session: + await session.execute( + text( + "INSERT INTO search_vector_chunks (" + "id, entity_id, project_id, chunk_key, chunk_text, source_hash, " + "entity_fingerprint, embedding_model, vector_index, embedding_status" + ") VALUES (" + ":id, :entity_id, :project_id, :chunk_key, 'text', 'hash', " + "'fingerprint', 'stub:model', 'sqlite-vec', :embedding_status)" + ), + manifest_rows, + ) + await session.execute( + text( + "INSERT INTO search_vector_embeddings " + "(rowid, project_id, embedding, source_hash) " + "VALUES (:rowid, :project_id, :embedding, 'hash')" + ), + vector_rows, + ) + await session.commit() + + max_progress_calls = 2_000 + progress_calls = 0 + + def interrupt_quadratic_plan() -> int: + nonlocal progress_calls + progress_calls += 1 + return int(progress_calls > max_progress_calls) + + async def set_progress_handler(handler: Callable[[], int] | None, steps: int) -> None: + async with engine.connect() as connection: + raw_connection = await connection.get_raw_connection() + driver_connection = raw_connection.driver_connection + assert driver_connection is not None + await driver_connection.set_progress_handler(handler, steps) + + await set_progress_handler(interrupt_quadratic_plan, 1_000) + try: + try: + await index.delete_orphans(test_project.id, []) + except SAOperationalError: + pytest.fail("sqlite-vec reconciliation exceeded its bounded SQLite operation budget") + finally: + await set_progress_handler(None, 0) + + async with db.scoped_session(session_maker) as session: + remaining = await session.execute( + text("SELECT rowid FROM search_vector_embeddings ORDER BY rowid") + ) + remaining_rowids = cast(list[int], remaining.scalars().all()) + + assert remaining_rowids == list(range(1, 2_001)) + assert progress_calls <= max_progress_calls