ringbuffer: more over-time aligners

Author: MichaHoffmann

ringbuffer/overtime.go

@@ -8,27 +8,33 @@ import (
 	"math"
 
 	"github.com/thanos-io/promql-engine/execution/telemetry"
+	"github.com/thanos-io/promql-engine/execution/warnings"
 	"github.com/thanos-io/promql-engine/query"
 
 	"github.com/prometheus/prometheus/model/histogram"
+	"github.com/prometheus/prometheus/promql/parser/posrange"
+	"github.com/prometheus/prometheus/util/annotations"
 )
 
-// TODO: this is mostly copied over from the rate aligner and not at all special to
-// count_over_time - we could of course generalize it to other over_time functions and we
-// will do that soon.
-// This is really just to get warmed up.
-// The idea here is that we tile the range into "stepRanges" and when we push a sample
-// and it happens to fall into a step range, we increment the state that we hold in a
-// corresponding "resultSamples" slice.
+// TODO: ... we should reuse execution/aggregation/accumulator here.. this is a bit janky
+// but it helps already
+type aggregator func(cur *stepState, next Value) bool
 
-// CountOverTimeBuffer is a Buffer which can calculate count_over_time for a series in a
+type stepState struct {
+	sample      Value
+	initialized bool
+	warning     error
+}
+
+// OverTimeBuffer is a Buffer which can calculate [agg]_over_time for a series in a
 // streaming manner, calculating the value incrementally for each step where the sample is used.
-type CountOverTimeBuffer struct {
+type OverTimeBuffer struct {
 	// stepRanges contain the bounds and number of samples for each evaluation step.
 	stepRanges []stepRange
-
-	// resultscounts contains the resulting state for each evaluation step.
-	resultCounts []float64
+	// stepStates contains the aggregation state for the corresponding stepRange
+	stepStates []stepState
+	// aggregator aggregates the current stepState with the new value
+	aggregator aggregator
 
 	// firstTimestamps contains the timestamp of the first sample for each evaluation step.
 	firstTimestamps []int64
@@ -39,8 +45,7 @@ type CountOverTimeBuffer struct {
 	step int64
 }
 
-// NewCountOverTime creates a new CountOverTimeBuffer.
-func NewCountOverTimeBuffer(opts query.Options, selectRange, offset int64) *CountOverTimeBuffer {
+func newOverTimeBuffer(opts query.Options, selectRange, offset int64, agg aggregator) *OverTimeBuffer {
 	var (
 		step     = max(1, opts.Step.Milliseconds())
 		numSteps = min(
@@ -49,38 +54,93 @@ func NewCountOverTimeBuffer(opts query.Options, selectRange, offset int64) *Coun
 		)
 
 		current         = opts.Start.UnixMilli()
-		resultCounts    = make([]float64, 0, numSteps)
 		firstTimestamps = make([]int64, 0, numSteps)
 		stepRanges      = make([]stepRange, 0, numSteps)
+		stepStates      = make([]stepState, 0, numSteps)
 	)
 	for range int(numSteps) {
 		var (
 			maxt = current - offset
 			mint = maxt - selectRange
 		)
 		stepRanges = append(stepRanges, stepRange{mint: mint, maxt: maxt})
-		resultCounts = append(resultCounts, 0.)
+		stepStates = append(stepStates, stepState{})
 		firstTimestamps = append(firstTimestamps, math.MaxInt64)
 
 		current += step
 	}
 
-	return &CountOverTimeBuffer{
+	return &OverTimeBuffer{
+		aggregator:      agg,
 		step:            step,
 		stepRanges:      stepRanges,
-		resultCounts:    resultCounts,
+		stepStates:      stepStates,
 		firstTimestamps: firstTimestamps,
 		lastTimestamp:   math.MinInt64,
 	}
 }
 
-func (r *CountOverTimeBuffer) SampleCount() int {
+// NewCountOverTimeBuffer creates a new OverTimeBuffer for the count_over_time function.
+func NewCountOverTimeBuffer(opts query.Options, selectRange, offset int64) *OverTimeBuffer {
+	return newOverTimeBuffer(opts, selectRange, offset, func(cur *stepState, _ Value) bool {
+		cur.sample.F += 1
+		return true
+	})
+}
+
+// NewMaxOverTimeBuffer creates a new OverTimeBuffer for the max_over_time function.
+func NewMaxOverTimeBuffer(opts query.Options, selectRange, offset int64) *OverTimeBuffer {
+	return newOverTimeBuffer(opts, selectRange, offset, func(cur *stepState, next Value) bool {
+		if next.H != nil {
+			if cur.initialized {
+				cur.warning = annotations.NewHistogramIgnoredInMixedRangeInfo("", posrange.PositionRange{})
+			}
+			return false
+		}
+		if !cur.initialized {
+			cur.sample.F = next.F
+			cur.initialized = true
+		} else if cur.sample.F < next.F || math.IsNaN(cur.sample.F) {
+			cur.sample.F = next.F
+		}
+		return true
+	})
+}
+
+// NewMinOverTime creates a new OverTimeBuffer for the min_over_time function.
+func NewMinOverTimeBuffer(opts query.Options, selectRange, offset int64) *OverTimeBuffer {
+	return newOverTimeBuffer(opts, selectRange, offset, func(cur *stepState, next Value) bool {
+		if next.H != nil {
+			if cur.initialized {
+				cur.warning = annotations.NewHistogramIgnoredInMixedRangeInfo("", posrange.PositionRange{})
+			}
+			return false
+		}
+		if !cur.initialized {
+			cur.sample.F = next.F
+			cur.initialized = true
+		} else if cur.sample.F > next.F || math.IsNaN(cur.sample.F) {
+			cur.sample.F = next.F
+		}
+		return true
+	})
+}
+
+// NewLastOverTime creates a new OverTimeBuffer for the last_over_time function.
+func NewLastOverTimeBuffer(opts query.Options, selectRange, offset int64) *OverTimeBuffer {
+	return newOverTimeBuffer(opts, selectRange, offset, func(cur *stepState, next Value) bool {
+		cur.sample = next
+		return true
+	})
+}
+
+func (r *OverTimeBuffer) SampleCount() int {
 	return r.stepRanges[0].sampleCount
 }
 
-func (r *CountOverTimeBuffer) MaxT() int64 { return r.lastTimestamp }
+func (r *OverTimeBuffer) MaxT() int64 { return r.lastTimestamp }
 
-func (r *CountOverTimeBuffer) Push(t int64, v Value) {
+func (r *OverTimeBuffer) Push(t int64, v Value) {
 	// Set the lastSample sample for the current evaluation step.
 	r.lastTimestamp = t
 
@@ -93,8 +153,10 @@ func (r *CountOverTimeBuffer) Push(t int64, v Value) {
 			r.stepRanges[i].sampleCount++
 		}
 
-		// Count the current sample in its range
-		r.resultCounts[i] += 1
+		// Aggregate the sample to the current step
+		if ok := r.aggregator(&(r.stepStates[i]), v); !ok {
+			continue
+		}
 
 		if fts := r.firstTimestamps[i]; t >= fts {
 			continue
@@ -103,7 +165,7 @@ func (r *CountOverTimeBuffer) Push(t int64, v Value) {
 	}
 }
 
-func (r *CountOverTimeBuffer) Reset(mint int64, evalt int64) {
+func (r *OverTimeBuffer) Reset(mint int64, evalt int64) {
 	if r.stepRanges[0].mint == mint {
 		return
 	}
@@ -113,22 +175,36 @@ func (r *CountOverTimeBuffer) Reset(mint int64, evalt int64) {
 		nextMint = r.stepRanges[lastSample].mint + r.step
 		nextMaxt = r.stepRanges[lastSample].maxt + r.step
 	)
+	nextStepRange := r.stepRanges[0]
 	copy(r.stepRanges, r.stepRanges[1:])
-	r.stepRanges[lastSample] = stepRange{mint: nextMint, maxt: nextMaxt}
-
-	copy(r.resultCounts, r.resultCounts[1:])
-	r.resultCounts[lastSample] = 0
+	r.stepRanges[lastSample] = nextStepRange
+	r.stepRanges[lastSample].mint = nextMint
+	r.stepRanges[lastSample].maxt = nextMaxt
+	r.stepRanges[lastSample].sampleCount = 0
+	r.stepRanges[lastSample].numSamples = 0
+
+	nextFirstState := r.stepStates[0]
+	copy(r.stepStates, r.stepStates[1:])
+	r.stepStates[lastSample] = nextFirstState
+	r.stepStates[lastSample].sample.H = nil
+	r.stepStates[lastSample].sample.F = 0
+	r.stepStates[lastSample].initialized = false
+	r.stepStates[lastSample].warning = nil
 
 	nextFirstTimestamp := r.firstTimestamps[0]
 	copy(r.firstTimestamps, r.firstTimestamps[1:])
 	r.firstTimestamps[lastSample] = nextFirstTimestamp
 	r.firstTimestamps[lastSample] = math.MaxInt64
 }
 
-func (r *CountOverTimeBuffer) Eval(ctx context.Context, _, _ float64, _ int64) (float64, *histogram.FloatHistogram, bool, error) {
+func (r *OverTimeBuffer) Eval(ctx context.Context, _, _ float64, _ int64) (float64, *histogram.FloatHistogram, bool, error) {
+	if r.stepStates[0].warning != nil {
+		warnings.AddToContext(r.stepStates[0].warning, ctx)
+	}
+
 	if r.firstTimestamps[0] == math.MaxInt64 {
 		return 0, nil, false, nil
 	}
 
-	return r.resultCounts[0], nil, true, nil
+	return r.stepStates[0].sample.F, r.stepStates[0].sample.H, true, nil
 }

storage/prometheus/matrix_selector.go

@@ -280,6 +280,12 @@ func (o *matrixSelector) newBuffer(ctx context.Context) ringbuffer.Buffer {
 		return ringbuffer.NewRateBuffer(ctx, *o.opts, false, false, o.selectRange, o.offset)
 	case "count_over_time":
 		return ringbuffer.NewCountOverTimeBuffer(*o.opts, o.selectRange, o.offset)
+	case "max_over_time":
+		return ringbuffer.NewMaxOverTimeBuffer(*o.opts, o.selectRange, o.offset)
+	case "min_over_time":
+		return ringbuffer.NewMinOverTimeBuffer(*o.opts, o.selectRange, o.offset)
+	case "last_over_time":
+		return ringbuffer.NewLastOverTimeBuffer(*o.opts, o.selectRange, o.offset)
 	}
 
 	if o.isExtFunction {