Fix: selection_data re-derives the filtered document set 5 times over (#13229)

* Fix: selection_data re-derives the filtered document set 5 times over

_get_selection_data_for_queryset() (powers ?include_selection_data=true on
the document list and search endpoints) computed document_count for
Correspondent/DocumentType/StoragePath/Tag/CustomField by embedding the
caller's full filtered queryset -- filters plus the permission check -- as
a subquery inside 5 separate Count(filter=Q(documents__in=queryset), ...)
calls. Each one re-evaluates that whole queryset from scratch.

Resolve the document ids once into a concrete list and reuse it across all
five annotations instead. For Tag/CustomField specifically (M2M via a
through-model table), also route through annotate_document_count_by_ids()
-- extracted from the tag/custom-field document_count fix (#13203) -- to
avoid the same Count(filter=Q(id__in=...), distinct=True)-on-an-M2M-relation
anti-pattern diagnosed there.

On a 400k-document/1,000-tag corpus, the correspondents portion alone
previously didn't finish within several minutes (killed twice while
investigating, including one run that left a zombie query still consuming
a CPU 30+ minutes later). All five queries together now complete in
~30-35s. Root-caused from a real report (paperless-ngx#13201) via a
different, already-fixed query (#13205) -- this one hasn't been reported
in the wild yet, found by auditing the same call path.

Depends on #13203 for annotate_document_count_by_ids().

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

* Address review feedback: drop unnecessary ordering before collecting ids

queryset passed into _get_selection_data_for_queryset() carries the
default/user-specified ordering, which is irrelevant once we're only
collecting a flat id list. Clearing it removes a pointless sort.

No measurable change in benchmarking at 400k documents -- the id-list
materialization/IN-clause cost still dominates -- but it's a free,
strictly-correct cleanup, not just noise-neutral in the other direction.

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

src/documents/permissions.py

@@ -224,51 +224,56 @@ def get_document_count_filter_for_user(user, related_name: str = "documents"):
     return Q(**{f"{related_name}__id__in": permitted_ids})
 
 
-def annotate_document_count_for_related_queryset(
+def annotate_document_count_by_ids(
     queryset: QuerySet[Any],
     through_model: Any,
     related_object_field: str,
+    document_ids: Any,
     target_field: str = "document_id",
-    user: User | None = None,
 ) -> QuerySet[Any]:
     """
-    Annotate a queryset with a permissions-aware document count for a relation
-    to Document that goes through an M2M/through-model table (e.g. Tag via
-    ``Document.tags.through``, or CustomField via ``CustomFieldInstance``).
+    Annotate a queryset with a document count for a relation to Document that
+    goes through an M2M/through-model table (e.g. Tag via
+    ``Document.tags.through``, or CustomField via ``CustomFieldInstance``),
+    for an explicit, already-resolved set of document ids.
 
     Counts are computed via a single, independent GROUP BY over the relation
-    table -- with the permission filter expressed as a plain ``WHERE`` rather
-    than an aggregate ``FILTER`` -- then injected via ``Case``/``When``. This
+    table -- with the id filter expressed as a plain ``WHERE`` rather than an
+    aggregate ``FILTER`` -- then injected via ``Case``/``When``. This
     deliberately avoids two slower alternatives found while building this:
 
     - A per-outer-row correlated subquery (one execution per row of the
       annotated queryset): fine at a handful of rows, catastrophic once the
       queryset has hundreds/thousands of rows.
-    - ``Count(..., filter=Q(id__in=permitted_ids), distinct=True)`` applied
+    - ``Count(..., filter=Q(id__in=document_ids), distinct=True)`` applied
       directly to the M2M relation: Postgres can fail to plan the ``id__in``
       check as a semi-join and instead re-checks subquery membership once per
       row of the (much larger) M2M join -- worse than the correlated subquery.
 
     Aggregation is restricted to rows whose ``related_object_field`` is one of
     ``queryset``'s pks, so passing a subset (e.g. a handful of tag descendants)
-    doesn't pay the cost of counting for every row matching the permission
-    filter.
+    doesn't pay the cost of counting for every row matching ``document_ids``.
 
     Args:
         queryset: base queryset to annotate (must contain pk)
         through_model: model representing the relation (e.g., Document.tags.through
                        or CustomFieldInstance)
         related_object_field: field on the relation pointing back to queryset pk
+        document_ids: the document ids to count against -- a concrete list/set,
+                       or a simple (already resolved) queryset of ids. Callers
+                       that need this filtered by a complex condition (e.g. a
+                       permission check) should resolve it to a concrete list
+                       first if the same ids will be reused across multiple
+                       calls, rather than passing the complex queryset itself
+                       into each -- see ``_get_selection_data_for_queryset``.
         target_field: field on the relation pointing to Document id
-        user: the user for whom to filter permitted document ids
     """
 
-    permitted_ids = _permitted_document_ids(user)
     counts = (
         through_model.objects.filter(
             **{
                 f"{related_object_field}__in": queryset.values("pk"),
-                f"{target_field}__in": permitted_ids,
+                f"{target_field}__in": document_ids,
             },
         )
         .values(related_object_field)
@@ -290,6 +295,27 @@ def annotate_document_count_for_related_queryset(
     )
 
 
+def annotate_document_count_for_related_queryset(
+    queryset: QuerySet[Any],
+    through_model: Any,
+    related_object_field: str,
+    target_field: str = "document_id",
+    user: User | None = None,
+) -> QuerySet[Any]:
+    """
+    Same as ``annotate_document_count_by_ids``, but resolves the document ids
+    from the given user's view permissions rather than taking them directly.
+    """
+
+    return annotate_document_count_by_ids(
+        queryset,
+        through_model=through_model,
+        related_object_field=related_object_field,
+        document_ids=_permitted_document_ids(user),
+        target_field=target_field,
+    )
+
+
 def get_objects_for_user_owner_aware(
     user: User | None,
     perms: str | list[str],

src/documents/tests/test_api_documents.py

@@ -1347,6 +1347,63 @@ class TestDocumentApi(DirectoriesMixin, ConsumeTaskMixin, APITestCase):
         self.assertEqual(selected_type["document_count"], 1)
         self.assertEqual(selected_storage_path["document_count"], 1)
 
+    def test_selection_data_document_counts_per_tag(self) -> None:
+        """
+        GIVEN:
+            - Multiple tags with different numbers of matching documents
+              within the filtered set, including one with no matches
+        WHEN:
+            - Requesting the document list with include_selection_data=true
+        THEN:
+            - Each tag's document_count reflects only documents in the
+              filtered set, not the instance-wide count
+        """
+        tag_a = Tag.objects.create(name="a")
+        tag_b = Tag.objects.create(name="b")
+        tag_unused = Tag.objects.create(name="unused")
+        custom_field = CustomField.objects.create(
+            name="cf1",
+            data_type=CustomField.FieldDataType.STRING,
+        )
+
+        doc1 = Document.objects.create(checksum="1", correspondent=None)
+        doc1.tags.add(tag_a)
+        doc2 = Document.objects.create(checksum="2")
+        doc2.tags.add(tag_a, tag_b)
+        doc3 = Document.objects.create(checksum="3")
+        doc3.tags.add(tag_b)
+        CustomFieldInstance.objects.create(
+            document=doc1,
+            field=custom_field,
+            value_text="x",
+        )
+
+        # Excluded from the filtered set entirely.
+        excluded = Document.objects.create(checksum="4")
+        excluded.tags.add(tag_a, tag_b, tag_unused)
+
+        response = self.client.get(
+            f"/api/documents/?id__in={doc1.id},{doc2.id},{doc3.id}"
+            "&include_selection_data=true",
+        )
+
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        selection_data = response.data["selection_data"]
+
+        counts_by_tag = {
+            item["id"]: item["document_count"]
+            for item in selection_data["selected_tags"]
+        }
+        self.assertEqual(counts_by_tag[tag_a.id], 2)
+        self.assertEqual(counts_by_tag[tag_b.id], 2)
+        self.assertEqual(counts_by_tag[tag_unused.id], 0)
+
+        counts_by_field = {
+            item["id"]: item["document_count"]
+            for item in selection_data["selected_custom_fields"]
+        }
+        self.assertEqual(counts_by_field[custom_field.id], 1)
+
     def test_statistics(self) -> None:
         doc1 = Document.objects.create(
             title="none1",

src/documents/views.py

@@ -169,6 +169,7 @@ from documents.permissions import PaperlessAdminPermissions
 from documents.permissions import PaperlessNotePermissions
 from documents.permissions import PaperlessObjectPermissions
 from documents.permissions import ViewDocumentsPermissions
+from documents.permissions import annotate_document_count_by_ids
 from documents.permissions import annotate_document_count_for_related_queryset
 from documents.permissions import get_document_count_filter_for_user
 from documents.permissions import get_objects_for_user_owner_aware
@@ -992,42 +993,54 @@ class DocumentViewSet(
     )
 
     def _get_selection_data_for_queryset(self, queryset):
+        # Resolve once instead of once per model below. `queryset` can carry an
+        # arbitrarily expensive WHERE clause (user filters plus the permission
+        # filter); re-embedding it as a subquery inside 5 separate Count(...)
+        # calls forces the database to re-evaluate that whole thing 5 times, and
+        # -- for FK relations especially -- can defeat semi-join planning
+        # entirely at scale. A concrete id list is cheap to reuse.
+        # order_by() drops the default/user ordering -- irrelevant for a plain
+        # id list, but left in place it forces a sort over the full filtered
+        # set before the ids can even be collected.
+        document_ids = list(queryset.order_by().values_list("pk", flat=True))
+
         correspondents = Correspondent.objects.annotate(
             document_count=Count(
                 "documents",
-                filter=Q(documents__in=queryset),
-                distinct=True,
-            ),
-        )
-        tags = Tag.objects.annotate(
-            document_count=Count(
-                "documents",
-                filter=Q(documents__in=queryset),
+                filter=Q(documents__id__in=document_ids),
                 distinct=True,
             ),
         )
         document_types = DocumentType.objects.annotate(
             document_count=Count(
                 "documents",
-                filter=Q(documents__in=queryset),
+                filter=Q(documents__id__in=document_ids),
                 distinct=True,
             ),
         )
         storage_paths = StoragePath.objects.annotate(
             document_count=Count(
                 "documents",
-                filter=Q(documents__in=queryset),
+                filter=Q(documents__id__in=document_ids),
                 distinct=True,
             ),
         )
-        custom_fields = CustomField.objects.annotate(
-            document_count=Count(
-                "fields__document",
-                filter=Q(fields__document__in=queryset),
-                distinct=True,
-            ),
+        # Tag and CustomField reach Document through an M2M/through-model table;
+        # a plain Count(filter=...) there is a much more expensive plan than the
+        # FK relations above once the bridge table is large -- see
+        # annotate_document_count_by_ids() for why.
+        tags = annotate_document_count_by_ids(
+            Tag.objects.all(),
+            through_model=Document.tags.through,
+            related_object_field="tag_id",
+            document_ids=document_ids,
+        )
+        custom_fields = annotate_document_count_by_ids(
+            CustomField.objects.all(),
+            through_model=CustomFieldInstance,
+            related_object_field="field_id",
+            document_ids=document_ids,
         )
-
         return {
             "selected_correspondents": [
                 {"id": t.id, "document_count": t.document_count} for t in correspondents