From 3e6579a4ecbbf25f83b7d84268ec43673ac44fe4 Mon Sep 17 00:00:00 2001 From: Trenton Holmes <797416+stumpylog@users.noreply.github.com> Date: Sun, 26 Jul 2026 20:30:04 -0700 Subject: [PATCH] Perf: point-delete vector store chunks by id instead of a document_id scan vec0 only gets an efficient lookup on a metadata column inside a KNN query; a plain document_id-filtered DELETE is a full table scan regardless of index size. Add a document_chunks side table (plain SQLite, real index) to look up chunk ids for a document, then delete each by its id primary key instead. Includes a v1 -> v2 migration to backfill document_chunks for indexes created before this existed. Co-Authored-By: Claude Sonnet 5 --- src/paperless_ai/tests/test_vector_store.py | 113 ++++++++++++-- src/paperless_ai/vector_store.py | 161 +++++++++++++++++--- 2 files changed, 242 insertions(+), 32 deletions(-) diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index 75cc0d17e..97c483633 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -262,6 +262,41 @@ class TestUpsert: assert ids == ["a3"] +class TestDocumentChunksIndex: + """document_chunks lets delete()/upsert_document() find a document's chunk + ids without a vec0 full table scan on document_id -- see + PaperlessSqliteVecVectorStore._delete_chunks_by_document_id.""" + + def _chunk_index_rows(self, store) -> list[tuple[str, str]]: + rows = store.client.execute( + "SELECT chunk_id, document_id FROM document_chunks ORDER BY chunk_id", + ).fetchall() + return [(r["chunk_id"], r["document_id"]) for r in rows] + + def test_add_populates_document_chunks(self, store) -> None: + store.add([make_node("a1", "1"), make_node("a2", "1"), make_node("b1", "2")]) + assert self._chunk_index_rows(store) == [ + ("a1", "1"), + ("a2", "1"), + ("b1", "2"), + ] + + def test_delete_clears_document_chunks(self, store) -> None: + store.add([make_node("a1", "1"), make_node("a2", "1"), make_node("b1", "2")]) + store.delete("1") + assert self._chunk_index_rows(store) == [("b1", "2")] + + def test_upsert_replaces_document_chunks(self, store) -> None: + store.add([make_node("a1", "1"), make_node("b1", "2")]) + store.upsert_document("1", [make_node("a2", "1")]) + assert self._chunk_index_rows(store) == [("a2", "1"), ("b1", "2")] + + def test_drop_table_clears_document_chunks(self, store) -> None: + store.add([make_node("a1", "1")]) + store.drop_table() + assert self._chunk_index_rows(store) == [] + + class TestMetadataCoercion: def test_none_metadata_values_become_empty_strings(self, store) -> None: node = make_node("a1", "1") @@ -381,6 +416,22 @@ class TestCompact: store.compact() store.compact(force=True) + def test_compact_preserves_document_chunks(self, store) -> None: + """document_chunks must survive the file-swap rebuild, or delete() + would silently stop finding chunk ids for anything indexed before a + compaction ran.""" + store.add([make_node("a1", "1"), make_node("b1", "2")]) + self._churn(store, 5) + store.compact(force=True) + + rows = store.client.execute( + "SELECT chunk_id, document_id FROM document_chunks WHERE document_id = '2'", + ).fetchall() + assert [(r["chunk_id"], r["document_id"]) for r in rows] == [("b1", "2")] + + store.delete("2") + assert "b1" not in _query(store, [0.0] * DIM, top_k=10).ids + def test_failed_compact_removes_temp_wal_and_shm( self, store, @@ -469,18 +520,22 @@ class TestMigrations: def test_reembed_migration_returns_true(self, store, tmp_path: Path) -> None: store.add([make_node("a1", "1")]) + # store was just created at the real (current) SCHEMA_VERSION, so the + # simulated pending migration must target one version past that -- + # not a hardcoded 1 -> 2, which would collide with the real + # document_chunks migration already registered in MIGRATIONS. + from paperless_ai import vector_store as vs_mod + + original = vs_mod.SCHEMA_VERSION migration = Migration( - from_version=1, - to_version=2, + from_version=original, + to_version=original + 1, kind="re-embed", description="test re-embed", ) MIGRATIONS.append(migration) try: - from paperless_ai import vector_store as vs_mod - - original = vs_mod.SCHEMA_VERSION - vs_mod.SCHEMA_VERSION = 2 + vs_mod.SCHEMA_VERSION = original + 1 result = store.check_and_run_migrations() finally: MIGRATIONS.remove(migration) @@ -537,26 +592,26 @@ class TestMigrations: ) dst.execute("COMMIT") + from paperless_ai import vector_store as vs_mod + + original = vs_mod.SCHEMA_VERSION migration = Migration( - from_version=1, - to_version=2, + from_version=original, + to_version=original + 1, kind="structural", description="test structural", apply=apply, ) MIGRATIONS.append(migration) try: - from paperless_ai import vector_store as vs_mod - - original = vs_mod.SCHEMA_VERSION - vs_mod.SCHEMA_VERSION = 2 + vs_mod.SCHEMA_VERSION = original + 1 result = store.check_and_run_migrations() finally: MIGRATIONS.remove(migration) vs_mod.SCHEMA_VERSION = original assert result is False - assert self._schema_version(store) == 2 + assert self._schema_version(store) == original + 1 ids = {n.node_id for n in store.get_nodes()} assert ids == {"a1", "b1"} @@ -566,6 +621,38 @@ class TestMigrations: store.compact(force=True) assert self._schema_version(store) == SCHEMA_VERSION + def test_v1_to_v2_migration_backfills_document_chunks(self, store) -> None: + """The real v1 -> v2 migration (registered in MIGRATIONS) must + backfill document_chunks for rows written before it existed, or + delete()/upsert_document() would find zero chunk ids for them and + leave their vec0 rows orphaned forever.""" + store.add([make_node("a1", "1"), make_node("a2", "1"), make_node("b1", "2")]) + # Simulate a pre-migration (schema v1) store: rows exist in the vec0 + # table, but document_chunks (added by the v1 -> v2 migration) has no + # entries for them, mirroring a real on-disk index created before + # this migration existed. + store.client.execute("DELETE FROM document_chunks") + store.client.execute( + "UPDATE index_meta SET value = '1' WHERE key = 'schema_version'", + ) + + assert store.check_and_run_migrations() is False # structural, not re-embed + assert self._schema_version(store) == SCHEMA_VERSION + + rows = store.client.execute( + "SELECT chunk_id, document_id FROM document_chunks ORDER BY chunk_id", + ).fetchall() + assert [(r["chunk_id"], r["document_id"]) for r in rows] == [ + ("a1", "1"), + ("a2", "1"), + ("b1", "2"), + ] + + # And delete() now actually removes rows for a document that + # predates the migration, instead of silently no-op'ing. + store.delete("1") + assert sorted(_query(store, [0.0] * DIM, top_k=10).ids) == ["b1"] + def test_stop_at_reembed_boundary(self, store) -> None: # Registry: structural v2, re-embed v3, structural v4. # Only v2 should apply; the re-embed boundary must stop execution diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index f6f333576..7b61e45e2 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -31,10 +31,20 @@ logger = logging.getLogger("paperless_ai.vector_store") DB_FILENAME = "llmindex.db" DEFAULT_TABLE_NAME = "documents" +_INSERT = ( + "INSERT INTO " + + DEFAULT_TABLE_NAME + + " (id, document_id, modified, node_content, embedding) VALUES (?, ?, ?, ?, ?)" +) + +_INSERT_CHUNK_INDEX = ( + "INSERT INTO document_chunks (chunk_id, document_id) VALUES (?, ?)" +) + # Current schema version. Written to index_meta at table creation and bumped # whenever a Migration is added to MIGRATIONS. check_and_run_migrations() uses # this to decide which migrations to run on an existing store. -SCHEMA_VERSION = 1 +SCHEMA_VERSION = 2 # compact(): rebuild when the cumulative rowid count exceeds this multiple of # the live row count. DELETEs on vec0 tables never reclaim space (upstream @@ -80,8 +90,10 @@ class Migration: ) -# Registry of all schema migrations in order. Empty at v1 -- this is the -# baseline. Add entries here (and bump SCHEMA_VERSION) when the schema changes. +# Registry of all schema migrations in order. Populated after the class body +# below, since v1 -> v2's apply() needs PaperlessSqliteVecVectorStore's own +# static methods. Add entries here (and bump SCHEMA_VERSION) when the schema +# changes. MIGRATIONS: list[Migration] = [] @@ -189,6 +201,18 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): conn.execute( "CREATE TABLE IF NOT EXISTS index_meta (key TEXT PRIMARY KEY, value TEXT)", ) + # vec0 metadata columns only get an efficient lookup path inside a KNN + # (MATCH) query; a plain `WHERE document_id = ?` is a full table scan + # regardless of index size. This plain, indexed table is how delete()/ + # upsert_document() find a document's chunk ids without that scan. + conn.execute( + "CREATE TABLE IF NOT EXISTS document_chunks " + "(chunk_id TEXT PRIMARY KEY, document_id TEXT NOT NULL)", + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_document_chunks_document_id " + "ON document_chunks (document_id)", + ) return conn @property @@ -259,6 +283,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): def drop_table(self) -> None: self._conn.execute("DROP TABLE IF EXISTS " + DEFAULT_TABLE_NAME) self._conn.execute("DELETE FROM index_meta") + self._conn.execute("DELETE FROM document_chunks") def stored_model_name(self) -> str | None: """Return the embedding model name recorded at table creation, or None.""" @@ -325,11 +350,40 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): _pack(node.get_embedding()), ) - _INSERT = ( - "INSERT INTO " - + DEFAULT_TABLE_NAME - + " (id, document_id, modified, node_content, embedding) VALUES (?, ?, ?, ?, ?)" - ) + def _index_chunks(self, rows: list[tuple[str, str, str, str, bytes]]) -> None: + """Record each row's (chunk_id, document_id) in the document_chunks + side table, kept in lockstep with every insert into the vec0 table.""" + self._conn.executemany( + _INSERT_CHUNK_INDEX, + [(row[0], row[1]) for row in rows], + ) + + def _delete_chunks_by_document_id(self, document_id: str) -> None: + """Delete all of a document's chunks via point-deletes on `id`. + + vec0 has no efficient lookup on the document_id metadata column + outside a KNN query (see _open_connection), so a plain + `DELETE ... WHERE document_id = ?` is a full table scan regardless of + index size. Looking the chunk ids up in document_chunks first (a real + indexed lookup) and deleting each by its `id` primary key instead + turns that scan into a handful of O(1) point deletes. + """ + chunk_ids = [ + row[0] + for row in self._conn.execute( + "SELECT chunk_id FROM document_chunks WHERE document_id = ?", + (str(document_id),), + ) + ] + for chunk_id in chunk_ids: + self._conn.execute( + "DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE id = ?", + (chunk_id,), + ) + self._conn.execute( + "DELETE FROM document_chunks WHERE document_id = ?", + (str(document_id),), + ) def _increment_total_inserts(self, count: int) -> None: """Increment the cumulative insert counter stored in index_meta. @@ -348,7 +402,8 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): rows = [self._row(node) for node in nodes] with self._transaction(): self._ensure_table(len(nodes[0].get_embedding())) - self._conn.executemany(self._INSERT, rows) + self._conn.executemany(_INSERT, rows) + self._index_chunks(rows) self._increment_total_inserts(len(rows)) return [node.node_id for node in nodes] @@ -365,22 +420,17 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): if nodes: self._ensure_table(len(nodes[0].get_embedding())) if self.table_exists(): - self._conn.execute( - "DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE document_id = ?", - (str(document_id),), - ) + self._delete_chunks_by_document_id(document_id) if rows: - self._conn.executemany(self._INSERT, rows) + self._conn.executemany(_INSERT, rows) + self._index_chunks(rows) self._increment_total_inserts(len(rows)) return [node.node_id for node in nodes] def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None: if self.table_exists(): with self._transaction(): - self._conn.execute( - "DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE document_id = ?", - (str(ref_doc_id),), - ) + self._delete_chunks_by_document_id(ref_doc_id) def _rows_to_nodes(self, rows: list[sqlite3.Row]) -> list[BaseNode]: nodes: list[BaseNode] = [] @@ -518,7 +568,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): # not cause an OOM during routine maintenance compactions. while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE): new_conn.executemany( - self._INSERT, + _INSERT, [ ( r["id"], @@ -530,6 +580,10 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): for r in batch ], ) + new_conn.executemany( + _INSERT_CHUNK_INDEX, + [(r["id"], r["document_id"]) for r in batch], + ) # Reset the cumulative counter: after compact, total_inserts == live. self._meta_set_on(new_conn, "total_inserts", str(live)) new_conn.execute("COMMIT") @@ -614,3 +668,72 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): raise new_conn.close() self._swap_in_compact(compact_path, db_path) + + +def _migrate_v1_to_v2_add_document_chunks( + src_conn: sqlite3.Connection, + dst_conn: sqlite3.Connection, + dim: int, +) -> None: + """v1 -> v2: backfill the document_chunks side table. + + document_chunks (see PaperlessSqliteVecVectorStore._open_connection) lets + delete()/upsert_document() find a document's chunk ids without a vec0 + full table scan on the document_id metadata column. Every row written + before this migration predates that table, so without backfilling, + deleting a pre-migration document would find zero chunk ids and leave its + vec0 rows permanently orphaned. This copies all live rows into the new + file -- the same approach compact() uses -- indexing each as it goes. + """ + PaperlessSqliteVecVectorStore._create_vec_table(dst_conn, dim) + PaperlessSqliteVecVectorStore._meta_set_on(dst_conn, "dim", str(dim)) + embed_model_row = src_conn.execute( + "SELECT value FROM index_meta WHERE key = 'embed_model'", + ).fetchone() + if embed_model_row is not None: + PaperlessSqliteVecVectorStore._meta_set_on( + dst_conn, + "embed_model", + embed_model_row["value"], + ) + + src_cursor = src_conn.execute( + "SELECT id, document_id, modified, node_content, embedding FROM " + + DEFAULT_TABLE_NAME, + ) + dst_conn.execute("BEGIN IMMEDIATE") + live = 0 + while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE): + dst_conn.executemany( + _INSERT, + [ + ( + r["id"], + r["document_id"], + r["modified"], + r["node_content"], + bytes(r["embedding"]), + ) + for r in batch + ], + ) + dst_conn.executemany( + _INSERT_CHUNK_INDEX, + [(r["id"], r["document_id"]) for r in batch], + ) + live += len(batch) + # This migration only ever copies live rows (like compact()), so the + # cumulative counter resets to match -- the new file has no bloat yet. + PaperlessSqliteVecVectorStore._meta_set_on(dst_conn, "total_inserts", str(live)) + dst_conn.execute("COMMIT") + + +MIGRATIONS.append( + Migration( + from_version=1, + to_version=2, + kind="structural", + description="add document_chunks side table for O(1) per-document deletes", + apply=_migrate_v1_to_v2_add_document_chunks, + ), +)