Add approx_topk and multi-variant expressions to LogQL parser

logql/v2/ast.go

@@ -1039,3 +1039,91 @@ func (v *VectorExpr) String() string {
 }
 
 func (v *VectorExpr) Walk(fn WalkFn) { fn(v) }
+
+type LogMultiVariantExpr struct {
+	defaultLogQLExpr //nolint:unused
+	logRange         *LogRangeQueryExpr
+	variants         []LogMetricSampleExpr
+	Expr
+}
+
+func newLogMultiVariantExpr(variants []LogMetricSampleExpr, logRange LogSelectorExpr) *LogMultiVariantExpr {
+	rangeExpr, ok := logRange.(*LogRangeQueryExpr)
+	if !ok {
+		rangeExpr = nil
+	}
+
+	return &LogMultiVariantExpr{
+		logRange: rangeExpr,
+		variants: variants,
+	}
+}
+
+func (e *LogMultiVariantExpr) LogRange() *LogRangeQueryExpr {
+	return e.logRange
+}
+
+func (e *LogMultiVariantExpr) Matchers() []*labels.Matcher {
+	if e.logRange != nil {
+		return e.logRange.Matchers()
+	}
+	return nil
+}
+
+func (e *LogMultiVariantExpr) Interval() time.Duration {
+	if e.logRange == nil {
+		return 0
+	}
+
+	duration, err := parseDuration(e.logRange.rng)
+	if err != nil {
+		return 0
+	}
+	return duration
+}
+
+func (e *LogMultiVariantExpr) Variants() []LogMetricSampleExpr {
+	return e.variants
+}
+
+func (e *LogMultiVariantExpr) AddVariant(v LogMetricSampleExpr) {
+	e.variants = append(e.variants, v)
+}
+
+func (e *LogMultiVariantExpr) SetVariant(i int, v LogMetricSampleExpr) error {
+	if i >= len(e.variants) {
+		return fmt.Errorf("variant index out of range")
+	}
+
+	e.variants[i] = v
+	return nil
+}
+
+func (e *LogMultiVariantExpr) Walk(fn WalkFn) {
+	fn(e)
+	if e.logRange != nil {
+		e.logRange.Walk(fn)
+	}
+
+	for _, v := range e.variants {
+		v.Walk(fn)
+	}
+}
+
+func (e *LogMultiVariantExpr) String() string {
+	var sb strings.Builder
+	sb.WriteString("variants(")
+	for i, v := range e.variants {
+		sb.WriteString(v.String())
+		if i+1 != len(e.variants) {
+			sb.WriteString(", ")
+		}
+	}
+	sb.WriteString(") of (")
+	sb.WriteString(e.logRange.String())
+	sb.WriteString(")")
+
+	return sb.String()
+}
+
+func (e *LogMultiVariantExpr) logQLExpr() {}

logql/v2/ast_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/prometheus/prometheus/model/labels"
+	"github.com/stretchr/testify/require"
 )
 
 func Test_AstWalker_SimpleCountExpr(t *testing.T) {
@@ -20,6 +21,12 @@ func Test_AstWalker_SimpleCountExpr(t *testing.T) {
 		}, {
 			input: `{first="value"} |= "baz" |= ip("8.8.8.8")`,
 			total: 2,
+		}, {
+			input: `variants(count_over_time({foo="bar"}[5m])) of ({foo="bar"}[5m])`,
+			total: 8,
+		}, {
+			input: `variants(count_over_time({foo="bar"}[5m]), rate({foo="bar"}[5m])) of ({foo="bar"}[5m])`,
+			total: 12,
 		},
 	}
 	for _, tc := range tc {
@@ -124,10 +131,24 @@ func Test_AstWalker_AppendMatcher(t *testing.T) {
 			input:  `(count_over_time({first="value"}[10h]))`,
 			output: `(count_over_time({first="value", second="next"}[10h]))`,
 		},
+		// approx_topkk tests
 		{
 			input:  `topk(25,(count_over_time({first="value"}[10h])))`,
 			output: `topk(25,(count_over_time({first="value", second="next"}[10h])))`,
 		},
+		{
+			input:  `approx_topk(25,(count_over_time({first="value"}[10h])))`,
+			output: `approx_topk(25,(count_over_time({first="value", second="next"}[10h])))`,
+		},
+		// variant tests
+		{
+			input:  `variants(count_over_time({job="foo"}[5m])) of ({job="foo"}[5m])`,
+			output: `variants(count_over_time({job="foo", second="next"}[5m])) of ({job="foo", second="next"}[5m])`,
+		},
+		{
+			input:  `variants(count_over_time({job="foo"}[5m]), bytes_over_time({job="foo"}[5m])) of ({job="foo"}[5m])`,
+			output: `variants(count_over_time({job="foo", second="next"}[5m]), bytes_over_time({job="foo", second="next"}[5m])) of ({job="foo", second="next"}[5m])`,
+		},
 	}
 	for _, tc := range tc {
 		expr, err := ParseExpr(tc.input)
@@ -220,3 +241,36 @@ func Test_AstWalker_AppendORMatcher(t *testing.T) {
 		}
 	}
 }
+
+func Test_VariantsExpr_String(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		expr string
+	}{
+		{
+			`variants(count_over_time({foo="bar"}[5m])) of ({foo="bar"}[5m])`,
+		},
+		{
+			`variants(count_over_time({baz="qux", foo=~"bar"}[5m]), bytes_over_time({baz="qux", foo=~"bar"}[5m])) of ({baz="qux", foo=~"bar"} | logfmt | this = "that"[5m])`,
+		},
+		{
+			`variants(count_over_time({baz="qux", foo!="bar"}[5m]),rate({baz="qux", foo!="bar"}[5m])) of ({baz="qux", foo!="bar"} |= "that" [5m])`,
+		},
+		{
+			`variants(sum by (app) (count_over_time({baz="qux", foo!="bar"}[5m])),rate({baz="qux", foo!="bar"}[5m])) of ({baz="qux", foo!="bar"} |= "that" [5m])`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.expr, func(t *testing.T) {
+			t.Parallel()
+			expr, err := ParseExpr(tt.expr)
+			require.NoError(t, err)
+
+			expr2, err := ParseExpr(expr.String())
+			require.Nil(t, err)
+
+			require.Equal(t, expr.String(), expr2.String())
+		})
+	}
+}

logql/v2/expr.y

@@ -23,6 +23,8 @@ import (
   LogStageLabels       LogLabelList
   LogParserLabels      LogLabelList
   LogFMTFlags          []string
+  LogMultiVariantExpr  *LogMultiVariantExpr
+  LogMetricExprs       []LogMetricSampleExpr
   Matcher              *labels.Matcher
   Matchers             []*labels.Matcher
   MetricOp             string
@@ -60,6 +62,8 @@ import (
 %type <LogStageLabels>       logStageLabels
 %type <LogParserLabels>      logParserLabels
 %type <LogFMTFlags>          logFMTFlags
+%type <LogMultiVariantExpr>  logMultiVariantExpr
+%type <LogMetricExprs>       logMetricExprs
 %type <Matcher>              matcher
 %type <Matchers>             matchers
 %type <MetricOp>             metricOp
@@ -70,11 +74,11 @@ import (
 %token  <str>      IDENTIFIER STRING RANGE NUMBER LOGFMTSTRICT LOGFMTKEEPEMPTY
 %token  <duration> DURATION
 %token  <val>      MATCHERS LABELS EQ RE NRE OPEN_BRACE CLOSE_BRACE OPEN_BRACKET CLOSE_BRACKET COMMA DOT
-                   OPEN_PARENTHESIS CLOSE_PARENTHESIS COUNT_OVER_TIME RATE RATE_COUNTER SUM AVG MAX MIN COUNT STDDEV STDVAR BOTTOMK TOPK SORT SORT_DESC
+                   OPEN_PARENTHESIS CLOSE_PARENTHESIS COUNT_OVER_TIME RATE RATE_COUNTER SUM AVG MAX MIN COUNT STDDEV STDVAR BOTTOMK TOPK APPROX_TOPK SORT SORT_DESC
                    BYTES_OVER_TIME BYTES_RATE BOOL JSON REGEXP LOGFMT PIPE_MATCH PIPE_EXACT PIPE_MATCH_PATTERN PIPE_NOT_MATCH_PATTERN PIPE LINE_FMT LABEL_FMT UNWRAP AVG_OVER_TIME SUM_OVER_TIME MIN_OVER_TIME
                    MAX_OVER_TIME STDVAR_OVER_TIME STDDEV_OVER_TIME QUANTILE_OVER_TIME FIRST_OVER_TIME LAST_OVER_TIME ABSENT_OVER_TIME
                    BY WITHOUT VECTOR LABEL_REPLACE IP UNPACK PATTERN OFFSET BYTES_CONV DURATION_CONV DURATION_SECONDS_CONV ON IGNORING GROUP_LEFT GROUP_RIGHT
-                   DECOLORIZE DROP KEEP
+                   DECOLORIZE DROP KEEP VARIANTS OF
 
 %left <binaryOp> OR
 %left <binaryOp> AND UNLESS
@@ -91,6 +95,7 @@ expr:
         |       logMetricExpr   { $$ = $1 }
         |       logBinaryOpExpr { $$ = $1 }
         |       logNumberExpr   { $$ = $1 }
+        |       logMultiVariantExpr { $$ = $1 }
         |       OPEN_PARENTHESIS expr CLOSE_PARENTHESIS { $$ = newParenthesisExpr($2) }
         ;
 
@@ -241,6 +246,15 @@ logFMTFlags:
         |       logFMTFlags logFMTFlags { $$ = mergeParserFlags($1, $2) }
         ;
 
+logMultiVariantExpr:
+                VARIANTS OPEN_PARENTHESIS logMetricExprs CLOSE_PARENTHESIS OF OPEN_PARENTHESIS logRangeQueryExpr CLOSE_PARENTHESIS { $$ = newLogMultiVariantExpr($3, $7) }
+        ;
+
+logMetricExprs:
+                logMetricExpr                      { $$ = []LogMetricSampleExpr{$1} }
+        |       logMetricExprs COMMA logMetricExpr { $$ = append($1, $3) }
+        ;
+
 binaryOpOptions:
                 { $$ = BinaryOpOptions{} }
         |       BOOL { $$ = BinaryOpOptions{ ReturnBool: true } }
@@ -389,6 +403,7 @@ metricOp:
         |       STDVAR             { $$ = VectorOpTypeStdvar     }
         |       BOTTOMK            { $$ = VectorOpTypeBottomK    }
         |       TOPK               { $$ = VectorOpTypeTopK       }
+        |       APPROX_TOPK        { $$ = VectorOpTypeApproxTopK }
         |       SORT               { $$ = VectorOpTypeSort       }
         |       SORT_DESC          { $$ = VectorOpTypeSortDesc   }
         |       VECTOR             { $$ = OpTypeVector           }