From fa01396dfdb81d798b531245ab0e703a197621a4 Mon Sep 17 00:00:00 2001 From: Trenton H <797416+stumpylog@users.noreply.github.com> Date: Thu, 23 Jul 2026 12:21:20 -0700 Subject: [PATCH] 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 * 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 --------- Co-authored-by: Claude Sonnet 5 --- src/documents/permissions.py | 52 +++++++++++++++------ src/documents/tests/test_api_documents.py | 57 +++++++++++++++++++++++ src/documents/views.py | 47 ++++++++++++------- 3 files changed, 126 insertions(+), 30 deletions(-) diff --git a/src/documents/permissions.py b/src/documents/permissions.py index 06e1e271a..adfaaec1a 100644 --- a/src/documents/permissions.py +++ b/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], diff --git a/src/documents/tests/test_api_documents.py b/src/documents/tests/test_api_documents.py index 23a150918..2e3e9cbfc 100644 --- a/src/documents/tests/test_api_documents.py +++ b/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", diff --git a/src/documents/views.py b/src/documents/views.py index 750b54248..64594a3e7 100644 --- a/src/documents/views.py +++ b/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