Perf: stream update_llm_index()'s document queryset instead of loading it whole

update_llm_index() built its documents queryset with
prefetch_related("tags", "notes", "custom_fields__field") and iterated
it directly -- no .iterator(), so a full rebuild pulled every Document
(including content, which can be megabytes each) plus every prefetch
cache into memory simultaneously.

Add _StreamedDocuments, a thin QuerySet wrapper matching the shape of
documents/search/_backend.py's _DocumentViewerStream: iterates via
.iterator(chunk_size=1000) so Django streams from a server-side cursor
(and, since 4.1, still runs prefetch_related per batch instead of all at
once), while __len__ still supports a real progress-bar total. Used for
both loop sites (the full rebuild and the modified-time-scoped update).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

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,57 @@ 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_uses_queryset_count(
+        self,
+        mocker: pytest_mock.MockerFixture,
+    ) -> None:
+        """
+        GIVEN:
+            - A mock queryset
+        WHEN:
+            - len() is called on a _StreamedDocuments wrapping it
+        THEN:
+            - The queryset's count() is used, not a materializing len()
+        """
+        mock_queryset = mocker.MagicMock()
+        mock_queryset.count.return_value = 42
+
+        assert len(indexing._StreamedDocuments(mock_queryset)) == 42
+        mock_queryset.count.assert_called_once_with()
+
+    def test_iter_streams_via_queryset_iterator(
+        self,
+        mocker: pytest_mock.MockerFixture,
+    ) -> None:
+        """
+        GIVEN:
+            - A mock queryset whose .iterator() yields documents
+        WHEN:
+            - A _StreamedDocuments wrapping it is iterated
+        THEN:
+            - .iterator() is called with the streaming chunk size (not
+              plain iteration, which would materialize prefetches for the
+              whole queryset instead of one batch at a time), and its
+              results are yielded through unchanged
+        """
+        mock_queryset = mocker.MagicMock()
+        mock_queryset.iterator.return_value = iter(["doc-1", "doc-2"])
+
+        result = list(indexing._StreamedDocuments(mock_queryset))
+
+        assert result == ["doc-1", "doc-2"]
+        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,