fix(eap): Map virtual context columns properly for topn queries (#90343)

Chart is empty for device.class top n queries.
This is because virtual context columns are currently not supported on
timeseries.
So we need to map it back to original column and values and back for
topn


[example](https://demo.sentry.io/explore/traces/?field=id&field=span.op&field=span.description&field=span.duration&field=transaction&field=timestamp&field=tags%5Bhttp.client.resolve_dns_ms%2Cnumber%5D&groupBy=device.class&mode=aggregate&project=4508968118321152&query=span.op%3Aapp.start.cold&sort=-count%28span.duration%29&statsPeriod=14d&table=span&visualize=%7B%22chartType%22%3A1%2C%22yAxes%22%3A%5B%22count%28span.duration%29%22%5D%7D)

Author: shruthilayaj

src/sentry/search/eap/resolver.py

@@ -403,7 +403,9 @@ def resolve_term(
 
         value = term.value.value
         if self.params.is_timeseries_request and context_definition is not None:
-            resolved_column, value = self.map_context_to_original_column(term, context_definition)
+            resolved_column, value = self.map_search_term_context_to_original_column(
+                term, context_definition
+            )
             context_definition = None
 
         if not isinstance(resolved_column.proto_definition, AttributeKey):
@@ -503,9 +505,8 @@ def resolve_term(
 
     def map_context_to_original_column(
         self,
-        term: event_search.SearchFilter,
         context_definition: VirtualColumnDefinition,
-    ) -> tuple[ResolvedAttribute, str | int | list[str]]:
+    ) -> ResolvedAttribute:
         """
         Time series request do not support virtual column contexts, so we have to remap the value back to the original column.
         (see https://github.com/getsentry/eap-planning/issues/236)
@@ -525,10 +526,30 @@ def map_context_to_original_column(
         if public_alias is None:
             raise InvalidSearchQuery(f"Cannot map {context.from_column_name} to a public alias")
 
-        value = term.value.value
         resolved_column, _ = self.resolve_column(public_alias)
+
         if not isinstance(resolved_column.proto_definition, AttributeKey):
-            raise ValueError(f"{term.key.name} is not valid search term")
+            raise ValueError(f"{resolved_column.public_alias} is not valid search term")
+
+        return resolved_column
+
+    def map_search_term_context_to_original_column(
+        self,
+        term: event_search.SearchFilter,
+        context_definition: VirtualColumnDefinition,
+    ) -> tuple[ResolvedAttribute, str | int | list[str]]:
+        """
+        Time series request do not support virtual column contexts, so we have to remap the value back to the original column.
+        (see https://github.com/getsentry/eap-planning/issues/236)
+        """
+        context = context_definition.constructor(self.params)
+        is_number_column = (
+            context.from_column_name in SPANS_INTERNAL_TO_PUBLIC_ALIAS_MAPPINGS["number"]
+        )
+
+        resolved_column = self.map_context_to_original_column(context_definition)
+
+        value = term.value.value
 
         inverse_value_map: dict[str, list[str]] = {}
         for key, val in context.value_map.items():

src/sentry/snuba/rpc_dataset_common.py

@@ -114,7 +114,17 @@ def get_timeseries_query(
     meta = search_resolver.resolve_meta(referrer=referrer, sampling_mode=sampling_mode)
     query, _, query_contexts = search_resolver.resolve_query(query_string)
     (functions, _) = search_resolver.resolve_functions(y_axes)
-    (groupbys, _) = search_resolver.resolve_attributes(groupby)
+    groupbys, groupby_contexts = search_resolver.resolve_attributes(groupby)
+
+    # Virtual context columns (VCCs) are currently only supported in TraceItemTable.
+    # Since they are not supported here - we map them manually back to the original
+    # column the virtual context column would have used.
+    for i, groupby_definition in enumerate(zip(groupbys, groupby_contexts)):
+        _, context = groupby_definition
+        if context is not None:
+            col = search_resolver.map_context_to_original_column(context)
+            groupbys[i] = col
+
     if extra_conditions is not None:
         if query is not None:
             query = TraceItemFilter(and_filter=AndFilter(filters=[query, extra_conditions]))

src/sentry/snuba/spans_rpc.py

@@ -207,9 +207,20 @@ def run_top_events_timeseries_query(
     change this"""
     """Make a table query first to get what we need to filter by"""
     rpc_dataset_common.validate_granularity(params)
-    search_resolver = get_resolver(params=params, config=config)
+
+    # Virtual context columns (VCCs) are currently only supported in TraceItemTable.
+    # For TopN queries, we want table and timeseries data to match.
+    # Here, we want to run the table request the the VCCs. SnubaParams has
+    # a property `is_timeseries_request` which resolves to true if granularity_secs is set.
+    # `is_timeseries_request` is used to evaluate if VCCs should be used.
+    # Unset granularity_secs, so this gets treated as a table request with
+    # the correct VCC.
+    table_query_params = params.copy()
+    table_query_params.granularity_secs = None
+    table_search_resolver = get_resolver(params=table_query_params, config=config)
+
     top_events = run_table_query(
-        params,
+        table_query_params,
         query_string,
         raw_groupby + y_axes,
         orderby,
@@ -218,10 +229,12 @@ def run_top_events_timeseries_query(
         referrer,
         config,
         sampling_mode,
-        search_resolver,
+        table_search_resolver,
     )
     if len(top_events["data"]) == 0:
         return {}
+
+    search_resolver = get_resolver(params=params, config=config)
     # Need to change the project slug columns to project.id because timeseries requests don't take virtual_column_contexts
     groupby_columns = [col for col in raw_groupby if not is_function(col)]
     groupby_columns_without_project = [
@@ -275,8 +288,22 @@ def run_top_events_timeseries_query(
                     int(groupby_attributes[resolved_groupby.internal_name])
                 ]
             else:
-                resolved_groupby, _ = search_resolver.resolve_attribute(col)
-                remapped_groupby[col] = groupby_attributes[resolved_groupby.internal_name]
+                resolved_groupby, context = search_resolver.resolve_attribute(col)
+
+                # Virtual context columns (VCCs) are currently only supported in TraceItemTable.
+                # Since timeseries run the query with the original column, we need to map
+                # them correctly so they map the table result. We need to map both the column name
+                # and the values.
+                if context is not None:
+                    resolved_groupby = search_resolver.map_context_to_original_column(context)
+
+                groupby_value = groupby_attributes[resolved_groupby.internal_name]
+                if context is not None:
+                    groupby_value = context.constructor(params).value_map[groupby_value]
+                    groupby_attributes[resolved_groupby.internal_name] = groupby_value
+
+                remapped_groupby[col] = groupby_value
+
         result_key = create_result_key(remapped_groupby, groupby_columns, {})
         map_result_key_to_timeseries[result_key].append(timeseries)
     final_result = {}

tests/snuba/api/endpoints/test_organization_events_span_indexed.py

@@ -2568,6 +2568,32 @@ def test_device_class_filter_unknown(self):
         assert data[0]["device.class"] == "Unknown"
         assert meta["dataset"] == self.dataset
 
+    def test_device_class_column(self):
+        self.store_spans(
+            [
+                self.create_span(
+                    {"sentry_tags": {"device.class": "1"}}, start_ts=self.ten_mins_ago
+                ),
+            ],
+            is_eap=self.is_eap,
+        )
+        response = self.do_request(
+            {
+                "field": ["device.class", "count()"],
+                "query": "",
+                "orderby": "count()",
+                "project": self.project.id,
+                "dataset": self.dataset,
+            }
+        )
+
+        assert response.status_code == 200, response.content
+        data = response.data["data"]
+        meta = response.data["meta"]
+        assert len(data) == 1
+        assert data[0]["device.class"] == "low"
+        assert meta["dataset"] == self.dataset
+
     def test_http_response_count(self):
         self.store_spans(
             [

tests/snuba/api/endpoints/test_organization_events_stats_span_indexed.py

@@ -3,6 +3,7 @@
 import pytest
 from django.urls import reverse
 
+from sentry.search.utils import DEVICE_CLASS
 from sentry.testutils.helpers.datetime import before_now
 from tests.snuba.api.endpoints.test_organization_events import OrganizationEventsEndpointTestBase
 from tests.snuba.api.endpoints.test_organization_events_span_indexed import KNOWN_PREFLIGHT_ID
@@ -1319,6 +1320,70 @@ def test_device_class_filter(self):
             for test in zip(event_counts, rows):
                 assert test[1][1][0]["count"] == test[0]
 
+    def test_device_class_top_events(self):
+        event_counts = [
+            ("low", 6),
+            ("medium", 0),
+            ("low", 6),
+            ("medium", 6),
+            ("low", 0),
+            ("medium", 3),
+        ]
+        spans = []
+        for hour, count in enumerate(event_counts):
+            spans.extend(
+                [
+                    self.create_span(
+                        {
+                            "description": "foo",
+                            "sentry_tags": {
+                                "status": "success",
+                                "device.class": (
+                                    list(DEVICE_CLASS["low"])[0]
+                                    if count[0] == "low"
+                                    else list(DEVICE_CLASS["medium"])[0]
+                                ),
+                            },
+                        },
+                        start_ts=self.day_ago + timedelta(hours=hour, minutes=minute),
+                    )
+                    for minute in range(count[1])
+                ],
+            )
+        self.store_spans(spans, is_eap=self.is_eap)
+
+        response = self._do_request(
+            data={
+                "start": self.day_ago,
+                "end": self.day_ago + timedelta(hours=6),
+                "interval": "1h",
+                "yAxis": "count()",
+                "field": ["device.class", "count()"],
+                "topEvents": 5,
+                "query": "",
+                "project": self.project.id,
+                "dataset": self.dataset,
+            },
+        )
+        assert response.status_code == 200, response.content
+        low = response.data["low"]["data"]
+        assert len(low) == 6
+
+        rows = low[0:6]
+        for i, test in enumerate(zip(event_counts, rows)):
+            test_data, row = test
+            test_count = test_data[1] if test_data[0] == "low" else 0.0
+            assert row[1][0]["count"] == test_count
+
+        medium = response.data["medium"]["data"]
+        assert len(medium) == 6
+
+        rows = medium[0:6]
+        for i, test in enumerate(zip(event_counts, rows)):
+            test_data, row = test
+            test_count = test_data[1] if test_data[0] == "medium" else 0.0
+            assert row[1][0]["count"] == test_count
+
     def test_top_events_filters_out_groupby_even_when_its_just_one_row(self):
         self.store_spans(
             [