From 84c1a8a0e25fa1e07b5baa2caa381633929809cc Mon Sep 17 00:00:00 2001 From: Paurush Garg Date: Mon, 13 Jul 2026 16:45:55 -0700 Subject: [PATCH] storage,logicalplan: lazy iterator allocation with batching for range functions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two changes that together fix ruler OOM from unbounded iterator allocation: 1. logicalplan: enable series batching for range vector functions The SelectorBatchSize optimizer previously disabled batching for ALL FunctionCall nodes. Range vector functions (rate, increase, present_over_time, etc.) are per-series independent — each series' result depends only on that series' own range data. This change only disables batching for instant vector functions (no matrix arg) while allowing it for range vector functions. 2. storage: defer iterator and buffer allocation to batch time Previously, loadSeries() called s.Iterator(nil) and newBuffer() for ALL series upfront. With high-cardinality tenants (4.5M series), this allocates 4.5M merge iterators simultaneously, causing 22GB+ heap and OOM kills on ruler pods. Move iterator and buffer allocation from loadSeries() to Next(), where they are initialized per-batch. With seriesBatchSize=256, only 256 merge iterators exist at any time instead of the full series count. Together: the optimizer enables batching for the OOM-causing query patterns, and the lazy allocation makes batching actually reduce peak memory from O(N) to O(batchSize). Signed-off-by: Paurush Garg --- engine/bench_test.go | 6 +- engine/enginefuzz_test.go | 7 ++- logicalplan/set_batch_size.go | 81 +++++++++++++++++++-------- logicalplan/set_batch_size_test.go | 42 +++++++++++++- storage/prometheus/matrix_selector.go | 16 +++++- 5 files changed, 121 insertions(+), 31 deletions(-) diff --git a/engine/bench_test.go b/engine/bench_test.go index 36306f6a..9b7e91a4 100644 --- a/engine/bench_test.go +++ b/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() diff --git a/engine/enginefuzz_test.go b/engine/enginefuzz_test.go index 3557612d..7add461f 100644 --- a/engine/enginefuzz_test.go +++ b/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( diff --git a/logicalplan/set_batch_size.go b/logicalplan/set_batch_size.go index ffc666f7..39a7e87f 100644 --- a/logicalplan/set_batch_size.go +++ b/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) } diff --git a/logicalplan/set_batch_size_test.go b/logicalplan/set_batch_size_test.go index 15733eac..67755d2d 100644 --- a/logicalplan/set_batch_size_test.go +++ b/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...) diff --git a/storage/prometheus/matrix_selector.go b/storage/prometheus/matrix_selector.go index 7f858ff1..967ece95 100644 --- a/storage/prometheus/matrix_selector.go +++ b/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 }