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 }