Use sync.Pool for mergeIterator batchesBuf

CHANGELOG.md

@@ -53,6 +53,7 @@
 * [ENHANCEMENT] Compactor: Reduce object storage GET calls when updating the bucket index by skipping re-reading parquet converter markers for blocks that already have a valid-version parquet entry in the previous index. #7669
 * [ENHANCEMENT] Upgrade Thanos and promql-engine to latest. #7740
 * [ENHANCEMENT] Ruler: Adjust ruler frontend decoder to not wrap query error messages with execution prefix, this makes error responses consistent between internal and external ruler paths. #7741
+* [ENHANCEMENT] Querier: Use sync.Pool for mergeIterator batchesBuf to reduce memory allocations during series iteration. #7765
 * [BUGFIX] Querier: Fix queryWithRetry and labelsWithRetry returning (nil, nil) on cancelled context by propagating ctx.Err(). #7370
 * [BUGFIX] Metrics Helper: Fix non-deterministic bucket order in merged histograms by sorting buckets after map iteration, matching Prometheus client library behavior. #7380
 * [BUGFIX] Distributor: Return HTTP 401 Unauthorized when tenant ID resolution fails in the Prometheus Remote Write 2.0 path. #7389

pkg/querier/batch/merge.go

@@ -3,23 +3,31 @@ package batch
 import (
 	"container/heap"
 	"sort"
+	"sync"
 
 	"github.com/prometheus/prometheus/tsdb/chunkenc"
 
 	promchunk "github.com/cortexproject/cortex/pkg/chunk"
 )
 
+var batchesBufPool = sync.Pool{
+	New: func() any {
+		buf := make(batchStream, 3)
+		return &buf
+	},
+}
+
 type mergeIterator struct {
 	its []*nonOverlappingIterator
 	h   iteratorHeap
 
 	// Store the current sorted batchStream
 	batches batchStream
 
-	// Buffers to merge in.
-	batchesBuf   batchStream
 	nextBatchBuf [1]promchunk.Batch
 
+	numPartitions int
+
 	currErr error
 }
 
@@ -32,9 +40,9 @@ func newMergeIterator(it iterator, cs []GenericChunk) *mergeIterator {
 		c = mIterator.Reset(len(css))
 	} else {
 		c = &mergeIterator{
-			h:          make(iteratorHeap, 0, len(css)),
-			batches:    make(batchStream, 0, len(css)),
-			batchesBuf: make(batchStream, len(css)),
+			h:             make(iteratorHeap, 0, len(css)),
+			batches:       make(batchStream, 0, len(css)),
+			numPartitions: len(css),
 		}
 	}
 
@@ -65,15 +73,7 @@ func (c *mergeIterator) Reset(size int) *mergeIterator {
 	c.its = c.its[:0]
 	c.h = c.h[:0]
 	c.batches = c.batches[:0]
-
-	if size > cap(c.batchesBuf) {
-		c.batchesBuf = make(batchStream, len(c.its))
-	} else {
-		c.batchesBuf = c.batchesBuf[:size]
-		for i := range size {
-			c.batchesBuf[i] = promchunk.Batch{}
-		}
-	}
+	c.numPartitions = size
 
 	for i := range len(c.nextBatchBuf) {
 		c.nextBatchBuf[i] = promchunk.Batch{}
@@ -141,12 +141,31 @@ func (c *mergeIterator) nextBatchEndTime() int64 {
 }
 
 func (c *mergeIterator) buildNextBatch(size int) chunkenc.ValueType {
+	if len(c.h) == 0 && len(c.batches) > 0 {
+		return c.batches[0].ValType
+	}
+	if len(c.h) == 0 {
+		return chunkenc.ValNone
+	}
+
+	bp := batchesBufPool.Get().(*batchStream)
+	batchesBuf := *bp
+	if cap(batchesBuf) < c.numPartitions {
+		batchesBuf = make(batchStream, c.numPartitions)
+	} else {
+		batchesBuf = batchesBuf[:c.numPartitions]
+		for i := range batchesBuf {
+			batchesBuf[i] = promchunk.Batch{}
+		}
+	}
+	defer batchesBufPool.Put(&batchesBuf)
+
 	// All we need to do is get enough batches that our first batch's last entry
 	// is before all iterators next entry.
 	for len(c.h) > 0 && (len(c.batches) == 0 || c.nextBatchEndTime() >= c.h[0].AtTime()) {
 		c.nextBatchBuf[0] = c.h[0].Batch()
-		c.batchesBuf = mergeStreams(c.batches, c.nextBatchBuf[:], c.batchesBuf, size)
-		c.batches = append(c.batches[:0], c.batchesBuf...)
+		batchesBuf = mergeStreams(c.batches, c.nextBatchBuf[:], batchesBuf, size)
+		c.batches = append(c.batches[:0], batchesBuf...)
 
 		if valType := c.h[0].Next(size); valType != chunkenc.ValNone {
 			heap.Fix(&c.h, 0)