diff --git a/src/documents/search/_backend.py b/src/documents/search/_backend.py index 7464de79a..227d0ccd7 100644 --- a/src/documents/search/_backend.py +++ b/src/documents/search/_backend.py @@ -39,6 +39,7 @@ from documents.search._tokenizer import ascii_fold from documents.search._tokenizer import autocomplete_tokens from documents.search._tokenizer import register_tokenizers from documents.utils import IterWrapper +from documents.utils import QuerySetStream from documents.utils import identity if TYPE_CHECKING: @@ -1035,14 +1036,15 @@ _EMPTY_VIEWER_GRANT: Final[ViewerGrant] = ViewerGrant( ) -class _DocumentViewerStream: +class _DocumentViewerStream(QuerySetStream["Document"]): """Yield document permission data while batch-loading grants. Viewer permissions are fetched in batches (see ``_bulk_get_viewer_permissions``), but documents are yielded individually so a progress bar wrapped around this stream advances per document rather than - jumping a whole chunk at a time. ``__len__`` lets the progress helper still - discover the total (it inspects ``QuerySet``/``Sized``). + jumping a whole chunk at a time. ``__len__`` (inherited from + ``QuerySetStream``) lets the progress helper still discover the total (it + inspects ``QuerySet``/``Sized``). The viewer and group ids travel with each document in the yielded pair rather than through a separate mutable attribute, so the pairing survives @@ -1051,18 +1053,11 @@ class _DocumentViewerStream: generator in lock-step. """ - def __init__(self, documents: QuerySet[Document], *, chunk_size: int) -> None: - self._documents = documents - self._chunk_size = chunk_size - - def __len__(self) -> int: - return self._documents.count() - def __iter__(self) -> Iterator[tuple[Document, ViewerGrant]]: # iterator(chunk_size=…) streams from a server-side cursor instead of # materialising the whole queryset in memory; since Django 4.1 it still # honours prefetch_related, running the prefetches one batch at a time. - documents = self._documents.iterator(chunk_size=self._chunk_size) + documents = self._queryset.iterator(chunk_size=self._chunk_size) for chunk in chunked(documents, self._chunk_size): grants_by_pk = _bulk_get_viewer_permissions([doc.pk for doc in chunk]) for doc in chunk: diff --git a/src/documents/tests/test_utils.py b/src/documents/tests/test_utils.py new file mode 100644 index 000000000..af45880ee --- /dev/null +++ b/src/documents/tests/test_utils.py @@ -0,0 +1,35 @@ +import pytest_mock + +from documents.utils import QuerySetStream + + +class TestQuerySetStream: + def test_len_and_iter_delegate_to_streaming_queryset_methods( + self, + mocker: pytest_mock.MockerFixture, + ) -> None: + """ + GIVEN: + - A mock queryset + WHEN: + - A QuerySetStream 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 the whole queryset, plus any prefetch + caches, into Django's own result cache at once) + """ + mock_queryset = mocker.MagicMock() + mock_queryset.count.return_value = 42 + mock_queryset.iterator.return_value = iter(["row-1", "row-2"]) + streamed = QuerySetStream(mock_queryset, chunk_size=1000) + + assert len(streamed) == 42 + assert list(streamed) == ["row-1", "row-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=1000) diff --git a/src/documents/utils.py b/src/documents/utils.py index 2ed6758dd..d6e505aa4 100644 --- a/src/documents/utils.py +++ b/src/documents/utils.py @@ -3,16 +3,24 @@ import logging import shutil from collections.abc import Callable from collections.abc import Iterable +from collections.abc import Iterator from os import utime from pathlib import Path from subprocess import CompletedProcess from subprocess import run +from typing import TYPE_CHECKING +from typing import Generic from typing import TypeVar from django.conf import settings from PIL import Image +if TYPE_CHECKING: + from django.db.models import Model + from django.db.models import QuerySet + _T = TypeVar("_T") +_M = TypeVar("_M", bound="Model") # A function that wraps an iterable — typically used to inject a progress bar. IterWrapper = Callable[[Iterable[_T]], Iterable[_T]] @@ -23,6 +31,40 @@ def identity(iterable: Iterable[_T]) -> Iterable[_T]: return iterable +class QuerySetStream(Generic[_M]): + """Stream a QuerySet via .iterator(chunk_size=...) instead of + materializing it (plus any prefetch caches) all at once, while still + supporting len() via count() so a progress bar wrapped around this + (e.g. via IterWrapper) shows a real total instead of falling back to + indeterminate. + + Plain QuerySet iteration (``for row in queryset:``) is not lazy: Django + fetches every matching row in one query and caches the fully-hydrated + result in the queryset's own ``_result_cache`` before yielding the + first item -- wrapping that in a progress bar or any other iterable + adapter doesn't change this, since none of them alter how the + underlying queryset produces items. ``.iterator(chunk_size=...)`` is + the specific Django API that bypasses ``_result_cache`` and streams + from a server-side cursor instead, discarding each chunk once consumed + (and, since Django 4.1, still honours ``prefetch_related``, running the + prefetches one batch at a time rather than for the whole queryset). + + Subclass to layer additional per-batch work on top (see + ``documents.search._backend._DocumentViewerStream``) by overriding + ``__iter__`` -- ``__len__`` and the constructor are inherited for free. + """ + + def __init__(self, queryset: "QuerySet[_M]", *, chunk_size: int) -> None: + self._queryset = queryset + self._chunk_size = chunk_size + + def __len__(self) -> int: + return self._queryset.count() + + def __iter__(self) -> Iterator[_M]: + return iter(self._queryset.iterator(chunk_size=self._chunk_size)) + + def _coerce_to_path( source: Path | str, dest: Path | str, diff --git a/src/paperless_ai/indexing.py b/src/paperless_ai/indexing.py index ff08d2aeb..5ee55c620 100644 --- a/src/paperless_ai/indexing.py +++ b/src/paperless_ai/indexing.py @@ -1,6 +1,5 @@ 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 @@ -15,6 +14,7 @@ from filelock import Timeout from documents.models import Document from documents.models import PaperlessTask from documents.utils import IterWrapper +from documents.utils import QuerySetStream from documents.utils import identity from paperless.config import AIConfig from paperless_ai.db import db_connection_released @@ -23,7 +23,6 @@ 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 @@ -35,35 +34,11 @@ 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. +# documents for a rebuild/update via QuerySetStream, 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 # concurrent workers can both observe no running task and both enqueue a @@ -416,7 +391,9 @@ def update_llm_index( if rebuild or not store.table_exists(): logger.info("Rebuilding LLM index.") store.drop_table() - for document in iter_wrapper(_StreamedDocuments(documents)): + for document in iter_wrapper( + QuerySetStream(documents, chunk_size=_INDEX_STREAM_CHUNK_SIZE), + ): nodes = build_document_node(document, chunk_size=chunk_size) _embed_nodes(nodes, embed_model) store.add(nodes) @@ -429,7 +406,9 @@ def update_llm_index( ) existing = store.get_modified_times() changed = 0 - for document in iter_wrapper(_StreamedDocuments(scoped_documents)): + for document in iter_wrapper( + QuerySetStream(scoped_documents, chunk_size=_INDEX_STREAM_CHUNK_SIZE), + ): doc_id = str(document.id) if existing.get(doc_id) == document.modified.isoformat(): continue diff --git a/src/paperless_ai/tests/test_ai_indexing.py b/src/paperless_ai/tests/test_ai_indexing.py index 7ae5477cb..667b43d06 100644 --- a/src/paperless_ai/tests/test_ai_indexing.py +++ b/src/paperless_ai/tests/test_ai_indexing.py @@ -186,45 +186,6 @@ 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,