fix(v2): query series labels from dataset index

pkg/experiment/query_backend/block_reader_test.go

@@ -3,6 +3,7 @@ package query_backend
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"os"
 	"path/filepath"
 	"slices"
@@ -179,3 +180,49 @@ func (s *testSuite) Test_QueryPprof_Metadata() {
 	s.Assert().Equal("cpu", p.StringTable[p.PeriodType.Type])
 	s.Assert().Equal("nanoseconds", p.StringTable[p.PeriodType.Unit])
 }
+
+func (s *testSuite) Test_DatasetIndex_SeriesLabels_GroupBy() {
+	selector := `{service_repository="https://github.com/grafana/pyroscope"}`
+	resp, err := s.reader.Invoke(s.ctx, &queryv1.InvokeRequest{
+		EndTime:       time.Now().UnixMilli(),
+		LabelSelector: selector,
+		QueryPlan:     s.plan,
+		Query: []*queryv1.Query{{
+			QueryType: queryv1.QueryType_QUERY_SERIES_LABELS,
+			SeriesLabels: &queryv1.SeriesLabelsQuery{
+				LabelNames: []string{"service_name", "__profile_type__"},
+			},
+		}},
+	})
+
+	s.Require().NoError(err)
+	s.Require().NotNil(resp)
+	s.Require().Len(resp.Reports, 1)
+
+	expected, err := os.ReadFile("testdata/fixtures/series_labels_by.json")
+	s.Require().NoError(err)
+	actual, _ := json.Marshal(resp.Reports[0].SeriesLabels)
+	s.Assert().JSONEq(string(expected), string(actual))
+}
+
+func (s *testSuite) Test_SeriesLabels() {
+	selector := `{service_name="pyroscope"}`
+	resp, err := s.reader.Invoke(s.ctx, &queryv1.InvokeRequest{
+		EndTime:       time.Now().UnixMilli(),
+		LabelSelector: selector,
+		QueryPlan:     s.plan,
+		Query: []*queryv1.Query{{
+			QueryType:    queryv1.QueryType_QUERY_SERIES_LABELS,
+			SeriesLabels: &queryv1.SeriesLabelsQuery{},
+		}},
+	})
+
+	s.Require().NoError(err)
+	s.Require().NotNil(resp)
+	s.Require().Len(resp.Reports, 1)
+
+	expected, err := os.ReadFile("testdata/fixtures/series_labels.json")
+	s.Require().NoError(err)
+	actual, _ := json.Marshal(resp.Reports[0].SeriesLabels)
+	s.Assert().JSONEq(string(expected), string(actual))
+}

pkg/experiment/query_backend/query_profile_entry.go

@@ -28,7 +28,7 @@ type ProfileEntry struct {
 func (e ProfileEntry) RowNumber() int64 { return e.RowNum }
 
 func profileEntryIterator(q *queryContext, groupBy ...string) (iter.Iterator[ProfileEntry], error) {
-	series, err := getSeriesLabels(q.ds.Index(), q.req.matchers, groupBy...)
+	series, err := getSeries(q.ds.Index(), q.req.matchers, groupBy...)
 	if err != nil {
 		return nil, err
 	}
@@ -62,35 +62,34 @@ func profileEntryIterator(q *queryContext, groupBy ...string) (iter.Iterator[Pro
 	return entries, nil
 }
 
-type seriesLabels struct {
+type series struct {
 	fingerprint model.Fingerprint
 	labels      phlaremodel.Labels
 }
 
-func getSeriesLabels(reader phlaredb.IndexReader, matchers []*labels.Matcher, by ...string) (map[uint32]seriesLabels, error) {
+func getSeries(reader phlaredb.IndexReader, matchers []*labels.Matcher, by ...string) (map[uint32]series, error) {
 	postings, err := getPostings(reader, matchers...)
 	if err != nil {
 		return nil, err
 	}
 	chunks := make([]index.ChunkMeta, 1)
-	series := make(map[uint32]seriesLabels)
+	s := make(map[uint32]series)
 	l := make(phlaremodel.Labels, 0, 6)
 	for postings.Next() {
 		fp, err := reader.SeriesBy(postings.At(), &l, &chunks, by...)
 		if err != nil {
 			return nil, err
 		}
-		_, ok := series[chunks[0].SeriesIndex]
+		_, ok := s[chunks[0].SeriesIndex]
 		if ok {
 			continue
 		}
-		series[chunks[0].SeriesIndex] = seriesLabels{
+		s[chunks[0].SeriesIndex] = series{
 			fingerprint: model.Fingerprint(fp),
 			labels:      l.Clone(),
 		}
 	}
-
-	return series, postings.Err()
+	return s, postings.Err()
 }
 
 func getPostings(reader phlaredb.IndexReader, matchers ...*labels.Matcher) (index.Postings, error) {

pkg/experiment/query_backend/query_series_labels.go

@@ -1,12 +1,18 @@
 package query_backend
 
 import (
+	"errors"
+	"slices"
 	"sync"
 
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/storage"
+
 	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
 	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
 	"github.com/grafana/pyroscope/pkg/experiment/block"
 	phlaremodel "github.com/grafana/pyroscope/pkg/model"
+	"github.com/grafana/pyroscope/pkg/phlaredb"
 )
 
 func init() {
@@ -20,16 +26,10 @@ func init() {
 }
 
 func querySeriesLabels(q *queryContext, query *queryv1.Query) (*queryv1.Report, error) {
-	m, err := getSeriesLabels(q.ds.Index(), q.req.matchers)
+	series, err := getSeriesLabels(q.ds.Index(), q.req.matchers, query.SeriesLabels.LabelNames...)
 	if err != nil {
 		return nil, err
 	}
-	series := make([]*typesv1.Labels, len(m))
-	var i int
-	for _, s := range m {
-		series[i] = &typesv1.Labels{Labels: s.labels}
-		i++
-	}
 	resp := &queryv1.Report{
 		SeriesLabels: &queryv1.SeriesLabelsReport{
 			Query:        query.SeriesLabels.CloneVT(),
@@ -67,3 +67,62 @@ func (a *seriesLabelsAggregator) build() *queryv1.Report {
 		},
 	}
 }
+
+func getSeriesLabels(reader phlaredb.IndexReader, matchers []*labels.Matcher, by ...string) ([]*typesv1.Labels, error) {
+	names, err := reader.LabelNames()
+	if err != nil {
+		return nil, err
+	}
+	if len(by) > 0 {
+		names = slices.DeleteFunc(names, func(n string) bool {
+			for j := 0; j < len(by); j++ {
+				if by[j] == n {
+					return false
+				}
+			}
+			return true
+		})
+		if len(names) == 0 {
+			return nil, nil
+		}
+	}
+
+	postings, err := getPostings(reader, matchers...)
+	if err != nil {
+		return nil, err
+	}
+
+	visited := make(map[uint64]struct{})
+	sets := make([]*typesv1.Labels, 0, 32)
+	ls := make(phlaremodel.Labels, len(names))
+	for i := range names {
+		ls[i] = new(typesv1.LabelPair)
+	}
+
+	for postings.Next() {
+		var j int
+		var v string
+		for i := range names {
+			v, err = reader.LabelValueFor(postings.At(), names[i])
+			switch {
+			case err == nil:
+			case errors.Is(err, storage.ErrNotFound):
+			default:
+				return nil, err
+			}
+			if v != "" {
+				ls[j].Name = names[i]
+				ls[j].Value = v
+				j++
+			}
+		}
+		set := ls[:j]
+		h := set.Hash()
+		if _, ok := visited[h]; !ok {
+			visited[h] = struct{}{}
+			sets = append(sets, &typesv1.Labels{Labels: set.Clone()})
+		}
+	}
+
+	return sets, postings.Err()
+}