chore: Make profileEntryIterator more configurable (#4621)

* chore: Make profileEntryIterator more configurable

* Implement profileIDSelector and address review comments

* Disable ID retrieval until it is used in #4615

Author: simonswine

pkg/phlaredb/schemas/v1/profiles.go

@@ -17,6 +17,7 @@ import (
 )
 
 const (
+	IDColumnName                  = "ID"
 	SeriesIndexColumnName         = "SeriesIndex"
 	TimeNanosColumnName           = "TimeNanos"
 	StacktracePartitionColumnName = "StacktracePartition"

pkg/querybackend/query_profile_entry.go

@@ -1,6 +1,10 @@
 package querybackend
 
 import (
+	"slices"
+	"strings"
+
+	"github.com/google/uuid"
 	"github.com/parquet-go/parquet-go"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/labels"
@@ -23,39 +27,198 @@ type ProfileEntry struct {
 	Fingerprint model.Fingerprint
 	Labels      phlaremodel.Labels
 	Partition   uint64
+	ID          string
 }
 
 func (e ProfileEntry) RowNumber() int64 { return e.RowNum }
 
-func profileEntryIterator(q *queryContext, groupBy ...string) (iter.Iterator[ProfileEntry], error) {
-	series, err := getSeries(q.ds.Index(), q.req.matchers, groupBy...)
+type profileIteratorOption struct {
+	iterator func(*iteratorOpts)
+	series   func(*seriesOpts)
+}
+
+func withAllLabels() profileIteratorOption {
+	return profileIteratorOption{
+		series: func(opts *seriesOpts) {
+			opts.allLabels = true
+		},
+	}
+}
+
+func withGroupByLabels(by ...string) profileIteratorOption {
+	return profileIteratorOption{
+		series: func(opts *seriesOpts) {
+			opts.groupBy = by
+		},
+	}
+}
+
+func withFetchPartition(v bool) profileIteratorOption {
+	return profileIteratorOption{
+		iterator: func(opts *iteratorOpts) {
+			opts.fetchPartition = v
+		},
+	}
+}
+
+func withFetchProfileIDs(v bool) profileIteratorOption {
+	return profileIteratorOption{
+		iterator: func(opts *iteratorOpts) {
+			opts.fetchProfileIDs = v
+		},
+	}
+}
+
+func withProfileIDSelector(ids ...string) (profileIteratorOption, error) {
+	// convert profile ids into uuids
+	uuids := make([]string, 0, len(ids))
+	for _, id := range ids {
+		u, err := uuid.Parse(id)
+		if err != nil {
+			return profileIteratorOption{}, err
+		}
+		uuids = append(uuids, string(u[:]))
+	}
+
+	return profileIteratorOption{
+		iterator: func(opts *iteratorOpts) {
+			opts.profileIDSelector = uuids
+		},
+	}, nil
+}
+
+type iteratorOpts struct {
+	profileIDSelector []string // this is a slice of the byte form of the UUID
+	fetchProfileIDs   bool
+	fetchPartition    bool
+}
+
+func iteratorOptsFromOptions(options []profileIteratorOption) iteratorOpts {
+	opts := iteratorOpts{
+		fetchPartition: true,
+	}
+	for _, f := range options {
+		if f.iterator != nil {
+			f.iterator(&opts)
+		}
+	}
+	return opts
+}
+
+type queryColumn struct {
+	name      string
+	predicate parquetquery.Predicate
+	priority  int
+}
+
+type queryColumns []queryColumn
+
+func (c queryColumns) names() []string {
+	result := make([]string, len(c))
+	for idx := range result {
+		result[idx] = c[idx].name
+	}
+	return result
+}
+
+func (c queryColumns) join(q *queryContext) parquetquery.Iterator {
+	var result parquetquery.Iterator
+
+	// sort columns by priority, without modifying queryColumn slice
+	order := make([]int, len(c))
+	for idx := range order {
+		order[idx] = idx
+	}
+	slices.SortFunc(order, func(a, b int) int {
+		if r := c[a].priority - c[b].priority; r != 0 {
+			return r
+		}
+		return strings.Compare(c[a].name, c[b].name)
+	})
+
+	for _, idx := range order {
+		it := q.ds.Profiles().Column(q.ctx, c[idx].name, c[idx].predicate)
+		if result == nil {
+			result = it
+			continue
+		}
+		result = parquetquery.NewBinaryJoinIterator(0,
+			result,
+			it,
+		)
+	}
+	return result
+}
+
+func profileEntryIterator(q *queryContext, options ...profileIteratorOption) (iter.Iterator[ProfileEntry], error) {
+	opts := iteratorOptsFromOptions(options)
+
+	series, err := getSeries(q.ds.Index(), q.req.matchers, options...)
 	if err != nil {
 		return nil, err
 	}
-	results := parquetquery.NewBinaryJoinIterator(0,
-		q.ds.Profiles().Column(q.ctx, "SeriesIndex", parquetquery.NewMapPredicate(series)),
-		q.ds.Profiles().Column(q.ctx, "TimeNanos", parquetquery.NewIntBetweenPredicate(q.req.startTime, q.req.endTime)),
-	)
-	results = parquetquery.NewBinaryJoinIterator(0, results,
-		q.ds.Profiles().Column(q.ctx, "StacktracePartition", nil),
-	)
 
-	buf := make([][]parquet.Value, 3)
+	columns := queryColumns{
+		{schemav1.SeriesIndexColumnName, parquetquery.NewMapPredicate(series), 10},
+		{schemav1.TimeNanosColumnName, parquetquery.NewIntBetweenPredicate(q.req.startTime, q.req.endTime), 15},
+	}
+	processor := []func([][]parquet.Value, *ProfileEntry){}
+
+	// fetch partition if requested
+	if opts.fetchPartition {
+		offset := len(columns)
+		columns = append(
+			columns,
+			queryColumn{schemav1.StacktracePartitionColumnName, nil, 20},
+		)
+		processor = append(processor, func(buf [][]parquet.Value, e *ProfileEntry) {
+			e.Partition = buf[offset][0].Uint64()
+		})
+	}
+	// fetch profile id if requested or part of the predicate
+	if opts.fetchProfileIDs || len(opts.profileIDSelector) > 0 {
+		var (
+			predicate parquetquery.Predicate
+			priority  = 20
+		)
+		if len(opts.profileIDSelector) > 0 {
+			predicate = parquetquery.NewStringInPredicate(opts.profileIDSelector)
+			priority = 5
+		}
+		offset := len(columns)
+		columns = append(
+			columns,
+			queryColumn{schemav1.IDColumnName, predicate, priority},
+		)
+		var u uuid.UUID
+		processor = append(processor, func(buf [][]parquet.Value, e *ProfileEntry) {
+			b := buf[offset][0].Bytes()
+			if len(b) != 16 {
+				return
+			}
+			copy(u[:], b)
+			e.ID = u.String()
+		})
+	}
+
+	buf := make([][]parquet.Value, len(columns))
+	columnNames := columns.names()
+
 	entries := iter.NewAsyncBatchIterator[*parquetquery.IteratorResult, ProfileEntry](
-		results, bigBatchSize,
+		columns.join(q), bigBatchSize,
 		func(r *parquetquery.IteratorResult) ProfileEntry {
-			buf = r.Columns(buf,
-				schemav1.SeriesIndexColumnName,
-				schemav1.TimeNanosColumnName,
-				schemav1.StacktracePartitionColumnName)
+			buf = r.Columns(buf, columnNames...)
 			x := series[buf[0][0].Uint32()]
-			return ProfileEntry{
+			e := ProfileEntry{
 				RowNum:      r.RowNumber[0],
 				Timestamp:   model.TimeFromUnixNano(buf[1][0].Int64()),
 				Fingerprint: x.fingerprint,
 				Labels:      x.labels,
-				Partition:   buf[2][0].Uint64(),
 			}
+			for _, proc := range processor {
+				proc(buf, &e)
+			}
+			return e
 		},
 		func([]ProfileEntry) {},
 	)
@@ -67,7 +230,19 @@ type series struct {
 	labels      phlaremodel.Labels
 }
 
-func getSeries(reader phlaredb.IndexReader, matchers []*labels.Matcher, by ...string) (map[uint32]series, error) {
+type seriesOpts struct {
+	allLabels bool // when this is true, groupBy is ignored
+	groupBy   []string
+}
+
+func getSeries(reader phlaredb.IndexReader, matchers []*labels.Matcher, options ...profileIteratorOption) (map[uint32]series, error) {
+	var opts seriesOpts
+	for _, f := range options {
+		if f.series != nil {
+			f.series(&opts)
+		}
+	}
+
 	postings, err := getPostings(reader, matchers...)
 	if err != nil {
 		return nil, err
@@ -76,7 +251,12 @@ func getSeries(reader phlaredb.IndexReader, matchers []*labels.Matcher, by ...st
 	s := make(map[uint32]series)
 	l := make(phlaremodel.Labels, 0, 6)
 	for postings.Next() {
-		fp, err := reader.SeriesBy(postings.At(), &l, &chunks, by...)
+		var fp uint64
+		if opts.allLabels {
+			fp, err = reader.Series(postings.At(), &l, &chunks)
+		} else {
+			fp, err = reader.SeriesBy(postings.At(), &l, &chunks, opts.groupBy...)
+		}
 		if err != nil {
 			return nil, err
 		}

pkg/querybackend/query_time_series.go

@@ -31,7 +31,22 @@ func init() {
 }
 
 func queryTimeSeries(q *queryContext, query *queryv1.Query) (r *queryv1.Report, err error) {
-	entries, err := profileEntryIterator(q, query.TimeSeries.GroupBy...)
+	opts := []profileIteratorOption{
+		withFetchPartition(false), // Partition data not needed, as we don't access stacktraces at all
+	}
+	exemplarsEnabled := false // TODO: This will be enabled as part of #4615
+	if exemplarsEnabled {
+		opts = append(opts,
+			withAllLabels(),
+			withFetchProfileIDs(true),
+		)
+	} else {
+		opts = append(opts,
+			withGroupByLabels(query.TimeSeries.GroupBy...),
+		)
+	}
+
+	entries, err := profileEntryIterator(q, opts...)
 	if err != nil {
 		return nil, err
 	}