feat: add exemplar support to timeseries responses for individual profile retrieval

pkg/frontend/readpath/queryfrontend/query_select_time_series.go

@@ -48,9 +48,10 @@ func (q *QueryFrontend) SelectSeries(
 		Query: []*queryv1.Query{{
 			QueryType: queryv1.QueryType_QUERY_TIME_SERIES,
 			TimeSeries: &queryv1.TimeSeriesQuery{
-				Step:    c.Msg.GetStep(),
-				GroupBy: c.Msg.GetGroupBy(),
-				Limit:   c.Msg.GetLimit(),
+				Step:         c.Msg.GetStep(),
+				GroupBy:      c.Msg.GetGroupBy(),
+				Limit:        c.Msg.GetLimit(),
+				ExemplarType: c.Msg.GetExemplarType(),
 			},
 		}},
 	})

pkg/model/exemplar_builder.go

@@ -0,0 +1,102 @@
+package model
+
+import (
+	"sort"
+
+	"github.com/prometheus/common/model"
+
+	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
+)
+
+// ExemplarBuilder builds exemplars for a single time series.
+type ExemplarBuilder struct {
+	labelSetIndex map[uint64]int
+	labelSets     []Labels
+	exemplars     []exemplar
+}
+
+type exemplar struct {
+	timestamp   int64
+	profileID   string
+	labelSetRef int
+	value       uint64
+}
+
+// NewExemplarBuilder creates a new ExemplarBuilder.
+func NewExemplarBuilder() *ExemplarBuilder {
+	return &ExemplarBuilder{
+		labelSetIndex: make(map[uint64]int),
+		labelSets:     make([]Labels, 0),
+		exemplars:     make([]exemplar, 0),
+	}
+}
+
+// Add adds an exemplar with its full labels.
+func (eb *ExemplarBuilder) Add(fp model.Fingerprint, labels Labels, ts int64, profileID string, value uint64) {
+	if profileID == "" {
+		return
+	}
+
+	labelSetIdx, exists := eb.labelSetIndex[uint64(fp)]
+	if !exists {
+		eb.labelSets = append(eb.labelSets, labels.Clone())
+		labelSetIdx = len(eb.labelSets) - 1
+		eb.labelSetIndex[uint64(fp)] = labelSetIdx
+	}
+
+	eb.exemplars = append(eb.exemplars, exemplar{
+		timestamp:   ts,
+		profileID:   profileID,
+		labelSetRef: labelSetIdx,
+		value:       value,
+	})
+}
+
+// Build returns the final exemplars, sorted and deduplicated.
+// Exemplars with the same (profileID, timestamp) are merged by intersecting their labels.
+func (eb *ExemplarBuilder) Build() []*typesv1.Exemplar {
+	if len(eb.exemplars) == 0 {
+		return nil
+	}
+
+	sort.Slice(eb.exemplars, func(i, j int) bool {
+		if eb.exemplars[i].timestamp != eb.exemplars[j].timestamp {
+			return eb.exemplars[i].timestamp < eb.exemplars[j].timestamp
+		}
+		return eb.exemplars[i].profileID < eb.exemplars[j].profileID
+	})
+
+	return eb.deduplicateAndIntersect()
+}
+
+// deduplicateAndIntersect merges exemplars with the same profileID, timestamp by intersecting their label sets.
+func (eb *ExemplarBuilder) deduplicateAndIntersect() []*typesv1.Exemplar {
+	result := make([]*typesv1.Exemplar, 0, len(eb.exemplars))
+
+	i := 0
+	for i < len(eb.exemplars) {
+		curr := eb.exemplars[i]
+
+		labelSetsToIntersect := []Labels{eb.labelSets[curr.labelSetRef]}
+		j := i + 1
+		for j < len(eb.exemplars) &&
+			eb.exemplars[j].profileID == curr.profileID &&
+			eb.exemplars[j].timestamp == curr.timestamp {
+			labelSetsToIntersect = append(labelSetsToIntersect, eb.labelSets[eb.exemplars[j].labelSetRef])
+			j++
+		}
+
+		finalLabels := IntersectAll(labelSetsToIntersect)
+
+		result = append(result, &typesv1.Exemplar{
+			Timestamp: curr.timestamp,
+			ProfileId: curr.profileID,
+			Value:     curr.value,
+			Labels:    finalLabels,
+		})
+
+		i = j
+	}
+
+	return result
+}

pkg/model/labels.go

@@ -584,3 +584,52 @@ func ServiceVersionFromLabels(lbls Labels) (ServiceVersion, bool) {
 		RootPath:   rootPath,
 	}, repo != "" || gitref != "" || rootPath != ""
 }
+
+// IntersectAll returns only the labels that are present in all label sets
+// with the same value. Used for merging exemplars with dynamic labels.
+// Reuses the existing Intersect method by iteratively intersecting all label sets.
+func IntersectAll(labelSets []Labels) Labels {
+	if len(labelSets) == 0 {
+		return nil
+	}
+	if len(labelSets) == 1 {
+		return labelSets[0]
+	}
+
+	result := labelSets[0].Clone()
+	for i := 1; i < len(labelSets); i++ {
+		result = result.Intersect(labelSets[i])
+		if len(result) == 0 {
+			return nil
+		}
+	}
+
+	if len(result) == 0 {
+		return nil
+	}
+	return result
+}
+
+// WithoutLabels returns a new Labels with the specified label names removed.
+func (ls Labels) WithoutLabels(names ...string) Labels {
+	if len(names) == 0 {
+		return ls
+	}
+
+	toRemove := make(Labels, 0, len(names))
+	for _, name := range names {
+		for _, l := range ls {
+			if l.Name == name {
+				toRemove = append(toRemove, l)
+				break
+			}
+		}
+	}
+
+	if len(toRemove) == 0 {
+		return ls
+	}
+
+	result := ls.Clone()
+	return result.Subtract(toRemove)
+}

pkg/model/labels_test.go

@@ -666,3 +666,168 @@ func TestLabels_Intersect(t *testing.T) {
 		})
 	}
 }
+
+func TestLabels_IntersectAll(t *testing.T) {
+	tests := []struct {
+		name      string
+		labelSets []Labels
+		expected  Labels
+	}{
+		{
+			name:      "Empty input",
+			labelSets: []Labels{},
+			expected:  nil,
+		},
+		{
+			name: "Single label set",
+			labelSets: []Labels{
+				{
+					{Name: "env", Value: "prod"},
+					{Name: "service", Value: "api"},
+				},
+			},
+			expected: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "service", Value: "api"},
+			},
+		},
+		{
+			name: "Two sets with full overlap",
+			labelSets: []Labels{
+				{
+					{Name: "env", Value: "prod"},
+					{Name: "service", Value: "api"},
+				},
+				{
+					{Name: "env", Value: "prod"},
+					{Name: "service", Value: "api"},
+				},
+			},
+			expected: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "service", Value: "api"},
+			},
+		},
+		{
+			name: "Two sets with partial overlap",
+			labelSets: []Labels{
+				{
+					{Name: "env", Value: "prod"},
+					{Name: "pod", Value: "pod-1"},
+					{Name: "region", Value: "us-east"},
+				},
+				{
+					{Name: "env", Value: "prod"},
+					{Name: "pod", Value: "pod-2"},
+					{Name: "region", Value: "us-east"},
+				},
+			},
+			expected: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "region", Value: "us-east"},
+			},
+		},
+		{
+			name: "No common labels",
+			labelSets: []Labels{
+				{
+					{Name: "env", Value: "prod"},
+				},
+				{
+					{Name: "region", Value: "us-east"},
+				},
+			},
+			expected: nil,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			result := IntersectAll(test.labelSets)
+			assert.Equal(t, test.expected, result)
+		})
+	}
+}
+
+func TestLabels_WithoutLabels(t *testing.T) {
+	tests := []struct {
+		name     string
+		labels   Labels
+		remove   []string
+		expected Labels
+	}{
+		{
+			name: "Remove nothing",
+			labels: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "service", Value: "api"},
+			},
+			remove: []string{},
+			expected: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "service", Value: "api"},
+			},
+		},
+		{
+			name: "Remove single label",
+			labels: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "pod", Value: "pod-1"},
+				{Name: "service", Value: "api"},
+			},
+			remove: []string{"pod"},
+			expected: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "service", Value: "api"},
+			},
+		},
+		{
+			name: "Remove multiple labels",
+			labels: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "pod", Value: "pod-1"},
+				{Name: "region", Value: "us-east"},
+				{Name: "service", Value: "api"},
+			},
+			remove: []string{"pod", "region"},
+			expected: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "service", Value: "api"},
+			},
+		},
+		{
+			name: "Remove non-existent label",
+			labels: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "service", Value: "api"},
+			},
+			remove: []string{"nonexistent"},
+			expected: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "service", Value: "api"},
+			},
+		},
+		{
+			name: "Remove all labels",
+			labels: Labels{
+				{Name: "env", Value: "prod"},
+				{Name: "service", Value: "api"},
+			},
+			remove:   []string{"env", "service"},
+			expected: Labels{},
+		},
+		{
+			name:     "Empty labels",
+			labels:   Labels{},
+			remove:   []string{"env"},
+			expected: Labels{},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			result := test.labels.WithoutLabels(test.remove...)
+			assert.Equal(t, test.expected, result)
+		})
+	}
+}