feat(mep): Add avg, count, percentile to MEP (#33465)

- This adds avg, count and percentile to the supported functions in MEP

src/sentry/search/events/constants.py

@@ -210,3 +210,11 @@ DRY_RUN_COLUMNS = {
     "transaction",
     "transaction.status",
 }
+METRIC_PERCENTILES = {
+    0.5,
+    0.75,
+    0.9,
+    0.95,
+    0.99,
+    1,
+}

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

@@ -76,6 +76,24 @@ class MetricsDatasetConfig(DatasetConfig):
                     snql_distribution=self._resolve_apdex_function,
                     default_result_type="number",
                 ),
+                fields.MetricsFunction(
+                    "avg",
+                    snql_distribution=lambda args, alias: Function(
+                        "avgIf",
+                        [
+                            Column("value"),
+                            Function(
+                                "equals",
+                                [
+                                    Column("metric_id"),
+                                    self.resolve_metric("transaction.duration"),
+                                ],
+                            ),
+                        ],
+                        alias,
+                    ),
+                    default_result_type="integer",
+                ),
                 fields.MetricsFunction(
                     "count_miserable",
                     required_args=[fields.MetricArg("column", allowed_columns=["user"])],
@@ -181,14 +199,22 @@ class MetricsDatasetConfig(DatasetConfig):
                         ),
                     ],
                     calculated_args=[resolve_metric_id],
-                    snql_distribution=lambda args, alias: Function(
-                        "maxIf",
-                        [
-                            Column("value"),
-                            Function("equals", [Column("metric_id"), args["metric_id"]]),
-                        ],
-                        alias,
-                    ),
+                    snql_distribution=lambda args, alias: self._resolve_percentile(args, alias, 1),
+                    default_result_type="duration",
+                ),
+                fields.MetricsFunction(
+                    "percentile",
+                    required_args=[
+                        fields.with_default(
+                            "transaction.duration",
+                            fields.MetricArg(
+                                "column", allowed_columns=constants.METRIC_DURATION_COLUMNS
+                            ),
+                        ),
+                        fields.NumberRange("percentile", 0, 1),
+                    ],
+                    calculated_args=[resolve_metric_id],
+                    snql_distribution=self._resolve_percentile,
                     default_result_type="duration",
                 ),
                 fields.MetricsFunction(
@@ -205,6 +231,24 @@ class MetricsDatasetConfig(DatasetConfig):
                     ),
                     default_result_type="integer",
                 ),
+                fields.MetricsFunction(
+                    "count",
+                    snql_distribution=lambda args, alias: Function(
+                        "countIf",
+                        [
+                            Column("value"),
+                            Function(
+                                "equals",
+                                [
+                                    Column("metric_id"),
+                                    self.resolve_metric("transaction.duration"),
+                                ],
+                            ),
+                        ],
+                        alias,
+                    ),
+                    default_result_type="integer",
+                ),
                 fields.MetricsFunction(
                     "count_web_vitals",
                     required_args=[
@@ -504,21 +548,36 @@ class MetricsDatasetConfig(DatasetConfig):
         self,
         args: Mapping[str, Union[str, Column, SelectType, int, float]],
         alias: str,
-        fixed_percentile: float,
+        fixed_percentile: Optional[float] = None,
     ) -> SelectType:
-        return Function(
-            "arrayElement",
-            [
-                Function(
-                    f"quantilesIf({fixed_percentile})",
-                    [
-                        Column("value"),
-                        Function("equals", [Column("metric_id"), args["metric_id"]]),
-                    ],
-                ),
-                1,
-            ],
-            alias,
+        if fixed_percentile is None:
+            fixed_percentile = args["percentile"]
+        if fixed_percentile not in constants.METRIC_PERCENTILES:
+            raise IncompatibleMetricsQuery("Custom quantile incompatible with metrics")
+        return (
+            Function(
+                "maxIf",
+                [
+                    Column("value"),
+                    Function("equals", [Column("metric_id"), args["metric_id"]]),
+                ],
+                alias,
+            )
+            if fixed_percentile == 1
+            else Function(
+                "arrayElement",
+                [
+                    Function(
+                        f"quantilesIf({fixed_percentile})",
+                        [
+                            Column("value"),
+                            Function("equals", [Column("metric_id"), args["metric_id"]]),
+                        ],
+                    ),
+                    1,
+                ],
+                alias,
+            )
         )
 
     def _key_transaction_filter_converter(self, search_filter: SearchFilter) -> Optional[WhereType]:

tests/sentry/search/events/test_builder.py

@@ -726,6 +726,66 @@ class MetricQueryBuilderTest(MetricBuilderBaseTest):
             ],
         )
 
+    def test_custom_percentile_throws_error(self):
+        with self.assertRaises(IncompatibleMetricsQuery):
+            MetricsQueryBuilder(
+                self.params,
+                "",
+                selected_columns=[
+                    "percentile(transaction.duration, 0.11)",
+                ],
+            )
+
+    def test_percentile_function(self):
+        self.maxDiff = None
+        query = MetricsQueryBuilder(
+            self.params,
+            "",
+            selected_columns=[
+                "percentile(transaction.duration, 0.75)",
+            ],
+        )
+        self.assertCountEqual(
+            query.where,
+            [
+                *self.default_conditions,
+                *_metric_conditions(
+                    self.organization.id,
+                    [
+                        "transaction.duration",
+                    ],
+                ),
+            ],
+        )
+        self.assertCountEqual(
+            query.distributions,
+            [
+                Function(
+                    "arrayElement",
+                    [
+                        Function(
+                            "quantilesIf(0.75)",
+                            [
+                                Column("value"),
+                                Function(
+                                    "equals",
+                                    [
+                                        Column("metric_id"),
+                                        indexer.resolve(
+                                            self.organization.id,
+                                            constants.METRICS_MAP["transaction.duration"],
+                                        ),
+                                    ],
+                                ),
+                            ],
+                        ),
+                        1,
+                    ],
+                    "percentile_transaction_duration_0_75",
+                )
+            ],
+        )
+
     def test_metric_condition_dedupe(self):
         org_id = 1
         query = MetricsQueryBuilder(
@@ -1202,6 +1262,48 @@ class MetricQueryBuilderTest(MetricBuilderBaseTest):
         assert data["tpm"] == 5 / ((self.end - self.start).total_seconds() / 60)
         assert data["tpm"] / 60 == data["tps"]
 
+    def test_count(self):
+        for _ in range(3):
+            self.store_metric(
+                150,
+                timestamp=self.start + datetime.timedelta(minutes=5),
+            )
+            self.store_metric(
+                50,
+                timestamp=self.start + datetime.timedelta(minutes=5),
+            )
+        query = MetricsQueryBuilder(
+            self.params,
+            "",
+            selected_columns=[
+                "count()",
+            ],
+        )
+        result = query.run_query("test_query")
+        data = result["data"][0]
+        assert data["count"] == 6
+
+    def test_avg(self):
+        for _ in range(3):
+            self.store_metric(
+                150,
+                timestamp=self.start + datetime.timedelta(minutes=5),
+            )
+            self.store_metric(
+                50,
+                timestamp=self.start + datetime.timedelta(minutes=5),
+            )
+        query = MetricsQueryBuilder(
+            self.params,
+            "",
+            selected_columns=[
+                "avg()",
+            ],
+        )
+        result = query.run_query("test_query")
+        data = result["data"][0]
+        assert data["avg"] == 100
+
     def test_failure_rate(self):
         for _ in range(3):
             self.store_metric(