feat(eap-spans): add perf score function (#87127)

Closes getsentry/team-visibility#37

1. Implements `performance_score(x)` function
2. Adds `attribute_resolver` to `AggregateDefinition` to handle more
complicated aggregates where we don't simply aggregate on the parameter.
3. Some refactoring to improve readability and logic

Author: DominikB2014

src/sentry/search/eap/columns.py

@@ -87,7 +87,7 @@ class ArgumentDefinition:
     # Sets the argument as an attribute, for custom functions like `http_response rate` we might have non-attribute parameters
     is_attribute: bool = True
     # Validator to check if the value is allowed for this argument
-    validator: Callable[[Any], Any] | None = None
+    validator: Callable[[str], bool] | None = None
     # Whether this argument is completely ignored, used for `count()`
     ignored: bool = False
 
@@ -238,6 +238,11 @@ def resolve(
 @dataclass(kw_only=True)
 class AggregateDefinition(FunctionDefinition):
     internal_function: Function.ValueType
+    """
+    An optional function that takes in the resolved argument and returns the attribute key to aggregate on.
+    If not provided, assumes the aggregate is on the first argument.
+    """
+    attribute_resolver: Callable[[ResolvedArguments], AttributeKey] | None = None
 
     def resolve(
         self,
@@ -249,11 +254,15 @@ def resolve(
             raise InvalidSearchQuery(
                 f"Aggregates expects exactly 1 argument, got {len(resolved_arguments)}"
             )
-        resolved_argument = None
+
+        resolved_attribute = None
+
         if len(resolved_arguments) == 1:
             if not isinstance(resolved_arguments[0], AttributeKey):
                 raise InvalidSearchQuery("Aggregates accept attribute keys only")
-            resolved_argument = resolved_arguments[0]
+            resolved_attribute = resolved_arguments[0]
+            if self.attribute_resolver is not None:
+                resolved_attribute = self.attribute_resolver(resolved_arguments)
 
         return ResolvedAggregate(
             public_alias=alias,
@@ -262,7 +271,7 @@ def resolve(
             internal_type=self.internal_type,
             processor=self.processor,
             extrapolation=self.extrapolation,
-            argument=resolved_argument,
+            argument=resolved_attribute,
         )
 
 

src/sentry/search/eap/resolver.py

@@ -730,43 +730,44 @@ def resolve_function(self, column: str, match: Match | None = None) -> tuple[
                 f"Invalid number of arguments for {function}, was expecting {len(function_definition.required_arguments)} arguments"
             )
 
-        for index, argument in enumerate(function_definition.arguments):
-            if argument.ignored:
+        for index, argument_definition in enumerate(function_definition.arguments):
+            if argument_definition.ignored:
                 continue
-            if argument.validator is not None:
-                if not argument.validator(arguments[index]):
-                    raise InvalidSearchQuery(
-                        f"{arguments[index]} is not a valid argument for {function}"
-                    )
 
             if index < len(arguments):
-                if argument.is_attribute:
-                    parsed_argument, _ = self.resolve_attribute(arguments[index])
+                argument = arguments[index]
+                if argument_definition.validator is not None:
+                    if not argument_definition.validator(argument):
+                        raise InvalidSearchQuery(
+                            f"{argument} is not a valid argument for {function}"
+                        )
+                if argument_definition.is_attribute:
+                    parsed_argument, _ = self.resolve_attribute(argument)
                 else:
-                    if argument.argument_types is None:
-                        parsed_args.append(arguments[index])  # assume it's a string
+                    if argument_definition.argument_types is None:
+                        parsed_args.append(argument)  # assume it's a string
                         continue
                     # TODO: we assume that the argument is only one type for now, and we only support string/integer
-                    for type in argument.argument_types:
+                    for type in argument_definition.argument_types:
                         if type == "integer":
-                            parsed_args.append(int(arguments[index]))
+                            parsed_args.append(int(argument))
                         else:
-                            parsed_args.append(arguments[index])
+                            parsed_args.append(argument)
                     continue
 
-            elif argument.default_arg:
-                parsed_argument, _ = self.resolve_attribute(argument.default_arg)
+            elif argument_definition.default_arg:
+                parsed_argument, _ = self.resolve_attribute(argument_definition.default_arg)
             else:
                 raise InvalidSearchQuery(
                     f"Invalid number of arguments for {function}, was expecting {len(function_definition.required_arguments)} arguments"
                 )
 
             if (
-                argument.argument_types is not None
-                and parsed_argument.search_type not in argument.argument_types
+                argument_definition.argument_types is not None
+                and parsed_argument.search_type not in argument_definition.argument_types
             ):
                 raise InvalidSearchQuery(
-                    f"{parsed_argument.public_alias} is invalid for parameter {index+1} in {function}. Its a {parsed_argument.search_type} type field, but it must be one of these types: {argument.argument_types}"
+                    f"{parsed_argument.public_alias} is invalid for parameter {index+1} in {function}. Its a {parsed_argument.search_type} type field, but it must be one of these types: {argument_definition.argument_types}"
                 )
             parsed_args.append(parsed_argument)
 

src/sentry/search/eap/spans/aggregates.py

@@ -66,14 +66,21 @@ def resolve_key_eq_value_filter(args: ResolvedArguments) -> tuple[AttributeKey,
 
 # TODO: We should eventually update the frontend to query the ratio column directly
 def resolve_count_scores(args: ResolvedArguments) -> tuple[AttributeKey, TraceItemFilter]:
-    score_column = cast(str, args[0])
-    ratio_column_name = score_column.replace("measurements.score", "score.ratio")
-    if ratio_column_name == "score.ratio.total":
-        ratio_column_name = "score.total"
-    attribute_key = AttributeKey(name=ratio_column_name, type=AttributeKey.TYPE_DOUBLE)
-    filter = TraceItemFilter(exists_filter=ExistsFilter(key=attribute_key))
+    score_attribute = cast(AttributeKey, args[0])
+    ratio_attribute = transform_vital_score_to_ratio([score_attribute])
+    filter = TraceItemFilter(exists_filter=ExistsFilter(key=ratio_attribute))
 
-    return (attribute_key, filter)
+    return (ratio_attribute, filter)
+
+
+def transform_vital_score_to_ratio(args: ResolvedArguments) -> AttributeKey:
+    score_attribute = cast(AttributeKey, args[0])
+    score_name = score_attribute.name
+
+    ratio_score_name = score_name.replace("score", "score.ratio")
+    if ratio_score_name == "score.ratio.total":
+        ratio_score_name = "score.total"
+    return AttributeKey(name=ratio_score_name, type=AttributeKey.TYPE_DOUBLE)
 
 
 SPAN_CONDITIONAL_AGGREGATE_DEFINITIONS = {
@@ -113,9 +120,14 @@ def resolve_count_scores(args: ResolvedArguments) -> tuple[AttributeKey, TraceIt
         default_search_type="integer",
         arguments=[
             ArgumentDefinition(
-                argument_types={"string"},
+                argument_types={
+                    "duration",
+                    "number",
+                    "percentage",
+                    *constants.SIZE_TYPE,
+                    *constants.DURATION_TYPE,
+                },
                 validator=literal_validator(WEB_VITALS_MEASUREMENTS),
-                is_attribute=False,
             )
         ],
         aggregate_resolver=resolve_count_scores,
@@ -355,4 +367,21 @@ def resolve_count_scores(args: ResolvedArguments) -> tuple[AttributeKey, TraceIt
             )
         ],
     ),
+    "performance_score": AggregateDefinition(
+        internal_function=Function.FUNCTION_AVG,
+        default_search_type="integer",
+        arguments=[
+            ArgumentDefinition(
+                argument_types={
+                    "duration",
+                    "number",
+                    "percentage",
+                    *constants.SIZE_TYPE,
+                    *constants.DURATION_TYPE,
+                },
+                validator=literal_validator(WEB_VITALS_MEASUREMENTS),
+            ),
+        ],
+        attribute_resolver=transform_vital_score_to_ratio,
+    ),
 }

tests/snuba/api/endpoints/test_organization_events_span_indexed.py

@@ -2981,6 +2981,54 @@ def test_count_scores(self):
         assert data[0]["count_scores(measurements.score.total)"] == 3
         assert meta["dataset"] == self.dataset
 
+    def test_performance_score(self):
+        self.store_spans(
+            [
+                self.create_span(
+                    {
+                        "measurements": {
+                            "score.ratio.lcp": {"value": 0.02},
+                        }
+                    },
+                    start_ts=self.ten_mins_ago,
+                ),
+                self.create_span(
+                    {
+                        "measurements": {
+                            "score.ratio.lcp": {"value": 0.08},
+                        }
+                    },
+                    start_ts=self.ten_mins_ago,
+                ),
+                self.create_span(
+                    {
+                        "measurements": {
+                            "score.ratio.lcp": {"value": 0.08},
+                        }
+                    },
+                    start_ts=self.ten_mins_ago,
+                ),
+            ],
+            is_eap=self.is_eap,
+        )
+
+        response = self.do_request(
+            {
+                "field": [
+                    "performance_score(measurements.score.lcp)",
+                ],
+                "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]["performance_score(measurements.score.lcp)"] == 0.06
+        assert meta["dataset"] == self.dataset
+
     @pytest.mark.skip(reason="replay id alias not migrated over")
     def test_replay_id(self):
         super().test_replay_id()