engine,execution: use logical matrix and vector (#373)

* engine,execution: use logical matrix and vector

As the last phase of creating a new plan we replace all Matrix and
Vector selectors by our logical ones. This makes it possible to not
handle them twice everywhere in the codebase. It will also facilitate
more general sharding of them later on since we can annotate them
directly with their sharding information.
We had to do matrix and vector selector in the same PR since the
distributed engine cannot handle logical nodes well (because they
generally change the String() representation of the node).
We also had to change the String() representation of the logical nodes
to be equivalent to the original PromQL expression they came from.

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

* plan: use custom string render method for tests

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

---------

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

Author: MichaHoffmann

engine/user_defined_test.go

@@ -66,10 +66,10 @@ type injectVectorSelector struct{}
 
 func (i injectVectorSelector) Optimize(plan parser.Expr, opts *query.Options) (parser.Expr, annotations.Annotations) {
 	logicalplan.TraverseBottomUp(nil, &plan, func(_, current *parser.Expr) bool {
-		switch (*current).(type) {
-		case *parser.VectorSelector:
+		switch t := (*current).(type) {
+		case *logicalplan.VectorSelector:
 			*current = &logicalVectorSelector{
-				VectorSelector: (*current).(*parser.VectorSelector),
+				VectorSelector: t,
 			}
 		}
 		return false
@@ -78,7 +78,7 @@ func (i injectVectorSelector) Optimize(plan parser.Expr, opts *query.Options) (p
 }
 
 type logicalVectorSelector struct {
-	*parser.VectorSelector
+	*logicalplan.VectorSelector
 }
 
 func (c logicalVectorSelector) MakeExecutionOperator(vectors *model.VectorPool, _ *engstore.SelectorPool, opts *query.Options, hints storage.SelectHints) (model.VectorOperator, error) {

execution/execution.go

@@ -19,7 +19,6 @@ package execution
 import (
 	"runtime"
 	"sort"
-	"time"
 
 	"github.com/efficientgo/core/errors"
 	"github.com/prometheus/prometheus/model/labels"
@@ -59,8 +58,8 @@ func newOperator(expr parser.Expr, storage *engstore.SelectorPool, opts *query.O
 	switch e := expr.(type) {
 	case *parser.NumberLiteral:
 		return scan.NewNumberLiteralSelector(model.NewVectorPool(opts.StepsBatch), opts, e.Val), nil
-	case *parser.VectorSelector, *logicalplan.VectorSelector:
-		return newVectorSelector(expr, storage, opts, hints)
+	case *logicalplan.VectorSelector:
+		return newVectorSelector(e, storage, opts, hints)
 	case *parser.Call:
 		return newCall(e, storage, opts, hints)
 	case *parser.AggregateExpr:
@@ -82,34 +81,18 @@ func newOperator(expr parser.Expr, storage *engstore.SelectorPool, opts *query.O
 	case logicalplan.UserDefinedExpr:
 		return e.MakeExecutionOperator(model.NewVectorPool(opts.StepsBatch), storage, opts, hints)
 	default:
-		return nil, errors.Wrapf(parse.ErrNotSupportedExpr, "got: %s", e)
+		return nil, errors.Wrapf(parse.ErrNotSupportedExpr, "got: %s (%T)", e, e)
 	}
 }
 
-func newVectorSelector(expr parser.Expr, storage *engstore.SelectorPool, opts *query.Options, hints storage.SelectHints) (model.VectorOperator, error) {
-	var (
-		batchsize int64
-		offset    time.Duration
-		selector  engstore.SeriesSelector
-	)
-	switch e := expr.(type) {
-	case *parser.VectorSelector:
-		start, end := getTimeRangesForVectorSelector(e, opts, 0)
-		hints.Start = start
-		hints.End = end
-		offset = e.Offset
-		batchsize = 0
-		selector = storage.GetSelector(start, end, opts.Step.Milliseconds(), e.LabelMatchers, hints)
-	case *logicalplan.VectorSelector:
-		start, end := getTimeRangesForVectorSelector(e.VectorSelector, opts, 0)
-		hints.Start = start
-		hints.End = end
-		offset = e.Offset
-		batchsize = e.BatchSize
-		selector = storage.GetFilteredSelector(start, end, opts.Step.Milliseconds(), e.LabelMatchers, e.Filters, hints)
-	default:
-		return nil, errors.Wrapf(parse.ErrNotSupportedExpr, "got: %s", e)
-	}
+func newVectorSelector(e *logicalplan.VectorSelector, storage *engstore.SelectorPool, opts *query.Options, hints storage.SelectHints) (model.VectorOperator, error) {
+	start, end := getTimeRangesForVectorSelector(e, opts, 0)
+	hints.Start = start
+	hints.End = end
+
+	offset := e.Offset
+	batchsize := e.BatchSize
+	selector := storage.GetFilteredSelector(start, end, opts.Step.Milliseconds(), e.LabelMatchers, e.Filters, hints)
 
 	numShards := runtime.GOMAXPROCS(0) / 2
 	if numShards < 1 {
@@ -138,19 +121,13 @@ func newCall(e *parser.Call, storage *engstore.SelectorPool, opts *query.Options
 	}
 	if e.Func.Name == "timestamp" {
 		switch arg := e.Args[0].(type) {
-		case *logicalplan.VectorSelector, *parser.VectorSelector:
+		case *logicalplan.VectorSelector:
 			// We push down the timestamp function into the scanner through the hints.
-			return newVectorSelector(e.Args[0], storage, opts, hints)
+			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 *parser.VectorSelector:
-				// Prometheus weirdness.
-				if vs.Timestamp != nil {
-					vs.OriginalOffset = 0
-				}
-				return newVectorSelector(vs, storage, opts, hints)
 			case *logicalplan.VectorSelector:
 				// Prometheus weirdness.
 				if vs.Timestamp != nil {
@@ -175,7 +152,7 @@ func newCall(e *parser.Call, storage *engstore.SelectorPool, opts *query.Options
 				return nil, parse.ErrNotImplemented
 			}
 			return newSubqueryFunction(e, t, storage, opts, hints)
-		case *parser.MatrixSelector:
+		case *logicalplan.MatrixSelector:
 			return newRangeVectorFunction(e, t, storage, opts, hints)
 		}
 	}
@@ -203,7 +180,7 @@ func newAbsentOverTimeOperator(call *parser.Call, storage *engstore.SelectorPool
 			Args: []parser.Expr{matrixCall},
 		}
 		return function.NewFunctionOperator(f, []model.VectorOperator{argOp}, opts.StepsBatch, opts)
-	case *parser.MatrixSelector:
+	case *logicalplan.MatrixSelector:
 		matrixCall := &parser.Call{
 			Func: &parser.Function{Name: "last_over_time"},
 			Args: call.Args,
@@ -221,9 +198,12 @@ func newAbsentOverTimeOperator(call *parser.Call, storage *engstore.SelectorPool
 		vs.LabelMatchers = append(vs.LabelMatchers, filters...)
 		f := &parser.Call{
 			Func: &parser.Function{Name: "absent"},
-			Args: []parser.Expr{&parser.MatrixSelector{
-				VectorSelector: vs,
-				Range:          arg.Range,
+			Args: []parser.Expr{&logicalplan.MatrixSelector{
+				MatrixSelector: &parser.MatrixSelector{
+					VectorSelector: vs,
+					Range:          arg.Range,
+				},
+				OriginalString: arg.String(),
 			}},
 		}
 		return function.NewFunctionOperator(f, []model.VectorOperator{argOp}, opts.StepsBatch, opts)
@@ -232,7 +212,7 @@ func newAbsentOverTimeOperator(call *parser.Call, storage *engstore.SelectorPool
 	}
 }
 
-func newRangeVectorFunction(e *parser.Call, t *parser.MatrixSelector, storage *engstore.SelectorPool, opts *query.Options, hints storage.SelectHints) (model.VectorOperator, error) {
+func newRangeVectorFunction(e *parser.Call, t *logicalplan.MatrixSelector, storage *engstore.SelectorPool, opts *query.Options, hints storage.SelectHints) (model.VectorOperator, error) {
 	// TODO(saswatamcode): Range vector result might need new operator
 	// before it can be non-nested. https://github.com/thanos-io/promql-engine/issues/39
 	batchSize, vs, filters, err := unpackVectorSelector(t)
@@ -466,7 +446,7 @@ func newRemoteExecution(e logicalplan.RemoteExecution, opts *query.Options, hint
 }
 
 // Copy from https://github.com/prometheus/prometheus/blob/v2.39.1/promql/engine.go#L791.
-func getTimeRangesForVectorSelector(n *parser.VectorSelector, opts *query.Options, evalRange int64) (int64, int64) {
+func getTimeRangesForVectorSelector(n *logicalplan.VectorSelector, opts *query.Options, evalRange int64) (int64, int64) {
 	start := opts.Start.UnixMilli()
 	end := opts.End.UnixMilli()
 	if n.Timestamp != nil {
@@ -493,12 +473,10 @@ func unwrapConstVal(e parser.Expr) (float64, error) {
 	return 0, errors.Wrap(parse.ErrNotSupportedExpr, "matrix selector argument must be a constant")
 }
 
-func unpackVectorSelector(t *parser.MatrixSelector) (int64, *parser.VectorSelector, []*labels.Matcher, error) {
+func unpackVectorSelector(t *logicalplan.MatrixSelector) (int64, *logicalplan.VectorSelector, []*labels.Matcher, error) {
 	switch t := t.VectorSelector.(type) {
-	case *parser.VectorSelector:
-		return 0, t, nil, nil
 	case *logicalplan.VectorSelector:
-		return t.BatchSize, t.VectorSelector, t.Filters, nil
+		return t.BatchSize, t, t.Filters, nil
 	default:
 		return 0, nil, nil, parse.ErrNotSupportedExpr
 	}

execution/function/absent.go

@@ -13,6 +13,7 @@ import (
 	"github.com/prometheus/prometheus/promql/parser"
 
 	"github.com/thanos-io/promql-engine/execution/model"
+	"github.com/thanos-io/promql-engine/logicalplan"
 )
 
 type absentOperator struct {
@@ -49,10 +50,10 @@ func (o *absentOperator) loadSeries() {
 		// https://github.com/prometheus/prometheus/blob/main/promql/functions.go#L1385
 		var lm []*labels.Matcher
 		switch n := o.funcExpr.Args[0].(type) {
-		case *parser.VectorSelector:
+		case *logicalplan.VectorSelector:
 			lm = n.LabelMatchers
-		case *parser.MatrixSelector:
-			lm = n.VectorSelector.(*parser.VectorSelector).LabelMatchers
+		case *logicalplan.MatrixSelector:
+			lm = n.VectorSelector.(*logicalplan.VectorSelector).LabelMatchers
 		default:
 			o.series = []labels.Labels{labels.EmptyLabels()}
 			return

execution/function/histogram.go

@@ -13,7 +13,6 @@ import (
 
 	"github.com/cespare/xxhash/v2"
 	"github.com/prometheus/prometheus/model/labels"
-
 	"github.com/prometheus/prometheus/promql/parser"
 
 	"github.com/thanos-io/promql-engine/execution/model"
@@ -174,8 +173,8 @@ func (o *histogramOperator) processInputSeries(vectors []model.StepVector) ([]mo
 		out = append(out, step)
 		o.vectorOp.GetPool().PutStepVector(vector)
 	}
-
 	o.vectorOp.GetPool().PutVectors(vectors)
+
 	return out, nil
 }
 

execution/scan/vector_selector.go

@@ -169,7 +169,6 @@ func (o *vectorSelector) Next(ctx context.Context) ([]model.StepVector, error) {
 		o.currentStep += o.step * int64(o.numSteps)
 		o.currentSeries = 0
 	}
-
 	return vectors, nil
 }
 

execution/step_invariant/step_invariant.go

@@ -14,6 +14,7 @@ import (
 	"github.com/prometheus/prometheus/promql/parser"
 
 	"github.com/thanos-io/promql-engine/execution/model"
+	"github.com/thanos-io/promql-engine/logicalplan"
 	"github.com/thanos-io/promql-engine/query"
 )
 
@@ -72,7 +73,7 @@ func NewStepInvariantOperator(
 	// We do not duplicate results for range selectors since result is a matrix
 	// with their unique timestamps which does not depend on the step.
 	switch expr.(type) {
-	case *parser.MatrixSelector, *parser.SubqueryExpr:
+	case *logicalplan.MatrixSelector, *parser.SubqueryExpr:
 		u.cacheResult = false
 	}
 	u.OperatorTelemetry = &model.NoopTelemetry{}

logicalplan/distribute.go

@@ -413,9 +413,9 @@ func calculateStartOffset(expr *parser.Expr, lookbackDelta time.Duration) time.D
 		switch n := (*node).(type) {
 		case *parser.SubqueryExpr:
 			selectRange += n.Range
-		case *parser.MatrixSelector:
+		case *MatrixSelector:
 			selectRange += n.Range
-		case *parser.VectorSelector:
+		case *VectorSelector:
 			offset = n.Offset
 		}
 	})
@@ -469,7 +469,7 @@ func matchesExternalLabelSet(expr parser.Expr, externalLabelSet []labels.Labels)
 	}
 	var selectorSet [][]*labels.Matcher
 	traverse(&expr, func(current *parser.Expr) {
-		vs, ok := (*current).(*parser.VectorSelector)
+		vs, ok := (*current).(*VectorSelector)
 		if ok {
 			selectorSet = append(selectorSet, vs.LabelMatchers)
 		}

logicalplan/distribute_test.go

@@ -551,7 +551,7 @@ sum(
 			plan := New(expr, &query.Options{Start: tcase.queryStart, End: tcase.queryEnd, Step: time.Minute})
 			optimizedPlan, _ := plan.Optimize(optimizers)
 			expectedPlan := cleanUp(replacements, tcase.expected)
-			testutil.Equals(t, expectedPlan, optimizedPlan.Expr().String())
+			testutil.Equals(t, expectedPlan, renderExprTree(optimizedPlan.Expr()))
 		})
 	}
 }

logicalplan/filter.go

@@ -33,7 +33,7 @@ func (m MergeSelectsOptimizer) Optimize(plan parser.Expr, _ *query.Options) (par
 
 func extractSelectors(selectors matcherHeap, expr parser.Expr) {
 	traverse(&expr, func(node *parser.Expr) {
-		e, ok := (*node).(*parser.VectorSelector)
+		e, ok := (*node).(*VectorSelector)
 		if !ok {
 			return
 		}
@@ -49,8 +49,6 @@ func replaceMatchers(selectors matcherHeap, expr *parser.Expr) {
 	traverse(expr, func(node *parser.Expr) {
 		var matchers []*labels.Matcher
 		switch e := (*node).(type) {
-		case *parser.VectorSelector:
-			matchers = e.LabelMatchers
 		case *VectorSelector:
 			matchers = e.LabelMatchers
 		default:
@@ -85,12 +83,6 @@ func replaceMatchers(selectors matcherHeap, expr *parser.Expr) {
 			}
 
 			switch e := (*node).(type) {
-			case *parser.VectorSelector:
-				e.LabelMatchers = replacement
-				*node = &VectorSelector{
-					VectorSelector: e,
-					Filters:        filters,
-				}
 			case *VectorSelector:
 				e.LabelMatchers = replacement
 				e.Filters = filters

logicalplan/merge_selects.go

@@ -47,7 +47,7 @@ func TestMergeSelects(t *testing.T) {
 
 			plan := New(expr, &query.Options{})
 			optimizedPlan, _ := plan.Optimize(optimizers)
-			testutil.Equals(t, tcase.expected, optimizedPlan.Expr().String())
+			testutil.Equals(t, tcase.expected, renderExprTree(optimizedPlan.Expr()))
 		})
 	}
 }

logicalplan/merge_selects_test.go

@@ -63,7 +63,7 @@ func (m PassthroughOptimizer) Optimize(plan parser.Expr, opts *query.Options) (p
 
 	matchingLabelsEngines := make([]api.RemoteEngine, 0, len(engines))
 	TraverseBottomUp(nil, &plan, func(parent, current *parser.Expr) (stop bool) {
-		if vs, ok := (*current).(*parser.VectorSelector); ok {
+		if vs, ok := (*current).(*VectorSelector); ok {
 			for _, e := range engines {
 				if !labelSetsMatch(vs.LabelMatchers, e.LabelSets()...) {
 					continue