Fix: Use prefetching to reduce query counts during classifier training (#14122)

src/documents/classifier.py

@@ -19,6 +19,7 @@ if TYPE_CHECKING:
 from django.conf import settings
 from django.core.cache import cache
 from django.core.cache import caches
+from django.db.models import Prefetch
 
 from documents.caching import CACHE_5_MINUTES
 from documents.caching import CACHE_50_MINUTES
@@ -28,6 +29,7 @@ from documents.caching import CLASSIFIER_VERSION_KEY
 from documents.caching import StoredLRUCache
 from documents.models import Document
 from documents.models import MatchingModel
+from documents.models import Tag
 from paperless.signed_pickle import SignedPickleError
 from paperless.signed_pickle import signed_pickle_dumps
 from paperless.signed_pickle import signed_pickle_loads
@@ -260,7 +262,17 @@ class DocumentClassifier:
                 tags__is_inbox_tag=True,
             )
             .select_related("document_type", "correspondent", "storage_path")
-            .prefetch_related("tags")
+            .prefetch_related(
+                Prefetch(
+                    "tags",
+                    queryset=Tag.objects.filter(
+                        matching_algorithm=MatchingModel.MATCH_AUTO,
+                    )
+                    .order_by("pk")
+                    .only("pk"),
+                    to_attr="auto_tags",
+                ),
+            )
             .order_by("pk")
         )
 
@@ -292,11 +304,7 @@ class DocumentClassifier:
             hasher.update(y.to_bytes(4, "little", signed=True))
             labels_correspondent.append(y)
 
-            tags: list[int] = list(
-                doc.tags.filter(matching_algorithm=MatchingModel.MATCH_AUTO)
-                .order_by("pk")
-                .values_list("pk", flat=True),
-            )
+            tags: list[int] = [tag.pk for tag in doc.auto_tags]
             for tag in tags:
                 hasher.update(tag.to_bytes(4, "little", signed=True))
             labels_tags.append(tags)

src/documents/tests/test_classifier.py

@@ -6,8 +6,11 @@ from unittest import mock
 import numpy as np
 import pytest
 from django.conf import settings
+from django.db import connection
 from django.test import TestCase
 from django.test import override_settings
+from django.test.utils import CaptureQueriesContext
+from pytest_mock import MockerFixture
 
 from documents.classifier import ClassifierModelCorruptError
 from documents.classifier import DocumentClassifier
@@ -20,6 +23,8 @@ from documents.models import DocumentType
 from documents.models import MatchingModel
 from documents.models import StoragePath
 from documents.models import Tag
+from documents.tests.factories import DocumentFactory
+from documents.tests.factories import TagFactory
 from documents.tests.utils import DirectoriesMixin
 from paperless.signed_pickle import HMAC_SIZE
 from paperless.signed_pickle import signed_pickle_dumps
@@ -1017,3 +1022,70 @@ def test_preprocess_content_nltk_load_fail(mocker) -> None:
         expected_preprocess_content = f.read().rstrip()
     result = classifier.preprocess_content(content)
     assert result == expected_preprocess_content
+
+
+@pytest.mark.django_db
+class TestClassifierTrainTagLabels:
+    @pytest.fixture(autouse=True)
+    def _simple_preprocess(self, mocker: MockerFixture) -> None:
+        mocker.patch.object(
+            DocumentClassifier,
+            "preprocess_content",
+            side_effect=dummy_preprocess,
+        )
+
+    @pytest.fixture
+    def auto_tags(self) -> list[Tag]:
+        return TagFactory.create_batch(2, matching_algorithm=MatchingModel.MATCH_AUTO)
+
+    def test_train_query_count_does_not_scale_with_documents(
+        self,
+        auto_tags: list[Tag],
+    ) -> None:
+        """
+        GIVEN:
+            - Documents with auto matching tags
+        WHEN:
+            - The classifier is trained, then more documents are added and it is
+              trained again
+        THEN:
+            - Both trainings run the same number of queries
+        """
+        for doc in DocumentFactory.create_batch(2):
+            doc.tags.set(auto_tags)
+
+        with CaptureQueriesContext(connection) as few_documents:
+            DocumentClassifier().train()
+
+        for doc in DocumentFactory.create_batch(6):
+            doc.tags.set(auto_tags)
+
+        with CaptureQueriesContext(connection) as more_documents:
+            DocumentClassifier().train()
+
+        assert len(more_documents) == len(few_documents)
+
+    def test_train_uses_only_auto_tags_as_labels(
+        self,
+        auto_tags: list[Tag],
+    ) -> None:
+        """
+        GIVEN:
+            - Documents with both auto matching and non auto matching tags
+        WHEN:
+            - The classifier is trained
+        THEN:
+            - Only the auto matching tags are used as tag labels
+        """
+        manual_tag = TagFactory(matching_algorithm=MatchingModel.MATCH_ANY)
+        first, second, third = DocumentFactory.create_batch(3)
+        first.tags.set([auto_tags[0], manual_tag])
+        second.tags.set([auto_tags[1], manual_tag])
+        third.tags.set([*auto_tags, manual_tag])
+
+        classifier = DocumentClassifier()
+        classifier.train()
+
+        assert list(classifier.tags_binarizer.classes_) == sorted(
+            tag.pk for tag in auto_tags
+        )