Merge pull request #12243 from leizor/leizor/prometheus/issues/11274

Optimize and test MemoizedSeriesIterator

promql/engine.go

@@ -1850,14 +1850,14 @@ func (ev *evaluator) vectorSelectorSingle(it *storage.MemoizedSeriesIterator, no
 		}
 	case chunkenc.ValFloat:
 		t, v = it.At()
-	case chunkenc.ValHistogram, chunkenc.ValFloatHistogram:
+	case chunkenc.ValFloatHistogram:
 		t, h = it.AtFloatHistogram()
 	default:
 		panic(fmt.Errorf("unknown value type %v", valueType))
 	}
 	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

@@ -21,6 +21,9 @@ import (
 )
 
 // MemoizedSeriesIterator wraps an iterator with a buffer to look back the previous element.
+//
+// This iterator regards integer histograms as float histograms; calls to Seek() will never return chunkenc.Histogram.
+// This iterator deliberately does not implement chunkenc.Iterator.
 type MemoizedSeriesIterator struct {
 	it    chunkenc.Iterator
 	delta int64
@@ -31,12 +34,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 +64,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.
@@ -83,8 +81,11 @@ func (b *MemoizedSeriesIterator) Seek(t int64) chunkenc.ValueType {
 		b.prevTime = math.MinInt64
 
 		b.valueType = b.it.Seek(t0)
-		if b.valueType == chunkenc.ValNone {
+		switch b.valueType {
+		case chunkenc.ValNone:
 			return chunkenc.ValNone
+		case chunkenc.ValHistogram:
+			b.valueType = chunkenc.ValFloatHistogram
 		}
 		b.lastTime = b.it.AtT()
 	}
@@ -100,7 +101,8 @@ func (b *MemoizedSeriesIterator) Seek(t int64) chunkenc.ValueType {
 	return chunkenc.ValNone
 }
 
-// Next advances the iterator to the next element.
+// Next advances the iterator to the next element. Note that this does not check whether the element being buffered is
+// within the time range of the current element and the duration of delta before.
 func (b *MemoizedSeriesIterator) Next() chunkenc.ValueType {
 	// Keep track of the previous element.
 	switch b.valueType {
@@ -108,15 +110,9 @@ 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:
+	case chunkenc.ValHistogram, chunkenc.ValFloatHistogram:
 		b.prevValue = 0
-		b.prevTime, b.prevHistogram = b.it.AtHistogram()
-		_, b.prevFloatHistogram = b.it.AtFloatHistogram()
-	case chunkenc.ValFloatHistogram:
-		b.prevValue = 0
-		b.prevHistogram = nil
 		b.prevTime, b.prevFloatHistogram = b.it.AtFloatHistogram()
 	}
 
@@ -124,6 +120,9 @@ func (b *MemoizedSeriesIterator) Next() chunkenc.ValueType {
 	if b.valueType != chunkenc.ValNone {
 		b.lastTime = b.it.AtT()
 	}
+	if b.valueType == chunkenc.ValHistogram {
+		b.valueType = chunkenc.ValFloatHistogram
+	}
 	return b.valueType
 }
 
@@ -132,21 +131,11 @@ func (b *MemoizedSeriesIterator) At() (int64, float64) {
 	return b.it.At()
 }
 
-// AtHistogram returns the current histogram element of the iterator.
-func (b *MemoizedSeriesIterator) AtHistogram() (int64, *histogram.Histogram) {
-	return b.it.AtHistogram()
-}
-
 // AtFloatHistogram returns the current float-histogram element of the iterator.
 func (b *MemoizedSeriesIterator) AtFloatHistogram() (int64, *histogram.FloatHistogram) {
 	return b.it.AtFloatHistogram()
 }
 
-// AtT returns the current timestamp of the iterator.
-func (b *MemoizedSeriesIterator) AtT() int64 {
-	return b.it.AtT()
-}
-
 // Err returns the last encountered error.
 func (b *MemoizedSeriesIterator) Err() error {
 	return b.it.Err()

storage/memoized_iterator_test.go

@@ -18,23 +18,34 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/tsdb/chunkenc"
+	"github.com/prometheus/prometheus/tsdb/tsdbutil"
 )
 
 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,31 +57,50 @@ 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},
+		hSample{t: 399, h: tsdbutil.GenerateTestHistogram(6)},
+		fSample{t: 400, f: 12},
 	}), 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)
-	prevSampleEq(1, 2, 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)
-	prevSampleEq(4, 5, 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")
+	// Seek to a histogram sample with a previous float sample.
-	sampleEq(101, 10)
+	require.Equal(t, it.Seek(102), chunkenc.ValFloatHistogram, "seek failed")
-	prevSampleEq(100, 9, true)
+	sampleEq(102, 10, tsdbutil.GenerateTestFloatHistogram(0))
+	prevSampleEq(101, 10, nil, true)
+
+	// Attempt to seek backwards (no-op).
+	require.Equal(t, it.Seek(50), chunkenc.ValFloatHistogram, "seek failed")
+	sampleEq(102, 10, tsdbutil.GenerateTestFloatHistogram(0))
+	prevSampleEq(101, 10, nil, true)
+
+	// Seek to a float histogram sample with a previous histogram sample.
+	require.Equal(t, it.Seek(104), chunkenc.ValFloatHistogram, "seek failed")
+	sampleEq(104, 0, tsdbutil.GenerateTestFloatHistogram(2))
+	prevSampleEq(103, 0, tsdbutil.GenerateTestFloatHistogram(1), true)
+
+	// Seek to a float sample with a previous float histogram sample.
+	require.Equal(t, chunkenc.ValFloat, it.Seek(300), "seek failed")
+	sampleEq(300, 11, nil)
+	prevSampleEq(299, 0, tsdbutil.GenerateTestFloatHistogram(5), true)
+
+	// Seek to a float sample with a previous histogram sample.
+	require.Equal(t, chunkenc.ValFloat, it.Seek(400), "seek failed")
+	sampleEq(400, 12, nil)
+	prevSampleEq(399, 0, tsdbutil.GenerateTestFloatHistogram(6), true)
 
-	require.Equal(t, it.Next(), chunkenc.ValNone, "next succeeded unexpectedly")
 	require.Equal(t, it.Seek(1024), chunkenc.ValNone, "seek succeeded unexpectedly")
 }