Optimize and test MemoizedSeriesIterator

Signed-off-by: Justin Lei <justin.lei@grafana.com>

promql/engine.go

@@ -1857,7 +1857,7 @@ func (ev *evaluator) vectorSelectorSingle(it *storage.MemoizedSeriesIterator, no
 	}
 	if valueType == chunkenc.ValNone || t > refTime {
 		var ok bool
-		t, v, _, h, ok = it.PeekPrev()
+		t, v, h, ok = it.PeekPrev()
 		if !ok || t < refTime-durationMilliseconds(ev.lookbackDelta) {
 			return 0, 0, nil, false
 		}

storage/memoized_iterator.go

@@ -31,12 +31,7 @@ type MemoizedSeriesIterator struct {
 	// Keep track of the previously returned value.
 	prevTime           int64
 	prevValue          float64
-	prevHistogram      *histogram.Histogram
 	prevFloatHistogram *histogram.FloatHistogram
-	// TODO(beorn7): MemoizedSeriesIterator is currently only used by the
-	// PromQL engine, which only works with FloatHistograms. For better
-	// performance, we could change MemoizedSeriesIterator to also only
-	// handle FloatHistograms.
 }
 
 // NewMemoizedEmptyIterator is like NewMemoizedIterator but it's initialised with an empty iterator.
@@ -66,11 +61,11 @@ func (b *MemoizedSeriesIterator) Reset(it chunkenc.Iterator) {
 
 // PeekPrev returns the previous element of the iterator. If there is none buffered,
 // ok is false.
-func (b *MemoizedSeriesIterator) PeekPrev() (t int64, v float64, h *histogram.Histogram, fh *histogram.FloatHistogram, ok bool) {
+func (b *MemoizedSeriesIterator) PeekPrev() (t int64, v float64, fh *histogram.FloatHistogram, ok bool) {
 	if b.prevTime == math.MinInt64 {
-		return 0, 0, nil, nil, false
+		return 0, 0, nil, false
 	}
-	return b.prevTime, b.prevValue, b.prevHistogram, b.prevFloatHistogram, true
+	return b.prevTime, b.prevValue, b.prevFloatHistogram, true
 }
 
 // Seek advances the iterator to the element at time t or greater.
@@ -108,15 +103,14 @@ func (b *MemoizedSeriesIterator) Next() chunkenc.ValueType {
 		return chunkenc.ValNone
 	case chunkenc.ValFloat:
 		b.prevTime, b.prevValue = b.it.At()
-		b.prevHistogram = nil
 		b.prevFloatHistogram = nil
 	case chunkenc.ValHistogram:
 		b.prevValue = 0
-		b.prevTime, b.prevHistogram = b.it.AtHistogram()
+		ts, h := b.it.AtHistogram()
-		_, b.prevFloatHistogram = b.it.AtFloatHistogram()
+		b.prevTime = ts
+		b.prevFloatHistogram = h.ToFloat()
 	case chunkenc.ValFloatHistogram:
 		b.prevValue = 0
-		b.prevHistogram = nil
 		b.prevTime, b.prevFloatHistogram = b.it.AtFloatHistogram()
 	}
 
@@ -132,13 +126,17 @@ func (b *MemoizedSeriesIterator) At() (int64, float64) {
 	return b.it.At()
 }
 
-// AtHistogram returns the current histogram element of the iterator.
+// AtHistogram is not supported by this iterator.
 func (b *MemoizedSeriesIterator) AtHistogram() (int64, *histogram.Histogram) {
-	return b.it.AtHistogram()
+	panic("AtHistogram is not supported by this iterator.")
 }
 
 // AtFloatHistogram returns the current float-histogram element of the iterator.
 func (b *MemoizedSeriesIterator) AtFloatHistogram() (int64, *histogram.FloatHistogram) {
+	if b.valueType == chunkenc.ValHistogram {
+		ts, h := b.it.AtHistogram()
+		return ts, h.ToFloat()
+	}
 	return b.it.AtFloatHistogram()
 }
 

storage/memoized_iterator_test.go

@@ -16,25 +16,36 @@ package storage
 import (
 	"testing"
 
-	"github.com/stretchr/testify/require"
+	"github.com/prometheus/prometheus/model/histogram"
-
 	"github.com/prometheus/prometheus/tsdb/chunkenc"
+	"github.com/prometheus/prometheus/tsdb/tsdbutil"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestMemoizedSeriesIterator(t *testing.T) {
-	// TODO(beorn7): Include histograms in testing.
 	var it *MemoizedSeriesIterator
 
-	sampleEq := func(ets int64, ev float64) {
+	sampleEq := func(ets int64, ev float64, efh *histogram.FloatHistogram) {
-		ts, v := it.At()
+		if efh == nil {
-		require.Equal(t, ets, ts, "timestamp mismatch")
+			ts, v := it.At()
-		require.Equal(t, ev, v, "value mismatch")
+			require.Equal(t, ets, ts, "timestamp mismatch")
+			require.Equal(t, ev, v, "value mismatch")
+		} else {
+			ts, fh := it.AtFloatHistogram()
+			require.Equal(t, ets, ts, "timestamp mismatch")
+			require.Equal(t, efh, fh, "histogram mismatch")
+		}
 	}
-	prevSampleEq := func(ets int64, ev float64, eok bool) {
+	prevSampleEq := func(ets int64, ev float64, efh *histogram.FloatHistogram, eok bool) {
-		ts, v, _, _, ok := it.PeekPrev()
+		ts, v, fh, ok := it.PeekPrev()
 		require.Equal(t, eok, ok, "exist mismatch")
 		require.Equal(t, ets, ts, "timestamp mismatch")
-		require.Equal(t, ev, v, "value mismatch")
+		if efh == nil {
+			require.Equal(t, ev, v, "value mismatch")
+		} else {
+			require.Equal(t, efh, fh, "histogram mismatch")
+		}
 	}
 
 	it = NewMemoizedIterator(NewListSeriesIterator(samples{
@@ -46,29 +57,60 @@ func TestMemoizedSeriesIterator(t *testing.T) {
 		fSample{t: 99, f: 8},
 		fSample{t: 100, f: 9},
 		fSample{t: 101, f: 10},
+		hSample{t: 102, h: tsdbutil.GenerateTestHistogram(0)},
+		hSample{t: 103, h: tsdbutil.GenerateTestHistogram(1)},
+		fhSample{t: 104, fh: tsdbutil.GenerateTestFloatHistogram(2)},
+		fhSample{t: 199, fh: tsdbutil.GenerateTestFloatHistogram(3)},
+		hSample{t: 200, h: tsdbutil.GenerateTestHistogram(4)},
+		fhSample{t: 299, fh: tsdbutil.GenerateTestFloatHistogram(5)},
+		fSample{t: 300, f: 11},
 	}), 2)
 
 	require.Equal(t, it.Seek(-123), chunkenc.ValFloat, "seek failed")
-	sampleEq(1, 2)
+	sampleEq(1, 2, nil)
-	prevSampleEq(0, 0, false)
+	prevSampleEq(0, 0, nil, false)
 
 	require.Equal(t, it.Next(), chunkenc.ValFloat, "next failed")
-	sampleEq(2, 3)
+	sampleEq(2, 3, nil)
-	prevSampleEq(1, 2, true)
+	prevSampleEq(1, 2, nil, true)
 
 	require.Equal(t, it.Next(), chunkenc.ValFloat, "next failed")
 	require.Equal(t, it.Next(), chunkenc.ValFloat, "next failed")
 	require.Equal(t, it.Next(), chunkenc.ValFloat, "next failed")
-	sampleEq(5, 6)
+	sampleEq(5, 6, nil)
-	prevSampleEq(4, 5, true)
+	prevSampleEq(4, 5, nil, true)
 
 	require.Equal(t, it.Seek(5), chunkenc.ValFloat, "seek failed")
-	sampleEq(5, 6)
+	sampleEq(5, 6, nil)
-	prevSampleEq(4, 5, true)
+	prevSampleEq(4, 5, nil, true)
 
 	require.Equal(t, it.Seek(101), chunkenc.ValFloat, "seek failed")
-	sampleEq(101, 10)
+	sampleEq(101, 10, nil)
-	prevSampleEq(100, 9, true)
+	prevSampleEq(100, 9, nil, true)
+
+	require.Equal(t, chunkenc.ValHistogram, it.Next(), "next failed")
+	sampleEq(102, 0, tsdbutil.GenerateTestFloatHistogram(0))
+	prevSampleEq(101, 10, nil, true)
+
+	require.Equal(t, chunkenc.ValHistogram, it.Next(), "next failed")
+	sampleEq(103, 0, tsdbutil.GenerateTestFloatHistogram(1))
+	prevSampleEq(102, 0, tsdbutil.GenerateTestFloatHistogram(0), true)
+
+	require.Equal(t, chunkenc.ValFloatHistogram, it.Seek(104), "seek failed")
+	sampleEq(104, 0, tsdbutil.GenerateTestFloatHistogram(2))
+	prevSampleEq(103, 0, tsdbutil.GenerateTestFloatHistogram(1), true)
+
+	require.Equal(t, chunkenc.ValFloatHistogram, it.Next(), "next failed")
+	sampleEq(199, 0, tsdbutil.GenerateTestFloatHistogram(3))
+	prevSampleEq(104, 0, tsdbutil.GenerateTestFloatHistogram(2), true)
+
+	require.Equal(t, chunkenc.ValFloatHistogram, it.Seek(280), "seek failed")
+	sampleEq(299, 0, tsdbutil.GenerateTestFloatHistogram(5))
+	prevSampleEq(0, 0, nil, false)
+
+	require.Equal(t, chunkenc.ValFloat, it.Next(), "next failed")
+	sampleEq(300, 11, nil)
+	prevSampleEq(299, 0, tsdbutil.GenerateTestFloatHistogram(5), true)
 
 	require.Equal(t, it.Next(), chunkenc.ValNone, "next succeeded unexpectedly")
 	require.Equal(t, it.Seek(1024), chunkenc.ValNone, "seek succeeded unexpectedly")