Repository: getsentry/sentry
Pull request: feat(indexer): Add bulk_reverse_resolve to avoid N+1 db queries.
This PR adds functionality for bulk conversion between indexes and strings in the String indexer .

fixes: https://github.com/getsentry/sentry/issues/51786

Problem statement:
Add `bulk_reverse_resolve` to StringIndexer implementation
### Problem Statement

When retrieving data from Snuba we need to convert back from indexes to strings, for this we use `StringIndexer.reverse_resolve`. 
Since we are doing this for every row we end up doing another query for each result row ( excluding cached results).
This is unreasonable and can be handled by batching multiple indexes and looking them up together.

### Solution Brainstorm

Proposed solution:

Implement a method:
```
def bulk_reverse_resolve(self, use_case_id: UseCaseID, org_id: int, ids: Collection[int]) -> List[Optional[str]]:
```
Which resolves all indexes at once.
The result shall have the same length as the ids parameter. Indexes that cannot be resolve shall have None return in their respective position.

### Product Area

Other

Hints from issue discussion:

Commit message: initial implementation of bulk_reverse_resolve

Patch to review:
diff --git a/src/sentry/sentry_metrics/indexer/__init__.py b/src/sentry/sentry_metrics/indexer/__init__.py
index 56e7298564c65b..a5fc956e93e31d 100644
--- a/src/sentry/sentry_metrics/indexer/__init__.py
+++ b/src/sentry/sentry_metrics/indexer/__init__.py
@@ -19,3 +19,4 @@
     resolve = StringIndexer().resolve
     reverse_resolve = StringIndexer().reverse_resolve
     resolve_shared_org = StringIndexer().resolve_shared_org
+    bulk_reverse_resolve = StringIndexer().bulk_reverse_resolve
diff --git a/src/sentry/sentry_metrics/indexer/base.py b/src/sentry/sentry_metrics/indexer/base.py
index ec9fc266fe2943..781ece726b5e07 100644
--- a/src/sentry/sentry_metrics/indexer/base.py
+++ b/src/sentry/sentry_metrics/indexer/base.py
@@ -6,6 +6,8 @@
 from typing import (
     Any,
     Callable,
+    Collection,
+    List,
     Mapping,
     MutableMapping,
     MutableSequence,
@@ -443,6 +445,7 @@ class StringIndexer(Service):
         "bulk_record",
         "resolve_shared_org",
         "reverse_shared_org_resolve",
+        "bulk_reverse_resolve",
     )
 
     def bulk_record(
@@ -506,6 +509,16 @@ def reverse_resolve(self, use_case_id: UseCaseID, org_id: int, id: int) -> Optio
         """
         raise NotImplementedError()
 
+    def bulk_reverse_resolve(
+        self, use_case_id: UseCaseID, org_id: int, ids: Collection[int]
+    ) -> List[Optional[str]]:
+        """Lookup the stored strings for multiple ids belonging to the same use-case and organisation.
+
+        Returns a list with the same length as the ids parameter, the resolved strings correspond with the
+        positions in the ids parameter. If an id is not resolved the corresponding position will be None.
+        """
+        raise NotImplementedError()
+
     def resolve_shared_org(self, string: str) -> Optional[int]:
         """
         Look up the index for a shared (cross organisation) string.
diff --git a/src/sentry/sentry_metrics/indexer/cache.py b/src/sentry/sentry_metrics/indexer/cache.py
index e7e34a0f1c009a..4d20154508253e 100644
--- a/src/sentry/sentry_metrics/indexer/cache.py
+++ b/src/sentry/sentry_metrics/indexer/cache.py
@@ -1,6 +1,6 @@
 import logging
 import random
-from typing import Mapping, MutableMapping, Optional, Sequence, Set
+from typing import Collection, List, Mapping, MutableMapping, Optional, Sequence, Set
 
 from django.conf import settings
 from django.core.cache import caches
@@ -177,6 +177,11 @@ def resolve(self, use_case_id: UseCaseID, org_id: int, string: str) -> Optional[
     def reverse_resolve(self, use_case_id: UseCaseID, org_id: int, id: int) -> Optional[str]:
         return self.indexer.reverse_resolve(use_case_id, org_id, id)
 
+    def bulk_reverse_resolve(
+        self, use_case_id: UseCaseID, org_id: int, ids: Collection[int]
+    ) -> List[Optional[str]]:
+        return self.indexer.bulk_reverse_resolve(use_case_id, org_id, ids)
+
     def resolve_shared_org(self, string: str) -> Optional[int]:
         raise NotImplementedError(
             "This class should not be used directly, use a wrapping class that derives from StaticStringIndexer"
diff --git a/src/sentry/sentry_metrics/indexer/mock.py b/src/sentry/sentry_metrics/indexer/mock.py
index f9e631a4293956..312e05f3c6b542 100644
--- a/src/sentry/sentry_metrics/indexer/mock.py
+++ b/src/sentry/sentry_metrics/indexer/mock.py
@@ -1,6 +1,6 @@
 import itertools
 from collections import defaultdict
-from typing import DefaultDict, Dict, Mapping, Optional, Set
+from typing import Collection, DefaultDict, Dict, List, Mapping, Optional, Set
 
 from sentry.sentry_metrics.indexer.base import (
     FetchType,
@@ -73,6 +73,15 @@ def resolve(self, use_case_id: UseCaseID, org_id: int, string: str) -> Optional[
     def reverse_resolve(self, use_case_id: UseCaseID, org_id: int, id: int) -> Optional[str]:
         return self._reverse.get(id)
 
+    def bulk_reverse_resolve(
+        self, use_case_id: UseCaseID, org_id: int, ids: Collection[int]
+    ) -> List[Optional[str]]:
+        # Performance is not an issue for this indexer, so we can fall back on reverse_resolve
+        ret_val = []
+        for ident in ids:
+            ret_val.append(self.reverse_resolve(use_case_id, org_id, ident))
+        return ret_val
+
     def _record(self, use_case_id: UseCaseID, org_id: OrgId, string: str) -> Optional[int]:
         index = self._strings[use_case_id][org_id][string]
         if index is not None:
diff --git a/src/sentry/sentry_metrics/indexer/postgres/postgres_v2.py b/src/sentry/sentry_metrics/indexer/postgres/postgres_v2.py
index ce04ba0dfbdbab..a9a838fa3c907a 100644
--- a/src/sentry/sentry_metrics/indexer/postgres/postgres_v2.py
+++ b/src/sentry/sentry_metrics/indexer/postgres/postgres_v2.py
@@ -1,7 +1,7 @@
 from functools import reduce
 from operator import or_
 from time import sleep
-from typing import Any, Collection, Mapping, Optional, Sequence, Set
+from typing import Any, Collection, List, Mapping, Optional, Sequence, Set
 
 import sentry_sdk
 from django.conf import settings
@@ -409,6 +409,27 @@ def reverse_resolve(self, use_case_id: UseCaseID, org_id: int, id: int) -> Optio
         string: str = obj.string
         return string
 
+    def bulk_reverse_resolve(
+        self, use_case_id: UseCaseID, org_id: int, ids: Collection[int]
+    ) -> List[Optional[str]]:
+        """Looks up the stored strings for the given ids.
+
+        Returns a list with the same length as ids parameter. The at position n the result contains
+        either the resolved string for ids[n] or None if ids[n] cannot be resolved.
+        """
+        metric_path_key = METRIC_PATH_MAPPING[use_case_id]
+        table = self._get_table_from_metric_path_key(metric_path_key)
+        try:
+            strings = table.objects.get_many_from_cache(id=id, use_replica=True)
+        except table.DoesNotExist:
+            return [None] * len(ids)
+
+        for obj in strings:
+            if obj:
+                assert obj.organization_id == org_id
+
+        return [obj.string if obj else None for obj in strings]
+
     def _get_metric_path_key(self, use_case_ids: Collection[UseCaseID]) -> UseCaseKey:
         metrics_paths = {METRIC_PATH_MAPPING[use_case_id] for use_case_id in use_case_ids}
         if len(metrics_paths) > 1:
diff --git a/src/sentry/sentry_metrics/indexer/strings.py b/src/sentry/sentry_metrics/indexer/strings.py
index 30549bb45c6980..e3b7de7858b213 100644
--- a/src/sentry/sentry_metrics/indexer/strings.py
+++ b/src/sentry/sentry_metrics/indexer/strings.py
@@ -1,4 +1,4 @@
-from typing import Mapping, Optional, Set
+from typing import Collection, List, Mapping, Optional, Set
 
 from sentry.sentry_metrics.indexer.base import (
     FetchType,
@@ -223,6 +223,34 @@ def reverse_resolve(self, use_case_id: UseCaseID, org_id: int, id: int) -> Optio
             return REVERSE_SHARED_STRINGS[id]
         return self.indexer.reverse_resolve(use_case_id, org_id, id)
 
+    def bulk_reverse_resolve(
+        self, use_case_id: UseCaseID, org_id: int, ids: Collection[int]
+    ) -> List[Optional[str]]:
+        ret_val: List[Optional[str]] = [None] * len(ids)
+        # keeps a mapping from the position in the request to the base indexer
+        # to the position in the global result (since shared strings are not going into the
+        # base indexer query)
+        unresolved_ids_map = []
+        unresolved_ids = []
+        for idx, ident in enumerate(ids):
+            if ident in REVERSE_SHARED_STRINGS:
+                # resolved the shared string
+                ret_val[idx] = REVERSE_SHARED_STRINGS[ident]
+            else:
+                # remember the position of the strings we need to resolve
+                unresolved_ids_map.append(idx)
+                unresolved_ids.append(ident)
+
+        # insert the strings resolved by the base indexer in the global result
+        if len(unresolved_ids) > 0:
+            resolved_strings = self.indexer.bulk_reverse_resolve(
+                use_case_id, org_id, unresolved_ids
+            )
+            for idx, resolved_str in enumerate(resolved_strings):
+                position = unresolved_ids_map[idx]
+                ret_val[position] = resolved_strings[idx]
+        return ret_val
+
     def resolve_shared_org(self, string: str) -> Optional[int]:
         if string in SHARED_STRINGS:
             return SHARED_STRINGS[string]


Code review comments:
diff_hunk: @@ -409,6 +409,27 @@ def reverse_resolve(self, use_case_id: UseCaseID, org_id: int, id: int) -> Optio
         string: str = obj.string
         return string
 
+    def bulk_reverse_resolve(
+        self, use_case_id: UseCaseID, org_id: int, ids: Collection[int]
+    ) -> List[Optional[str]]:
+        """Looks up the stored strings for the given ids.
+
+        Returns a list with the same length as ids parameter. The at position n the result contains
+        either the resolved string for ids[n] or None if ids[n] cannot be resolved.
+        """
+        metric_path_key = METRIC_PATH_MAPPING[use_case_id]
+        table = self._get_table_from_metric_path_key(metric_path_key)
+        try:
+            strings = table.objects.get_many_from_cache(id=id, use_replica=True)
+        except table.DoesNotExist:
+            return [None] * len(ids)
+
+        for obj in strings:
+            if obj:
+                assert obj.organization_id == org_id
+
+        return [obj.string if obj else None for obj in strings]
path: src/sentry/sentry_metrics/indexer/postgres/postgres_v2.py
text: @user1:
How do you know here that the ordering of strings from the query will match the `ids`?

@author:
You are right... initially I thought it is the only way it makes sense but since they return objects ( and we can see what id corresponds to each element my assumption is wrong, will fix.

@user1:
It would make sense if the id is monotonically increasing and the input we get is sorted in ascending order, in that case we can assume that the db rows are returned in the order of insertion but I don't think we should make such an assumption.

@author:
yep, but there's no need for this.

diff_hunk: @@ -409,6 +409,27 @@ def reverse_resolve(self, use_case_id: UseCaseID, org_id: int, id: int) -> Optio
         string: str = obj.string
         return string
 
+    def bulk_reverse_resolve(
+        self, use_case_id: UseCaseID, org_id: int, ids: Collection[int]
+    ) -> List[Optional[str]]:
+        """Looks up the stored strings for the given ids.
+
+        Returns a list with the same length as ids parameter. The at position n the result contains
+        either the resolved string for ids[n] or None if ids[n] cannot be resolved.
+        """
+        metric_path_key = METRIC_PATH_MAPPING[use_case_id]
+        table = self._get_table_from_metric_path_key(metric_path_key)
+        try:
+            strings = table.objects.get_many_from_cache(id=id, use_replica=True)
path: src/sentry/sentry_metrics/indexer/postgres/postgres_v2.py
text: @user1:
I can't seem to find where this `id` is created, maybe you meant:
```suggestion
            strings = table.objects.get_many_from_cache(id__in=ids, use_replica=True)
```

@author:
yep that's wrong (copy&paste mistake) it should be:
`get_many_from_cache(ids)`

@user1:
If we want all the elements whose id match `ids` shouldn't we use `id__in=ids`?

@author:
no, `get_many_from_cache` is more specialised, it doesn't allow normal django parameters.
It only gets a list of keys and optionally then name of the key column ( by default the primary key i.e. `"pk"`).
It also doesn't allow usage of a replica db for the query(which is unfortunate).