From cc76174387a068cf9b1bba73c922ca89ef491ed3 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Mon, 17 Jan 2022 15:27:00 +0200
Subject: [PATCH] app/vmselect/promql: properly keep metric names when
 optimized path is used for aggregate function calculations

For example, `sum(rate(...) keep_metric_names) by (__name__)` didn't leave the original metric name because of this issue.

This is a follup-up for 1bdc71d917f9d586aadff90b7a5ef9fd830633bb

Udates https://github.com/VictoriaMetrics/VictoriaMetrics/issues/949
---
 app/vmselect/promql/eval.go      | 12 ++++++++++
 app/vmselect/promql/exec_test.go | 39 ++++++++++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/app/vmselect/promql/eval.go b/app/vmselect/promql/eval.go
index 48375d074..57f64005f 100644
--- a/app/vmselect/promql/eval.go
+++ b/app/vmselect/promql/eval.go
@@ -496,6 +496,9 @@ func getRollupExprArg(arg metricsql.Expr) *metricsql.RollupExpr {
 	return &reNew
 }
 
+// expr may contain:
+// - rollupFunc(m) if iafc is nil
+// - aggrFunc(rollupFunc(m)) if iafc isn't nil
 func evalRollupFunc(ec *EvalConfig, funcName string, rf rollupFunc, expr metricsql.Expr, re *metricsql.RollupExpr, iafc *incrementalAggrFuncContext) ([]*timeseries, error) {
 	if re.At == nil {
 		return evalRollupFuncWithoutAt(ec, funcName, rf, expr, re, iafc)
@@ -641,6 +644,15 @@ func evalRollupFuncWithSubquery(ec *EvalConfig, funcName string, rf rollupFunc,
 }
 
 func getKeepMetricNames(expr metricsql.Expr) bool {
+	if ae, ok := expr.(*metricsql.AggrFuncExpr); ok {
+		// Extract rollupFunc(...) from aggrFunc(rollupFunc(...)).
+		// This case is possible when optimized aggrFunc calculations are used
+		// such as `sum(rate(...))`
+		if len(ae.Args) != 1 {
+			return false
+		}
+		expr = ae.Args[0]
+	}
 	if fe, ok := expr.(*metricsql.FuncExpr); ok {
 		return fe.KeepMetricNames
 	}
diff --git a/app/vmselect/promql/exec_test.go b/app/vmselect/promql/exec_test.go
index b64def380..4c02a3710 100644
--- a/app/vmselect/promql/exec_test.go
+++ b/app/vmselect/promql/exec_test.go
@@ -1029,6 +1029,17 @@ func TestExecSuccess(t *testing.T) {
 		resultExpected := []netstorage.Result{r}
 		f(q, resultExpected)
 	})
+	t.Run("time() @ end() offset 10m", func(t *testing.T) {
+		t.Parallel()
+		q := `time() @ end() offset 10m`
+		r := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{1400, 1400, 1400, 1400, 1400, 1400},
+			Timestamps: timestampsExpected,
+		}
+		resultExpected := []netstorage.Result{r}
+		f(q, resultExpected)
+	})
 	t.Run("time() @ (end()-10m)", func(t *testing.T) {
 		t.Parallel()
 		q := `time() @ (end()-10m)`
@@ -6229,18 +6240,42 @@ func TestExecSuccess(t *testing.T) {
 	})
 	t.Run(`rate(time())`, func(t *testing.T) {
 		t.Parallel()
-		q := `rate(alias(time(), "foo"))`
+		q := `rate(label_set(alias(time(), "foo"), "x", "y"))`
 		r := netstorage.Result{
 			MetricName: metricNameExpected,
 			Values:     []float64{1, 1, 1, 1, 1, 1},
 			Timestamps: timestampsExpected,
 		}
+		r.MetricName.Tags = []storage.Tag{
+			{
+				Key:   []byte("x"),
+				Value: []byte("y"),
+			},
+		}
 		resultExpected := []netstorage.Result{r}
 		f(q, resultExpected)
 	})
 	t.Run(`rate(time()) keep_metric_names`, func(t *testing.T) {
 		t.Parallel()
-		q := `rate(alias(time(), "foo")) keep_metric_names`
+		q := `rate(label_set(alias(time(), "foo"), "x", "y")) keep_metric_names`
+		r := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{1, 1, 1, 1, 1, 1},
+			Timestamps: timestampsExpected,
+		}
+		r.MetricName.MetricGroup = []byte("foo")
+		r.MetricName.Tags = []storage.Tag{
+			{
+				Key:   []byte("x"),
+				Value: []byte("y"),
+			},
+		}
+		resultExpected := []netstorage.Result{r}
+		f(q, resultExpected)
+	})
+	t.Run(`sum(rate(time()) keep_metric_names) by (__name__)`, func(t *testing.T) {
+		t.Parallel()
+		q := `sum(rate(label_set(alias(time(), "foo"), "x", "y")) keep_metric_names) by (__name__)`
 		r := netstorage.Result{
 			MetricName: metricNameExpected,
 			Values:     []float64{1, 1, 1, 1, 1, 1},