Add unit test for chunkIterator.Seek() when the current batch can s…

…atisfy the sought timestamp (#9040)

* Fix typo in comment

* Add test for seeking before the current batch of a `chunkIterator`

* Clarify documentation

* Address PR feedback: clarify comments

pkg/querier/batch/batch.go

@@ -38,7 +38,10 @@ func (c GenericChunk) Iterator(reuse chunk.Iterator) chunk.Iterator {
 
 // iterator iterates over batches.
 type iterator interface {
-	// Seek to the batch with sample at (or after) time t.
+	// Seek advances the iterator forward to batch containing the sample at or after the given timestamp.
+	// If the current batch contains a sample at or after the given timestamp, then Seek retains the current batch.
+	//
+	// The batch's Index is advanced to point to the sample at or after the given timestamp.
 	Seek(t int64, size int) chunkenc.ValueType
 
 	// Next moves to the next batch.

pkg/querier/batch/chunk.go

@@ -32,8 +32,6 @@ func (i *chunkIterator) reset(chunk GenericChunk) {
 	i.batch.Index = 0
 }
 
-// Seek advances the iterator forward to the value at or after
-// the given timestamp.
 func (i *chunkIterator) Seek(t int64, size int) chunkenc.ValueType {
 	// We assume seeks only care about a specific window; if this chunk doesn't
 	// contain samples in that window, we can shortcut.

pkg/querier/batch/chunk_test.go

@@ -222,3 +222,40 @@ func (i *mockIterator) Batch(_ int, valueType chunkenc.ValueType, _ *zeropool.Po
 func (i *mockIterator) Err() error {
 	return nil
 }
+
+func TestChunkIterator_SeekBeforeCurrentBatch(t *testing.T) {
+	chunkTimestamps := []int64{50, 60, 70, 80, 90, 100}
+
+	ch := chunkenc.NewXORChunk()
+	app, err := ch.Appender()
+	require.NoError(t, err)
+	for _, ts := range chunkTimestamps {
+		app.Append(ts, float64(ts))
+	}
+
+	genericChunk := NewGenericChunk(chunkTimestamps[0], chunkTimestamps[len(chunkTimestamps)-1], func(reuse chunk.Iterator) chunk.Iterator {
+		chk, err := chunk.NewForEncoding(chunk.PrometheusXorChunk)
+		require.NoError(t, err)
+		require.NoError(t, chk.UnmarshalFromBuf(ch.Bytes()))
+
+		// We should never need to reset this iterator, as the Seek call below should be satisfiable by the initial batch.
+		return &chunkIteratorThatForbidsFindAtOrAfter{chk.NewIterator(reuse)}
+	})
+
+	it := &chunkIterator{}
+	it.reset(genericChunk)
+
+	require.Equal(t, chunkenc.ValFloat, it.Next(2))
+	require.Equal(t, int64(50), it.AtTime())
+
+	require.Equal(t, chunkenc.ValFloat, it.Seek(45, 1))
+	require.Equal(t, int64(50), it.AtTime())
+}
+
+type chunkIteratorThatForbidsFindAtOrAfter struct {
+	chunk.Iterator
+}
+
+func (it *chunkIteratorThatForbidsFindAtOrAfter) FindAtOrAfter(model.Time) chunkenc.ValueType {
+	panic("FindAtOrAfter should never be called")
+}

pkg/querier/batch/stream.go

@@ -13,7 +13,7 @@ import (
 	"github.com/grafana/mimir/pkg/storage/chunk"
 )
 
-// batchStream deals with iteratoring through multiple, non-overlapping batches,
+// batchStream deals with iterating through multiple, non-overlapping batches,
 // and building new slices of non-overlapping batches.  Designed to be used
 // without allocations.
 type batchStream struct {