Prometheus: Reduce memory allocations in QueryData.processExemplars (#63373)

* tsdb/prometheus/querydata: Reduce memory allocations in QueryData.processExemplars

---------

Signed-off-by: Arve Knudsen <arve.knudsen@gmail.com>

pkg/tsdb/prometheus/models/result.go

@@ -47,7 +47,8 @@ func (r ResultType) String() string {
 
 type Exemplar struct {
 	SeriesLabels map[string]string
-	Labels       map[string]string
+	Fields       []*data.Field
+	RowIdx       int
 	Value        float64
 	Timestamp    time.Time
 }

pkg/tsdb/prometheus/querydata/exemplar/framer.go

@@ -55,22 +55,29 @@ func (f *Framer) Frames() (data.Frames, error) {
 	valueField := data.NewField(data.TimeSeriesValueFieldName, nil, make([]float64, 0, len(exemplars)))
 	exemplarFrame.Fields = append(exemplarFrame.Fields, timeField, valueField)
 	labelNames := f.labelTracker.GetNames()
+	exemplarLabels := make(map[string]string, len(labelNames))
 	for _, labelName := range labelNames {
 		exemplarFrame.Fields = append(exemplarFrame.Fields, data.NewField(labelName, nil, make([]string, 0, len(exemplars))))
 	}
 
 	// add the sampled exemplars to the new exemplar frame
 	for _, b := range exemplars {
+		// Fill labels map with default values
+		for _, n := range labelNames {
+			exemplarLabels[n] = b.SeriesLabels[n]
+		}
+		// Enter corresponding label values from exemplar fields
+		for _, bf := range b.Fields {
+			if _, exists := exemplarLabels[bf.Name]; exists {
+				exemplarLabels[bf.Name] = bf.CopyAt(b.RowIdx).(string)
+			}
+		}
+
 		timeField.Append(b.Timestamp)
 		valueField.Append(b.Value)
 		for i, labelName := range labelNames {
-			labelValue, ok := b.Labels[labelName]
-			if !ok {
-				// if the label is not present in the exemplar labels, then use the series label
-				labelValue = b.SeriesLabels[labelName]
-			}
 			colIdx := i + 2 // +2 to skip time and value fields
-			exemplarFrame.Fields[colIdx].Append(labelValue)
+			exemplarFrame.Fields[colIdx].Append(exemplarLabels[labelName])
 		}
 	}
 

pkg/tsdb/prometheus/querydata/exemplar/labels.go

@@ -1,11 +1,16 @@
 package exemplar
 
-import "sort"
+import (
+	"sort"
+
+	"github.com/grafana/grafana-plugin-sdk-go/data"
+)
 
 var _ LabelTracker = (*labelTracker)(nil)
 
 type LabelTracker interface {
 	Add(labels map[string]string)
+	AddFields(fields []*data.Field)
 	GetNames() []string
 }
 
@@ -23,9 +28,14 @@ func NewLabelTracker() LabelTracker {
 // so that they can be used to build the label fields in the exemplar frame
 func (l *labelTracker) Add(labels map[string]string) {
 	for k := range labels {
-		if _, ok := l.labelSet[k]; !ok {
-			l.labelSet[k] = struct{}{}
-		}
+		l.labelSet[k] = struct{}{}
+	}
+}
+
+// AddFields saves field names so that they can be used to build the label fields in the exemplar frame.
+func (l *labelTracker) AddFields(fields []*data.Field) {
+	for _, f := range fields {
+		l.labelSet[f.Name] = struct{}{}
 	}
 }
 

pkg/tsdb/prometheus/querydata/framing_bench_test.go

@@ -48,8 +48,11 @@ func BenchmarkExemplarJson(b *testing.B) {
 			Body:       io.NopCloser(bytes.NewReader(responseBytes)),
 		}
 		tCtx.httpProvider.setResponse(&res)
-		_, err := tCtx.queryData.Execute(context.Background(), query)
+		resp, err := tCtx.queryData.Execute(context.Background(), query)
 		require.NoError(b, err)
+		for _, r := range resp.Responses {
+			require.NoError(b, r.Error)
+		}
 	}
 }
 

pkg/tsdb/prometheus/querydata/request.go

@@ -99,7 +99,7 @@ func (s *QueryData) Execute(ctx context.Context, req *backend.QueryDataRequest)
 		}
 		r := s.fetch(ctx, s.client, query, req.Headers)
 		if r == nil {
-			s.log.FromContext(ctx).Debug("Received nilresponse from runQuery", "query", query.Expr)
+			s.log.FromContext(ctx).Debug("Received nil response from runQuery", "query", query.Expr)
 			continue
 		}
 		result.Responses[q.RefID] = *r

pkg/tsdb/prometheus/querydata/response.go

@@ -78,14 +78,15 @@ func (s *QueryData) processExemplars(q *models.Query, dr backend.DataResponse) b
 
 		seriesLabels := getSeriesLabels(frame)
 		labelTracker.Add(seriesLabels)
+		labelTracker.AddFields(frame.Fields[2:])
 		for rowIdx := 0; rowIdx < frame.Fields[0].Len(); rowIdx++ {
-			row := frame.RowCopy(rowIdx)
-			labels := getLabels(frame, row)
-			labelTracker.Add(labels)
+			ts := frame.CopyAt(0, rowIdx).(time.Time)
+			val := frame.CopyAt(1, rowIdx).(float64)
 			ex := models.Exemplar{
-				Labels:       labels,
-				Value:        row[1].(float64),
-				Timestamp:    row[0].(time.Time),
+				RowIdx:       rowIdx,
+				Fields:       frame.Fields[2:],
+				Value:        val,
+				Timestamp:    ts,
 				SeriesLabels: seriesLabels,
 			}
 			sampler.Add(ex)
@@ -200,11 +201,3 @@ func getSeriesLabels(frame *data.Frame) data.Labels {
 	// series labels are stored on the value field (index 1)
 	return frame.Fields[1].Labels.Copy()
 }
-
-func getLabels(frame *data.Frame, row []interface{}) map[string]string {
-	labels := make(map[string]string)
-	for i := 2; i < len(row); i++ {
-		labels[frame.Fields[i].Name] = row[i].(string)
-	}
-	return labels
-}