From 35a15e8f04b3b5716da73c0f4faf1c2be34b79a0 Mon Sep 17 00:00:00 2001
From: Filip Petkovski <filip.petkovsky@gmail.com>
Date: Thu, 23 Nov 2023 15:09:17 +0100
Subject: [PATCH] Add benchmark for native histograms (#13160)

* Add benchmark for native histograms

This commit adds a PromQL benchmark for queries on native histograms.

Signed-off-by: Filip Petkovski <filip.petkovsky@gmail.com>
---
 promql/bench_test.go | 95 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 95 insertions(+)

diff --git a/promql/bench_test.go b/promql/bench_test.go
index 8e443b5a6..13eba3714 100644
--- a/promql/bench_test.go
+++ b/promql/bench_test.go
@@ -21,9 +21,11 @@ import (
 	"testing"
 	"time"
 
+	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/storage"
+	"github.com/prometheus/prometheus/tsdb/tsdbutil"
 	"github.com/prometheus/prometheus/util/teststorage"
 )
 
@@ -269,6 +271,99 @@ func BenchmarkRangeQuery(b *testing.B) {
 	}
 }
 
+func BenchmarkNativeHistograms(b *testing.B) {
+	testStorage := teststorage.New(b)
+	defer testStorage.Close()
+
+	app := testStorage.Appender(context.TODO())
+	if err := generateNativeHistogramSeries(app, 3000); err != nil {
+		b.Fatal(err)
+	}
+	if err := app.Commit(); err != nil {
+		b.Fatal(err)
+	}
+
+	start := time.Unix(0, 0)
+	end := start.Add(2 * time.Hour)
+	step := time.Second * 30
+
+	cases := []struct {
+		name  string
+		query string
+	}{
+		{
+			name:  "sum",
+			query: "sum(native_histogram_series)",
+		},
+		{
+			name:  "sum rate",
+			query: "sum(rate(native_histogram_series[1m]))",
+		},
+	}
+
+	opts := EngineOpts{
+		Logger:               nil,
+		Reg:                  nil,
+		MaxSamples:           50000000,
+		Timeout:              100 * time.Second,
+		EnableAtModifier:     true,
+		EnableNegativeOffset: true,
+	}
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for _, tc := range cases {
+		b.Run(tc.name, func(b *testing.B) {
+			ng := NewEngine(opts)
+			for i := 0; i < b.N; i++ {
+				qry, err := ng.NewRangeQuery(context.Background(), testStorage, nil, tc.query, start, end, step)
+				if err != nil {
+					b.Fatal(err)
+				}
+				if result := qry.Exec(context.Background()); result.Err != nil {
+					b.Fatal(result.Err)
+				}
+			}
+		})
+	}
+}
+
+func generateNativeHistogramSeries(app storage.Appender, numSeries int) error {
+	commonLabels := []string{labels.MetricName, "native_histogram_series", "foo", "bar"}
+	series := make([][]*histogram.Histogram, numSeries)
+	for i := range series {
+		series[i] = tsdbutil.GenerateTestHistograms(2000)
+	}
+	higherSchemaHist := &histogram.Histogram{
+		Schema: 3,
+		PositiveSpans: []histogram.Span{
+			{Offset: -5, Length: 2}, // -5 -4
+			{Offset: 2, Length: 3},  // -1 0 1
+			{Offset: 2, Length: 2},  // 4 5
+		},
+		PositiveBuckets: []int64{1, 2, -2, 1, -1, 0, 3},
+		Count:           13,
+	}
+	for sid, histograms := range series {
+		seriesLabels := labels.FromStrings(append(commonLabels, "h", strconv.Itoa(sid))...)
+		for i := range histograms {
+			ts := time.Unix(int64(i*15), 0).UnixMilli()
+			if i == 0 {
+				// Inject a histogram with a higher schema.
+				if _, err := app.AppendHistogram(0, seriesLabels, ts, higherSchemaHist, nil); err != nil {
+					return err
+				}
+			}
+			if _, err := app.AppendHistogram(0, seriesLabels, ts, histograms[i], nil); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
 func BenchmarkParser(b *testing.B) {
 	cases := []string{
 		"a",