execution,plan: move timestamp handling to plan (#397)

* move timestamp handling to plan

Signed-off-by: Michael Hoffmann <[email protected]>

Author: MichaHoffmann

engine/distributed_test.go

@@ -244,6 +244,8 @@ func TestDistributedAggregations(t *testing.T) {
 			query:      `max_over_time(min_over_time(sum(bar)[15s:15s])[15s:15s])`,
 		},
 		{name: "subquery over distributed binary expression", query: `max_over_time((bar/bar)[30s:15s])`},
+		{name: "timestamp", query: `timestamp(bar)`},
+		{name: "timestamp - step invariant", query: `timestamp(bar @ 6000)`},
 	}
 
 	lookbackDeltas := []time.Duration{0, 30 * time.Second, 5 * time.Minute}

execution/execution.go

@@ -60,15 +60,12 @@ func newOperator(expr parser.Expr, storage *engstore.SelectorPool, opts *query.O
 	case *parser.NumberLiteral:
 		return scan.NewNumberLiteralSelector(model.NewVectorPool(opts.StepsBatch), opts, e.Val), nil
 	case *logicalplan.VectorSelector:
-		op, err := newVectorSelector(e, storage, opts, hints)
-		if err != nil {
-			return nil, err
-		}
+		op := newVectorSelector(e, storage, opts, hints)
 		if e.SelectTimestamp {
 			// we will drop the __name__ label here, so we need to check for duplicate labels
 			return exchange.NewDuplicateLabelCheck(op, opts), nil
 		}
-		return op, err
+		return op, nil
 	case *parser.Call:
 		op, err := newCall(e, storage, opts, hints)
 		if err != nil {
@@ -110,7 +107,7 @@ func newOperator(expr parser.Expr, storage *engstore.SelectorPool, opts *query.O
 	}
 }
 
-func newVectorSelector(e *logicalplan.VectorSelector, storage *engstore.SelectorPool, opts *query.Options, hints storage.SelectHints) (model.VectorOperator, error) {
+func newVectorSelector(e *logicalplan.VectorSelector, storage *engstore.SelectorPool, opts *query.Options, hints storage.SelectHints) model.VectorOperator {
 	start, end := getTimeRangesForVectorSelector(e, opts, 0)
 	hints.Start = start
 	hints.End = end
@@ -132,7 +129,7 @@ func newVectorSelector(e *logicalplan.VectorSelector, storage *engstore.Selector
 		operators = append(operators, exchange.NewConcurrent(operator, 2, opts))
 	}
 
-	return exchange.NewCoalesce(model.NewVectorPool(opts.StepsBatch), opts, batchsize*int64(numShards), operators...), nil
+	return exchange.NewCoalesce(model.NewVectorPool(opts.StepsBatch), opts, batchsize*int64(numShards), operators...)
 }
 
 func newCall(e *parser.Call, storage *engstore.SelectorPool, opts *query.Options, hints storage.SelectHints) (model.VectorOperator, error) {
@@ -143,28 +140,6 @@ func newCall(e *parser.Call, storage *engstore.SelectorPool, opts *query.Options
 	if e.Func.Name == "absent_over_time" {
 		return newAbsentOverTimeOperator(e, storage, opts, hints)
 	}
-	if e.Func.Name == "timestamp" {
-		switch arg := e.Args[0].(type) {
-		case *logicalplan.VectorSelector:
-			arg.SelectTimestamp = true
-			return newVectorSelector(arg, storage, opts, hints)
-		case *parser.StepInvariantExpr:
-			// Step invariant expressions on vector selectors need to be unwrapped so that we
-			// can return the original timestamp rather than the step invariant one.
-			switch vs := arg.Expr.(type) {
-			case *logicalplan.VectorSelector:
-				// Prometheus weirdness.
-				if vs.Timestamp != nil {
-					vs.OriginalOffset = 0
-				}
-				vs.SelectTimestamp = true
-				return newVectorSelector(vs, storage, opts, hints)
-			}
-			return newInstantVectorFunction(e, storage, opts, hints)
-		}
-		return newInstantVectorFunction(e, storage, opts, hints)
-	}
-
 	// TODO(saswatamcode): Range vector result might need new operator
 	// before it can be non-nested. https://github.com/thanos-io/promql-engine/issues/39
 	for i := range e.Args {

logicalplan/plan.go

@@ -48,8 +48,7 @@ func New(expr parser.Expr, opts *query.Options) Plan {
 	setOffsetForAtModifier(opts.Start.UnixMilli(), expr)
 	setOffsetForInnerSubqueries(expr, opts)
 
-	// * the engine handles sorting at the presentation layer
-	// * parens are just annoying and getting rid of them doesnt change the query
+	// the engine handles sorting at the presentation layer
 	expr = trimSorts(expr)
 
 	// replace scanners by our logical nodes
@@ -68,8 +67,10 @@ func (p *plan) Optimize(optimizers []Optimizer) (Plan, annotations.Annotations)
 		p.expr, a = o.Optimize(p.expr, p.opts)
 		annos.Merge(a)
 	}
+	// parens are just annoying and getting rid of them doesnt change the query
+	expr := trimParens(p.expr)
 
-	return &plan{expr: trimParens(p.expr), opts: p.opts}, *annos
+	return &plan{expr: expr, opts: p.opts}, *annos
 }
 
 func (p *plan) Expr() parser.Expr {
@@ -174,7 +175,6 @@ func TraverseBottomUp(parent *parser.Expr, current *parser.Expr, transform func(
 		}
 		return transform(parent, current)
 	}
-
 	return true
 }
 
@@ -185,6 +185,23 @@ func replaceSelectors(plan parser.Expr) parser.Expr {
 			*current = &MatrixSelector{MatrixSelector: t, OriginalString: t.String()}
 		case *parser.VectorSelector:
 			*current = &VectorSelector{VectorSelector: t}
+		case *parser.Call:
+			if t.Func.Name != "timestamp" {
+				return
+			}
+			switch v := unwrapParens(t.Args[0]).(type) {
+			case *parser.VectorSelector:
+				*current = &VectorSelector{VectorSelector: v, SelectTimestamp: true}
+			case *parser.StepInvariantExpr:
+				vs, ok := unwrapParens(v.Expr).(*parser.VectorSelector)
+				if ok {
+					// Prometheus weirdness
+					if vs.Timestamp != nil {
+						vs.OriginalOffset = 0
+					}
+					*current = &VectorSelector{VectorSelector: vs, SelectTimestamp: true}
+				}
+			}
 		}
 	})
 	return plan
@@ -355,6 +372,15 @@ func newStepInvariantExpr(expr parser.Expr) parser.Expr {
 	return &parser.StepInvariantExpr{Expr: expr}
 }
 
+func unwrapParens(expr parser.Expr) parser.Expr {
+	switch t := expr.(type) {
+	case *parser.ParenExpr:
+		return unwrapParens(t.Expr)
+	default:
+		return t
+	}
+}
+
 // Copy from https://github.com/prometheus/prometheus/blob/v2.39.1/promql/engine.go#L2658.
 func setOffsetForAtModifier(evalTime int64, expr parser.Expr) {
 	getOffset := func(ts *int64, originalOffset time.Duration, path []parser.Node) time.Duration {
@@ -439,6 +465,11 @@ type VectorSelector struct {
 }
 
 func (f VectorSelector) String() string {
+	if f.SelectTimestamp {
+		// If we pushed down timestamp into the vector selector we need to render the proper
+		// PromQL again.
+		return fmt.Sprintf("timestamp(%s)", f.VectorSelector.String())
+	}
 	return f.VectorSelector.String()
 }