Repository: getsentry/snuba
Pull request #7153: feat(eap): Transform an order by on timestamp to the full sort key order to improve performance 
Author: login: xurui-c
Created: 2025-05-02T18:00:52Z | Merged: 2025-05-13T04:46:35Z

redoing https://github.com/getsentry/snuba/pull/7149

only add all the sort key columns to order by if the GROUP BY is empty or else we'd get error DB::Exception: Column blah is not under aggregate function and not in GROUP BY



Diff:
diff --git a/snuba/web/rpc/v1/resolvers/R_eap_items/resolver_trace_item_table.py b/snuba/web/rpc/v1/resolvers/R_eap_items/resolver_trace_item_table.py
index cbb0ea96047..b57af93a1f2 100644
--- a/snuba/web/rpc/v1/resolvers/R_eap_items/resolver_trace_item_table.py
+++ b/snuba/web/rpc/v1/resolvers/R_eap_items/resolver_trace_item_table.py
@@ -1,5 +1,5 @@
 from dataclasses import replace
-from typing import Sequence
+from typing import List, Sequence
 
 from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import (
     AggregationComparisonFilter,
@@ -118,9 +118,35 @@ def aggregation_filter_to_expression(
 
 
 def _convert_order_by(
+    groupby: List[Expression],
     order_by: Sequence[TraceItemTableRequest.OrderBy],
     request_meta: RequestMeta,
 ) -> Sequence[OrderBy]:
+    if len(order_by) == 1:
+        order = order_by[0]
+        if order.column.key.name == "sentry.timestamp" and len(groupby) == 0:
+            direction = (
+                OrderByDirection.DESC if order.descending else OrderByDirection.ASC
+            )
+            return [
+                OrderBy(
+                    direction=direction,
+                    expression=snuba_column("organization_id"),
+                ),
+                OrderBy(
+                    direction=direction,
+                    expression=snuba_column("project_id"),
+                ),
+                OrderBy(
+                    direction=direction,
+                    expression=snuba_column("item_type"),
+                ),
+                OrderBy(
+                    direction=direction,
+                    expression=snuba_column("timestamp"),
+                ),
+            ]
+
     res: list[OrderBy] = []
     for x in order_by:
         direction = OrderByDirection.DESC if x.descending else OrderByDirection.ASC
@@ -270,6 +296,9 @@ def build_query(request: TraceItemTableRequest) -> Query:
     item_type_conds = [
         f.equals(snuba_column("item_type"), request.meta.trace_item_type)
     ]
+    groupby = [
+        attribute_key_to_expression_eap_items(attr_key) for attr_key in request.group_by
+    ]
     res = Query(
         from_clause=entity,
         selected_columns=selected_columns,
@@ -281,11 +310,12 @@ def build_query(request: TraceItemTableRequest) -> Query:
             ),
             *item_type_conds,
         ),
-        order_by=_convert_order_by(request.order_by, request.meta),
-        groupby=[
-            attribute_key_to_expression_eap_items(attr_key)
-            for attr_key in request.group_by
-        ],
+        order_by=_convert_order_by(
+            groupby,
+            request.order_by,
+            request.meta,
+        ),
+        groupby=groupby,
         # Only support offset page tokens for now
         offset=request.page_token.offset,
         # protobuf sets limit to 0 by default if it is not set,
diff --git a/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py b/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py
index 0444d0dffce..7efacc7cbb2 100644
--- a/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py
+++ b/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py
@@ -1135,6 +1135,33 @@ def test_order_by_aggregation(self, setup_teardown: Any) -> None:
         measurements = [v.val_double for v in response.column_values[1].results]
         assert sorted(measurements) == measurements
 
+    def test_order_by_does_not_error_if_groupby_exists(self) -> None:
+        message = TraceItemTableRequest(
+            meta=RequestMeta(
+                project_ids=[1, 2, 3],
+                trace_item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN,
+            ),
+            columns=[
+                Column(
+                    key=AttributeKey(
+                        type=AttributeKey.TYPE_STRING, name="sentry.timestamp"
+                    )
+                ),
+                Column(
+                    aggregation=AttributeAggregation(
+                        aggregate=Function.FUNCTION_COUNT,
+                        key=AttributeKey(
+                            type=AttributeKey.TYPE_STRING, name="sentry.timestamp"
+                        ),
+                    )
+                ),
+            ],
+            group_by=[
+                AttributeKey(type=AttributeKey.TYPE_STRING, name="sentry.timestamp")
+            ],
+        )
+        EndpointTraceItemTable().execute(message)
+
     def test_aggregation_on_attribute_column_backward_compat(
         self,
         setup_teardown: Any,


Reviews:
user: login: volokluev
body: please add a test to make sure this doesn't reproduce the bug
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-02T20:50:57Z
user: login: volokluev
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-02T22:07:27Z
user: login: xurui-c
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-02T23:07:37Z
user: login: volokluev
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-08T18:35:15Z

Review comments:
diff_hunk: @@ -1235,6 +1238,71 @@ def test_order_by_aggregation(self, setup_teardown: Any) -> None:
         measurements = [v.val_double for v in response.column_values[1].results]
         assert sorted(measurements) == measurements
 
+    def test_order_by_full_sort_key_if_no_groupby(self) -> None:
path: tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py
user: login: volokluev
body: this is testing implementation details too much. the problem that started this task was that [an exception was being thrown ](https://sentry.sentry.io/issues/6581772260/?project=300688&query=is%3Aunresolved%20issue.priority%3A%5Bhigh%2C%20medium%5D%20transaction%3A%5BEndpointTraceItemAttributeNames__v1%2CEndpointTimeSeries__v1%2CEndpointTraceItemTable__v1%2CEndpointGetTrace__v1%5D%20%21title%3A%22%5C%22BadSnubaRPCRequestException%2A%22&referrer=issue-stream&sort=date&stream_index=12) during certain conditions. We want to make sure we don't repeat that mistake. that's all
created_at: 2025-05-02T22:07:27Z
author_association: MEMBER
diff_hunk: @@ -1235,6 +1238,71 @@ def test_order_by_aggregation(self, setup_teardown: Any) -> None:
         measurements = [v.val_double for v in response.column_values[1].results]
         assert sorted(measurements) == measurements
 
+    def test_order_by_full_sort_key_if_no_groupby(self) -> None:
path: tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py
user: login: xurui-c
body: done
created_at: 2025-05-02T23:07:36Z
author_association: MEMBER

PR comments:
user: login: codecov[bot]
created_at: 2025-05-02T18:16:55Z
body: ### :x: 1 Tests Failed:
| Tests completed | Failed | Passed | Skipped |
|---|---|---|---|
| 2714 | 1 | 2713 | 11 |
<details><summary>View the top 1 failed test(s) by shortest run time</summary>

> 
> ```python
> tests.web.rpc.test_base::test_trim_time_range
> ```
> 
> <details><summary>Stack Traces | 0.505s run time</summary>
> 
> > 
> > ```python
> > Traceback (most recent call last):
> >   File ".../web/rpc/test_base.py", line 199, in test_trim_time_range
> >     response = EndpointTraceItemTable().execute(message)
> >                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> >   File ".../web/rpc/__init__.py", line 187, in execute
> >     return self.__after_execute(in_msg, out, error)
> >            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> >   File ".../web/rpc/__init__.py", line 252, in __after_execute
> >     raise error
> >   File ".../web/rpc/__init__.py", line 156, in execute
> >     out = self._execute(in_msg)
> >           ^^^^^^^^^^^^^^^^^^^^^
> >   File ".../rpc/v1/endpoint_trace_item_table.py", line 136, in _execute
> >     return resolver.resolve(in_msg)
> >            ^^^^^^^^^^^^^^^^^^^^^^^^
> >   File ".../resolvers/R_eap_spans/resolver_trace_item_table.py", line 19, in resolve
> >     return ResolverTraceItemTableEAPItems().resolve(
> >            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> >   File ".../resolvers/R_eap_items/resolver_trace_item_table.py", line 361, in resolve
> >     res = run_query_to_correct_tier(
> >           ^^^^^^^^^^^^^^^^^^^^^^^^^^
> >   File ".../utils/metrics/util.py", line 75, in wrapper
> >     return func(*args, **kwargs)
> >            ^^^^^^^^^^^^^^^^^^^^^
> >   File ".../R_eap_items/storage_routing/sampling_in_storage_util.py", line 36, in run_query_to_correct_tier
> >     return selected_strategy.run_query_to_correct_tier(routing_context)
> >            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> >   File ".../utils/metrics/util.py", line 75, in wrapper
> >     return func(*args, **kwargs)
> >            ^^^^^^^^^^^^^^^^^^^^^
> >   File ".../storage_routing/routing_strategies/storage_routing.py", line 385, in run_query_to_correct_tier
> >     output = self._run_query(routing_context)
> >              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> >   File ".../storage_routing/routing_strategies/storage_routing.py", line 199, in _run_query
> >     snuba_request = self._build_snuba_request(routing_context)
> >                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> >   File ".../storage_routing/routing_strategies/storage_routing.py", line 183, in _build_snuba_request
> >     query=routing_context.build_query(request),
> >           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> >   File ".../resolvers/R_eap_items/resolver_trace_item_table.py", line 322, in build_query
> >     order_by=_convert_order_by(set(groupby), request.order_by, request.meta),
> >                                ^^^^^^^^^^^^
> > TypeError: unhashable type: 'list'
> > ```
> 
> </details>

</details>

To view more test analytics, go to the [Test Analytics Dashboard](https://app.codecov.io/gh/getsentry/snuba/tests/rachel%2Fperf)
<sub>📋 Got 3 mins? [Take this short survey](https://forms.gle/BpocVj23nhr2Y45G7) to help us improve Test Analytics.</sub>
author_association: NONE
user: login: volokluev
created_at: 2025-05-02T19:59:30Z
body: This won't work, or rather won't do anything. The columns are never in the groupby because they are not selected. what you want to do is only apply this optimization if there is *no* groupby
author_association: MEMBER