Performance: Use server side iterators during LLM index updating (#13430)

src/paperless_ai/indexing.py

@@ -1,5 +1,6 @@
 import logging
 from collections.abc import Iterable
+from collections.abc import Iterator
 from contextlib import contextmanager
 from datetime import timedelta
 from typing import TYPE_CHECKING
@@ -22,6 +23,7 @@ from paperless_ai.embedding import get_configured_model_name
 from paperless_ai.embedding import get_embedding_model
 
 if TYPE_CHECKING:
+    from django.db.models import QuerySet
     from llama_index.core.schema import BaseNode
 
     from paperless_ai.vector_store import PaperlessSqliteVecVectorStore
@@ -32,6 +34,35 @@ logger = logging.getLogger("paperless_ai.indexing")
 RAG_NUM_OUTPUT = 512
 RAG_CHUNK_OVERLAP = 200
 
+# update_llm_index(): row count per .iterator() batch when streaming
+# documents for a rebuild/update, matching _DocumentViewerStream's chunk
+# size in documents/search/_backend.py.
+_INDEX_STREAM_CHUNK_SIZE = 1000
+
+
+class _StreamedDocuments:
+    """A thin QuerySet wrapper that streams via ``.iterator()`` instead of
+    materializing every row (plus its ``content`` and prefetch caches) into
+    memory at once, while still supporting ``len()`` so ``iter_wrapper``'s
+    progress bar shows a real total instead of falling back to indeterminate.
+    Same shape as ``documents/search/_backend.py``'s ``_DocumentViewerStream``,
+    just without that class's extra per-batch permission lookup -- nothing
+    here needs one.
+    """
+
+    def __init__(self, documents: "QuerySet[Document]") -> None:
+        self._documents = documents
+
+    def __len__(self) -> int:
+        return self._documents.count()
+
+    def __iter__(self) -> Iterator[Document]:
+        # iterator(chunk_size=...) streams from a server-side cursor instead
+        # of materializing the whole queryset in memory; since Django 4.1 it
+        # still honours prefetch_related, running the prefetches one batch
+        # at a time.
+        return iter(self._documents.iterator(chunk_size=_INDEX_STREAM_CHUNK_SIZE))
+
 
 def queue_llm_index_update_if_needed(*, rebuild: bool, reason: str) -> bool:
     # NOTE: The check-then-enqueue sequence below is non-atomic (TOCTOU): two
@@ -385,7 +416,7 @@ def update_llm_index(
         if rebuild or not store.table_exists():
             logger.info("Rebuilding LLM index.")
             store.drop_table()
-            for document in iter_wrapper(documents):
+            for document in iter_wrapper(_StreamedDocuments(documents)):
                 nodes = build_document_node(document, chunk_size=chunk_size)
                 _embed_nodes(nodes, embed_model)
                 store.add(nodes)
@@ -398,7 +429,7 @@ def update_llm_index(
             )
             existing = store.get_modified_times()
             changed = 0
-            for document in iter_wrapper(scoped_documents):
+            for document in iter_wrapper(_StreamedDocuments(scoped_documents)):
                 doc_id = str(document.id)
                 if existing.get(doc_id) == document.modified.isoformat():
                     continue

src/paperless_ai/tests/test_ai_indexing.py

@@ -186,6 +186,45 @@ def test_truncate_embedding_query_returns_single_chunk() -> None:
     assert "word199" not in result
 
 
+class TestStreamedDocuments:
+    """_StreamedDocuments streams via .iterator() instead of materializing
+    the whole queryset (plus its content and prefetch caches) in memory at
+    once, while still supporting len() so a progress bar wrapped around it
+    shows a real total.
+    """
+
+    def test_len_and_iter_delegate_to_streaming_queryset_methods(
+        self,
+        mocker: pytest_mock.MockerFixture,
+    ) -> None:
+        """
+        GIVEN:
+            - A mock queryset
+        WHEN:
+            - A _StreamedDocuments wrapping it is measured and iterated
+        THEN:
+            - len() uses count() (not a materializing len()), and iteration
+              uses .iterator(chunk_size=...) (not plain iteration, which
+              would materialize prefetches for the whole queryset at once)
+        """
+        mock_queryset = mocker.MagicMock()
+        mock_queryset.count.return_value = 42
+        mock_queryset.iterator.return_value = iter(["doc-1", "doc-2"])
+        streamed = indexing._StreamedDocuments(mock_queryset)
+
+        assert len(streamed) == 42
+        assert list(streamed) == ["doc-1", "doc-2"]
+        # count.call_count isn't asserted exactly: list()'s own size-hint
+        # optimization calls len(streamed) again internally, on top of the
+        # explicit len() call above -- both legitimately delegate to
+        # count(), so only the delegation itself (not the call count) is
+        # the thing being verified here.
+        mock_queryset.count.assert_called_with()
+        mock_queryset.iterator.assert_called_once_with(
+            chunk_size=indexing._INDEX_STREAM_CHUNK_SIZE,
+        )
+
+
 @pytest.mark.django_db
 def test_update_llm_index(
     temp_llm_index_dir: Path,