storage,logicalplan: lazy iterator allocation with batching for range…

engine/bench_test.go

@@ -736,7 +736,8 @@ func BenchmarkInstantQuery(b *testing.B) {
 		b.Run(tc.name, func(b *testing.B) {
 			b.Run("new_engine", func(b *testing.B) {
 				ng := engine.New(engine.Opts{
-					EngineOpts: promql.EngineOpts{Timeout: 100 * time.Second},
+					EngineOpts:        promql.EngineOpts{Timeout: 100 * time.Second},
+					SelectorBatchSize: 256,
 				})
 				b.ResetTimer()
 				b.ReportAllocs()
@@ -777,7 +778,8 @@ func BenchmarkInstantQuery(b *testing.B) {
 			})
 			b.Run("new_engine", func(b *testing.B) {
 				ng := engine.New(engine.Opts{
-					EngineOpts: promql.EngineOpts{Timeout: 100 * time.Second},
+					EngineOpts:        promql.EngineOpts{Timeout: 100 * time.Second},
+					SelectorBatchSize: 256,
 				})
 				b.ResetTimer()
 				b.ReportAllocs()

engine/enginefuzz_test.go

@@ -177,7 +177,7 @@ func FuzzEnginePromQLSmithRangeQuery(f *testing.F) {
 		}
 		ps := promqlsmith.New(rnd, seriesSet, psOpts...)
 
-		newEngine := engine.New(engine.Opts{EngineOpts: opts, EnableAnalysis: true})
+		newEngine := engine.New(engine.Opts{EngineOpts: opts, EnableAnalysis: true, SelectorBatchSize: 256})
 		oldEngine := promql.NewEngine(opts)
 
 		var (
@@ -263,6 +263,7 @@ func FuzzEnginePromQLSmithInstantQuery(f *testing.F) {
 			EngineOpts:        opts,
 			LogicalOptimizers: logicalplan.AllOptimizers,
 			EnableAnalysis:    true,
+			SelectorBatchSize: 256,
 		})
 		oldEngine := promql.NewEngine(opts)
 
@@ -481,7 +482,7 @@ func FuzzNativeHistogramQuery(f *testing.F) {
 
 		rnd := rand.New(rand.NewSource(seed))
 		ps := promqlsmith.New(rnd, seriesSet, psOpts...)
-		newEngine := engine.New(engine.Opts{EngineOpts: opts, EnableAnalysis: true})
+		newEngine := engine.New(engine.Opts{EngineOpts: opts, EnableAnalysis: true, SelectorBatchSize: 256})
 		oldEngine := promql.NewEngine(opts)
 
 		instantCases := make([]*testCase, 0, testRuns/2)
@@ -696,7 +697,7 @@ func FuzzDistributedEngineQuery(f *testing.F) {
 		end := time.Unix(int64(endTS), 0)
 		interval := time.Duration(intervalSeconds) * time.Second
 
-		engineOpts := engine.Opts{EngineOpts: opts}
+		engineOpts := engine.Opts{EngineOpts: opts, SelectorBatchSize: 256}
 
 		// Create remote engines for each partition
 		remoteEngine1 := engine.NewRemoteEngine(

logicalplan/set_batch_size.go

@@ -17,32 +17,67 @@ type SelectorBatchSize struct {
 }
 
 // Optimize configures the batch size of selector based on the query plan.
-// If any aggregate is present in the plan, the batch size is set to the configured value.
-// The two exceptions where this cannot be done is if the aggregate is quantile, or
-// when a binary expression precedes the aggregate.
+// It recursively traverses the plan and sets BatchSize on VectorSelectors
+// that can safely produce results in batches. Batching is enabled when an
+// Aggregation is encountered and propagated down through nodes that are
+// per-series independent (range functions, set operations, scalar-vector
+// arithmetic). Batching is disabled by nodes that require cross-series
+// visibility (group_left/group_right, histogram_quantile, topk, etc.).
 func (m SelectorBatchSize) Optimize(plan Node, _ *query.Options) (Node, annotations.Annotations) {
-	canBatch := false
-	Traverse(&plan, func(current *Node) {
-		switch e := (*current).(type) {
-		case *FunctionCall:
-			//TODO: calls can reduce the labelset of the input; think histogram_quantile reducing
-			// multiple "le" labels into one output. We cannot handle this in batching. Revisit
-			// what is safe here.
-			canBatch = false
-		case *Binary:
+	m.setBatchSize(&plan, false)
+	return plan, nil
+}
+
+func (m SelectorBatchSize) setBatchSize(node *Node, canBatch bool) {
+	switch e := (*node).(type) {
+	case *Aggregation:
+		if e.Op == parser.QUANTILE || e.Op == parser.TOPK || e.Op == parser.BOTTOMK || e.Op == parser.LIMITK || e.Op == parser.LIMIT_RATIO {
 			canBatch = false
-		case *Aggregation:
-			if e.Op == parser.QUANTILE || e.Op == parser.TOPK || e.Op == parser.BOTTOMK || e.Op == parser.LIMITK || e.Op == parser.LIMIT_RATIO {
-				canBatch = false
-				return
-			}
+		} else {
 			canBatch = true
-		case *VectorSelector:
-			if canBatch {
-				e.BatchSize = m.Size
-			}
+		}
+	case *FunctionCall:
+		// Range vector functions (rate, increase, present_over_time, etc.) are safe for
+		// batching because each output series depends only on that series' own range data.
+		// The function receives a single series' samples over the [range] window and produces
+		// one output value — no cross-series state is needed.
+		//
+		// Instant vector functions (e.g., histogram_quantile) may reduce or combine series
+		// (e.g., merging multiple "le" labels into one output), requiring full series visibility.
+		// We disable batching for those.
+		if !isRangeVectorFunction(e) {
 			canBatch = false
 		}
-	})
-	return plan, nil
+	case *Binary:
+		if isManyToOneOrOneToMany(e) {
+			canBatch = false
+		}
+	case *VectorSelector:
+		if canBatch {
+			e.BatchSize = m.Size
+		}
+		return
+	}
+
+	for _, child := range (*node).Children() {
+		m.setBatchSize(child, canBatch)
+	}
+}
+
+func isRangeVectorFunction(f *FunctionCall) bool {
+	for _, arg := range f.Args {
+		switch arg.(type) {
+		case *MatrixSelector, *Subquery:
+			return true
+		}
+	}
+	return false
+}
+
+// isManyToOneOrOneToMany returns true for binary operations with group_left/group_right
+// matching, which require full series visibility for correct join behavior.
+// All other binary operations (set ops, one-to-one, scalar-vector) are safe with batching.
+func isManyToOneOrOneToMany(b *Binary) bool {
+	return b.VectorMatching != nil &&
+		(b.VectorMatching.Card == parser.CardManyToOne || b.VectorMatching.Card == parser.CardOneToMany)
 }

logicalplan/set_batch_size_test.go

@@ -46,7 +46,7 @@ func TestSetBatchSize(t *testing.T) {
 		{
 			name:     "aggregation of binary expression",
 			expr:     `max by (pod) (metric_a / metric_b)`,
-			expected: `max by (pod) (metric_a / metric_b)`,
+			expected: `max by (pod) (metric_a[batch=10] / metric_b[batch=10])`,
 		},
 		{
 			name:     "binary operation of aggregations",
@@ -88,6 +88,46 @@ func TestSetBatchSize(t *testing.T) {
 			expr:     `histogram_quantile(scalar(max(quantile)), http_requests_total)`,
 			expected: `histogram_quantile(scalar(max(quantile[batch=10])), http_requests_total)`,
 		},
+		{
+			name:     "aggregation of range vector function (rate)",
+			expr:     `sum(rate(http_requests_total[5m]))`,
+			expected: `sum(rate(http_requests_total[batch=10][5m0s]))`,
+		},
+		{
+			name:     "aggregation of range vector function (increase)",
+			expr:     `sum(increase(http_requests_total[5m]))`,
+			expected: `sum(increase(http_requests_total[batch=10][5m0s]))`,
+		},
+		{
+			name:     "aggregation of range vector function (present_over_time)",
+			expr:     `sum(present_over_time(http_requests_total[1h]))`,
+			expected: `sum(present_over_time(http_requests_total[batch=10][1h0m0s]))`,
+		},
+		{
+			name:     "nested aggregation with range vector function",
+			expr:     `max by (pod) (sum by (pod) (rate(http_requests_total[5m])))`,
+			expected: `max by (pod) (sum by (pod) (rate(http_requests_total[batch=10][5m0s])))`,
+		},
+		{
+			name:     "histogram_quantile does not allow batching",
+			expr:     `sum(histogram_quantile(0.9, rate(http_requests_total[5m])))`,
+			expected: `sum(histogram_quantile(0.9, rate(http_requests_total[5m0s])))`,
+		},
+		{
+			name:     "or operator allows batching",
+			expr:     `sum(increase(http_requests_total[5m]) or present_over_time(http_requests_total[1h]))`,
+			expected: `sum(increase(http_requests_total[batch=10][5m0s]) or present_over_time(http_requests_total[batch=10][1h0m0s]))`,
+		},
+		{
+			name:     "and operator allows batching",
+			expr:     `sum(rate(http_requests_total[5m])) and sum(rate(http_requests_total[5m]))`,
+			expected: `sum(rate(http_requests_total[batch=10][5m0s])) and sum(rate(http_requests_total[batch=10][5m0s]))`,
+		},
+		{
+			name:     "group_left disables batching",
+			expr:     `max by (pod) (metric_a * on (pod) group_left (namespace) metric_b)`,
+			expected: `max by (pod) (metric_a * on (pod) group_left (namespace) metric_b)`,
+		},
 	}
 
 	optimizers := append([]Optimizer{SelectorBatchSize{Size: 10}}, DefaultOptimizers...)

storage/prometheus/matrix_selector.go

@@ -34,6 +34,7 @@ type matrixScanner struct {
 
 	buffer     ringbuffer.Buffer
 	iterator   chunkenc.Iterator
+	rawSeries  SignedSeries
 	lastSample ringbuffer.Sample
 }
 
@@ -171,6 +172,17 @@ func (o *matrixSelector) Next(ctx context.Context, buf []model.StepVector) (int,
 	ts = o.currentStep
 	firstSeries := o.currentSeries
 	batchSamplesDelta := 0
+
+	lastInBatch := min(firstSeries+o.seriesBatchSize, int64(len(o.scanners)))
+	// Initialize iterators lazily per-batch.
+	// TODO: reuse the iterator created for the previous scanner.
+	for i := firstSeries; i < lastInBatch; i++ {
+		if o.scanners[i].iterator == nil {
+			o.scanners[i].iterator = o.scanners[i].rawSeries.Iterator(nil)
+			o.scanners[i].buffer = o.newBuffer(ctx)
+		}
+	}
+
 	for ; o.currentSeries-firstSeries < o.seriesBatchSize && o.currentSeries < int64(len(o.scanners)); o.currentSeries++ {
 		var (
 			scanner  = &o.scanners[o.currentSeries]
@@ -217,6 +229,7 @@ func (o *matrixSelector) Next(ctx context.Context, buf []model.StepVector) (int,
 
 		if o.opts.IsInstantQuery() {
 			scanner.buffer.Reset(math.MaxInt64, 0)
+			scanner.iterator = nil
 			batchSamplesDelta -= sampleCountAfter
 		}
 
@@ -268,9 +281,8 @@ func (o *matrixSelector) loadSeries(ctx context.Context) error {
 				labels:     lbls,
 				metricName: origLbls.Get(labels.MetricName),
 				signature:  s.Signature,
-				iterator:   s.Iterator(nil),
+				rawSeries:  s,
 				lastSample: ringbuffer.Sample{T: math.MinInt64},
-				buffer:     o.newBuffer(ctx),
 			}
 			o.series[i] = lbls
 		}