feat(starfish): Treat percentiles as an entity (#53068)

- This treats percentiles as their own entity in the query framework,
the intent is so that we'll run cheaper distributions first when
possible, then only compute the percentiles for a smaller set after
- How this works is that if you have a query with sum(duration) and
p50(duration) grouping by span.description ordering by sum(duration)
1. A query with span.description + sum(duration) is run first, we get 50
of those results ordered by sum(duration)
2. Those 50 results are then used as a filter so we only get the
p50(duration) of up to 50 descriptions

src/sentry/search/events/builder/metrics.py

@@ -65,6 +65,7 @@ class MetricsQueryBuilder(QueryBuilder):
         self.distributions: List[CurriedFunction] = []
         self.sets: List[CurriedFunction] = []
         self.counters: List[CurriedFunction] = []
+        self.percentiles: List[CurriedFunction] = []
         self.metric_ids: Set[int] = set()
         self.allow_metric_aggregates = allow_metric_aggregates
         self._indexer_cache: Dict[str, Optional[int]] = {}
@@ -417,6 +418,13 @@ class MetricsQueryBuilder(QueryBuilder):
                 # Still add to aggregates so groupby is correct
                 self.aggregates.append(resolved_function)
             return resolved_function
+        if snql_function.snql_percentile is not None:
+            resolved_function = snql_function.snql_percentile(arguments, alias)
+            if not resolve_only:
+                self.percentiles.append(resolved_function)
+                # Still add to aggregates so groupby is correct
+                self.aggregates.append(resolved_function)
+            return resolved_function
         if snql_function.snql_set is not None:
             resolved_function = snql_function.snql_set(arguments, alias)
             if not resolve_only:
@@ -673,16 +681,19 @@ class MetricsQueryBuilder(QueryBuilder):
                 functions=self.sets,
                 entity=Entity(f"{prefix}metrics_sets", sample=self.sample_rate),
             ),
+            # Percentiles are a part of distributions but they're expensive, treat them as their own entity so we'll run
+            # a query with the cheap distributions first then only get page_size quantiles
+            "percentiles": QueryFramework(
+                orderby=[],
+                having=[],
+                functions=self.percentiles,
+                entity=Entity(f"{prefix}metrics_distributions", sample=self.sample_rate),
+            ),
         }
         primary = None
         # if orderby spans more than one table, the query isn't possible with metrics
         for orderby in self.orderby:
-            if orderby.exp in self.distributions:
-                query_framework["distribution"].orderby.append(orderby)
-                if primary not in [None, "distribution"]:
-                    raise IncompatibleMetricsQuery("Can't order across tables")
-                primary = "distribution"
-            elif orderby.exp in self.sets:
+            if orderby.exp in self.sets:
                 query_framework["set"].orderby.append(orderby)
                 if primary not in [None, "set"]:
                     raise IncompatibleMetricsQuery("Can't order across tables")
@@ -692,6 +703,16 @@ class MetricsQueryBuilder(QueryBuilder):
                 if primary not in [None, "counter"]:
                     raise IncompatibleMetricsQuery("Can't order across tables")
                 primary = "counter"
+            elif orderby.exp in self.distributions:
+                query_framework["distribution"].orderby.append(orderby)
+                if primary not in [None, "distribution"]:
+                    raise IncompatibleMetricsQuery("Can't order across tables")
+                primary = "distribution"
+            elif orderby.exp in self.percentiles:
+                query_framework["percentiles"].orderby.append(orderby)
+                if primary not in [None, "percentiles"]:
+                    raise IncompatibleMetricsQuery("Can't order across tables")
+                primary = "percentiles"
             else:
                 # An orderby that isn't on a function add it to all of them
                 for framework in query_framework.values():
@@ -699,14 +720,7 @@ class MetricsQueryBuilder(QueryBuilder):
 
         having_entity: Optional[str] = None
         for condition in self.flattened_having:
-            if condition.lhs in self.distributions:
-                if having_entity is None:
-                    having_entity = "distribution"
-                elif having_entity != "distribution":
-                    raise IncompatibleMetricsQuery(
-                        "Can only have aggregate conditions on one entity"
-                    )
-            elif condition.lhs in self.sets:
+            if condition.lhs in self.sets:
                 if having_entity is None:
                     having_entity = "set"
                 elif having_entity != "set":
@@ -720,6 +734,20 @@ class MetricsQueryBuilder(QueryBuilder):
                     raise IncompatibleMetricsQuery(
                         "Can only have aggregate conditions on one entity"
                     )
+            elif condition.lhs in self.distributions:
+                if having_entity is None:
+                    having_entity = "distribution"
+                elif having_entity != "distribution":
+                    raise IncompatibleMetricsQuery(
+                        "Can only have aggregate conditions on one entity"
+                    )
+            elif condition.lhs in self.percentiles:
+                if having_entity is None:
+                    having_entity = "percentiles"
+                elif having_entity != "percentiles":
+                    raise IncompatibleMetricsQuery(
+                        "Can only have aggregate conditions on one entity"
+                    )
 
         if primary is not None and having_entity is not None and having_entity != primary:
             raise IncompatibleMetricsQuery(
@@ -730,12 +758,14 @@ class MetricsQueryBuilder(QueryBuilder):
         if primary is None:
             if having_entity is not None:
                 primary = having_entity
-            elif len(self.distributions) > 0:
-                primary = "distribution"
             elif len(self.counters) > 0:
                 primary = "counter"
             elif len(self.sets) > 0:
                 primary = "set"
+            elif len(self.distributions) > 0:
+                primary = "distribution"
+            elif len(self.percentiles) > 0:
+                primary = "percentiles"
             else:
                 raise IncompatibleMetricsQuery("Need at least one function")
 

src/sentry/search/events/datasets/spans_metrics.py

@@ -145,7 +145,7 @@ class SpansMetricsDatasetConfig(DatasetConfig):
                         fields.NumberRange("percentile", 0, 1),
                     ],
                     calculated_args=[resolve_metric_id],
-                    snql_distribution=function_aliases.resolve_metrics_percentile,
+                    snql_percentile=function_aliases.resolve_metrics_percentile,
                     result_type_fn=self.reflective_result_type(),
                     default_result_type="duration",
                 ),
@@ -162,7 +162,7 @@ class SpansMetricsDatasetConfig(DatasetConfig):
                         ),
                     ],
                     calculated_args=[resolve_metric_id],
-                    snql_distribution=lambda args, alias: function_aliases.resolve_metrics_percentile(
+                    snql_percentile=lambda args, alias: function_aliases.resolve_metrics_percentile(
                         args=args, alias=alias, fixed_percentile=0.50
                     ),
                     default_result_type="duration",
@@ -180,7 +180,7 @@ class SpansMetricsDatasetConfig(DatasetConfig):
                         ),
                     ],
                     calculated_args=[resolve_metric_id],
-                    snql_distribution=lambda args, alias: function_aliases.resolve_metrics_percentile(
+                    snql_percentile=lambda args, alias: function_aliases.resolve_metrics_percentile(
                         args=args, alias=alias, fixed_percentile=0.75
                     ),
                     default_result_type="duration",
@@ -198,21 +198,11 @@ class SpansMetricsDatasetConfig(DatasetConfig):
                         ),
                     ],
                     calculated_args=[resolve_metric_id],
-                    snql_distribution=lambda args, alias: function_aliases.resolve_metrics_percentile(
+                    snql_percentile=lambda args, alias: function_aliases.resolve_metrics_percentile(
                         args=args, alias=alias, fixed_percentile=0.95
                     ),
                     default_result_type="duration",
                 ),
-                fields.MetricsFunction(
-                    "time_spent_percentage",
-                    optional_args=[
-                        fields.with_default(
-                            "app", fields.SnQLStringArg("scope", allowed_strings=["app", "local"])
-                        )
-                    ],
-                    snql_distribution=self._resolve_time_spent_percentage,
-                    default_result_type="percentage",
-                ),
                 fields.MetricsFunction(
                     "p99",
                     optional_args=[
@@ -226,7 +216,7 @@ class SpansMetricsDatasetConfig(DatasetConfig):
                         ),
                     ],
                     calculated_args=[resolve_metric_id],
-                    snql_distribution=lambda args, alias: function_aliases.resolve_metrics_percentile(
+                    snql_percentile=lambda args, alias: function_aliases.resolve_metrics_percentile(
                         args=args, alias=alias, fixed_percentile=0.99
                     ),
                     default_result_type="duration",
@@ -244,11 +234,21 @@ class SpansMetricsDatasetConfig(DatasetConfig):
                         ),
                     ],
                     calculated_args=[resolve_metric_id],
-                    snql_distribution=lambda args, alias: function_aliases.resolve_metrics_percentile(
+                    snql_percentile=lambda args, alias: function_aliases.resolve_metrics_percentile(
                         args=args, alias=alias, fixed_percentile=1
                     ),
                     default_result_type="duration",
                 ),
+                fields.MetricsFunction(
+                    "time_spent_percentage",
+                    optional_args=[
+                        fields.with_default(
+                            "app", fields.SnQLStringArg("scope", allowed_strings=["app", "local"])
+                        )
+                    ],
+                    snql_distribution=self._resolve_time_spent_percentage,
+                    default_result_type="percentage",
+                ),
                 fields.MetricsFunction(
                     "http_error_rate",
                     snql_distribution=lambda args, alias: self.builder.resolve_division(
@@ -288,7 +288,7 @@ class SpansMetricsDatasetConfig(DatasetConfig):
                         fields.SnQLDateArg("middle"),
                     ],
                     calculated_args=[resolve_metric_id],
-                    snql_distribution=lambda args, alias: function_aliases.resolve_metrics_percentile(
+                    snql_percentile=lambda args, alias: function_aliases.resolve_metrics_percentile(
                         args=args,
                         alias=alias,
                         fixed_percentile=args["percentile"],
@@ -315,7 +315,7 @@ class SpansMetricsDatasetConfig(DatasetConfig):
                         fields.NumberRange("percentile", 0, 1),
                     ],
                     calculated_args=[resolve_metric_id],
-                    snql_distribution=self._resolve_percentile_percent_change,
+                    snql_percentile=self._resolve_percentile_percent_change,
                     default_result_type="percent_change",
                 ),
                 fields.MetricsFunction(

src/sentry/search/events/fields.py

@@ -2132,6 +2132,7 @@ class MetricsFunction(SnQLFunction):
 
     def __init__(self, *args, **kwargs) -> None:
         self.snql_distribution = kwargs.pop("snql_distribution", None)
+        self.snql_percentile = kwargs.pop("snql_percentile", None)
         self.snql_set = kwargs.pop("snql_set", None)
         self.snql_counter = kwargs.pop("snql_counter", None)
         self.snql_metric_layer = kwargs.pop("snql_metric_layer", None)
@@ -2150,6 +2151,7 @@ class MetricsFunction(SnQLFunction):
                     self.snql_distribution is not None,
                     self.snql_set is not None,
                     self.snql_counter is not None,
+                    self.snql_percentile is not None,
                     self.snql_column is not None,
                     self.snql_metric_layer is not None,
                 ]