add support for @ modifier

It is available via RollupExpr.At field.

Contrary to PromQL, MetricsQL accepts arbitrary expression as `@` modifier.
Also MetricsQL accepts `@` modifier at any place in the query.

Reference: prometheus/prometheus#10121 and https://prometheus.io/docs/prometheus/latest/querying/basics/#modifier

Updates VictoriaMetrics/VictoriaMetrics#1348

lexer.go

@@ -83,7 +83,7 @@ again:
 		}
 		lex.sTail = s[n+1:]
 		goto again
-	case '{', '}', '[', ']', '(', ')', ',':
+	case '{', '}', '[', ']', '(', ')', ',', '@':
 		token = s[:1]
 		goto tokenFoundLabel
 	}

optimizer.go

@@ -70,6 +70,11 @@ func optimizeBinaryOpArgs(be *BinaryOpExpr) *BinaryOpExpr {
 }
 
 func getMetricExprForOptimization(e Expr) *MetricExpr {
+	re, ok := e.(*RollupExpr)
+	if ok {
+		// Try optimizing the inner expression in RollupExpr.
+		return getMetricExprForOptimization(re.Expr)
+	}
 	me, ok := e.(*MetricExpr)
 	if ok {
 		// Ordinary metric expression, i.e. `foo{bar="baz"}`
@@ -95,17 +100,12 @@ func getMetricExprForOptimization(e Expr) *MetricExpr {
 		return nil
 	}
 	if IsRollupFunc(fe.Name) {
-		for _, arg := range fe.Args {
-			re, ok := arg.(*RollupExpr)
-			if !ok {
-				continue
-			}
-			if me, ok := re.Expr.(*MetricExpr); ok {
-				// rollup_func(foo{bar="baz"}[d])
-				return me
-			}
+		argIdx := GetRollupArgIdx(fe)
+		if argIdx >= len(fe.Args) {
+			return nil
 		}
-		return nil
+		arg := fe.Args[argIdx]
+		return getMetricExprForOptimization(arg)
 	}
 	if IsTransformFunc(fe.Name) {
 		switch strings.ToLower(fe.Name) {

optimizer_test.go

@@ -51,10 +51,17 @@ func TestOptimize(t *testing.T) {
 	f(`ABSENT(foo{bar="baz"}) + sqrt(a{z=~"c"})`, `ABSENT(foo{bar="baz"}) + sqrt(a{z=~"c"})`)
 
 	// rollup funcs
-	f(`RATE(foo[5m]) / rate(baz{a="b"}) + increase(x{y="z"} offset 5i)`, `(RATE(foo[5m]) / rate(baz{a="b"})) + increase(x{y="z"} offset 5i)`)
+	f(`RATE(foo[5m]) / rate(baz{a="b"}) + increase(x{y="z"} offset 5i)`, `(RATE(foo{a="b"}[5m]) / rate(baz{a="b"})) + increase(x{y="z"} offset 5i)`)
+	f(`sum(rate(foo[5m])) / rate(baz{a="b"})`, `sum(rate(foo[5m])) / rate(baz{a="b"})`)
+	f(`rate({__name__="foo"}) + rate({__name__="bar",x="y"}) - rate({__name__=~"baz"})`, `(rate(foo{x="y"}) + rate(bar{x="y"})) - rate({__name__=~"baz"})`)
+
+	// @ modifier
+	f(`foo @ end() + bar{baz="a"}`, `foo{baz="a"} @ end() + bar{baz="a"}`)
+	f(`sum(foo @ end()) + bar{baz="a"}`, `sum(foo @ end()) + bar{baz="a"}`)
 
 	// subqueries
-	f(`rate(avg_over_time(foo[5m:])) + bar{baz="a"}`, `rate(avg_over_time(foo[5m:])) + bar{baz="a"}`)
+	f(`rate(avg_over_time(foo[5m:])) + bar{baz="a"}`, `rate(avg_over_time(foo{baz="a"}[5m:])) + bar{baz="a"}`)
+	f(`rate(sum(foo[5m:])) + bar{baz="a"}`, `rate(sum(foo[5m:])) + bar{baz="a"}`)
 
 	// binary ops with constants or scalars
 	f(`100 * foo / bar{baz="a"}`, `(100 * foo{baz="a"}) / bar{baz="a"}`)

parser.go

@@ -104,6 +104,9 @@ func mustParseWithArgExpr(s string) *withArgExpr {
 func removeParensExpr(e Expr) Expr {
 	if re, ok := e.(*RollupExpr); ok {
 		re.Expr = removeParensExpr(re.Expr)
+		if re.At != nil {
+			re.At = removeParensExpr(re.At)
+		}
 		return re
 	}
 	if be, ok := e.(*BinaryOpExpr); ok {
@@ -415,13 +418,17 @@ func (p *parser) parseSingleExpr() (Expr, error) {
 	if err != nil {
 		return nil, err
 	}
-	if p.lex.Token != "[" && !isOffset(p.lex.Token) {
+	if !isRollupStartToken(p.lex.Token) {
 		// There is no rollup expression.
 		return e, nil
 	}
 	return p.parseRollupExpr(e)
 }
 
+func isRollupStartToken(token string) bool {
+	return token == "[" || token == "@" || isOffset(token)
+}
+
 func (p *parser) parseSingleExprWithoutRollupSuffix() (Expr, error) {
 	if isPositiveDuration(p.lex.Token) {
 		return p.parsePositiveDuration()
@@ -1268,6 +1275,20 @@ func (p *parser) parseWindowAndStep() (*DurationExpr, *DurationExpr, bool, error
 	return window, step, inheritStep, nil
 }
 
+func (p *parser) parseAtExpr() (Expr, error) {
+	if p.lex.Token != "@" {
+		return nil, fmt.Errorf(`unexpected token %q; want "@"`, p.lex.Token)
+	}
+	if err := p.lex.Next(); err != nil {
+		return nil, err
+	}
+	e, err := p.parseSingleExprWithoutRollupSuffix()
+	if err != nil {
+		return nil, fmt.Errorf("cannot parse `@` expresion: %w", err)
+	}
+	return e, nil
+}
+
 func (p *parser) parseOffset() (*DurationExpr, error) {
 	if !isOffset(p.lex.Token) {
 		return nil, fmt.Errorf(`offset: unexpected token %q; want "offset"`, p.lex.Token)
@@ -1374,11 +1395,11 @@ func (p *parser) parseIdentExpr() (Expr, error) {
 			return p.parseAggrFuncExpr()
 		}
 		return p.parseFuncExpr()
-	case "{", "[", ")", ",":
+	case "{", "[", ")", ",", "@":
 		p.lex.Prev()
 		return p.parseMetricExpr()
 	default:
-		return nil, fmt.Errorf(`identExpr: unexpected token %q; want "(", "{", "[", ")", ","`, p.lex.Token)
+		return nil, fmt.Errorf(`identExpr: unexpected token %q; want "(", "{", "[", ")", "," or "@"`, p.lex.Token)
 	}
 }
 
@@ -1417,15 +1438,34 @@ func (p *parser) parseRollupExpr(arg Expr) (Expr, error) {
 		re.Window = window
 		re.Step = step
 		re.InheritStep = inheritStep
-		if !isOffset(p.lex.Token) {
+		if !isOffset(p.lex.Token) && p.lex.Token != "@" {
 			return &re, nil
 		}
 	}
-	offset, err := p.parseOffset()
-	if err != nil {
-		return nil, err
+	if p.lex.Token == "@" {
+		at, err := p.parseAtExpr()
+		if err != nil {
+			return nil, err
+		}
+		re.At = at
+	}
+	if isOffset(p.lex.Token) {
+		offset, err := p.parseOffset()
+		if err != nil {
+			return nil, err
+		}
+		re.Offset = offset
+	}
+	if p.lex.Token == "@" {
+		if re.At != nil {
+			return nil, fmt.Errorf("duplicate `@` token")
+		}
+		at, err := p.parseAtExpr()
+		if err != nil {
+			return nil, err
+		}
+		re.At = at
 	}
-	re.Offset = offset
 	return &re, nil
 }
 
@@ -1677,6 +1717,12 @@ type RollupExpr struct {
 	// If set to true, then `foo[1h:]` would print the same
 	// instead of `foo[1h]`.
 	InheritStep bool
+
+	// At contains an optional expression after `@` modifier.
+	//
+	// For example, `foo @ end()` or `bar[5m] @ 12345`
+	// See https://prometheus.io/docs/prometheus/latest/querying/basics/#modifier
+	At Expr
 }
 
 // ForSubquery returns true if re represents subquery.
@@ -1720,6 +1766,17 @@ func (re *RollupExpr) AppendString(dst []byte) []byte {
 		dst = append(dst, " offset "...)
 		dst = re.Offset.AppendString(dst)
 	}
+	if re.At != nil {
+		dst = append(dst, " @ "...)
+		_, needAtParens := re.At.(*BinaryOpExpr)
+		if needAtParens {
+			dst = append(dst, '(')
+		}
+		dst = re.At.AppendString(dst)
+		if needAtParens {
+			dst = append(dst, ')')
+		}
+	}
 	return dst
 }
 

parser_test.go

@@ -71,6 +71,20 @@ func TestParseSuccess(t *testing.T) {
 	same(`metric{foo="bar", b="sdfsdf"}[2.34:5.6] offset 3600.5`)
 	same(`metric{foo="bar", b="sdfsdf"}[234:56] offset -3600`)
 	another(`  metric  {  foo  = "bar"  }  [  2d ]   offset   10h  `, `metric{foo="bar"}[2d] offset 10h`)
+	// @ modifier
+	// See https://prometheus.io/docs/prometheus/latest/querying/basics/#modifier
+	same(`foo @ 123.45`)
+	same(`foo\@ @ 123.45`)
+	same(`{foo=~"bar"} @ end()`)
+	same(`foo{bar="baz"} @ start()`)
+	same(`foo{bar="baz"}[5m] @ 12345`)
+	same(`foo{bar="baz"}[5m:4s] offset 5m @ (end() - 3.5m)`)
+	another(`foo{bar="baz"}[5m:4s] @ (end() - 3.5m) offset 2.4h`, `foo{bar="baz"}[5m:4s] offset 2.4h @ (end() - 3.5m)`)
+	another(`foo @ start() + (bar offset 3m @ end()) / baz OFFSET -5m`, `foo @ start() + (bar offset 3m @ end() / baz offset -5m)`)
+	same(`sum(foo) @ start() + rate(bar @ (end() - 5m))`)
+	another(`time() @ (start())`, `time() @ start()`)
+	another(`time() @ (start()+(2))`, `time() @ (start() + 2)`)
+	same(`time() @ (end() - 10m)`)
 	// metric name matching keywords
 	same("rate")
 	same("RATE")
@@ -513,6 +527,17 @@ func TestParseError(t *testing.T) {
 	f(`m{x=y/5}`)
 	f(`m{x=y+5}`)
 
+	// Invalid @ modifier
+	f(`@`)
+	f(`foo @`)
+	f(`foo @ ! `)
+	f(`foo @ @`)
+	f(`foo @ offset 5m`)
+	f(`foo @ [5m]`)
+	f(`foo offset @ 5m`)
+	f(`foo @ 123 offset 5m @ 456`)
+	f(`foo offset 5m @`)
+
 	// Invalid regexp
 	f(`foo{bar=~"x["}`)
 	f(`foo{bar=~"x("}`)

rollup.go

@@ -84,3 +84,23 @@ func IsRollupFunc(funcName string) bool {
 	s := strings.ToLower(funcName)
 	return rollupFuncs[s]
 }
+
+// GetRollupArgIdx returns the argument index for the given fe, which accepts the rollup argument.
+//
+// -1 is returned if fe isn't a rollup function.
+func GetRollupArgIdx(fe *FuncExpr) int {
+	funcName := fe.Name
+	funcName = strings.ToLower(funcName)
+	if !rollupFuncs[funcName] {
+		return -1
+	}
+	switch funcName {
+	case "quantile_over_time", "aggr_over_time",
+		"hoeffding_bound_lower", "hoeffding_bound_upper":
+		return 1
+	case "quantiles_over_time":
+		return len(fe.Args) - 1
+	default:
+		return 0
+	}
+}