perf(machinery): improve built-in machinery performance

- add focused index with the exact conditions we use
- cap limit of fetched results to avoid streaming too much data from the
  database
- do cheaper exact lookup with 100% similarity
- order matches by similarity

Author: nijel

docs/changes.rst

@@ -17,6 +17,7 @@ Weblate 5.17
 * Improved error messages in some of the :ref:`api` endpoints.
 * Improved performance of project and category search result pages with very large match sets.
 * :envvar:`WEBLATE_COMMIT_PENDING_HOURS` is now available in Docker container.
+* Improved performance of built-in Weblate machinery lookups.
 
 .. rubric:: Bug fixes
 

weblate/machinery/tests.py

@@ -21,7 +21,7 @@
 from botocore.stub import ANY, Stubber
 from django.core.management import call_command
 from django.core.management.base import CommandError
-from django.test import TestCase
+from django.test import SimpleTestCase, TestCase
 from django.urls import reverse
 from google.api_core import exceptions as google_api_exceptions
 from google.cloud.translate import (
@@ -2972,6 +2972,20 @@ def test_exists(self) -> None:
         results = machine.translate(unit, self.user)
         self.assertNotEqual(results, [])
 
+    @patch("weblate.machinery.weblatetm.adjust_similarity_threshold")
+    def test_matches_still_probe_fuzzy_lookup(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_called_once_with(0.98)
+
 
 class CyrTranslitTranslationTest(ViewTestCase, BaseMachineTranslationTest):
     ENGLISH = "sr@latin"
@@ -3318,6 +3332,113 @@ def test_configure_invalid(self) -> None:
         )
 
 
+class WeblateTranslationLookupTest(SimpleTestCase):
+    @patch("weblate.machinery.weblatetm.Unit.objects")
+    @patch("weblate.machinery.weblatetm.Translation.objects")
+    def test_get_base_queryset_uses_translation_subquery(
+        self, translation_objects, unit_objects
+    ) -> None:
+        machine = WeblateTranslation({})
+        user = MagicMock()
+        translations_using = MagicMock()
+        translations = MagicMock()
+        filtered_translations = MagicMock()
+        translation_ids = MagicMock()
+        units_using = MagicMock()
+        queryset = MagicMock()
+
+        translation_objects.using.return_value = translations_using
+        translations_using.all.return_value = translations
+        translations.filter_access.return_value = filtered_translations
+        filtered_translations.filter.return_value = translation_ids
+        translation_ids.values.return_value = "translation-subquery"
+        unit_objects.using.return_value = units_using
+        units_using.filter.return_value = queryset
+
+        result = machine.get_base_queryset(user, "en", "cs")
+
+        self.assertEqual(result, queryset)
+        translation_objects.using.assert_called_once_with("default")
+        translations.filter_access.assert_called_once_with(user)
+        filtered_translations.filter.assert_called_once_with(
+            component__source_language="en",
+            language="cs",
+        )
+        translation_ids.values.assert_called_once_with("id")
+        unit_objects.using.assert_called_once_with("default")
+        units_using.filter.assert_called_once_with(
+            state__gte=STATE_TRANSLATED,
+            translation_id__in="translation-subquery",
+        )
+
+    @patch("weblate.machinery.weblatetm.adjust_similarity_threshold")
+    def test_get_matching_units_uses_fuzzy_lookup(self, adjust_threshold) -> None:
+        machine = WeblateTranslation({})
+        base = MagicMock()
+        queryset = MagicMock()
+        annotated_queryset = MagicMock()
+        ordered_queryset = MagicMock()
+        fuzzy_match = MagicMock(pk=1)
+        base.filter.return_value = queryset
+        queryset.annotate.return_value = annotated_queryset
+        annotated_queryset.order_by.return_value = ordered_queryset
+
+        with patch.object(
+            machine, "prepare_queryset", return_value=[fuzzy_match]
+        ) as prepare_queryset:
+            results = machine.get_matching_units(base, "Hello", 75)
+
+        self.assertEqual(results, [fuzzy_match])
+        base.filter.assert_called_once_with(source__trgm_search="Hello")
+        queryset.annotate.assert_called_once()
+        annotated_queryset.order_by.assert_called_once_with("-match_similarity", "pk")
+        prepare_queryset.assert_called_once_with(ordered_queryset)
+        adjust_threshold.assert_called_once_with(0.98)
+
+    @patch("weblate.machinery.weblatetm.adjust_similarity_threshold")
+    def test_get_matching_units_orders_short_queries_before_slicing(
+        self, adjust_threshold
+    ) -> None:
+        machine = WeblateTranslation({})
+        base = MagicMock()
+        short_queryset = MagicMock()
+        fuzzy_match = MagicMock(pk=1)
+
+        with (
+            patch.object(
+                machine, "get_short_query_matches", return_value=short_queryset
+            ) as get_short_query_matches,
+            patch.object(machine, "prepare_queryset", return_value=[fuzzy_match]),
+        ):
+            results = machine.get_matching_units(base, "id", 75)
+
+        self.assertEqual(results, [fuzzy_match])
+        get_short_query_matches.assert_called_once_with(base, "id")
+        adjust_threshold.assert_called_once_with(0.98)
+
+    @patch("weblate.machinery.weblatetm.adjust_similarity_threshold")
+    def test_get_matching_units_uses_exact_lookup_at_full_threshold(
+        self, adjust_threshold
+    ) -> None:
+        machine = WeblateTranslation({})
+        base = MagicMock()
+        queryset = MagicMock()
+        ordered_queryset = MagicMock()
+        exact_match = MagicMock(pk=1)
+        base.filter.return_value = queryset
+        queryset.order_by.return_value = ordered_queryset
+
+        with patch.object(
+            machine, "prepare_queryset", return_value=[exact_match]
+        ) as prepare_queryset:
+            results = machine.get_matching_units(base, "Hello", 100)
+
+        self.assertEqual(results, [exact_match])
+        queryset.order_by.assert_called_once_with("pk")
+        prepare_queryset.assert_called_once_with(ordered_queryset)
+        adjust_threshold.assert_not_called()
+
+
 class CommandTest(FixtureTestCase):
     """Test for management commands."""
 

weblate/machinery/weblatetm.py

@@ -4,14 +4,15 @@
 
 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
-from django.db.models.functions import MD5, Lower
+from django.contrib.postgres.search import TrigramSimilarity
+from django.db.models import Case, IntegerField, Value, When
+from django.db.models.functions import MD5, Length, Lower
 
-from weblate.trans.models import Unit
-from weblate.utils.db import adjust_similarity_threshold
+from weblate.trans.models import Translation, Unit
+from weblate.utils.db import adjust_similarity_threshold, use_trgm_fallback
 from weblate.utils.state import STATE_TRANSLATED
 
 from .base import InternalMachineTranslation
@@ -28,6 +29,7 @@ class WeblateTranslation(InternalMachineTranslation):
     cache_translations = True
     # Cache results for 1 hour to avoid frequent database hits
     cache_expiry = 3600
+    candidate_limit = 50
 
     def download_translations(
         self,
@@ -39,42 +41,12 @@ def download_translations(
         threshold: int = 10,
     ) -> DownloadTranslations:
         """Download list of possible translations from a service."""
-        # Filter based on user access
-        base = Unit.objects.filter_access(user) if user else Unit.objects.all()
-
-        # Use memory_db for the query in case it exists. This is supposed
-        # to be a read-only replica for offloading expensive translation
-        # queries.
-        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__trgm_search"] = text
-        else:
-            # Utilize PostgreSQL index
-            lookup["source__lower__md5"] = MD5(Lower(Value(text)))
-            lookup["source"] = text
-
-        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()
+        matching_units = self.get_matching_units(
+            self.get_base_queryset(user, source_language, target_language),
+            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 +63,58 @@ def download_translations(
                 "origin_url": munit.get_absolute_url(),
                 "source": source,
             }
+
+    def get_base_queryset(self, user, source_language, target_language):
+        alias = "memory_db" if "memory_db" in settings.DATABASES else "default"
+
+        translations = Translation.objects.using(alias).all()
+        if user is not None:
+            translations = translations.filter_access(user)
+
+        translation_ids = translations.filter(
+            component__source_language=source_language,
+            language=target_language,
+        ).values("id")
+
+        return Unit.objects.using(alias).filter(
+            state__gte=STATE_TRANSLATED,
+            translation_id__in=translation_ids,
+        )
+
+    def get_matching_units(self, base, text: str, threshold: int):
+        if threshold < 100:
+            adjust_similarity_threshold(0.98)
+            if use_trgm_fallback(text):
+                queryset = self.get_short_query_matches(base, text)
+            else:
+                queryset = base.filter(source__trgm_search=text).annotate(
+                    match_similarity=TrigramSimilarity("source", text)
+                )
+                queryset = queryset.order_by("-match_similarity", "pk")
+        else:
+            queryset = base.filter(
+                source__lower__md5=MD5(Lower(Value(text))),
+                source=text,
+            ).order_by("pk")
+
+        return list(self.prepare_queryset(queryset)[: self.candidate_limit])
+
+    def prepare_queryset(self, queryset):
+        return queryset.exclude(target="").prefetch()
+
+    def get_short_query_matches(self, base, text: str):
+        max_source_length = max(len(text) + 4, len(text) * 2, 8)
+        return (
+            base.filter(source__icontains=text)
+            .annotate(
+                short_query_rank=Case(
+                    When(source__iexact=text, then=Value(0)),
+                    When(source__istartswith=text, then=Value(1)),
+                    default=Value(2),
+                    output_field=IntegerField(),
+                ),
+                source_length=Length("source"),
+            )
+            .filter(source_length__lte=max_source_length)
+            .order_by("short_query_rank", "source_length", "pk")
+        )

weblate/trans/migrations/0068_unit_source_tm_index.py

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

weblate/trans/models/unit.py

@@ -530,6 +530,11 @@ class Meta:
                 models.F("translation"),
                 name="unit_explanation_fulltext",
             ),
+            postgres_indexes.GinIndex(
+                postgres_indexes.OpClass(models.F("source"), name="gin_trgm_ops"),
+                condition=Q(state__gte=STATE_TRANSLATED) & ~Q(target=""),
+                name="trans_unit_source_tm_idx",
+            ),
         ]
 
     def __str__(self) -> str:

weblate/utils/db.py

@@ -55,6 +55,10 @@ def count_alnum(string):
     return sum(map(str.isalnum, string))
 
 
+def use_trgm_fallback(string: str) -> bool:
+    return count_alnum(string) <= 3
+
+
 class PostgreSQLFallbackLookupMixin(Lookup):
     """
     Mixin to block PostgreSQL from using trigram index.
@@ -79,7 +83,7 @@ def process_lhs(self, compiler, connection, lhs=None):
 
 class PostgreSQLFallbackLookup(PostgreSQLFallbackLookupMixin, PatternLookup):
     def __init__(self, lhs, rhs) -> None:
-        self._needs_fallback = isinstance(rhs, str) and count_alnum(rhs) <= 3
+        self._needs_fallback = isinstance(rhs, str) and use_trgm_fallback(rhs)
         super().__init__(lhs, rhs)