Prometheus: Add Exemplar sampling for streaming parser (#56049) (#56571)

(cherry picked from commit 152c7f149a)

Co-authored-by: Todd Treece <360020+toddtreece@users.noreply.github.com>

pkg/tsdb/prometheus/querydata/response.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"sort"
 	"strings"
+	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/data"
@@ -39,6 +40,7 @@ func (s *QueryData) parseResponse(ctx context.Context, q *models.Query, res *htt
 		}
 	}
 
+	r = processExemplars(q, r)
 	return r, nil
 }
 
@@ -67,7 +69,7 @@ func addMetadataToWideFrame(q *models.Query, frame *data.Frame) {
 	}
 	frame.Fields[0].Config = &data.FieldConfig{Interval: float64(q.Step.Milliseconds())}
 	for _, f := range frame.Fields {
-		if f.Name != data.TimeSeriesTimeFieldName {
+		if f.Type() == data.FieldTypeFloat64 || f.Type() == data.FieldTypeNullableFloat64 {
 			f.Name = getName(q, f)
 		}
 	}
@@ -132,3 +134,95 @@ func getName(q *models.Query, field *data.Field) string {
 
 	return legend
 }
+
+func processExemplars(q *models.Query, dr *backend.DataResponse) *backend.DataResponse {
+	sampler := newExemplarSampler()
+
+	// we are moving from a multi-frame response returned
+	// by the converter to a single exemplar frame,
+	// so we need to build a new frame array with the
+	// old exemplar frames filtered out
+	frames := []*data.Frame{}
+
+	// the new exemplar frame will be a single frame in long format
+	// with a timestamp, metric value, and one or more label fields
+	exemplarFrame := data.NewFrame("exemplar")
+
+	for _, frame := range dr.Frames {
+		// we don't need to process non-exemplar frames
+		// so they can be added to the response
+		if !isExemplarFrame(frame) {
+			frames = append(frames, frame)
+			continue
+		}
+
+		// copy the frame metadata to the new exemplar frame
+		exemplarFrame.Meta = frame.Meta
+		exemplarFrame.RefID = frame.RefID
+		frame.Meta.Type = data.FrameTypeTimeSeriesMany
+
+		step := time.Duration(frame.Fields[0].Config.Interval) * time.Millisecond
+		seriesLabels := getSeriesLabels(frame)
+		for rowIdx := 0; rowIdx < frame.Fields[0].Len(); rowIdx++ {
+			row := frame.RowCopy(rowIdx)
+			ts := row[0].(time.Time)
+			val := row[1].(float64)
+			labels := getLabels(frame, row)
+			sampler.update(step, ts, val, seriesLabels, labels)
+		}
+	}
+
+	exemplars := sampler.getSampledExemplars()
+	if len(exemplars) == 0 {
+		return dr
+	}
+
+	// init the fields for the new exemplar frame
+	timeField := data.NewField(data.TimeSeriesTimeFieldName, nil, make([]time.Time, 0, len(exemplars)))
+	valueField := data.NewField(data.TimeSeriesValueFieldName, nil, make([]float64, 0, len(exemplars)))
+	exemplarFrame.Fields = append(exemplarFrame.Fields, timeField, valueField)
+	labelNames := sampler.getLabelNames()
+	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 {
+		timeField.Append(b.ts)
+		valueField.Append(b.val)
+		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)
+		}
+	}
+
+	frames = append(frames, exemplarFrame)
+
+	return &backend.DataResponse{
+		Frames: frames,
+		Error:  dr.Error,
+	}
+}
+
+func isExemplarFrame(frame *data.Frame) bool {
+	rt := models.ResultTypeFromFrame(frame)
+	return rt == models.ResultTypeExemplar
+}
+
+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
+}