Batch search selection data queries

Co-Authored-By: Codex <267193182+codex@users.noreply.github.com>

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,
                 ),
             ):

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(

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