Fix correctness issue in predict_linear with step invariant

engine/engine_test.go

@@ -245,6 +245,15 @@ func TestQueriesAgainstOldEngine(t *testing.T) {
 			    http_requests_total{pod="nginx-2", route="/"} -108+173.00x40`,
 			query: `predict_linear(http_requests_total{route="/"}[1h:1m] offset 1m, 60)`,
 		},
+		{
+			name: "predict_linear with step invariant",
+			load: `load 30s
+			    http_requests_total{pod="nginx-1"} 1+1.1x10
+			    http_requests_total{pod="nginx-2"} 2+2.3x50`,
+			query: `predict_linear({__name__="http_requests_total",pod!~"nginx-1"}[5m] @ start(), -0.37690610678629094)`,
+			end:   time.Unix(600, 0),
+			start: time.Unix(300, 0),
+		},
 		{
 			name: "duplicate label fuzz",
 			load: `load 30s

logicalplan/logical_nodes.go

@@ -115,6 +115,7 @@ func (f *VectorSelector) ReturnType() parser.ValueType { return parser.ValueType
 type MatrixSelector struct {
 	VectorSelector *VectorSelector `json:"-"`
 	Range          time.Duration
+	Timestamp      *int64
 
 	// Needed because this operator is used in the distributed mode
 	OriginalString string

logicalplan/plan.go

@@ -248,7 +248,13 @@ func replacePrometheusNodes(plan parser.Expr) Node {
 		// parents of matrixselector nodes are always expected to be functions, not step invariant
 		// operators.
 		if m, ok := t.Expr.(*parser.MatrixSelector); ok {
-			return replacePrometheusNodes(m)
+			n := replacePrometheusNodes(m)
+			if ms, ok := n.(*MatrixSelector); ok {
+				// For cases such as predict_linear(metric[5m] @start() , 0.3), we mark the
+				// matrix selector to be invariant. The scanner will only scan points once in such cases.
+				ms.Timestamp = ms.VectorSelector.Timestamp
+			}
+			return n
 		}
 		return &StepInvariantExpr{Expr: replacePrometheusNodes(t.Expr)}
 	case *parser.MatrixSelector:

ringbuffer/generic.go

@@ -63,6 +63,10 @@ func (r *GenericRingBuffer) ReadIntoLast(f func(*Sample)) {
 	f(&r.items[len(r.items)-1])
 }
 
+func (r *GenericRingBuffer) SetStep(step int64) {
+	r.currentStep = step
+}
+
 // Push adds a new sample to the buffer.
 func (r *GenericRingBuffer) Push(t int64, v Value) {
 	n := len(r.items)

ringbuffer/rate.go

@@ -18,6 +18,7 @@ type Buffer interface {
 	MaxT() int64
 	Push(t int64, v Value)
 	Reset(mint int64, evalt int64)
+	SetStep(int64)
 	Eval(ctx context.Context, _, _ float64, _ *int64) (float64, *histogram.FloatHistogram, bool, error)
 	ReadIntoLast(f func(*Sample))
 }
@@ -96,6 +97,8 @@ func (r *RateBuffer) Len() int { return r.stepRanges[0].numSamples }
 
 func (r *RateBuffer) MaxT() int64 { return r.last.T }
 
+func (r *RateBuffer) SetStep(step int64) { r.step = step }
+
 func (r *RateBuffer) Push(t int64, v Value) {
 	// Detect resets and store the current and previous sample so that
 	// the rate is properly adjusted.

storage/prometheus/matrix_selector.go

@@ -36,6 +36,8 @@ type matrixScanner struct {
 	iterator         chunkenc.Iterator
 	lastSample       ringbuffer.Sample
 	metricAppearedTs *int64
+	timestamp        *int64
+	once             sync.Once
 }
 
 type matrixSelector struct {
@@ -57,6 +59,7 @@ type matrixSelector struct {
 	numSteps      int
 	mint          int64
 	maxt          int64
+	timestamp     *int64
 	step          int64
 	selectRange   int64
 	offset        int64
@@ -85,6 +88,7 @@ func NewMatrixSelector(
 	functionName string,
 	arg float64,
 	arg2 float64,
+	timestamp *int64,
 	opts *query.Options,
 	selectRange, offset time.Duration,
 	batchSize int64,
@@ -108,6 +112,7 @@ func NewMatrixSelector(
 		mint:          opts.Start.UnixMilli(),
 		maxt:          opts.End.UnixMilli(),
 		step:          opts.Step.Milliseconds(),
+		timestamp:     timestamp,
 		isExtFunction: parse.IsExtFunction(functionName),
 
 		selectRange:     selectRange.Milliseconds(),
@@ -250,6 +255,7 @@ func (o *matrixSelector) loadSeries(ctx context.Context) error {
 				iterator:   s.Iterator(nil),
 				lastSample: ringbuffer.Sample{T: math.MinInt64},
 				buffer:     o.newBuffer(ctx),
+				timestamp:  o.timestamp,
 			}
 			o.series[i] = lbls
 		}
@@ -301,6 +307,23 @@ func (o *matrixSelector) String() string {
 	return fmt.Sprintf("[matrixSelector] {%v}[%s] %v mod %v", o.storage.Matchers(), r, o.shard, o.numShards)
 }
 
+func (m *matrixScanner) selectPoints(
+	mint, maxt, evalt int64,
+	fh *histogram.FloatHistogram,
+	isExtFunction bool,
+) error {
+	if m.timestamp != nil {
+		// If timestamp is set, then the points should be selected only once.
+		var err error
+		m.once.Do(func() {
+			err = m.selectPointsInternal(mint, maxt, *m.timestamp, fh, isExtFunction)
+		})
+		m.buffer.SetStep(evalt)
+		return err
+	}
+	return m.selectPointsInternal(mint, maxt, evalt, fh, isExtFunction)
+}
+
 // matrixIterSlice populates a matrix vector covering the requested range for a
 // single time series, with points retrieved from an iterator.
 //
@@ -310,7 +333,7 @@ func (o *matrixSelector) String() string {
 // into the [mint, maxt] range are retained; only points with later timestamps
 // are populated from the iterator.
 // TODO(fpetkovski): Add max samples limit.
-func (m *matrixScanner) selectPoints(
+func (m *matrixScanner) selectPointsInternal(
 	mint, maxt, evalt int64,
 	fh *histogram.FloatHistogram,
 	isExtFunction bool,

storage/prometheus/scanners.go

@@ -134,6 +134,7 @@ func (p Scanners) NewMatrixSelector(
 			call.Func.Name,
 			arg,
 			arg2,
+			logicalNode.Timestamp,
 			opts,
 			logicalNode.Range,
 			vs.Offset,