From efb1461e362ddf032bef45556c630535f4c0a9b8 Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Fri, 6 Nov 2026 08:37:12 -0800 Subject: [PATCH] Batch search selection data queries Co-Authored-By: Codex <267193182+codex@users.noreply.github.com> --- src/documents/tests/test_api_documents.py | 14 +-- src/documents/tests/test_api_search.py | 60 +++++----- src/documents/views.py | 127 +++++++++++----------- 3 files changed, 99 insertions(+), 102 deletions(-) diff --git a/src/documents/tests/test_api_documents.py b/src/documents/tests/test_api_documents.py index 0007b1205b..28ae4dbe0c 100644 --- a/src/documents/tests/test_api_documents.py +++ b/src/documents/tests/test_api_documents.py @@ -1766,8 +1766,8 @@ class TestDocumentApi(DirectoriesMixin, ConsumeTaskMixin, APITestCase): include_selection_data=true THEN: - Counts reflect only the shared document, not every document the - owner has with that correspondent or document type, however the - link-table lookups are batched + owner has with that correspondent or document type, regardless of + the batch size """ owner = User.objects.create_user(username="owner") other = User.objects.create_user(username="other") @@ -1789,10 +1789,10 @@ class TestDocumentApi(DirectoriesMixin, ConsumeTaskMixin, APITestCase): for batch_size in (10_000, 1): with ( - self.subTest(link_batch_size=batch_size), + self.subTest(batch_size=batch_size), mock.patch.object( DocumentViewSet, - "selection_data_link_batch_size", + "selection_data_batch_size", batch_size, ), ): @@ -1823,7 +1823,7 @@ class TestDocumentApi(DirectoriesMixin, ConsumeTaskMixin, APITestCase): - Requesting the document list with include_selection_data=true THEN: - The document is counted once per correspondent, not once per - matching tag, however the link-table lookups are batched + matching tag, regardless of the batch size """ correspondent = Correspondent.objects.create(name="c1") tag_a = Tag.objects.create(name="a") @@ -1835,10 +1835,10 @@ class TestDocumentApi(DirectoriesMixin, ConsumeTaskMixin, APITestCase): for batch_size in (10_000, 1): with ( - self.subTest(link_batch_size=batch_size), + self.subTest(batch_size=batch_size), mock.patch.object( DocumentViewSet, - "selection_data_link_batch_size", + "selection_data_batch_size", batch_size, ), ): diff --git a/src/documents/tests/test_api_search.py b/src/documents/tests/test_api_search.py index 9c7ff158df..ebcd20280b 100644 --- a/src/documents/tests/test_api_search.py +++ b/src/documents/tests/test_api_search.py @@ -25,6 +25,7 @@ from documents.models import Tag from documents.models import Workflow from documents.search import get_backend from documents.search import reset_backend +from documents.views import DocumentViewSet from paperless_mail.models import MailAccount from paperless_mail.models import MailRule from paperless_testing.dirs import DirectoriesMixin @@ -345,38 +346,39 @@ class TestDocumentSearchApi(DirectoriesMixin, APITestCase): storage_path = StoragePath.objects.create(name="sp1") tag = Tag.objects.create(name="tag") - matching_doc = Document.objects.create( - title="bank statement", - content="bank content", - checksum="A", - correspondent=correspondent, - document_type=doc_type, - storage_path=storage_path, - ) - matching_doc.tags.add(tag) + for checksum, title in ( + ("A", "bank statement"), + ("B", "bank account"), + ("C", "credit card"), + ): + doc = Document.objects.create( + title=title, + checksum=checksum, + correspondent=correspondent, + document_type=doc_type, + storage_path=storage_path, + ) + doc.tags.add(tag) + get_backend().add_or_update(doc) - get_backend().add_or_update(matching_doc) - - response = self.client.get( - "/api/documents/?query=bank&include_selection_data=true", - ) + with mock.patch.object(DocumentViewSet, "selection_data_batch_size", 1): + response = self.client.get( + "/api/documents/?query=bank&page_size=1&include_selection_data=true", + ) self.assertEqual(response.status_code, status.HTTP_200_OK) - self.assertIn("selection_data", response.data) - - selected_correspondent = next( - item - for item in response.data["selection_data"]["selected_correspondents"] - if item["id"] == correspondent.id - ) - selected_tag = next( - item - for item in response.data["selection_data"]["selected_tags"] - if item["id"] == tag.id - ) - - self.assertEqual(selected_correspondent["document_count"], 1) - self.assertEqual(selected_tag["document_count"], 1) + self.assertEqual(response.data["count"], 2) + self.assertEqual(len(response.data["results"]), 1) + for key, pk in ( + ("selected_correspondents", correspondent.pk), + ("selected_document_types", doc_type.pk), + ("selected_storage_paths", storage_path.pk), + ("selected_tags", tag.pk), + ): + self.assertEqual( + response.data["selection_data"][key], + [{"id": pk, "document_count": 2}], + ) def test_search_custom_field_ordering(self) -> None: custom_field = CustomField.objects.create( diff --git a/src/documents/views.py b/src/documents/views.py index b529ba14cb..c7754d1770 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -11,6 +11,7 @@ from collections.abc import Callable from datetime import datetime from datetime import timedelta from http import HTTPStatus +from itertools import islice from pathlib import Path from time import mktime from time import sleep @@ -1022,83 +1023,80 @@ class DocumentViewSet( "custom_field_", ) - # Matching document ids are looked up in the link tables in batches of this - # size: one literal IN (...) over every match is a multi-MB statement per - # relation and, on SQLite, exceeds the bound-parameter limit. - selection_data_link_batch_size = 10_000 + # Bound the IN clauses used for search results and through-table lookups. + selection_data_batch_size = 10_000 - def _get_selection_data_for_queryset(self, queryset): - # The filtered queryset is evaluated once, as the page's own filter is, - # and never wrapped in a subquery: re-deriving a join or a HAVING - # (shared_by__id) per statement is very slow on some databases, and - # regrouping the queryset by another column would miscount. - # order_by() drops the default/user ordering, which only costs a sort. - correspondents: Counter = Counter() - document_types: Counter = Counter() - storage_paths: Counter = Counter() - document_ids: list[int] = [] - seen: set[int] = set() - rows = queryset.order_by().values_list( - "pk", - "correspondent_id", - "document_type_id", - "storage_path_id", + def _get_selection_data_for_queryset( + self, + queryset: QuerySet[Document], + *, + search_ids: list[int] | None = None, + ) -> dict[str, list[dict[str, int]]]: + # Keep permission and filter predicates at document granularity. + # Reusing the filtered queryset as a subquery for each count is slow. + rows = ( + queryset.order_by() + .prefetch_related(None) + .values_list( + "pk", + "correspondent_id", + "document_type_id", + "storage_path_id", + ) ) - for pk, correspondent, document_type, storage_path in rows: - # A filter can join, repeating a document row. - if pk in seen: + batch_size = self.selection_data_batch_size + + def document_rows(): + if search_ids is None: + yield from rows.iterator() + else: + for start in range(0, len(search_ids), batch_size): + yield from rows.filter( + pk__in=search_ids[start : start + batch_size], + ).iterator() + + correspondents: Counter[int | None] = Counter() + document_types: Counter[int | None] = Counter() + storage_paths: Counter[int | None] = Counter() + document_ids: set[int] = set() + for pk, correspondent, document_type, storage_path in document_rows(): + # Filters that join multi-valued relations can repeat a document. + if pk in document_ids: continue - seen.add(pk) - document_ids.append(pk) - # Documents without a value count under None, which selection() - # never looks up. + document_ids.add(pk) correspondents[correspondent] += 1 document_types[document_type] += 1 storage_paths[storage_path] += 1 - def link_table_counts(link_model, field_name: str) -> Counter: - result: Counter = Counter() - batch_size = self.selection_data_link_batch_size - for start in range(0, len(document_ids), batch_size): - result.update( + tags: Counter[int | None] = Counter() + custom_fields: Counter[int | None] = Counter() + ids = iter(document_ids) + while batch := list(islice(ids, batch_size)): + for model, field, counts in ( + (Document.tags.through, "tag_id", tags), + (CustomFieldInstance, "field_id", custom_fields), + ): + counts.update( dict( - link_model.objects.filter( - document_id__in=document_ids[start : start + batch_size], - ) - .values_list(field_name) + model.objects.filter(document_id__in=batch) + .order_by() + .values_list(field) .annotate(document_count=Count("document_id")), ), ) - return result - def selection(model, counts: Counter) -> list[dict[str, int]]: - # Every object is listed, with 0 for those with no matching document. + def selection(model: type[Model], counts: Counter[int | None]): return [ - {"id": pk, "document_count": counts.get(pk, 0)} - for pk in model.objects.values_list("pk", flat=True) + {"id": pk, "document_count": counts[pk]} + for pk in model._default_manager.values_list("pk", flat=True) ] return { - "selected_correspondents": selection( - Correspondent, - correspondents, - ), - "selected_tags": selection( - Tag, - link_table_counts(Document.tags.through, "tag_id"), - ), - "selected_document_types": selection( - DocumentType, - document_types, - ), - "selected_storage_paths": selection( - StoragePath, - storage_paths, - ), - "selected_custom_fields": selection( - CustomField, - link_table_counts(CustomFieldInstance, "field_id"), - ), + "selected_correspondents": selection(Correspondent, correspondents), + "selected_tags": selection(Tag, tags), + "selected_document_types": selection(DocumentType, document_types), + "selected_storage_paths": selection(StoragePath, storage_paths), + "selected_custom_fields": selection(CustomField, custom_fields), } @classmethod @@ -2719,13 +2717,10 @@ class UnifiedSearchViewSet(DocumentViewSet): if get_boolean( str(request.query_params.get("include_selection_data", "false")), ): - # NOTE: pk__in=ordered_ids generates a large SQL IN clause - # for big result sets. Acceptable today but may need a temp - # table or chunked approach if selection_data becomes slow - # at scale (tens of thousands of matching documents). response.data["selection_data"] = ( self._get_selection_data_for_queryset( - filtered_qs.filter(pk__in=result.ordered_ids), + filtered_qs, + search_ids=result.ordered_ids, ) ) return response