diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index 1c4dbf2f94..9981b89a74 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -20,8 +20,10 @@ from paperless_ai.tables import DocumentChunksTable from paperless_ai.tables import DocumentMetaTable from paperless_ai.vector_store import _MAX_IN_VALUES from paperless_ai.vector_store import DB_FILENAME +from paperless_ai.vector_store import DB_PAGE_SIZE from paperless_ai.vector_store import DEFAULT_TABLE_NAME from paperless_ai.vector_store import SCHEMA_VERSION +from paperless_ai.vector_store import VEC0_CHUNK_SIZE from paperless_ai.vector_store import PaperlessSqliteVecVectorStore from paperless_ai.vector_store import _build_where from paperless_ai.vector_store import _pack @@ -1128,3 +1130,127 @@ class TestV1ToV2Migration: assert "DocumentChunksTable.create" not in source assert "DocumentMetaTable.create" not in source assert "_create_vec_table(" not in source + + +LEGACY_PAGE_SIZE = 4096 + + +def vec_table_sql(db_path: Path) -> str: + conn = sqlite3.connect(db_path) + try: + return conn.execute( + "SELECT sql FROM sqlite_master WHERE name = ?", + (DEFAULT_TABLE_NAME,), + ).fetchone()[0] + finally: + conn.close() + + +def page_size_of(db_path: Path) -> int: + conn = sqlite3.connect(db_path) + try: + return conn.execute("PRAGMA page_size").fetchone()[0] + finally: + conn.close() + + +def create_legacy_vec_table(conn: sqlite3.Connection, dim: int) -> None: + """The table definition from before chunk_size was set explicitly.""" + conn.execute( + f"CREATE VIRTUAL TABLE {DEFAULT_TABLE_NAME} USING vec0(" + "id TEXT PRIMARY KEY, document_id INTEGER, +node_content TEXT, " + f"embedding float[{dim}] distance_metric=cosine)", + ) + + +@pytest.fixture +def legacy_page_size_db(tmp_path: Path) -> Path: + """An existing index file created before the page size was set.""" + db_path = tmp_path / DB_FILENAME + conn = sqlite3.connect(db_path) + conn.execute(f"PRAGMA page_size={LEGACY_PAGE_SIZE}") + conn.execute("CREATE TABLE _seed(x)") + conn.execute("DROP TABLE _seed") + conn.close() + assert page_size_of(db_path) == LEGACY_PAGE_SIZE + return db_path + + +class TestVecChunkSize: + def test_new_table_uses_configured_chunk_size( + self, + store: PaperlessSqliteVecVectorStore, + tmp_path: Path, + ) -> None: + """ + GIVEN an empty store + WHEN the first document is indexed and the vec0 table is created + THEN the table definition carries the configured chunk size + """ + store.upsert_document(1, [make_node("1-a", 1)]) + + assert f"chunk_size={VEC0_CHUNK_SIZE}" in vec_table_sql(tmp_path / DB_FILENAME) + + def test_compact_moves_legacy_table_to_configured_chunk_size( + self, + store: PaperlessSqliteVecVectorStore, + tmp_path: Path, + mocker: MockerFixture, + ) -> None: + """ + GIVEN an index whose vec0 table was created with the default chunk size + WHEN the index is compacted + THEN the rebuilt table uses the configured chunk size + AND no indexed chunks are lost + """ + legacy = mocker.patch.object( + PaperlessSqliteVecVectorStore, + "_create_vec_table", + staticmethod(create_legacy_vec_table), + ) + store.upsert_document(1, [make_node("1-a", 1), make_node("1-b", 1)]) + store.upsert_document(2, [make_node("2-a", 2)]) + mocker.stop(legacy) + assert "chunk_size" not in vec_table_sql(tmp_path / DB_FILENAME) + + store.compact(force=True) + + assert f"chunk_size={VEC0_CHUNK_SIZE}" in vec_table_sql(tmp_path / DB_FILENAME) + assert DocumentChunksTable.count(store.client) == 3 + + +class TestIndexPageSize: + def test_new_index_file_uses_configured_page_size( + self, + store: PaperlessSqliteVecVectorStore, + tmp_path: Path, + ) -> None: + """ + GIVEN no existing index file + WHEN a store is opened and a document is indexed + THEN the database file uses the configured page size + """ + store.upsert_document(1, [make_node("1-a", 1)]) + + assert page_size_of(tmp_path / DB_FILENAME) == DB_PAGE_SIZE + + def test_compact_moves_legacy_file_to_configured_page_size( + self, + legacy_page_size_db: Path, + ) -> None: + """ + GIVEN an existing index file created with a different page size + WHEN the index is compacted + THEN the rebuilt file uses the configured page size + AND no indexed chunks are lost + """ + with PaperlessSqliteVecVectorStore( + uri=str(legacy_page_size_db.parent), + ) as store: + store.upsert_document(1, [make_node("1-a", 1), make_node("1-b", 1)]) + assert page_size_of(legacy_page_size_db) == LEGACY_PAGE_SIZE + + store.compact(force=True) + + assert page_size_of(legacy_page_size_db) == DB_PAGE_SIZE + assert DocumentChunksTable.count(store.client) == 2 diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index c7675261d6..7b48934247 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -47,6 +47,20 @@ SCHEMA_VERSION = 2 # a rebuild copies the live rows into a fresh table. COMPACT_BLOAT_RATIO = 2.0 +# Rows per vec0 vector chunk (vec0's default is 1024, which makes each chunk a +# multi-megabyte blob). On spinning disks, inserts and re-indexing were dozens +# of times faster with small chunks, at no cost in query speed or file size. +# Only applies to tables created from here on; an existing index picks it up +# the next time it is rebuilt (compact or rebuild). +VEC0_CHUNK_SIZE = 32 + +# SQLite page size for the index file, in bytes. With the small vec0 chunks above, +# 64 KiB pages made cold scans about 35% faster and warm ones about 45% faster, +# at the cost of a larger write-ahead log (up to about 70 MB). Only takes effect +# on a newly created file, so like the chunk size it reaches an existing index +# when it is rebuilt. +DB_PAGE_SIZE = 65536 + # Number of rows fetched/copied per batch whenever this module streams rows # instead of materializing them all at once, keeping memory bounded regardless # of index size -- used by compact()'s rebuild, m0001_v1_to_v2's migration @@ -222,6 +236,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): conn.enable_load_extension(True) # noqa: FBT003 sqlite_vec.load(conn) conn.enable_load_extension(False) # noqa: FBT003 + conn.execute(f"PRAGMA page_size={DB_PAGE_SIZE}") conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA synchronous=NORMAL") IndexMetaTable.create(conn) @@ -329,7 +344,8 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): + " +node_content TEXT," + " embedding float[" + str(int(dim)) - + "] distance_metric=cosine" + + "] distance_metric=cosine, chunk_size=" + + str(VEC0_CHUNK_SIZE) + ")", )