feat: improve performance of first_over_time and last_over_time queries by sharding them

pkg/logql/downstream.go

@@ -200,6 +200,35 @@ func (e *QuantileSketchMergeExpr) Walk(f syntax.WalkFn) {
 	}
 }
 
+// Step 3
+type MergeFirstOverTimeExpr struct {
+	syntax.SampleExpr
+	downstreams []DownstreamSampleExpr
+}
+
+func (e MergeFirstOverTimeExpr) String() string {
+	var sb strings.Builder
+	for i, d := range e.downstreams {
+		if i >= defaultMaxDepth {
+			break
+		}
+
+		if i > 0 {
+			sb.WriteString(" ++ ")
+		}
+
+		sb.WriteString(d.String())
+	}
+	return fmt.Sprintf("MergeFirstOverTime<%s>", sb.String())
+}
+
+func (e *MergeFirstOverTimeExpr) Walk(f syntax.WalkFn) {
+	f(e)
+	for _, d := range e.downstreams {
+		d.Walk(f)
+	}
+}
+
 type Shards []astmapper.ShardAnnotation
 
 func (xs Shards) Encode() (encoded []string) {
@@ -394,6 +423,43 @@ func (ev *DownstreamEvaluator) NewStepEvaluator(
 		}
 		inner := NewQuantileSketchMatrixStepEvaluator(matrix, params)
 		return NewQuantileSketchVectorStepEvaluator(inner, *e.quantile), nil
+	// Step 2:
+	case *MergeFirstOverTimeExpr:
+		queries := make([]DownstreamQuery, len(e.downstreams))
+
+		for i, d := range e.downstreams {
+			qry := DownstreamQuery{
+				Params: ParamsWithExpressionOverride{
+					Params:             params,
+					ExpressionOverride: d.SampleExpr,
+				},
+			}
+			if shard := d.shard; shard != nil {
+				qry.Params = ParamsWithShardsOverride{
+					Params:         qry.Params,
+					ShardsOverride: Shards{*shard}.Encode(),
+				}
+			}
+			queries[i] = qry
+		}
+
+		results, err := ev.Downstream(ctx, queries)
+		if err != nil {
+			return nil, err
+		}
+
+		xs := make([]promql.Matrix, 0, len(queries))
+		for _, res := range results {
+
+			switch data := res.Data.(type) {
+			case promql.Matrix:
+				xs = append(xs, data)
+			default:
+				return nil, fmt.Errorf("unexpected type (%s) uncoercible to StepEvaluator", data.Type())
+			}
+		}
+
+		return NewMergeFirstOverTimeStepEvaluator(params, xs), nil
 
 	default:
 		return ev.defaultEvaluator.NewStepEvaluator(ctx, nextEvFactory, e, params)

pkg/logql/downstream_test.go

@@ -63,6 +63,9 @@ func TestMappingEquivalence(t *testing.T) {
 			`,
 			false,
 		},
+		// Step 1:
+		{`first_over_time({a=~".+"} | logfmt | unwrap value [1s])`, false},
+		{`first_over_time({a=~".+"} | logfmt | unwrap value [1s]) by (a)`, false},
 		// topk prefers already-seen values in tiebreakers. Since the test data generates
 		// the same log lines for each series & the resulting promql.Vectors aren't deterministically
 		// sorted by labels, we don't expect this to pass.
@@ -132,7 +135,7 @@ func TestMappingEquivalenceSketches(t *testing.T) {
 		query         string
 		realtiveError float64
 	}{
-		{`quantile_over_time(0.70, {a=~".+"} | logfmt | unwrap value [1s]) by (a)`, 0.03},
+		{`quantile_over_time(0.70, {a=~".+"} | logfmt | unwrap value [1s]) by (a)`, 0.05},
 		{`quantile_over_time(0.99, {a=~".+"} | logfmt | unwrap value [1s]) by (a)`, 0.02},
 	} {
 		q := NewMockQuerier(

pkg/logql/engine.go

@@ -351,7 +351,7 @@ func (q *query) evalSample(ctx context.Context, expr syntax.SampleExpr) (promql_
 	}
 	defer util.LogErrorWithContext(ctx, "closing SampleExpr", stepEvaluator.Close)
 
-	next, ts, r := stepEvaluator.Next()
+	next, _, r := stepEvaluator.Next()
 	if stepEvaluator.Error() != nil {
 		return nil, stepEvaluator.Error()
 	}
@@ -361,7 +361,7 @@ func (q *query) evalSample(ctx context.Context, expr syntax.SampleExpr) (promql_
 		case SampleVector:
 			maxSeriesCapture := func(id string) int { return q.limits.MaxQuerySeries(ctx, id) }
 			maxSeries := validation.SmallestPositiveIntPerTenant(tenantIDs, maxSeriesCapture)
-			return q.JoinSampleVector(next, ts, vec, stepEvaluator, maxSeries)
+			return q.JoinSampleVector(next, vec, stepEvaluator, maxSeries)
 		case ProbabilisticQuantileVector:
 			return JoinQuantileSketchVector(next, vec, stepEvaluator, q.params)
 		default:
@@ -371,7 +371,7 @@ func (q *query) evalSample(ctx context.Context, expr syntax.SampleExpr) (promql_
 	return nil, nil
 }
 
-func (q *query) JoinSampleVector(next bool, ts int64, r StepResult, stepEvaluator StepEvaluator, maxSeries int) (promql_parser.Value, error) {
+func (q *query) JoinSampleVector(next bool, r StepResult, stepEvaluator StepEvaluator, maxSeries int) (promql_parser.Value, error) {
 
 	seriesIndex := map[uint64]*promql.Series{}
 
@@ -419,15 +419,16 @@ func (q *query) JoinSampleVector(next bool, ts int64, r StepResult, stepEvaluato
 				seriesIndex[hash] = series
 			}
 			series.Floats = append(series.Floats, promql.FPoint{
-				T: ts,
+				//T: ts,
+				T: p.T,
 				F: p.F,
 			})
 		}
 		// as we slowly build the full query for each steps, make sure we don't go over the limit of unique series.
 		if len(seriesIndex) > maxSeries {
 			return nil, logqlmodel.NewSeriesLimitError(maxSeries)
 		}
-		next, ts, r = stepEvaluator.Next()
+		next, _, r = stepEvaluator.Next()
 		if stepEvaluator.Error() != nil {
 			return nil, stepEvaluator.Error()
 		}

pkg/logql/evaluator.go

@@ -556,6 +556,18 @@ func newRangeAggEvaluator(
 		return &QuantileSketchStepEvaluator{
 			iter: iter,
 		}, nil
+	// Step 6
+	case syntax.OpRangeTypeFirstWithTimestamp:
+		iter := newFirstWithTimestampIterator(
+			it,
+			expr.Left.Interval.Nanoseconds(),
+			q.Step().Nanoseconds(),
+			q.Start().UnixNano(), q.End().UnixNano(), o.Nanoseconds(),
+		)
+
+		return &RangeVectorEvaluator{
+			iter: iter,
+		}, nil
 	default:
 		iter, err := newRangeVectorIterator(
 			it, expr,

pkg/logql/explain.go

@@ -57,3 +57,16 @@ func (e *BinOpStepEvaluator) Explain(parent Node) {
 func (i *VectorIterator) Explain(parent Node) {
 	parent.Childf("%f vectorIterator", i.val)
 }
+
+func (e *QuantileSketchVectorStepEvaluator) Explain(parent Node) {
+	b := parent.Child("QuantileSketchVector")
+	e.inner.Explain(b)
+}
+
+func (e *firstOverTimeStepEvaluator) Explain(parent Node) {
+	parent.Child("MergeFirstOverTime")
+}
+
+func (EmptyEvaluator) Explain(parent Node) {
+	parent.Child("Empty")
+}

pkg/logql/first_over_time.go

@@ -0,0 +1,191 @@
+package logql
+
+import (
+	"math"
+	"time"
+
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/promql"
+
+	"github.com/grafana/loki/pkg/iter"
+)
+
+func newFirstWithTimestampIterator(
+	it iter.PeekingSampleIterator,
+	selRange, step, start, end, offset int64) RangeVectorIterator {
+	inner := &batchRangeVectorIterator{
+		iter:     it,
+		step:     step,
+		end:      end,
+		selRange: selRange,
+		metrics:  map[string]labels.Labels{},
+		window:   map[string]*promql.Series{},
+		agg:      nil,
+		current:  start - step, // first loop iteration will set it to start
+		offset:   offset,
+	}
+	return &firstWithTimestampBatchRangeVectorIterator{
+		batchRangeVectorIterator: inner,
+	}
+}
+
+type firstWithTimestampBatchRangeVectorIterator struct {
+	*batchRangeVectorIterator
+	at []promql.Sample
+}
+
+// Step 7
+func (r *firstWithTimestampBatchRangeVectorIterator) At() (int64, StepResult) {
+	if r.at == nil {
+		r.at = make([]promql.Sample, 0, len(r.window))
+	}
+	r.at = r.at[:0]
+	// convert ts from nano to milli seconds as the iterator work with nanoseconds
+	ts := r.current/1e+6 + r.offset/1e+6
+	for _, series := range r.window {
+		s := r.agg(series.Floats)
+		r.at = append(r.at, promql.Sample{
+			F:      s.F,
+			T:      s.T / int64(time.Millisecond),
+			Metric: series.Metric,
+		})
+	}
+	return ts, SampleVector(r.at)
+}
+
+func (r *firstWithTimestampBatchRangeVectorIterator) agg(samples []promql.FPoint) promql.FPoint {
+	if len(samples) == 0 {
+		return promql.FPoint{F: math.NaN(), T: 0}
+	}
+	return samples[0]
+}
+
+// Step 8
+type firstOverTimeStepEvaluator struct {
+	start, end, ts time.Time
+	step           time.Duration
+	matrices       []promql.Matrix
+	streamVec      map[int64]int
+}
+
+func NewMergeFirstOverTimeStepEvaluator(params Params, m []promql.Matrix) StepEvaluator {
+	if len(m) == 0 {
+		return EmptyEvaluator{}
+	}
+
+	var (
+		start = params.Start()
+		end   = params.End()
+		step  = params.Step()
+	)
+
+	index := make(map[int64]int, 0)
+	for i, series := range m[1] {
+		index[int64(series.Metric.Hash())] = i
+	}
+
+	return &firstOverTimeStepEvaluator{
+		start:     start,
+		end:       end,
+		ts:        start.Add(-step), // will be corrected on first Next() call
+		step:      step,
+		matrices:  m,
+		streamVec: index,
+	}
+}
+
+func (e *firstOverTimeStepEvaluator) Next() (bool, int64, StepResult) {
+
+	var (
+		vec promql.Vector
+		ok  bool
+	)
+
+	// TODO: build index metric to vec pos
+
+	e.ts = e.ts.Add(e.step)
+	if e.ts.After(e.end) {
+		return false, 0, nil
+	}
+	ts := e.ts.UnixNano() / int64(time.Millisecond)
+
+	// Process first result
+	// len(e.matrices) >= 1 was check during creation
+	for s, series := range e.matrices[0] {
+		if len(series.Floats) == 0 || !e.inRange(series.Floats[0].T, ts) {
+			continue
+		}
+
+		vec = append(vec, promql.Sample{
+			Metric: series.Metric,
+			T:      series.Floats[0].T,
+			F:      series.Floats[0].F,
+		})
+
+		e.pop(0, s)
+	}
+
+	if len(e.matrices) == 1 {
+		return ok, ts, SampleVector(vec)
+	}
+
+	if len(vec) == 0 {
+		return e.hasNext(), ts, SampleVector(vec)
+	}
+
+	// Merge other results
+	for i, m := range e.matrices[1:] {
+		// TODO: verify length and same labels/metric
+		for j, series := range m {
+
+			if len(series.Floats) == 0 || !e.inRange(series.Floats[0].T, ts) {
+				continue
+			}
+
+			// Merge
+			if vec[j].T > series.Floats[0].T {
+				vec[j].F = series.Floats[0].F
+				vec[j].T = series.Floats[0].T
+			}
+
+			// We've omitted the first matrix. That's why +1.
+			e.pop(i+1, j)
+		}
+	}
+
+	// Align vector timestamps with step
+	for i := range vec {
+		vec[i].T = ts
+	}
+
+	return true, ts, SampleVector(vec)
+}
+
+func (e *firstOverTimeStepEvaluator) pop(r, s int) {
+	if len(e.matrices[r][s].Floats) <= 1 {
+		e.matrices[r][s].Floats = nil
+		return
+	}
+	e.matrices[r][s].Floats = e.matrices[r][s].Floats[1:]
+}
+
+func (e *firstOverTimeStepEvaluator) inRange(t, ts int64) bool {
+	previous := ts - e.step.Milliseconds()
+	return previous <= t && t < ts
+}
+
+func (e *firstOverTimeStepEvaluator) hasNext() bool {
+	for _, m := range e.matrices {
+		for _, s := range m {
+			if len(s.Floats) != 0 {
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
+func (*firstOverTimeStepEvaluator) Close() error { return nil }
+
+func (*firstOverTimeStepEvaluator) Error() error { return nil }

pkg/logql/optimize.go

@@ -8,7 +8,7 @@ func optimizeSampleExpr(expr syntax.SampleExpr) (syntax.SampleExpr, error) {
 	// we skip sharding AST for now, it's not easy to clone them since they are not part of the language.
 	expr.Walk(func(e syntax.Expr) {
 		switch e.(type) {
-		case *ConcatSampleExpr, *DownstreamSampleExpr, *QuantileSketchEvalExpr, *QuantileSketchMergeExpr:
+		case *ConcatSampleExpr, *DownstreamSampleExpr, *QuantileSketchEvalExpr, *QuantileSketchMergeExpr, *MergeFirstOverTimeExpr:
 			skip = true
 			return
 		}

pkg/logql/quantile_over_time_sketch.go

@@ -454,8 +454,3 @@ func (e *QuantileSketchVectorStepEvaluator) Next() (bool, int64, StepResult) {
 func (*QuantileSketchVectorStepEvaluator) Close() error { return nil }
 
 func (*QuantileSketchVectorStepEvaluator) Error() error { return nil }
-
-func (e *QuantileSketchVectorStepEvaluator) Explain(parent Node) {
-	b := parent.Child("QuantileSketchVector")
-	e.inner.Explain(b)
-}