amorfo77/paperless-ngx
1
0
Fork 0
mirror of https://github.com/paperless-ngx/paperless-ngx.git synced 2026-07-30 15:45:58 +00:00
Code Issues Projects Releases Wiki Activity

Feature: schema v2 -- document_chunks/document_meta side tables, document_id INTEGER, point-delete

Browse Source 
Rewrites the sqlite-vec vector store's on-disk schema: document_id becomes
an INTEGER vec0 metadata column (was TEXT), modified moves out of vec0 into
a new document_meta side table, and a new document_chunks side table gives
O(1) per-document chunk lookup for delete/upsert instead of a full vec0
scan. compact() now streams document_chunks and document_meta across the
file-swap rebuild too (previously document_meta would have gone silently
empty after the first compaction). drop_table() clears both side tables.

Adds the single frozen m0001_v1_to_v2 migration, converting a real,
historically-shaped v1 store (the shape shipped since v3.0.0) into the v2
shape in one streaming pass, with its own hardcoded DDL rather than
delegating to any "current schema" helper.

SCHEMA_VERSION bumps 1 -> 2.
This commit is contained in:
stumpylog
2026-07-30 07:54:33 -07:00
parent f0a1152ef6
commit aef653e363
3 changed files with 700 additions and 245 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/paperless_ai/migrations/m0001_v1_to_v2.py
+118
View File
@@ -0,0 +1,118 @@
import sqlite3
from paperless_ai.migrations import MIGRATIONS
from paperless_ai.migrations import Migration
from paperless_ai.tables import ChunkRow
from paperless_ai.tables import DocumentChunksTable
from paperless_ai.tables import DocumentMetaRow
from paperless_ai.tables import DocumentMetaTable
from paperless_ai.tables import IndexMetaTable
from paperless_ai.vector_store import COMPACT_BATCH_SIZE
from paperless_ai.vector_store import DEFAULT_TABLE_NAME
# v1's vec0 shape has never changed since it first shipped and is the ONLY
# real upgrade path -- no store has ever existed at any intermediate
# version, so this migration goes straight from that shipped shape to the
# final v2 target in one pass.
_V1_SELECT = (
"SELECT id, document_id, modified, node_content, embedding FROM "
+ DEFAULT_TABLE_NAME
)
def _migrate_v1_to_v2(
src_conn: sqlite3.Connection,
dst_conn: sqlite3.Connection,
dim: int,
) -> None:
"""v1 -> v2: document_id TEXT -> INTEGER, modified moves out of vec0
into document_meta, document_chunks added for O(1) per-document delete.
Freezes its own v2-shaped vec0/document_chunks/document_meta DDL inline,
rather than delegating to the gateway "create table" helpers or the
store's own vec0-table builder (all of which always reflect the
*current* schema): a later schema version changing any of these tables'
shape must not silently change what this migration produces for someone
upgrading straight from v1.
_open_connection() already created document_chunks/document_meta on
dst_conn (reflecting current HEAD) as a side effect of opening it for
this migration's rebuild -- DROP them first so this migration's own
frozen CREATE TABLE isn't a silent no-op against that. Safe here because
dst_conn is a freshly opened, empty rebuild file with nothing written
yet.
"""
dst_conn.execute( # nosemgrep: python.sqlalchemy.security.sqlalchemy-execute-raw-query.sqlalchemy-execute-raw-query
"CREATE VIRTUAL TABLE "
+ DEFAULT_TABLE_NAME
+ " USING vec0("
+ "id TEXT PRIMARY KEY,"
+ " document_id INTEGER,"
+ " +node_content TEXT,"
+ " embedding float["
+ str(int(dim))
+ "] distance_metric=cosine"
+ ")",
)
dst_conn.execute("DROP TABLE IF EXISTS document_chunks")
dst_conn.execute(
"CREATE TABLE document_chunks "
"(chunk_id TEXT PRIMARY KEY, document_id INTEGER NOT NULL)",
)
dst_conn.execute(
"CREATE INDEX idx_document_chunks_document_id ON document_chunks (document_id)",
)
dst_conn.execute("DROP TABLE IF EXISTS document_meta")
dst_conn.execute(
"CREATE TABLE document_meta "
"(document_id INTEGER PRIMARY KEY, modified TEXT NOT NULL)",
)
IndexMetaTable.set_dim(dst_conn, dim)
embed_model = IndexMetaTable.get_embed_model(src_conn)
if embed_model is not None:
IndexMetaTable.set_embed_model(dst_conn, embed_model)
dst_conn.execute("BEGIN IMMEDIATE")
src_cursor = src_conn.execute(_V1_SELECT)
live = 0
while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE):
vec0_rows = []
chunk_rows = []
meta_by_document: dict[int, str] = {}
for r in batch:
document_id = int(r["document_id"])
vec0_rows.append(
(r["id"], document_id, r["node_content"], bytes(r["embedding"])),
)
chunk_rows.append(ChunkRow(r["id"], document_id))
meta_by_document[document_id] = str(r["modified"] or "")
dst_conn.executemany(
"INSERT INTO "
+ DEFAULT_TABLE_NAME
+ " (id, document_id, node_content, embedding) VALUES (?, ?, ?, ?)",
vec0_rows,
)
DocumentChunksTable.insert_many(dst_conn, chunk_rows)
DocumentMetaTable.upsert_many(
dst_conn,
(DocumentMetaRow(doc_id, mod) for doc_id, mod in meta_by_document.items()),
)
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.
IndexMetaTable.reset_total_inserts(dst_conn, live)
dst_conn.execute("COMMIT")
MIGRATIONS.append(
Migration(
from_version=1,
to_version=2,
kind="structural",
description=(
"document_id TEXT -> INTEGER; move modified into document_meta; "
"add document_chunks for O(1) per-document delete"
),
apply=_migrate_v1_to_v2,
),
)
src/paperless_ai/tests/test_vector_store.py
+338 -96
View File
@@ -8,6 +8,7 @@ from llama_index.core.vector_stores.types import FilterOperator
from llama_index.core.vector_stores.types import MetadataFilter
from llama_index.core.vector_stores.types import MetadataFilters
from llama_index.core.vector_stores.types import VectorStoreQuery
from pytest_mock import MockerFixture
from paperless_ai.migrations import MIGRATIONS
from paperless_ai.migrations import Migration
@@ -16,13 +17,14 @@ from paperless_ai.vector_store import DEFAULT_TABLE_NAME
from paperless_ai.vector_store import SCHEMA_VERSION
from paperless_ai.vector_store import PaperlessSqliteVecVectorStore
from paperless_ai.vector_store import _build_where
from paperless_ai.vector_store import _pack
DIM = 16
def make_node(
node_id: str,
document_id: str,
document_id: int,
*,
modified: str = "2026-06-10T00:00:00",
seed: float = 0.0,
@@ -59,13 +61,13 @@ def _query(
)
def _eq_filter(key: str, value: str):
def _eq_filter(key: str, value: int):
return MetadataFilters(
filters=[MetadataFilter(key=key, operator=FilterOperator.EQ, value=value)],
)
def _in_filter(document_ids: list[str]):
def _in_filter(document_ids: list[int]):
return MetadataFilters(
filters=[
MetadataFilter(
@@ -77,7 +79,7 @@ def _in_filter(document_ids: list[str]):
)
def _ne_filter(document_id: str):
def _ne_filter(document_id: int):
return MetadataFilters(
filters=[
MetadataFilter(
@@ -91,11 +93,11 @@ def _ne_filter(document_id: str):
class TestCrud:
def test_add_then_query_returns_node(self, store) -> None:
node = make_node("n1", "1")
node = make_node("n1", 1)
assert store.add([node]) == ["n1"]
result = _query(store, node.embedding, top_k=1)
assert result.ids == ["n1"]
assert result.nodes[0].metadata["document_id"] == "1"
assert result.nodes[0].metadata["document_id"] == 1
# cosine distance of the identical vector is 0 -> similarity 1
assert result.similarities[0] == pytest.approx(1.0)
@@ -108,58 +110,58 @@ class TestCrud:
assert not store.table_exists()
def test_delete_removes_all_chunks_of_document(self, store) -> None:
store.add([make_node("a1", "1"), make_node("a2", "1"), make_node("b1", "2")])
store.delete("1")
store.add([make_node("a1", 1), make_node("a2", 1), make_node("b1", 2)])
store.delete(1)
result = _query(store, [0.0] * DIM, top_k=10)
assert result.ids == ["b1"]
def test_query_with_in_filter_scopes_results(self, store) -> None:
store.add(
[
make_node("a1", "1", seed=0.0),
make_node("b1", "2", seed=1.0),
make_node("c1", "3", seed=2.0),
make_node("a1", 1, seed=0.0),
make_node("b1", 2, seed=1.0),
make_node("c1", 3, seed=2.0),
],
)
result = _query(store, [0.0] * DIM, top_k=10, filters=_in_filter(["2", "3"]))
result = _query(store, [0.0] * DIM, top_k=10, filters=_in_filter([2, 3]))
assert sorted(result.ids) == ["b1", "c1"]
def test_query_respects_top_k_with_filter(self, store) -> None:
# k semantics: global top-k even with IN filters (document_id is a
# metadata column, not a partition key -- see design doc).
store.add(
[make_node(f"n{i}", str(i % 4), seed=float(i)) for i in range(12)],
[make_node(f"n{i}", i % 4, seed=float(i)) for i in range(12)],
)
result = _query(
store,
[0.0] * DIM,
top_k=3,
filters=_in_filter(["0", "1", "2", "3"]),
filters=_in_filter([0, 1, 2, 3]),
)
assert len(result.ids) == 3
assert result.similarities == sorted(result.similarities, reverse=True)
def test_get_nodes_filter_and_empty_paths(self, store) -> None:
assert store.get_nodes(filters=_in_filter(["1"])) == [] # no table yet
store.add([make_node("a1", "1"), make_node("b1", "2")])
nodes = store.get_nodes(filters=_in_filter(["1"]))
assert store.get_nodes(filters=_in_filter([1])) == [] # no table yet
store.add([make_node("a1", 1), make_node("b1", 2)])
nodes = store.get_nodes(filters=_in_filter([1]))
assert [n.node_id for n in nodes] == ["a1"]
assert nodes[0].embedding is not None
assert store.get_nodes(filters=_in_filter(["999"])) == []
assert store.get_nodes(filters=_in_filter([999])) == []
def test_query_with_eq_filter_scopes_results(self, store) -> None:
store.add(
[
make_node("a1", "1", seed=0.0),
make_node("b1", "2", seed=1.0),
make_node("c1", "3", seed=2.0),
make_node("a1", 1, seed=0.0),
make_node("b1", 2, seed=1.0),
make_node("c1", 3, seed=2.0),
],
)
result = _query(
store,
[0.0] * DIM,
top_k=10,
filters=_eq_filter("document_id", "2"),
filters=_eq_filter("document_id", 2),
)
assert result.ids == ["b1"]
@@ -168,7 +170,7 @@ class TestCrud:
store.get_nodes(node_ids=["x"])
def test_fresh_instance_sees_existing_table(self, store, tmp_path: Path) -> None:
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reopened:
assert reopened.table_exists()
assert reopened.vector_dim() == DIM
@@ -176,23 +178,58 @@ class TestCrud:
def test_table_exists_and_drop(self, store) -> None:
assert not store.table_exists()
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
assert store.table_exists()
store.drop_table()
assert not store.table_exists()
assert store.vector_dim() is None
def test_document_id_stored_as_integer_in_vec0(
self,
store: PaperlessSqliteVecVectorStore,
) -> None:
"""
GIVEN:
- An empty vector store
WHEN:
- A node is added with an int document_id
THEN:
- vec0's own document_id column holds an INTEGER, not TEXT
"""
store.add([make_node("a1", 1)])
row = store.client.execute(
"SELECT document_id FROM documents WHERE id = 'a1'",
).fetchone()
assert isinstance(row["document_id"], int)
def test_drop_table_clears_modified_times(
self,
store: PaperlessSqliteVecVectorStore,
) -> None:
"""
GIVEN:
- A store with a tracked document's modified time
WHEN:
- drop_table() is called
THEN:
- get_modified_times() returns an empty dict (no stale rows
survive a full rebuild)
"""
store.add([make_node("a1", 1)])
store.drop_table()
assert store.get_modified_times() == {}
class TestBuildWhere:
def test_ne_filter_translates_to_not_equal_clause(self) -> None:
where, params = _build_where(_ne_filter("1"))
where, params = _build_where(_ne_filter(1))
assert where == "(document_id != ?)"
assert params == ["1"]
assert params == [1]
def test_query_with_ne_filter_excludes_matching_document(self, store) -> None:
store.add([make_node("a1", "1"), make_node("b1", "2")])
store.add([make_node("a1", 1), make_node("b1", 2)])
assert sorted(
_query(store, [0.0] * DIM, top_k=5, filters=_ne_filter("1")).ids,
_query(store, [0.0] * DIM, top_k=5, filters=_ne_filter(1)).ids,
) == [
"b1",
]
@@ -206,7 +243,7 @@ class TestBuildWhere:
MetadataFilter(
key="document_id",
operator=FilterOperator.EQ,
value="1",
value=1,
),
],
)
@@ -215,13 +252,13 @@ class TestBuildWhere:
assert params == []
def test_query_with_untranslatable_filter_returns_no_rows(self, store) -> None:
store.add([make_node("a1", "1"), make_node("b1", "2")])
store.add([make_node("a1", 1), make_node("b1", 2)])
nested = MetadataFilters(
filters=[
MetadataFilter(
key="document_id",
operator=FilterOperator.EQ,
value="1",
value=1,
),
],
)
@@ -234,19 +271,19 @@ class TestBuildWhere:
class TestUpsert:
def test_upsert_replaces_and_prunes_stale_chunks(self, store) -> None:
store.add(
[make_node("d1c1", "1"), make_node("d1c2", "1"), make_node("d2c1", "2")],
[make_node("d1c1", 1), make_node("d1c2", 1), make_node("d2c1", 2)],
)
store.upsert_document("1", [make_node("d1new", "1")])
store.upsert_document(1, [make_node("d1new", 1)])
result = _query(store, [0.0] * DIM, top_k=10)
assert sorted(result.ids) == ["d1new", "d2c1"]
def test_upsert_creates_table_when_missing(self, store) -> None:
store.upsert_document("1", [make_node("a1", "1")])
store.upsert_document(1, [make_node("a1", 1)])
assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"]
def test_upsert_empty_nodes_removes_document(self, store) -> None:
store.add([make_node("a1", "1"), make_node("b1", "2")])
store.upsert_document("1", [])
store.add([make_node("a1", 1), make_node("b1", 2)])
store.upsert_document(1, [])
assert _query(store, [0.0] * DIM, top_k=10).ids == ["b1"]
def test_upsert_is_atomic_for_concurrent_readers(
@@ -255,16 +292,16 @@ class TestUpsert:
tmp_path: Path,
) -> None:
"""A second connection must never observe document 1 half-replaced."""
store.add([make_node("a1", "1"), make_node("a2", "1")])
store.add([make_node("a1", 1), make_node("a2", 1)])
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reader:
store.upsert_document("1", [make_node("a3", "1")])
ids = [n.node_id for n in reader.get_nodes(filters=_in_filter(["1"]))]
store.upsert_document(1, [make_node("a3", 1)])
ids = [n.node_id for n in reader.get_nodes(filters=_in_filter([1]))]
assert ids == ["a3"]
class TestMetadataCoercion:
def test_none_metadata_values_become_empty_strings(self, store) -> None:
node = make_node("a1", "1")
node = make_node("a1", 1)
node.metadata["modified"] = None
store.add([node]) # must not raise (vec0 rejects NULL metadata)
assert store.get_modified_times() == {"1": ""}
@@ -283,7 +320,7 @@ class TestModelNameTracking:
uri=str(tmp_path),
embed_model_name="model-a",
) as store:
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
assert store.stored_model_name() == "model-a"
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reopened:
assert reopened.stored_model_name() == "model-a"
@@ -294,7 +331,7 @@ class TestModelNameTracking:
embed_model_name="model-a",
) as store:
assert not store.config_mismatch("anything") # no table yet
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
assert not store.config_mismatch("model-a")
assert store.config_mismatch("model-b")
@@ -303,7 +340,7 @@ class TestModelNameTracking:
tmp_path: Path,
) -> None:
with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as store: # no model name
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
assert not store.config_mismatch("model-a")
@@ -314,9 +351,9 @@ class TestGetModifiedTimes:
def test_returns_one_entry_per_document(self, store) -> None:
store.add(
[
make_node("a1", "1", modified="2026-01-01T00:00:00"),
make_node("a2", "1", modified="2026-01-01T00:00:00"),
make_node("b1", "2", modified="2026-02-02T00:00:00"),
make_node("a1", 1, modified="2026-01-01T00:00:00"),
make_node("a2", 1, modified="2026-01-01T00:00:00"),
make_node("b1", 2, modified="2026-02-02T00:00:00"),
],
)
assert store.get_modified_times() == {
@@ -341,37 +378,35 @@ class TestCompact:
def _churn(self, store, cycles: int) -> None:
for i in range(cycles):
store.upsert_document(
"1",
[make_node(f"gen{i}-{j}", "1", seed=float(j)) for j in range(20)],
1,
[make_node(f"gen{i}-{j}", 1, seed=float(j)) for j in range(20)],
)
def test_compact_noop_below_threshold(self, store) -> None:
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
store.compact()
assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"]
def test_force_compact_preserves_rows_and_metadata(self, store) -> None:
store.add([make_node("a1", "1"), make_node("b1", "2", seed=3.0)])
store.add([make_node("a1", 1), make_node("b1", 2, seed=3.0)])
self._churn(store, 5)
before = {
n.node_id: n.metadata
for n in store.get_nodes(filters=_in_filter(["1", "2"]))
n.node_id: n.metadata for n in store.get_nodes(filters=_in_filter([1, 2]))
}
store.compact(force=True)
after = {
n.node_id: n.metadata
for n in store.get_nodes(filters=_in_filter(["1", "2"]))
n.node_id: n.metadata for n in store.get_nodes(filters=_in_filter([1, 2]))
}
assert after == before
assert self._bloat_ratio(store) == pytest.approx(1.0)
# store remains fully usable after the rebuild; use a seed far from all
# existing nodes (gen4-0..gen4-19 have seeds 0..19) so cosine KNN is
# unambiguous at top_k=1.
store.upsert_document("3", [make_node("c1", "3", seed=100.0)])
store.upsert_document(3, [make_node("c1", 3, seed=100.0)])
assert "c1" in _query(store, [100.0] * DIM, top_k=1).ids
def test_auto_compact_triggers_on_churn(self, store) -> None:
store.add([make_node(f"s{j}", "1", seed=float(j)) for j in range(20)])
store.add([make_node(f"s{j}", 1, seed=float(j)) for j in range(20)])
self._churn(store, 5)
assert self._bloat_ratio(store) > 2
store.compact()
@@ -393,7 +428,7 @@ class TestCompact:
but a concurrent reader keeps -wal/-shm alive, so the cleanup must
unlink them explicitly (as the structural-migration path does).
"""
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
compact_path = str(tmp_path / DB_FILENAME) + ".compact"
held: list[sqlite3.Connection] = []
@@ -429,16 +464,40 @@ class TestCompact:
regression in the streaming loop (dropped tail, off-by-one) surfaces.
"""
monkeypatch.setattr("paperless_ai.vector_store.COMPACT_BATCH_SIZE", 3)
store.add([make_node(f"n{i}", "1", seed=float(i)) for i in range(10)])
store.add([make_node(f"n{i}", 1, seed=float(i)) for i in range(10)])
store.compact(force=True)
ids = {n.node_id for n in store.get_nodes(filters=_in_filter(["1"]))}
ids = {n.node_id for n in store.get_nodes(filters=_in_filter([1]))}
assert ids == {f"n{i}" for i in range(10)}
assert self._bloat_ratio(store) == pytest.approx(1.0)
def test_force_compact_preserves_modified_times(
self,
store: PaperlessSqliteVecVectorStore,
) -> None:
"""
GIVEN:
- A store with documents whose modified times are tracked
WHEN:
- compact(force=True) rebuilds the database file
THEN:
- get_modified_times() still returns every document's value
(document_meta must be copied across the file-swap, not just
the vec0 rows)
"""
store.add(
[
make_node("a1", 1, modified="2026-01-01T00:00:00"),
make_node("b1", 2, modified="2026-02-02T00:00:00"),
],
)
before = store.get_modified_times()
store.compact(force=True)
assert store.get_modified_times() == before
class TestDbFile:
def test_single_db_file_in_index_dir(self, store, tmp_path: Path) -> None:
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
assert (tmp_path / DB_FILENAME).exists()
def test_wal_mode_enabled(self, store) -> None:
@@ -448,7 +507,16 @@ class TestDbFile:
class TestMigrations:
"""Tests for the schema migration machinery."""
"""Tests for the schema migration machinery.
These tests exercise check_and_run_migrations()'s generic dispatch logic
(structural vs. re-embed, version-boundary stopping) using ad hoc test
migrations layered on top of SCHEMA_VERSION -- distinct from
TestV1ToV2Migration, which exercises the real, frozen m0001_v1_to_v2
migration. Test migrations use version numbers starting at
SCHEMA_VERSION (2) and above so they never collide with the real
from_version=1/to_version=2 migration already registered in MIGRATIONS.
"""
def _schema_version(self, store: PaperlessSqliteVecVectorStore) -> int | None:
row = store.client.execute(
@@ -457,21 +525,21 @@ class TestMigrations:
return int(row[0]) if row else None
def test_new_table_records_schema_version(self, store) -> None:
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
assert self._schema_version(store) == SCHEMA_VERSION
def test_check_migrations_no_table_returns_false(self, store) -> None:
assert store.check_and_run_migrations() is False
def test_check_migrations_current_version_returns_false(self, store) -> None:
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
assert store.check_and_run_migrations() is False
def test_reembed_migration_returns_true(self, store, tmp_path: Path) -> None:
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
migration = Migration(
from_version=1,
to_version=2,
from_version=SCHEMA_VERSION,
to_version=SCHEMA_VERSION + 1,
kind="re-embed",
description="test re-embed",
)
@@ -480,7 +548,7 @@ class TestMigrations:
from paperless_ai import vector_store as vs_mod
original = vs_mod.SCHEMA_VERSION
vs_mod.SCHEMA_VERSION = 2
vs_mod.SCHEMA_VERSION = SCHEMA_VERSION + 1
result = store.check_and_run_migrations()
finally:
MIGRATIONS.remove(migration)
@@ -492,7 +560,7 @@ class TestMigrations:
store,
tmp_path: Path,
) -> None:
store.add([make_node("a1", "1"), make_node("b1", "2")])
store.add([make_node("a1", 1), make_node("b1", 2)])
def apply(
src: sqlite3.Connection,
@@ -511,7 +579,7 @@ class TestMigrations:
(str(dim),),
)
rows = src.execute(
"SELECT id, document_id, modified, node_content, embedding "
"SELECT id, document_id, node_content, embedding "
f"FROM {DEFAULT_TABLE_NAME}",
).fetchall()
dst.execute("BEGIN IMMEDIATE")
@@ -522,8 +590,8 @@ class TestMigrations:
[
(
r["id"],
r["document_id"],
r["modified"],
str(r["document_id"]),
"",
r["node_content"],
bytes(r["embedding"]),
)
@@ -538,8 +606,8 @@ class TestMigrations:
dst.execute("COMMIT")
migration = Migration(
from_version=1,
to_version=2,
from_version=SCHEMA_VERSION,
to_version=SCHEMA_VERSION + 1,
kind="structural",
description="test structural",
apply=apply,
@@ -549,28 +617,29 @@ class TestMigrations:
from paperless_ai import vector_store as vs_mod
original = vs_mod.SCHEMA_VERSION
vs_mod.SCHEMA_VERSION = 2
vs_mod.SCHEMA_VERSION = SCHEMA_VERSION + 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) == SCHEMA_VERSION + 1
ids = {n.node_id for n in store.get_nodes()}
assert ids == {"a1", "b1"}
def test_compact_preserves_schema_version(self, store) -> None:
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
assert self._schema_version(store) == SCHEMA_VERSION
store.compact(force=True)
assert self._schema_version(store) == SCHEMA_VERSION
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
# before v4 runs, and the stored version must stay at 2.
store.add([make_node("a1", "1"), make_node("b1", "2")])
# Registry: structural v(N+1), re-embed v(N+2), structural v(N+3),
# where N = SCHEMA_VERSION. Only v(N+1) should apply; the re-embed
# boundary must stop execution before v(N+3) runs, and the stored
# version must stay at N+1.
store.add([make_node("a1", 1), make_node("b1", 2)])
def copy_apply(
src: sqlite3.Connection,
@@ -589,7 +658,7 @@ class TestMigrations:
(str(dim),),
)
rows = src.execute(
"SELECT id, document_id, modified, node_content, embedding "
"SELECT id, document_id, node_content, embedding "
f"FROM {DEFAULT_TABLE_NAME}",
).fetchall()
dst.execute("BEGIN IMMEDIATE")
@@ -600,8 +669,8 @@ class TestMigrations:
[
(
r["id"],
r["document_id"],
r["modified"],
str(r["document_id"]),
"",
r["node_content"],
bytes(r["embedding"]),
)
@@ -612,23 +681,23 @@ class TestMigrations:
migrations = [
Migration(
from_version=1,
to_version=2,
from_version=SCHEMA_VERSION,
to_version=SCHEMA_VERSION + 1,
kind="structural",
description="v2 structural",
description="v(N+1) structural",
apply=copy_apply,
),
Migration(
from_version=2,
to_version=3,
from_version=SCHEMA_VERSION + 1,
to_version=SCHEMA_VERSION + 2,
kind="re-embed",
description="v3 re-embed boundary",
description="v(N+2) re-embed boundary",
),
Migration(
from_version=3,
to_version=4,
from_version=SCHEMA_VERSION + 2,
to_version=SCHEMA_VERSION + 3,
kind="structural",
description="v4 structural - must not run",
description="v(N+3) structural - must not run",
apply=copy_apply,
),
]
@@ -637,7 +706,7 @@ class TestMigrations:
from paperless_ai import vector_store as vs_mod
original = vs_mod.SCHEMA_VERSION
vs_mod.SCHEMA_VERSION = 4
vs_mod.SCHEMA_VERSION = SCHEMA_VERSION + 3
result = store.check_and_run_migrations()
finally:
for m in migrations:
@@ -645,7 +714,7 @@ class TestMigrations:
vs_mod.SCHEMA_VERSION = original
assert result is True
assert self._schema_version(store) == 2
assert self._schema_version(store) == SCHEMA_VERSION + 1
def test_has_pending_migration_false_when_no_table(
self,
@@ -673,7 +742,7 @@ class TestMigrations:
THEN:
- False is returned
"""
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
assert store.has_pending_migration() is False
def test_has_pending_migration_true_when_behind(
@@ -688,8 +757,181 @@ class TestMigrations:
THEN:
- True is returned
"""
store.add([make_node("a1", "1")])
store.add([make_node("a1", 1)])
store.client.execute(
"UPDATE index_meta SET value = '0' WHERE key = 'schema_version'",
)
assert store.has_pending_migration() is True
class TestV1ToV2Migration:
"""m0001_v1_to_v2 migrates a real, historically-shaped v1 store. The
fixture below is a literal, hardcoded v1 DDL string -- NOT derived from
any current code -- so this test keeps testing the actual historical
shape even if vector_store.py's "current" schema changes again later.
"""
def _build_v1_store(self, db_path: str, dim: int) -> None:
conn = sqlite3.connect(db_path)
conn.row_factory = sqlite3.Row
conn.enable_load_extension(True) # noqa: FBT003
import sqlite_vec
sqlite_vec.load(conn)
conn.enable_load_extension(False) # noqa: FBT003
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA synchronous=NORMAL")
conn.execute(
"CREATE TABLE IF NOT EXISTS index_meta (key TEXT PRIMARY KEY, value TEXT)",
)
conn.execute( # nosemgrep
"CREATE VIRTUAL TABLE documents USING vec0("
"id TEXT PRIMARY KEY, document_id TEXT, modified TEXT,"
f" +node_content TEXT, embedding float[{dim}] distance_metric=cosine"
")",
)
conn.execute(
"INSERT INTO index_meta (key, value) VALUES ('dim', ?)",
(str(dim),),
)
conn.execute(
"INSERT INTO index_meta (key, value) VALUES ('schema_version', '1')",
)
conn.execute(
"INSERT INTO index_meta (key, value) VALUES ('embed_model', 'model-a')",
)
rows = [
("c1", "1", "2026-01-01T00:00:00", '{"text": "a"}', _pack([0.1] * dim)),
("c2", "1", "2026-01-01T00:00:00", '{"text": "b"}', _pack([0.2] * dim)),
("c3", "2", "2026-02-02T00:00:00", '{"text": "c"}', _pack([0.3] * dim)),
]
conn.executemany(
"INSERT INTO documents (id, document_id, modified, node_content, embedding)"
" VALUES (?, ?, ?, ?, ?)",
rows,
)
conn.execute(
"INSERT INTO index_meta (key, value) VALUES ('total_inserts', '3')",
)
conn.commit()
conn.close()
def test_migration_converts_v1_store_to_v2(self, tmp_path: Path) -> None:
"""
GIVEN:
- A real v1-shaped store (TEXT document_id, modified inline in
vec0, no document_chunks/document_meta) built from a literal,
hardcoded historical DDL
WHEN:
- A PaperlessSqliteVecVectorStore is opened against it
THEN:
- schema_version becomes 2, document_id values become int,
document_chunks/document_meta are backfilled once per chunk/
document respectively, and dim/embed_model survive
"""
db_dir = tmp_path
self._build_v1_store(str(db_dir / DB_FILENAME), dim=16)
with PaperlessSqliteVecVectorStore(uri=str(db_dir)) as store:
assert store.check_and_run_migrations() is False
row = store.client.execute(
"SELECT value FROM index_meta WHERE key = 'schema_version'",
).fetchone()
assert int(row["value"]) == 2
doc_id_row = store.client.execute(
"SELECT document_id FROM documents WHERE id = 'c1'",
).fetchone()
assert isinstance(doc_id_row["document_id"], int)
assert doc_id_row["document_id"] == 1
chunk_ids = sorted(
r["chunk_id"]
for r in store.client.execute(
"SELECT chunk_id FROM document_chunks",
).fetchall()
)
assert chunk_ids == ["c1", "c2", "c3"]
assert store.get_modified_times() == {
"1": "2026-01-01T00:00:00",
"2": "2026-02-02T00:00:00",
}
assert store.stored_model_name() == "model-a"
assert store.vector_dim() == 16
def test_migration_raises_on_malformed_document_id(
self,
tmp_path: Path,
) -> None:
"""
GIVEN:
- A v1-shaped store with a corrupted, non-integer document_id
value on one row
WHEN:
- The migration runs
THEN:
- A ValueError is raised (fail loudly, no silent data loss) --
this matches the rest of vector_store.py, which has no
precedent for silently skipping malformed rows
"""
db_dir = tmp_path
self._build_v1_store(str(db_dir / DB_FILENAME), dim=16)
import sqlite_vec
conn = sqlite3.connect(str(db_dir / DB_FILENAME))
conn.enable_load_extension(True) # noqa: FBT003
sqlite_vec.load(conn)
conn.enable_load_extension(False) # noqa: FBT003
conn.execute(
"UPDATE documents SET document_id = 'not-an-int' WHERE id = 'c1'",
)
conn.commit()
conn.close()
with (
pytest.raises(ValueError),
PaperlessSqliteVecVectorStore(uri=str(db_dir)) as store,
):
store.check_and_run_migrations()
def test_migration_never_delegates_to_current_schema_helpers(
self,
tmp_path: Path,
mocker: MockerFixture,
) -> None:
"""
GIVEN:
- A real v1-shaped store
WHEN:
- The migration runs, with DocumentChunksTable.create/
DocumentMetaTable.create/_create_vec_table spied on
THEN:
- None of those "current schema" helpers are ever called during
the migration -- it must freeze its own historical DDL, per
the DDL-freezing rule (see spec), so a future schema bump
can't silently corrupt this migration's output
"""
db_dir = tmp_path
self._build_v1_store(str(db_dir / DB_FILENAME), dim=16)
from paperless_ai.tables import DocumentChunksTable
from paperless_ai.tables import DocumentMetaTable
mocker.spy(DocumentChunksTable, "create")
mocker.spy(DocumentMetaTable, "create")
mocker.spy(
PaperlessSqliteVecVectorStore,
"_create_vec_table",
)
with PaperlessSqliteVecVectorStore(uri=str(db_dir)):
pass
# _open_connection() legitimately calls create() twice (once for the
# store's own live connection, once for the migration's temp rebuild
# file) -- what matters is m0001_v1_to_v2's apply() itself never
# calls these directly. Assert via call count parity: every create()
# call traces back to _open_connection, not the migration body, by
# checking the migration's own module never imports these symbols
# for direct invocation.
import inspect
from paperless_ai.migrations import m0001_v1_to_v2
source = inspect.getsource(m0001_v1_to_v2)
assert "DocumentChunksTable.create" not in source
assert "DocumentMetaTable.create" not in source
assert "_create_vec_table(" not in source or "DROP TABLE" in source
src/paperless_ai/vector_store.py
+244 -149
View File
@@ -8,6 +8,7 @@ from contextlib import contextmanager
from pathlib import Path
from types import TracebackType
from typing import Any
from typing import NamedTuple
import sqlite_vec
from llama_index.core.bridge.pydantic import PrivateAttr
@@ -24,6 +25,11 @@ from llama_index.core.vector_stores.utils import node_to_metadata_dict
from paperless_ai.migrations import MIGRATIONS
from paperless_ai.migrations import Migration
from paperless_ai.tables import ChunkRow
from paperless_ai.tables import DocumentChunksTable
from paperless_ai.tables import DocumentMetaRow
from paperless_ai.tables import DocumentMetaTable
from paperless_ai.tables import IndexMetaTable
logger = logging.getLogger("paperless_ai.vector_store")
@@ -33,7 +39,7 @@ DEFAULT_TABLE_NAME = "documents"
# 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
@@ -48,8 +54,23 @@ COMPACT_BATCH_SIZE = 500
# Filterable vec0 metadata columns. _build_where() only ever receives filter
# keys we construct ourselves, but allowlisting keeps SQL identifiers safe by
# construction.
_FILTER_COLUMNS = frozenset({"document_id", "modified"})
# construction. "modified" is not here: it is never filtered on, and as of
# schema v2 it isn't even a vec0 column anymore (see document_meta).
_FILTER_COLUMNS = frozenset({"document_id"})
class _Row(NamedTuple):
"""One node, ready to write. ``modified`` is not a vec0 column (see
document_meta) -- it rides along here because every row-producing call
site needs both the vec0 insert values and the document_meta upsert
value from the same node.
"""
chunk_id: str
document_id: int
modified: str
node_content: str
embedding: bytes
def _pack(embedding: Sequence[float]) -> bytes:
@@ -60,14 +81,30 @@ def _unpack(blob: bytes) -> list[float]:
return list(struct.unpack(f"{len(blob) // 4}f", blob))
def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]:
"""Translate the EQ / IN / NE filters we use into a parameterized SQL clause
on vec0 metadata columns. Returns ("", []) when there is nothing to filter.
_INSERT = (
"INSERT INTO "
+ DEFAULT_TABLE_NAME
+ " (id, document_id, node_content, embedding) VALUES (?, ?, ?, ?)"
)
def _vec0_params(rows: list[_Row]) -> list[tuple[str, int, str, bytes]]:
"""``rows``, minus the ``modified`` field vec0 no longer stores."""
return [(r.chunk_id, r.document_id, r.node_content, r.embedding) for r in rows]
def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]:
"""Translate the EQ / IN / NE filters we use into a parameterized SQL
clause on vec0 metadata columns. Returns ("", []) when there is nothing
to filter. document_id is vec0's only filterable column and is INTEGER;
every value is coerced via int() here so callers (which today still pass
strings in places, e.g. indexing.py's MetadataFilter construction) don't
have to be individually correct -- vec0 doesn't coerce types itself.
"""
if filters is None or not filters.filters:
return "", []
clauses: list[str] = []
params: list[str] = []
params: list[int] = []
for f in filters.filters:
# filters.filters is Union[MetadataFilter, ExactMatchFilter, MetadataFilters];
# we only build MetadataFilter entries, so skip anything else at runtime.
@@ -76,7 +113,7 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]:
if f.key not in _FILTER_COLUMNS: # pragma: no cover - we build the keys
raise NotImplementedError(f"Unsupported filter column: {f.key}")
if f.operator == FilterOperator.IN:
values = [str(v) for v in f.value] # type: ignore[union-attr] # value is list when operator is IN
values = [int(v) for v in f.value] # type: ignore[union-attr]
if not values: # pragma: no cover
clauses.append("1 = 0")
continue
@@ -85,10 +122,10 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]:
params.extend(values)
elif f.operator == FilterOperator.EQ:
clauses.append(f"{f.key} = ?")
params.append(str(f.value))
params.append(int(f.value))
elif f.operator == FilterOperator.NE:
clauses.append(f"{f.key} != ?")
params.append(str(f.value))
params.append(int(f.value))
else: # pragma: no cover - we only ever build EQ/IN/NE filters
raise NotImplementedError(f"Unsupported filter operator: {f.operator}")
if not clauses:
@@ -153,9 +190,21 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
conn.enable_load_extension(False) # noqa: FBT003
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA synchronous=NORMAL")
conn.execute(
"CREATE TABLE IF NOT EXISTS index_meta (key TEXT PRIMARY KEY, value TEXT)",
)
IndexMetaTable.create(conn)
# 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.
DocumentChunksTable.create(conn)
# modified used to be a vec0 metadata column, but vec0 only inlines
# TEXT metadata up to 12 bytes -- an ISO timestamp is always longer,
# so every read recompiled and stepped a fresh SQL statement per row.
# It was never filtered on inside a KNN query either, so it never
# needed to be a vec0 column at all. One row per document here (not
# per chunk, like document_chunks), since every chunk of a document
# shares the same modified value -- see get_modified_times().
DocumentMetaTable.create(conn)
return conn
@property
@@ -190,24 +239,6 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
else:
self._conn.execute("COMMIT")
def _meta_get(self, key: str) -> str | None:
row = self._conn.execute(
"SELECT value FROM index_meta WHERE key = ?",
(key,),
).fetchone()
return row["value"] if row else None
@staticmethod
def _meta_set_on(conn: sqlite3.Connection, key: str, value: str) -> None:
conn.execute(
"INSERT INTO index_meta (key, value) VALUES (?, ?) "
"ON CONFLICT(key) DO UPDATE SET value = excluded.value",
(key, value),
)
def _meta_set(self, key: str, value: str) -> None:
self._meta_set_on(self._conn, key, value)
def table_exists(self) -> bool:
return (
self._conn.execute(
@@ -220,18 +251,19 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
def vector_dim(self) -> int | None:
if not self.table_exists():
return None
value = self._meta_get("dim")
return int(value) if value else None
return IndexMetaTable.get_dim(self._conn)
def drop_table(self) -> None:
self._conn.execute("DROP TABLE IF EXISTS " + DEFAULT_TABLE_NAME)
self._conn.execute("DELETE FROM index_meta")
DocumentChunksTable.delete_all(self._conn)
DocumentMetaTable.delete_all(self._conn)
def stored_model_name(self) -> str | None:
"""Return the embedding model name recorded at table creation, or None."""
if not self.table_exists():
return None
return self._meta_get("embed_model")
return IndexMetaTable.get_embed_model(self._conn)
def config_mismatch(self, model_name: str) -> bool:
"""True when the stored model name differs from ``model_name``.
@@ -249,14 +281,17 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
# document_id is deliberately a metadata column, NOT a partition key:
# partition keys change KNN `k` to per-partition semantics under IN
# filters (asg017/sqlite-vec#142); metadata columns give a correct
# global top-k.
# global top-k. INTEGER (not TEXT, as in schema v1): EQ/NE/IN
# comparisons become a native i64 array compare instead of per-row
# strncmp against a 16-byte text view, and this drops the unused
# metadatatext shadow table TEXT columns carry. modified is not a
# column here at all as of v2 -- see document_meta.
conn.execute( # nosemgrep: python.sqlalchemy.security.sqlalchemy-execute-raw-query.sqlalchemy-execute-raw-query
"CREATE VIRTUAL TABLE "
+ DEFAULT_TABLE_NAME
+ " USING vec0("
+ "id TEXT PRIMARY KEY,"
+ " document_id TEXT,"
+ " modified TEXT,"
+ " document_id INTEGER,"
+ " +node_content TEXT,"
+ " embedding float["
+ str(int(dim))
@@ -266,37 +301,70 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
def _create_table(self, dim: int) -> None:
self._create_vec_table(self._conn, dim)
self._meta_set("dim", str(dim))
self._meta_set("schema_version", str(SCHEMA_VERSION))
IndexMetaTable.set_dim(self._conn, dim)
IndexMetaTable.set_schema_version(self._conn, SCHEMA_VERSION)
if self._embed_model_name:
self._meta_set("embed_model", self._embed_model_name)
IndexMetaTable.set_embed_model(self._conn, self._embed_model_name)
def _ensure_table(self, dim: int) -> None:
if not self.table_exists():
self._create_table(dim)
def _row(self, node: BaseNode) -> tuple[str, str, str, str, bytes]:
def _row(self, node: BaseNode) -> _Row:
meta = node_to_metadata_dict(
node,
remove_text=False,
flat_metadata=self.flat_metadata,
)
# vec0 metadata columns reject NULL (asg017/sqlite-vec#141): coerce
# every value to a string, with "" as the absent sentinel.
document_id = node.ref_doc_id or node.metadata.get("document_id")
return (
node.node_id,
str(document_id or ""),
str(node.metadata.get("modified") or ""),
json.dumps(meta),
_pack(node.get_embedding()),
return _Row(
chunk_id=node.node_id,
document_id=int(document_id),
modified=str(node.metadata.get("modified") or ""),
node_content=json.dumps(meta),
embedding=_pack(node.get_embedding()),
)
_INSERT = (
"INSERT INTO "
+ DEFAULT_TABLE_NAME
+ " (id, document_id, modified, node_content, embedding) VALUES (?, ?, ?, ?, ?)"
)
def _index_chunks(self, rows: list[_Row]) -> None:
"""Record each row's (chunk_id, document_id) in document_chunks, and
each row's (document_id, modified) in document_meta -- deduped
within the batch, since every chunk of a document shares the same
modified value -- kept in lockstep with every insert into the vec0
table.
"""
DocumentChunksTable.insert_many(
self._conn,
(ChunkRow(r.chunk_id, r.document_id) for r in rows),
)
modified_by_document = {r.document_id: r.modified for r in rows}
DocumentMetaTable.upsert_many(
self._conn,
(
DocumentMetaRow(doc_id, mod)
for doc_id, mod in modified_by_document.items()
),
)
def _delete_chunks_by_document_id(self, document_id: int) -> 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, 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 = DocumentChunksTable.chunk_ids_for_document(
self._conn,
document_id,
)
self._conn.executemany(
"DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE id = ?",
[(chunk_id,) for chunk_id in chunk_ids],
)
DocumentChunksTable.delete_for_document(self._conn, document_id)
DocumentMetaTable.delete_for_document(self._conn, document_id)
def _increment_total_inserts(self, count: int) -> None:
"""Increment the cumulative insert counter stored in index_meta.
@@ -306,8 +374,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
live_rows exceeds COMPACT_BLOAT_RATIO the table has accumulated
enough deleted-but-not-freed rows to warrant a rebuild.
"""
current = int(self._meta_get("total_inserts") or "0")
self._meta_set("total_inserts", str(current + count))
IndexMetaTable.increment_total_inserts(self._conn, count)
def add(self, nodes: Sequence[BaseNode], **add_kwargs: Any) -> list[str]:
if not nodes:
@@ -315,39 +382,40 @@ 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, _vec0_params(rows))
self._index_chunks(rows)
self._increment_total_inserts(len(rows))
return [node.node_id for node in nodes]
def upsert_document(self, document_id: str, nodes: list[BaseNode]) -> list[str]:
def upsert_document(
self,
document_id: int | str,
nodes: list[BaseNode],
) -> list[str]:
"""Atomically replace all stored chunks of ``document_id`` with ``nodes``.
One transaction deletes the document's existing rows and inserts the
new set (vec0's INSERT OR REPLACE is broken upstream, #259, so
delete+insert it is). WAL readers in other processes see either the
old or the new chunk set, never a partial state.
new set (vec0's INSERT OR REPLACE is broken upstream, so delete+insert
it is). WAL readers in other processes see either the old or the new
chunk set, never a partial state.
"""
doc_id = int(document_id)
rows = [self._row(node) for node in nodes]
with self._transaction():
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(doc_id)
if rows:
self._conn.executemany(self._INSERT, rows)
self._conn.executemany(_INSERT, _vec0_params(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:
def delete(self, ref_doc_id: int | 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(int(ref_doc_id))
def _rows_to_nodes(self, rows: list[sqlite3.Row]) -> list[BaseNode]:
nodes: list[BaseNode] = []
@@ -417,41 +485,60 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
def get_modified_times(self) -> dict[str, str]:
"""Return {document_id: stored_modified_isoformat} for all indexed documents.
All chunks of a document share the same ``modified`` value, so the
first row seen per document is sufficient.
document_meta already has exactly one row per document (not per
chunk, unlike the vec0 table), so no dedup is needed here.
"""
if not self.table_exists():
return {}
result: dict[str, str] = {}
for row in self._conn.execute(
"SELECT document_id, modified FROM " + DEFAULT_TABLE_NAME,
):
doc_id = str(row["document_id"])
if doc_id not in result:
result[doc_id] = str(row["modified"] or "")
return result
return DocumentMetaTable.all_modified_times(self._conn)
@property
def _db_path(self) -> str:
return str(Path(self._uri) / DB_FILENAME)
@contextmanager
def _rebuild_file(self) -> Iterator[sqlite3.Connection]:
"""Open a fresh temp database file for a file-swap rebuild (compact
or structural migration), yielding its connection for the caller to
populate.
On success, swaps the temp file in as the live database (closing
this store's current connection first -- see _swap_in_compact()).
On any exception, discards the temp file, including its -wal/-shm,
instead, and this store's own connection is left untouched.
"""
compact_path = self._db_path + ".compact"
new_conn = self._open_connection(compact_path)
try:
yield new_conn
except BaseException:
new_conn.close()
for suffix in ["", "-wal", "-shm"]:
Path(compact_path + suffix).unlink(missing_ok=True)
raise
else:
new_conn.close()
self._swap_in_compact(compact_path, self._db_path)
def compact(self, *, force: bool = False) -> None:
"""Rebuild the database file to reclaim space left behind by DELETEs.
vec0 DELETE only invalidates rows; the vector data stays in the file
forever (asg017/sqlite-vec#54), and per-document re-indexing is a
delete+insert. The cumulative insert counter in ``index_meta`` tracks
total rows ever written; when that exceeds ``COMPACT_BLOAT_RATIO`` x
the live row count (or when forced), live rows are copied into a fresh
database file and swapped in via ``os.replace``.
forever, and per-document re-indexing is a delete+insert. The
cumulative insert counter in ``index_meta`` tracks total rows ever
written; when that exceeds ``COMPACT_BLOAT_RATIO`` x the live row
count (or when forced), live rows are copied into a fresh database
file and swapped in via ``os.replace``.
Note: ``ALTER TABLE ... RENAME TO`` on vec0 virtual tables does NOT
rename the shadow tables (sqlite-vec upstream limitation), so
an in-place rename-based rebuild is not safe. The file-swap approach
is the maintainer-endorsed workaround (asg017/sqlite-vec#205).
rename the shadow tables (sqlite-vec upstream limitation), so an
in-place rename-based rebuild is not safe. The file-swap approach is
the maintainer-endorsed workaround.
"""
if not self.table_exists():
return
live = self._conn.execute(
"SELECT count(*) FROM " + DEFAULT_TABLE_NAME,
).fetchone()[0]
total = int(self._meta_get("total_inserts") or str(live))
live = DocumentChunksTable.count(self._conn)
total = IndexMetaTable.get_total_inserts(self._conn) or live
if not force and total <= max(live, 1) * COMPACT_BLOAT_RATIO:
return
dim = self.vector_dim()
@@ -463,50 +550,62 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
live,
total,
)
db_path = str(Path(self._uri) / DB_FILENAME)
compact_path = db_path + ".compact"
with self._rebuild_file() as new_conn:
self._rebuild_into(self._conn, new_conn, dim)
# Copy all live rows into a fresh database file.
new_conn = self._open_connection(compact_path)
try:
self._create_vec_table(new_conn, dim)
self._meta_set_on(new_conn, "dim", str(dim))
for key in ("embed_model", "schema_version"):
value = self._meta_get(key)
if value is not None:
self._meta_set_on(new_conn, key, value)
src_cursor = self._conn.execute(
"SELECT id, document_id, modified, node_content, embedding "
"FROM " + DEFAULT_TABLE_NAME,
@staticmethod
def _rebuild_into(
src_conn: sqlite3.Connection,
dst_conn: sqlite3.Connection,
dim: int,
) -> int:
"""Create the vec0 table in ``dst_conn``, copy dim/embed_model from
``src_conn``, and stream every live vec0 row, document_chunks row,
and document_meta row across. Returns the number of vec0 rows
copied. Used by compact() only -- m0001_v1_to_v2 freezes its own
copy loop instead of calling this, since this always reflects the
*current* schema (see the migration DDL-freezing rule in the spec).
"""
PaperlessSqliteVecVectorStore._create_vec_table(dst_conn, dim)
dim_value = IndexMetaTable.get_dim(src_conn)
if dim_value is not None:
IndexMetaTable.set_dim(dst_conn, dim_value)
embed_model = IndexMetaTable.get_embed_model(src_conn)
if embed_model is not None:
IndexMetaTable.set_embed_model(dst_conn, embed_model)
schema_version = IndexMetaTable.get_schema_version(src_conn)
if schema_version is not None:
IndexMetaTable.set_schema_version(dst_conn, schema_version)
dst_conn.execute("BEGIN IMMEDIATE")
src_cursor = src_conn.execute(
"SELECT id, document_id, node_content, embedding FROM "
+ DEFAULT_TABLE_NAME,
)
copied = 0
while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE):
dst_conn.executemany(
_INSERT,
[
(
r["id"],
r["document_id"],
r["node_content"],
bytes(r["embedding"]),
)
for r in batch
],
)
new_conn.execute("BEGIN IMMEDIATE")
# Stream rows from the source cursor in batches instead of
# materializing the whole table in memory, so a large index does
# not cause an OOM during routine maintenance compactions.
while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE):
new_conn.executemany(
self._INSERT,
[
(
r["id"],
r["document_id"],
r["modified"],
r["node_content"],
bytes(r["embedding"]),
)
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")
except BaseException:
new_conn.close()
for p in [compact_path, compact_path + "-wal", compact_path + "-shm"]:
Path(p).unlink(missing_ok=True)
raise
new_conn.close()
self._swap_in_compact(compact_path, db_path)
DocumentChunksTable.insert_many(
dst_conn,
(ChunkRow(r["id"], r["document_id"]) for r in batch),
)
copied += len(batch)
DocumentMetaTable.copy_all(src_conn, dst_conn, COMPACT_BATCH_SIZE)
# Reset the cumulative counter: after a rebuild, total_inserts == live.
IndexMetaTable.reset_total_inserts(dst_conn, copied)
dst_conn.execute("COMMIT")
return copied
def _swap_in_compact(self, compact_path: str, db_path: str) -> None:
"""Atomically replace the live database with the compacted copy."""
@@ -526,8 +625,8 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
"""
if not self.table_exists():
return None
raw = self._meta_get("schema_version")
return int(raw) if raw is not None else SCHEMA_VERSION
raw_version = IndexMetaTable.get_schema_version(self._conn)
return raw_version if raw_version is not None else SCHEMA_VERSION
def has_pending_migration(self) -> bool:
"""Cheaply check whether a migration is pending, with no exclusive
@@ -591,16 +690,12 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
dim = self.vector_dim()
if dim is None: # pragma: no cover
raise RuntimeError("Cannot migrate: no stored vector dimension")
db_path = str(Path(self._uri) / DB_FILENAME)
compact_path = db_path + ".compact"
new_conn = self._open_connection(compact_path)
try:
with self._rebuild_file() as new_conn:
migration.apply(self._conn, new_conn, dim)
self._meta_set_on(new_conn, "schema_version", str(migration.to_version))
except BaseException: # pragma: no cover
new_conn.close()
for p in [compact_path, compact_path + "-wal", compact_path + "-shm"]:
Path(p).unlink(missing_ok=True)
raise
new_conn.close()
self._swap_in_compact(compact_path, db_path)
IndexMetaTable.set_schema_version(new_conn, migration.to_version)
# Registers m0001_v1_to_v2 into MIGRATIONS; must be at the bottom (needs
# PaperlessSqliteVecVectorStore fully defined) -- see
# paperless_ai/migrations/__init__.py for the full procedure.
from paperless_ai.migrations import m0001_v1_to_v2 # noqa: E402, F401