promql: benchmark for join queries with more labels (#17130)

* promql: benchmark for join queries with more labels

Signed-off-by: Linas Medziunas <linas.medziunas@gmail.com>

* Add benchmark case for GROUP_LEFT

Signed-off-by: Linas Medziunas <linas.medziunas@gmail.com>

* Address PR feedback

Signed-off-by: Linas Medziunas <linas.medziunas@gmail.com>

---------

Signed-off-by: Linas Medziunas <linas.medziunas@gmail.com>

promql/bench_test.go

@@ -16,15 +16,18 @@ package promql_test
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"strconv"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/model/timestamp"
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/promql/promqltest"
@@ -94,6 +97,57 @@ func setupRangeQueryTestData(stor *teststorage.TestStorage, _ *promql.Engine, in
 	return nil
 }
 
+// setupJoinQueryTestData creates numInstances series for two metrics ("rpc_request_success_total" and
+// "rpc_request_error_total") that can be joined on "instance" label for benchmarking join performance.
+func setupJoinQueryTestData(stor *teststorage.TestStorage, _ *promql.Engine, interval, numIntervals, numInstances int) error {
+	ctx := context.Background()
+
+	commonLabels := labels.FromStrings(
+		"environment", "staging",
+		"cluster", "test-kubernetes-cluster",
+		"namespace", "test-kubernetes-namespace",
+		"job", "worker",
+		"rpc_method", "fetch-my-data-from-this-service",
+		"domain", "test-domain",
+	)
+	builder := labels.NewBuilder(commonLabels)
+
+	rnd := rand.New(rand.NewSource(0)) // Fixed seed for deterministic results.
+
+	var metrics []labels.Labels
+	for range numInstances {
+		instance, err := uuid.NewRandomFromReader(rnd)
+		if err != nil {
+			return err
+		}
+		builder.Set("instance", instance.String())
+
+		builder.Set("__name__", "rpc_request_success_total")
+		metrics = append(metrics, builder.Labels())
+
+		builder.Set("__name__", "rpc_request_error_total")
+		metrics = append(metrics, builder.Labels())
+	}
+
+	refs := make([]storage.SeriesRef, len(metrics))
+
+	for s := range numIntervals {
+		a := stor.Appender(context.Background())
+		ts := int64(s * interval)
+		for i, metric := range metrics {
+			ref, _ := a.Append(refs[i], metric, ts, float64(s)+float64(i)/float64(len(metrics)))
+			refs[i] = ref
+		}
+		if err := a.Commit(); err != nil {
+			return err
+		}
+	}
+
+	stor.ForceHeadMMap() // Ensure we have at most one head chunk for every series.
+	stor.Compact(ctx)
+	return nil
+}
+
 type benchCase struct {
 	expr  string
 	steps int
@@ -326,6 +380,69 @@ func BenchmarkRangeQuery(b *testing.B) {
 	}
 }
 
+func BenchmarkJoinQuery(b *testing.B) {
+	stor := teststorage.New(b)
+	stor.DisableCompactions() // Don't want auto-compaction disrupting timings.
+	defer stor.Close()
+
+	opts := promql.EngineOpts{
+		Logger:     nil,
+		Reg:        nil,
+		MaxSamples: 50000000,
+		Timeout:    100 * time.Second,
+	}
+	engine := promqltest.NewTestEngineWithOpts(b, opts)
+
+	const interval = 10000 // 10s interval.
+
+	// A day of data plus 10k steps.
+	numIntervals := 8640 + 10000
+
+	require.NoError(b, setupJoinQueryTestData(stor, engine, interval, numIntervals, 1000))
+
+	for _, c := range []benchCase{
+		{
+			expr:  `rpc_request_success_total + rpc_request_error_total`,
+			steps: 10000,
+		},
+		{
+			expr:  `rpc_request_success_total + ON (job, instance) GROUP_LEFT rpc_request_error_total`,
+			steps: 10000,
+		},
+		{
+			expr:  `rpc_request_success_total AND rpc_request_error_total{instance=~"0.*"}`, // 0.* keeps 1/16 of UUID values
+			steps: 10000,
+		},
+		{
+			expr:  `rpc_request_success_total OR rpc_request_error_total{instance=~"0.*"}`, // 0.* keeps 1/16 of UUID values
+			steps: 10000,
+		},
+		{
+			expr:  `rpc_request_success_total UNLESS rpc_request_error_total{instance=~"0.*"}`, // 0.* keeps 1/16 of UUID values
+			steps: 10000,
+		},
+	} {
+		name := fmt.Sprintf("expr=%s/steps=%d", c.expr, c.steps)
+		b.Run(name, func(b *testing.B) {
+			ctx := context.Background()
+			b.ReportAllocs()
+			for range b.N {
+				qry, err := engine.NewRangeQuery(
+					ctx, stor, nil, c.expr,
+					timestamp.Time(int64((numIntervals-c.steps)*10_000)),
+					timestamp.Time(int64(numIntervals*10_000)),
+					time.Second*10)
+				require.NoError(b, err)
+
+				res := qry.Exec(ctx)
+				require.NoError(b, res.Err)
+
+				qry.Close()
+			}
+		})
+	}
+}
+
 func BenchmarkNativeHistograms(b *testing.B) {
 	testStorage := teststorage.New(b)
 	defer testStorage.Close()