From 602eb69edfe87f90e8a1a2844168c946f83aa3f6 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Fri, 5 Apr 2024 11:37:55 +0100
Subject: [PATCH] promql: refactor: extract function nextSample

With sub-function nextValues which we shall use shortly.

Signed-off-by: Bryan Boreham <bjboreham@gmail.com>
---
 promql/engine.go | 43 ++++++++++++++++++++++++++-----------------
 1 file changed, 26 insertions(+), 17 deletions(-)

diff --git a/promql/engine.go b/promql/engine.go
index f915c3480..9e0a6b17c 100644
--- a/promql/engine.go
+++ b/promql/engine.go
@@ -2721,25 +2721,11 @@ func (ev *evaluator) aggregation(e *parser.AggregateExpr, q float64, inputMatrix
 		}
 	}
 
-	for si, series := range inputMatrix {
-		var s Sample
-
-		switch {
-		case len(series.Floats) > 0 && series.Floats[0].T == enh.Ts:
-			s = Sample{Metric: series.Metric, F: series.Floats[0].F, T: enh.Ts}
-			// Move input vectors forward so we don't have to re-scan the same
-			// past points at the next step.
-			inputMatrix[si].Floats = series.Floats[1:]
-		case len(series.Histograms) > 0 && series.Histograms[0].T == enh.Ts:
-			s = Sample{Metric: series.Metric, H: series.Histograms[0].H, T: enh.Ts}
-			inputMatrix[si].Histograms = series.Histograms[1:]
-		default:
+	for si := range inputMatrix {
+		s, ok := ev.nextSample(enh.Ts, inputMatrix, si)
+		if !ok {
 			continue
 		}
-		ev.currentSamples++
-		if ev.currentSamples > ev.maxSamples {
-			ev.error(ErrTooManySamples(env))
-		}
 
 		group := orderedResult[seriesToResult[si]]
 		// Initialize this group if it's the first time we've seen it.
@@ -3058,6 +3044,29 @@ func addToSeries(ss *Series, ts int64, f float64, h *histogram.FloatHistogram, n
 	}
 }
 
+func (ev *evaluator) nextValues(ts int64, series *Series) (f float64, h *histogram.FloatHistogram, b bool) {
+	switch {
+	case len(series.Floats) > 0 && series.Floats[0].T == ts:
+		f = series.Floats[0].F
+		series.Floats = series.Floats[1:] // Move input vectors forward
+	case len(series.Histograms) > 0 && series.Histograms[0].T == ts:
+		h = series.Histograms[0].H
+		series.Histograms = series.Histograms[1:]
+	default:
+		return f, h, false
+	}
+	return f, h, true
+}
+
+func (ev *evaluator) nextSample(ts int64, inputMatrix Matrix, si int) (Sample, bool) {
+	f, h, ok := ev.nextValues(ts, &inputMatrix[si])
+	ev.currentSamples++
+	if ev.currentSamples > ev.maxSamples {
+		ev.error(ErrTooManySamples(env))
+	}
+	return Sample{Metric: inputMatrix[si].Metric, F: f, H: h, T: ts}, ok
+}
+
 // groupingKey builds and returns the grouping key for the given metric and
 // grouping labels.
 func generateGroupingKey(metric labels.Labels, grouping []string, without bool, buf []byte) (uint64, []byte) {