feat(starfish): Add avg to starfish (#53456)

- This adds the average function to both metrics and indexed datasets
getsentry · Jul 24, 2023 · f0df914 · f0df914
1 parent 473d9bf
commit f0df914
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 0 deletions.
diff --git a/src/sentry/search/events/datasets/spans_indexed.py b/src/sentry/search/events/datasets/spans_indexed.py
@@ -73,6 +73,16 @@ def function_converter(self) -> Mapping[str, SnQLFunction]:
                     result_type_fn=self.reflective_result_type(),
                     default_result_type="duration",
                 ),
+                SnQLFunction(
+                    "avg",
+                    optional_args=[
+                        with_default("span.duration", NumericColumn("column", spans=True)),
+                    ],
+                    snql_aggregate=lambda args, alias: Function("avg", [args["column"]], alias),
+                    result_type_fn=self.reflective_result_type(),
+                    default_result_type="duration",
+                    redundant_grouping=True,
+                ),
                 SnQLFunction(
                     "percentile",
                     required_args=[

diff --git a/src/sentry/search/events/datasets/spans_metrics.py b/src/sentry/search/events/datasets/spans_metrics.py
@@ -133,6 +133,28 @@ def function_converter(self) -> Mapping[str, fields.MetricsFunction]:
                     ),
                     default_result_type="duration",
                 ),
+                fields.MetricsFunction(
+                    "avg",
+                    optional_args=[
+                        fields.with_default(
+                            "span.self_time",
+                            fields.MetricArg(
+                                "column", allowed_columns=constants.SPAN_METRIC_DURATION_COLUMNS
+                            ),
+                        ),
+                    ],
+                    calculated_args=[resolve_metric_id],
+                    snql_percentile=lambda args, alias: Function(
+                        "avgIf",
+                        [
+                            Column("value"),
+                            Function("equals", [Column("metric_id"), args["metric_id"]]),
+                        ],
+                        alias,
+                    ),
+                    result_type_fn=self.reflective_result_type(),
+                    default_result_type="duration",
+                ),
                 fields.MetricsFunction(
                     "percentile",
                     required_args=[

diff --git a/tests/snuba/api/endpoints/test_organization_events_span_metrics.py b/tests/snuba/api/endpoints/test_organization_events_span_metrics.py
@@ -172,6 +172,27 @@ def test_p50(self):
         assert data[0]["p50()"] == 1
         assert meta["dataset"] == "spansMetrics"
 
+    def test_avg(self):
+        self.store_span_metric(
+            1,
+            internal_metric=constants.SELF_TIME_LIGHT,
+            timestamp=self.min_ago,
+        )
+        response = self.do_request(
+            {
+                "field": ["avg()"],
+                "query": "",
+                "project": self.project.id,
+                "dataset": "spansMetrics",
+            }
+        )
+        assert response.status_code == 200, response.content
+        data = response.data["data"]
+        meta = response.data["meta"]
+        assert len(data) == 1
+        assert data[0]["avg()"] == 1
+        assert meta["dataset"] == "spansMetrics"
+
     def test_eps(self):
         for _ in range(6):
             self.store_span_metric(