Reuse prom serieset from Thanos (#5844)

* reuse prom serieset from Thanos

Signed-off-by: Ben Ye <[email protected]>

* fix lint

Signed-off-by: Ben Ye <[email protected]>

---------

Signed-off-by: Ben Ye <[email protected]>

Author: yeya24

pkg/querier/block.go

@@ -1,19 +1,8 @@
 package querier
 
 import (
-	"math"
-	"sort"
-
-	"github.com/pkg/errors"
 	"github.com/prometheus/prometheus/model/labels"
-	"github.com/prometheus/prometheus/storage"
-	"github.com/prometheus/prometheus/tsdb/chunkenc"
-	"github.com/prometheus/prometheus/util/annotations"
-	"github.com/thanos-io/thanos/pkg/store/labelpb"
 	"github.com/thanos-io/thanos/pkg/store/storepb"
-
-	"github.com/cortexproject/cortex/pkg/querier/iterators"
-	"github.com/cortexproject/cortex/pkg/querier/series"
 )
 
 func convertMatchersToLabelMatcher(matchers []*labels.Matcher) []storepb.LabelMatcher {
@@ -40,180 +29,28 @@ func convertMatchersToLabelMatcher(matchers []*labels.Matcher) []storepb.LabelMa
 	return converted
 }
 
-// Implementation of storage.SeriesSet, based on individual responses from store client.
-type blockQuerierSeriesSet struct {
-	series   []*storepb.Series
-	warnings annotations.Annotations
-
-	// next response to process
-	next int
-
-	currSeries storage.Series
+// storeSeriesSet implements a storepb SeriesSet against a list of storepb.Series.
+type storeSeriesSet struct {
+	series []*storepb.Series
+	i      int
 }
 
-func (bqss *blockQuerierSeriesSet) Next() bool {
-	bqss.currSeries = nil
+func newStoreSeriesSet(s []*storepb.Series) *storeSeriesSet {
+	return &storeSeriesSet{series: s, i: -1}
+}
 
-	if bqss.next >= len(bqss.series) {
+func (s *storeSeriesSet) Next() bool {
+	if s.i >= len(s.series)-1 {
 		return false
 	}
-
-	currLabels := labelpb.ZLabelsToPromLabels(bqss.series[bqss.next].Labels)
-	currChunks := bqss.series[bqss.next].Chunks
-
-	bqss.next++
-
-	// Merge chunks for current series. Chunks may come in multiple responses, but as soon
-	// as the response has chunks for a new series, we can stop searching. Series are sorted.
-	// See documentation for StoreClient.Series call for details.
-	for bqss.next < len(bqss.series) && labels.Compare(currLabels, labelpb.ZLabelsToPromLabels(bqss.series[bqss.next].Labels)) == 0 {
-		currChunks = append(currChunks, bqss.series[bqss.next].Chunks...)
-		bqss.next++
-	}
-
-	bqss.currSeries = newBlockQuerierSeries(currLabels, currChunks)
+	s.i++
 	return true
 }
 
-func (bqss *blockQuerierSeriesSet) At() storage.Series {
-	return bqss.currSeries
-}
-
-func (bqss *blockQuerierSeriesSet) Err() error {
+func (*storeSeriesSet) Err() error {
 	return nil
 }
 
-func (bqss *blockQuerierSeriesSet) Warnings() annotations.Annotations {
-	return bqss.warnings
-}
-
-// newBlockQuerierSeries makes a new blockQuerierSeries. Input labels must be already sorted by name.
-func newBlockQuerierSeries(lbls []labels.Label, chunks []storepb.AggrChunk) *blockQuerierSeries {
-	sort.Slice(chunks, func(i, j int) bool {
-		return chunks[i].MinTime < chunks[j].MinTime
-	})
-
-	return &blockQuerierSeries{labels: lbls, chunks: chunks}
-}
-
-type blockQuerierSeries struct {
-	labels labels.Labels
-	chunks []storepb.AggrChunk
-}
-
-func (bqs *blockQuerierSeries) Labels() labels.Labels {
-	return bqs.labels
-}
-
-func (bqs *blockQuerierSeries) Iterator(chunkenc.Iterator) chunkenc.Iterator {
-	if len(bqs.chunks) == 0 {
-		// should not happen in practice, but we have a unit test for it
-		return series.NewErrIterator(errors.New("no chunks"))
-	}
-
-	its := make([]chunkenc.Iterator, 0, len(bqs.chunks))
-
-	for _, c := range bqs.chunks {
-		// Ignore if the current chunk is not XOR chunk.
-		if c.Raw == nil {
-			continue
-		}
-		ch, err := chunkenc.FromData(chunkenc.EncXOR, c.Raw.Data)
-		if err != nil {
-			return series.NewErrIterator(errors.Wrapf(err, "failed to initialize chunk from XOR encoded raw data (series: %v min time: %d max time: %d)", bqs.Labels(), c.MinTime, c.MaxTime))
-		}
-
-		it := ch.Iterator(nil)
-		its = append(its, it)
-	}
-
-	return iterators.NewCompatibleChunksIterator(newBlockQuerierSeriesIterator(bqs.Labels(), its))
-}
-
-func newBlockQuerierSeriesIterator(labels labels.Labels, its []chunkenc.Iterator) *blockQuerierSeriesIterator {
-	return &blockQuerierSeriesIterator{labels: labels, iterators: its, lastT: math.MinInt64}
-}
-
-// blockQuerierSeriesIterator implements a series iterator on top
-// of a list of time-sorted, non-overlapping chunks.
-type blockQuerierSeriesIterator struct {
-	// only used for error reporting
-	labels labels.Labels
-
-	iterators []chunkenc.Iterator
-	i         int
-	lastT     int64
-}
-
-func (it *blockQuerierSeriesIterator) Seek(t int64) bool {
-	// We generally expect the chunks already to be cut down
-	// to the range we are interested in. There's not much to be gained from
-	// hopping across chunks so we just call next until we reach t.
-	for {
-		ct, _ := it.At()
-		if ct >= t {
-			return true
-		}
-		if !it.Next() {
-			return false
-		}
-	}
-}
-
-func (it *blockQuerierSeriesIterator) At() (int64, float64) {
-	if it.i >= len(it.iterators) {
-		return 0, 0
-	}
-
-	t, v := it.iterators[it.i].At()
-	it.lastT = t
-	return t, v
-}
-
-func (it *blockQuerierSeriesIterator) Next() bool {
-	if it.i >= len(it.iterators) {
-		return false
-	}
-
-	if it.iterators[it.i].Next() != chunkenc.ValNone {
-		return true
-	}
-	if it.iterators[it.i].Err() != nil {
-		return false
-	}
-
-	for {
-		it.i++
-
-		if it.i >= len(it.iterators) {
-			return false
-		}
-
-		// we must advance iterator first, to see if it has any samples.
-		// Seek will call At() as its first operation.
-		if it.iterators[it.i].Next() == chunkenc.ValNone {
-			if it.iterators[it.i].Err() != nil {
-				return false
-			}
-
-			// Found empty iterator without error, skip it.
-			continue
-		}
-
-		// Chunks are guaranteed to be ordered but not generally guaranteed to not overlap.
-		// We must ensure to skip any overlapping range between adjacent chunks.
-		return it.Seek(it.lastT + 1)
-	}
-}
-
-func (it *blockQuerierSeriesIterator) Err() error {
-	if it.i >= len(it.iterators) {
-		return nil
-	}
-
-	err := it.iterators[it.i].Err()
-	if err != nil {
-		return errors.Wrapf(err, "cannot iterate chunk for series: %v", it.labels)
-	}
-	return nil
+func (s *storeSeriesSet) At() (labels.Labels, []storepb.AggrChunk) {
+	return s.series[s.i].PromLabels(), s.series[s.i].Chunks
 }