chore: add multi-variant query support to LogQL (#16196)

pkg/logql/engine.go

@@ -323,6 +323,8 @@ func (q *query) Eval(ctx context.Context) (promql_parser.Value, error) {
 		defer util.LogErrorWithContext(ctx, "closing iterator", itr.Close)
 		streams, err := readStreams(itr, q.params.Limit(), q.params.Direction(), q.params.Interval())
 		return streams, err
+	case syntax.VariantsExpr:
+		return nil, logqlmodel.ErrVariantsDisabled
 	default:
 		return nil, fmt.Errorf("unexpected type (%T): cannot evaluate", e)
 	}

pkg/logql/syntax/ast.go

@@ -66,6 +66,7 @@ func (LogfmtExpressionParserExpr) isExpr() {}
 func (LogRangeExpr) isExpr()               {}
 func (OffsetExpr) isExpr()                 {}
 func (UnwrapExpr) isExpr()                 {}
+func (MultiVariantExpr) isExpr()           {}
 
 // LogSelectorExpr is a expression filtering and returning logs.
 type LogSelectorExpr interface {
@@ -1318,6 +1319,10 @@ const (
 
 	// probabilistic aggregations
 	OpTypeApproxTopK = "approx_topk"
+
+	// variants
+	OpVariants = "variants"
+	VariantsOf = "of"
 )
 
 func IsComparisonOperator(op string) bool {
@@ -2412,3 +2417,169 @@ func groupingReducesLabels(grp *Grouping) bool {
 
 	return false
 }
+
+// VariantsExpr is a LogQL expression that can produce multiple streams, defined by a set of variants,
+// over a single log selector.
+//
+//sumtype:decl
+type VariantsExpr interface {
+	LogRange() *LogRangeExpr
+	Matchers() []*labels.Matcher
+	Variants() []SampleExpr
+	SetVariant(i int, e SampleExpr) error
+	Interval() time.Duration
+	Offset() time.Duration
+	Extractors() ([]SampleExtractor, error)
+	Expr
+}
+
+type MultiVariantExpr struct {
+	logRange *LogRangeExpr
+	variants []SampleExpr
+	err      error
+}
+
+func NewMultiVariantExpr(
+	logRange *LogRangeExpr,
+	variants []SampleExpr,
+) MultiVariantExpr {
+	return MultiVariantExpr{
+		logRange: logRange,
+		variants: variants,
+	}
+}
+
+func (m *MultiVariantExpr) LogRange() *LogRangeExpr {
+	return m.logRange
+}
+
+func (m *MultiVariantExpr) SetLogSelector(e *LogRangeExpr) {
+	m.logRange = e
+}
+
+func (m *MultiVariantExpr) Matchers() []*labels.Matcher {
+	return m.logRange.Left.Matchers()
+}
+
+func (m *MultiVariantExpr) Interval() time.Duration {
+	return m.logRange.Interval
+}
+
+func (m *MultiVariantExpr) Offset() time.Duration {
+	return m.logRange.Offset
+}
+
+func (m *MultiVariantExpr) Variants() []SampleExpr {
+	return m.variants
+}
+
+func (m *MultiVariantExpr) AddVariant(v SampleExpr) {
+	m.variants = append(m.variants, v)
+}
+
+func (m *MultiVariantExpr) SetVariant(i int, v SampleExpr) error {
+	if i >= len(m.variants) {
+		return fmt.Errorf("variant index out of range")
+	}
+
+	m.variants[i] = v
+	return nil
+}
+
+func (m *MultiVariantExpr) Shardable(topLevel bool) bool {
+	if !m.logRange.Shardable(topLevel) {
+		return false
+	}
+
+	for _, v := range m.variants {
+		if !v.Shardable(topLevel) {
+			return false
+		}
+	}
+
+	return true
+}
+
+func (m *MultiVariantExpr) Walk(f WalkFn) {
+	f(m)
+
+	if m.logRange != nil {
+		m.logRange.Walk(f)
+	}
+
+	for _, v := range m.variants {
+		v.Walk(f)
+	}
+}
+
+func (m *MultiVariantExpr) String() string {
+	var sb strings.Builder
+	sb.WriteString(OpVariants)
+	sb.WriteString("(")
+	for i, v := range m.variants {
+		sb.WriteString(v.String())
+		if i+1 != len(m.variants) {
+			sb.WriteString(", ")
+		}
+	}
+	sb.WriteString(") ")
+
+	sb.WriteString(VariantsOf)
+	sb.WriteString(" (")
+	sb.WriteString(m.logRange.String())
+	sb.WriteString(")")
+
+	return sb.String()
+}
+
+func (m *MultiVariantExpr) Accept(v RootVisitor) {
+	v.VisitVariants(m)
+}
+
+// Pretty prettyfies any LogQL expression at given `level` of the whole LogQL query.
+func (m *MultiVariantExpr) Pretty(level int) string {
+	s := Indent(level)
+
+	s += OpVariants + "(\n"
+
+	variants := make([]string, 0, len(m.variants))
+	for _, v := range m.variants {
+		variants = append(variants, v.Pretty(level+1))
+	}
+
+	for i, v := range variants {
+		s += v
+		// LogQL doesn't allow `,` at the end of last argument.
+		if i < len(variants)-1 {
+			s += ","
+		}
+		s += "\n"
+	}
+
+	s += Indent(level) + ") of (\n"
+	s += m.logRange.Pretty(level + 1)
+	s += Indent(level) + "\n)"
+
+	return s
+}
+
+func (m *MultiVariantExpr) Extractors() ([]log.SampleExtractor, error) {
+	extractors := make([]log.SampleExtractor, 0, len(m.variants))
+	for _, v := range m.variants {
+		e, err := v.Extractor()
+		if err != nil {
+			return nil, err
+		}
+
+		extractors = append(extractors, e)
+	}
+
+	return extractors, nil
+}
+
+func newVariantsExpr(variants []SampleExpr, logRange *LogRangeExpr) VariantsExpr {
+	return &MultiVariantExpr{
+		variants: variants,
+		logRange: logRange,
+	}
+}

pkg/logql/syntax/ast_test.go

@@ -2,6 +2,7 @@ package syntax
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 	"time"
 
@@ -11,6 +12,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/loki/v3/pkg/logql/log"
+	"github.com/grafana/loki/v3/pkg/logqlmodel"
 )
 
 var labelBar, _ = ParseLabels("{app=\"bar\"}")
@@ -1157,3 +1159,126 @@ func TestCombineFilters(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)
+
+			AssertExpressions(t, expr, expr2)
+		})
+	}
+}
+
+func Test_VariantsExpr_Pretty(t *testing.T) {
+	tests := []struct {
+		expr   string
+		pretty string
+	}{
+		{`variants(count_over_time({foo="bar"}[5m])) of ({foo="bar"}[5m])`, `
+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]),
+  bytes_over_time({baz="qux", foo=~"bar"}[5m])
+) of (
+  {baz="qux", foo=~"bar"} | logfmt | this="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)
+
+			require.Equal(t, strings.TrimSpace(tt.pretty), strings.TrimSpace(expr.Pretty(0)))
+		})
+	}
+}
+
+func Test_MultiVariantExpr_Extractors(t *testing.T) {
+	emptyExpr := &MultiVariantExpr{
+		variants: []SampleExpr{},
+		logRange: &LogRangeExpr{},
+	}
+	validAgg := &RangeAggregationExpr{
+		Operation: OpRangeTypeCount,
+		Left: &LogRangeExpr{
+			Interval: time.Second,
+			Left: &MatchersExpr{
+				Mts: []*labels.Matcher{
+					mustNewMatcher(labels.MatchEqual, "foo", "bar"),
+				},
+			},
+		},
+	}
+	errorAgg := &RangeAggregationExpr{
+		err: logqlmodel.NewParseError("test error", 0, 0),
+	}
+	t.Run("should return empty slice for no variants", func(t *testing.T) {
+		extractors, err := emptyExpr.Extractors()
+		require.NoError(t, err)
+		require.Empty(t, extractors)
+	})
+
+	t.Run("should return extractors for all variants", func(t *testing.T) {
+		expr := &MultiVariantExpr{
+			variants: []SampleExpr{validAgg, validAgg}, // Two identical variants for simplicity
+			logRange: &LogRangeExpr{},
+		}
+
+		extractors, err := expr.Extractors()
+		require.NoError(t, err)
+		require.Len(t, extractors, 2)
+	})
+
+	t.Run("should propagate extractor errors", func(t *testing.T) {
+		expr := &MultiVariantExpr{
+			variants: []SampleExpr{errorAgg},
+			logRange: &LogRangeExpr{},
+		}
+
+		extractors, err := expr.Extractors()
+		require.Error(t, err)
+		require.Nil(t, extractors)
+	})
+
+	t.Run("should handle mixed valid and invalid variants", func(t *testing.T) {
+		expr := &MultiVariantExpr{
+			variants: []SampleExpr{validAgg, errorAgg},
+			logRange: &LogRangeExpr{},
+		}
+
+		extractors, err := expr.Extractors()
+		require.Error(t, err)
+		require.Nil(t, extractors)
+	})
+}

pkg/logql/syntax/clone.go

@@ -300,3 +300,16 @@ func (v *cloneVisitor) VisitLogfmtParser(e *LogfmtParserExpr) {
 		KeepEmpty: e.KeepEmpty,
 	}
 }
+
+func (v *cloneVisitor) VisitVariants(e *MultiVariantExpr) {
+	copied := &MultiVariantExpr{
+		logRange: MustClone[*LogRangeExpr](e.logRange),
+		variants: make([]SampleExpr, len(e.variants)),
+	}
+
+	for i, v := range e.variants {
+		copied.variants[i] = MustClone[SampleExpr](v)
+	}
+
+	v.cloned = copied
+}

pkg/logql/syntax/clone_test.go

@@ -64,6 +64,15 @@ func TestClone(t *testing.T) {
 		"true filter": {
 			query: `{ foo = "bar" } | foo =~".*"`,
 		},
+		"multiple variants": {
+			query: `variants(bytes_over_time({foo="bar"}[5m]), count_over_time({foo="bar"}[5m])) of ({foo="bar"}[5m])`,
+		},
+		"multiple variants with aggregation": {
+			query: `variants(sum by (app) (bytes_over_time({foo="bar"}[5m])), count_over_time({foo="bar"}[5m])) of ({foo="bar"}[5m])`,
+		},
+		"multiple variants with filters": {
+			query: `variants(bytes_over_time({foo="bar"}[5m]), count_over_time({foo="bar"}[5m])) of ({foo="bar"} | logfmt[5m])`,
+		},
 	}
 
 	for name, test := range tests {

pkg/logql/syntax/lex.go

@@ -80,6 +80,10 @@ var tokens = map[string]int{
 
 	// keep labels
 	OpKeep: KEEP,
+
+	// variants
+	OpVariants: VARIANTS,
+	VariantsOf: OF,
 }
 
 var parserFlags = map[string]struct{}{