amorfo77/paperless-ngx
1
0
Fork 0
mirror of https://github.com/paperless-ngx/paperless-ngx.git synced 2026-07-28 06:44:57 +00:00
Code Issues Projects Releases Wiki Activity

Perf: point-delete vector store chunks by id instead of a document_id scan

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
Trenton Holmes
2026-07-26 20:31:15 -07:00
co-authored by Claude Sonnet 5
parent f4e7a2e9e4
commit 3e6579a4ec
2 changed files with 242 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/paperless_ai/tests/test_vector_store.py
+100 -13
View File
@@ -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
src/paperless_ai/vector_store.py
+142 -19
View File
@@ -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,
),
)