perf(machinery): improve built-in machinery queries

This is AI-assisted approach:

- do an exact lookup first and skip trigram lookup in case some matches
  are found (is there problem with skipping other matches in this case?)
- create split indexes and split SQL query for translation memory, this
  seems to work better than forcing PostgreSQL to do the complex
  filtering on top of trigram

Author: nijel

docs/changes.rst

@@ -14,6 +14,7 @@ Weblate 5.17
 * Improved logic for adding monolingual plurals in :doc:`/formats/gettext`.
 * Improved error messages in some of the :ref:`api` endpoints.
 * :envvar:`WEBLATE_COMMIT_PENDING_HOURS` is now available in Docker container.
+* Improved performance of Weblate translation memory and built-in machinery lookups.
 
 .. rubric:: Bug fixes
 

weblate/machinery/tests.py

@@ -2934,6 +2934,20 @@ def test_exists(self) -> None:
         results = machine.translate(unit, self.user)
         self.assertNotEqual(results, [])
 
+    @patch("weblate.machinery.weblatetm.adjust_similarity_threshold")
+    def test_exact_match_skips_trigram_threshold(self, adjust_threshold) -> None:
+        unit = Unit.objects.filter(translation__language_code="cs")[0]
+        other = unit.translation.unit_set.exclude(pk=unit.pk)[0]
+        other.source = unit.source
+        other.target = "Preklad"
+        other.state = STATE_TRANSLATED
+        other.save()
+
+        machine = WeblateTranslation({})
+        machine.translate(unit, self.user)
+
+        adjust_threshold.assert_not_called()
+
 
 class CyrTranslitTranslationTest(ViewTestCase, BaseMachineTranslationTest):
     ENGLISH = "sr@latin"

weblate/machinery/weblatetm.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from django.conf import settings
 from django.db.models import Value
@@ -28,6 +28,7 @@ class WeblateTranslation(InternalMachineTranslation):
     cache_translations = True
     # Cache results for 1 hour to avoid frequent database hits
     cache_expiry = 3600
+    candidate_limit = 100
 
     def download_translations(
         self,
@@ -48,33 +49,16 @@ def download_translations(
         if "memory_db" in settings.DATABASES:
             base = base.using("memory_db")
 
-        # Build search query
-        lookup: dict[str, Any] = {}
-        if threshold < 100:
-            # Full text search
-            lookup["source__search"] = text
-        else:
-            # Utilize PostgreSQL index
-            lookup["source__lower__md5"] = MD5(Lower(Value(text)))
-            lookup["source"] = text
-
-        matching_units = (
+        matching_units = self.get_matching_units(
             base.filter(
                 translation__component__source_language=source_language,
                 translation__language=target_language,
                 state__gte=STATE_TRANSLATED,
-                **lookup,
-            )
-            .exclude(
-                # The read-only strings can be possibly blank
-                target__lower__md5=MD5(Lower(Value("")))
-            )
-            .prefetch()
+            ),
+            text,
+            threshold,
         )
 
-        # We want only close matches here
-        adjust_similarity_threshold(0.98)
-
         for munit in matching_units:
             source = munit.source_string
             if "forbidden" in munit.all_flags:
@@ -91,3 +75,34 @@ def download_translations(
                 "origin_url": munit.get_absolute_url(),
                 "source": source,
             }
+
+    def get_matching_units(self, base, text: str, threshold: int):
+        # Exact matches are common and are much cheaper than trigram search.
+        exact_matches = list(
+            self.prepare_queryset(
+                base.filter(
+                    source__lower__md5=MD5(Lower(Value(text))),
+                    source=text,
+                )
+            )[: self.candidate_limit]
+        )
+        if exact_matches or threshold >= 100:
+            return exact_matches
+
+        # We want only close matches here.
+        adjust_similarity_threshold(0.98)
+
+        return list(
+            self.prepare_queryset(base.filter(source__search=text))[
+                : self.candidate_limit
+            ]
+        )
+
+    def prepare_queryset(self, queryset):
+        return queryset.exclude(target="").select_related(
+            "source_unit",
+            "translation__language",
+            "translation__plural",
+            "translation__component",
+            "translation__component__project",
+        )

weblate/memory/migrations/0005_memory_lookup_indexes.py

@@ -0,0 +1,53 @@
+# Copyright © Michal Čihař <michal@weblate.org>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from django.db import migrations, models
+from django.db.models import Q
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("memory", "0004_memory_status_and_context"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="memory",
+            index=models.Index(
+                "source_language",
+                "target_language",
+                condition=Q(from_file=True),
+                name="memory_file_lang_idx",
+            ),
+        ),
+        migrations.AddIndex(
+            model_name="memory",
+            index=models.Index(
+                "project",
+                "source_language",
+                "target_language",
+                condition=Q(project__isnull=False, shared=False),
+                name="memory_project_lang_idx",
+            ),
+        ),
+        migrations.AddIndex(
+            model_name="memory",
+            index=models.Index(
+                "user",
+                "source_language",
+                "target_language",
+                condition=Q(user__isnull=False, shared=False),
+                name="memory_user_lang_idx",
+            ),
+        ),
+        migrations.AddIndex(
+            model_name="memory",
+            index=models.Index(
+                "source_language",
+                "target_language",
+                condition=Q(shared=True),
+                name="memory_shared_lang_idx",
+            ),
+        ),
+    ]

weblate/memory/models.py

@@ -167,39 +167,94 @@ def lookup(
         use_shared,
         threshold: int = 75,
     ):
+        exact_results = self.lookup_scopes(
+            source_language=source_language,
+            target_language=target_language,
+            text=text,
+            user=user,
+            project=project,
+            use_shared=use_shared,
+            exact=True,
+        )
+        if exact_results:
+            return exact_results
+
         # Adjust similarity based on string length to get more relevant matches
         # for long strings
         similarity_threshold = self.threshold_to_similarity(text, threshold)
 
-        results = self.none()
+        results: list[Memory] = []
 
         while len(results) == 0 and similarity_threshold > MIN_SIMILARITY_THRESHOLD:
             # Change PostgreSQL similarity threshold configuration
             adjust_similarity_threshold(similarity_threshold)
 
             # Actual database query
-            results = (
-                self.prefetch_project()
-                .filter_type(
-                    # Type filtering
-                    user=user,
-                    project=project,
-                    use_shared=use_shared,
-                    from_file=True,
-                )
-                .filter(
-                    # Full-text search on source
-                    source__search=text,
-                    # Language filtering
-                    source_language=source_language,
-                    target_language=target_language,
-                )[:50]
+            results = self.lookup_scopes(
+                source_language=source_language,
+                target_language=target_language,
+                text=text,
+                user=user,
+                project=project,
+                use_shared=use_shared,
             )
             # Decrease threshold in case no matches were found
             similarity_threshold -= 0.05
 
         return results
 
+    def lookup_scopes(
+        self,
+        *,
+        source_language,
+        target_language,
+        text: str,
+        user,
+        project,
+        use_shared: bool,
+        exact: bool = False,
+        limit: int = 50,
+    ) -> list[Memory]:
+        base = self.prefetch_project()
+        common_filters = {
+            "source_language": source_language,
+            "target_language": target_language,
+        }
+        if exact:
+            common_filters["source"] = text
+        else:
+            common_filters["source__search"] = text
+
+        results: list[Memory] = []
+        seen: set[int] = set()
+
+        for query in self.iter_lookup_scope_queries(
+            user=user,
+            project=project,
+            use_shared=use_shared,
+        ):
+            for memory in base.filter(query, **common_filters)[:limit]:
+                if memory.pk in seen:
+                    continue
+                seen.add(memory.pk)
+                results.append(memory)
+                if len(results) >= limit:
+                    return results
+
+        return results
+
+    def iter_lookup_scope_queries(
+        self, *, user, project, use_shared: bool
+    ) -> tuple[Q, ...]:
+        queries: list[Q] = [Q(from_file=True)]
+        if project is not None:
+            queries.append(Q(project=project, shared=False))
+        if use_shared:
+            queries.append(Q(shared=True))
+        if user is not None:
+            queries.append(Q(user=user, shared=False))
+        return tuple(queries)
+
     def prefetch_lang(self):
         return self.prefetch_related("source_language", "target_language")
 
@@ -601,6 +656,32 @@ class Meta:
                 condition=Q(from_file=True),
                 name="memory_from_file",
             ),
+            models.Index(
+                "source_language",
+                "target_language",
+                condition=Q(from_file=True),
+                name="memory_file_lang_idx",
+            ),
+            models.Index(
+                "project",
+                "source_language",
+                "target_language",
+                condition=Q(project__isnull=False, shared=False),
+                name="memory_project_lang_idx",
+            ),
+            models.Index(
+                "user",
+                "source_language",
+                "target_language",
+                condition=Q(user__isnull=False, shared=False),
+                name="memory_user_lang_idx",
+            ),
+            models.Index(
+                "source_language",
+                "target_language",
+                condition=Q(shared=True),
+                name="memory_shared_lang_idx",
+            ),
         ]
 
     def __str__(self) -> str:

weblate/memory/tests.py

@@ -8,6 +8,7 @@
 import tempfile
 from io import StringIO
 from typing import Any
+from unittest.mock import patch
 
 from django.core.management import call_command
 from django.core.management.base import CommandError
@@ -93,6 +94,16 @@ def test_machine_batch(self) -> None:
         del machinery["origin"]
         self.assertEqual(machinery, {"quality": [100], "translation": ["Ahoj"]})
 
+    @patch("weblate.memory.models.adjust_similarity_threshold")
+    def test_exact_match_skips_similarity_adjustment(self, adjust_threshold) -> None:
+        add_document()
+        unit = self.get_unit()
+
+        machine_translation = WeblateMemory({})
+        machine_translation.search(unit, "Hello", None)
+
+        adjust_threshold.assert_not_called()
+
     def test_machine_plurals(self) -> None:
         unit = self.get_unit("Orangutan has %d banana.\n")
 

weblate/trans/migrations/0067_unit_source_tm_index.py

@@ -0,0 +1,25 @@
+# Copyright © Michal Čihař <michal@weblate.org>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from django.contrib.postgres.indexes import GinIndex
+from django.db import migrations
+from django.db.models import Q
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("trans", "0066_alter_variant_key"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="unit",
+            index=GinIndex(
+                fields=["source"],
+                opclasses=["gin_trgm_ops"],
+                condition=Q(state__gte=20) & ~Q(target=""),
+                name="trans_unit_source_tm_idx",
+            ),
+        ),
+    ]

weblate/trans/models/unit.py

@@ -11,6 +11,7 @@
 
 import sentry_sdk
 from django.conf import settings
+from django.contrib.postgres.indexes import GinIndex
 from django.core.cache import cache
 from django.db import Error as DjangoDatabaseError
 from django.db import models, transaction
@@ -499,6 +500,12 @@ class Meta:
             models.Index(
                 MD5(Lower("context")), "translation", name="trans_unit_context_md5"
             ),
+            GinIndex(
+                fields=["source"],
+                opclasses=["gin_trgm_ops"],
+                condition=Q(state__gte=STATE_TRANSLATED) & ~Q(target=""),
+                name="trans_unit_source_tm_idx",
+            ),
         ]
 
     def __str__(self) -> str: