Merge pull request #16896 from bboreham/wrap-less

[PERF] PromQL: Move more work to preprocessing step

promql/engine.go

@@ -1670,12 +1670,8 @@ func (ev *evaluator) eval(ctx context.Context, expr parser.Expr) (parser.Value,
 		sortedGrouping := e.Grouping
 		slices.Sort(sortedGrouping)
 
-		unwrapParenExpr(&e.Param)
-		param := unwrapStepInvariantExpr(e.Param)
-		unwrapParenExpr(&param)
-
 		if e.Op == parser.COUNT_VALUES {
-			valueLabel := param.(*parser.StringLiteral)
+			valueLabel := e.Param.(*parser.StringLiteral)
 			if !model.LabelName(valueLabel.Val).IsValid() {
 				ev.errorf("invalid label name %s", valueLabel)
 			}
@@ -1690,8 +1686,8 @@ func (ev *evaluator) eval(ctx context.Context, expr parser.Expr) (parser.Value,
 
 		var warnings annotations.Annotations
 		originalNumSamples := ev.currentSamples
-		// param is the number k for topk/bottomk, or q for quantile.
+		// e.Param is the number k for topk/bottomk, or q for quantile.
-		fp, ws := newFParams(ctx, ev, param)
+		fp, ws := newFParams(ctx, ev, e.Param)
 		warnings.Merge(ws)
 		// Now fetch the data to be aggregated.
 		val, ws := ev.eval(ctx, e.Expr)
@@ -1711,9 +1707,7 @@ func (ev *evaluator) eval(ctx context.Context, expr parser.Expr) (parser.Value,
 			// Matrix evaluation always returns the evaluation time,
 			// so this function needs special handling when given
 			// a vector selector.
-			unwrapParenExpr(&e.Args[0])
 			arg := unwrapStepInvariantExpr(e.Args[0])
-			unwrapParenExpr(&arg)
 			vs, ok := arg.(*parser.VectorSelector)
 			if ok {
 				return ev.rangeEvalTimestampFunctionOverVectorSelector(ctx, vs, call, e)
@@ -1727,9 +1721,7 @@ func (ev *evaluator) eval(ctx context.Context, expr parser.Expr) (parser.Value,
 			warnings       annotations.Annotations
 		)
 		for i := range e.Args {
-			unwrapParenExpr(&e.Args[i])
+			a := e.Args[i]
-			a := unwrapStepInvariantExpr(e.Args[i])
-			unwrapParenExpr(&a)
 			if _, ok := a.(*parser.MatrixSelector); ok {
 				matrixArgIndex = i
 				matrixArg = true
@@ -1780,9 +1772,7 @@ func (ev *evaluator) eval(ctx context.Context, expr parser.Expr) (parser.Value,
 			}
 		}
 
-		unwrapParenExpr(&e.Args[matrixArgIndex])
+		arg := e.Args[matrixArgIndex]
-		arg := unwrapStepInvariantExpr(e.Args[matrixArgIndex])
-		unwrapParenExpr(&arg)
 		sel := arg.(*parser.MatrixSelector)
 		selVS := sel.VectorSelector.(*parser.VectorSelector)
 
@@ -2108,11 +2098,6 @@ func (ev *evaluator) eval(ctx context.Context, expr parser.Expr) (parser.Value,
 		ev.samplesStats.IncrementSamplesAtTimestamp(ev.endTimestamp, newEv.samplesStats.TotalSamples)
 		return res, ws
 	case *parser.StepInvariantExpr:
-		switch ce := e.Expr.(type) {
-		case *parser.StringLiteral, *parser.NumberLiteral:
-			return ev.eval(ctx, ce)
-		}
-
 		newEv := &evaluator{
 			startTimestamp:           ev.startTimestamp,
 			endTimestamp:             ev.startTimestamp, // Always a single evaluation.
@@ -3725,8 +3710,8 @@ func unwrapStepInvariantExpr(e parser.Expr) parser.Expr {
 }
 
 // PreprocessExpr wraps all possible step invariant parts of the given expression with
-// StepInvariantExpr. It also resolves the preprocessors and evaluates duration expressions
+// StepInvariantExpr. It also resolves the preprocessors, evaluates duration expressions
-// into their numeric values.
+// into their numeric values and removes superfluous parenthesis on parameters to functions and aggregations.
 func PreprocessExpr(expr parser.Expr, start, end time.Time, step time.Duration) (parser.Expr, error) {
 	detectHistogramStatsDecoding(expr)
 
@@ -3734,18 +3719,20 @@ func PreprocessExpr(expr parser.Expr, start, end time.Time, step time.Duration)
 		return nil, err
 	}
 
-	isStepInvariant := preprocessExprHelper(expr, start, end)
+	_, shouldWrap := preprocessExprHelper(expr, start, end)
-	if isStepInvariant {
+	if shouldWrap {
 		return newStepInvariantExpr(expr), nil
 	}
 	return expr, nil
 }
 
-// preprocessExprHelper wraps the child nodes of the expression
+// preprocessExprHelper wraps child nodes of expr with a StepInvariantExpr,
-// with a StepInvariantExpr wherever it's step invariant. The returned boolean is true if the
+// at the highest level within the tree that is step-invariant.
-// passed expression qualifies to be wrapped by StepInvariantExpr.
+// Also resolves start() and end() on selector and subquery nodes.
-// It also resolves the preprocessors.
+// Also remove superfluous parenthesis on parameters to functions and aggregations.
-func preprocessExprHelper(expr parser.Expr, start, end time.Time) bool {
+// Return isStepInvariant is true when the whole subexpression is step invariant.
+// Return shoudlWrap is false for cases like MatrixSelector and StringLiteral that never need to be wrapped.
+func preprocessExprHelper(expr parser.Expr, start, end time.Time) (isStepInvariant, shouldWrap bool) {
 	switch n := expr.(type) {
 	case *parser.VectorSelector:
 		switch n.StartOrEnd {
@@ -3754,49 +3741,56 @@ func preprocessExprHelper(expr parser.Expr, start, end time.Time) bool {
 		case parser.END:
 			n.Timestamp = makeInt64Pointer(timestamp.FromTime(end))
 		}
-		return n.Timestamp != nil
+		return n.Timestamp != nil, n.Timestamp != nil
 
 	case *parser.AggregateExpr:
+		unwrapParenExpr(&n.Expr)
+		unwrapParenExpr(&n.Param)
 		return preprocessExprHelper(n.Expr, start, end)
 
 	case *parser.BinaryExpr:
-		isInvariant1, isInvariant2 := preprocessExprHelper(n.LHS, start, end), preprocessExprHelper(n.RHS, start, end)
+		isInvariantLHS, shouldWrapLHS := preprocessExprHelper(n.LHS, start, end)
-		if isInvariant1 && isInvariant2 {
+		isInvariantRHS, shouldWrapRHS := preprocessExprHelper(n.RHS, start, end)
-			return true
+		if isInvariantLHS && isInvariantRHS {
+			return true, true
 		}
 
-		if isInvariant1 {
+		if shouldWrapLHS {
 			n.LHS = newStepInvariantExpr(n.LHS)
 		}
-		if isInvariant2 {
+		if shouldWrapRHS {
 			n.RHS = newStepInvariantExpr(n.RHS)
 		}
 
-		return false
+		return false, false
 
 	case *parser.Call:
 		_, ok := AtModifierUnsafeFunctions[n.Func.Name]
 		isStepInvariant := !ok
-		isStepInvariantSlice := make([]bool, len(n.Args))
+		shouldWrap := make([]bool, len(n.Args))
 		for i := range n.Args {
-			isStepInvariantSlice[i] = preprocessExprHelper(n.Args[i], start, end)
+			unwrapParenExpr(&n.Args[i])
-			isStepInvariant = isStepInvariant && isStepInvariantSlice[i]
+			var argIsStepInvariant bool
+			argIsStepInvariant, shouldWrap[i] = preprocessExprHelper(n.Args[i], start, end)
+			isStepInvariant = isStepInvariant && argIsStepInvariant
 		}
 
 		if isStepInvariant {
 			// The function and all arguments are step invariant.
-			return true
+			return true, true
 		}
 
-		for i, isi := range isStepInvariantSlice {
+		for i, isi := range shouldWrap {
 			if isi {
 				n.Args[i] = newStepInvariantExpr(n.Args[i])
 			}
 		}
-		return false
+		return false, false
 
 	case *parser.MatrixSelector:
-		return preprocessExprHelper(n.VectorSelector, start, end)
+		// We don't need to wrap a MatrixSelector because functions over range vectors evaluate those directly.
+		isStepInvariant, _ := preprocessExprHelper(n.VectorSelector, start, end)
+		return isStepInvariant, false
 
 	case *parser.SubqueryExpr:
 		// Since we adjust offset for the @ modifier evaluation,
@@ -3804,7 +3798,7 @@ func preprocessExprHelper(expr parser.Expr, start, end time.Time) bool {
 		// Hence we wrap the inside of subquery irrespective of
 		// @ on subquery (given it is also step invariant) so that
 		// it is evaluated only once w.r.t. the start time of subquery.
-		isInvariant := preprocessExprHelper(n.Expr, start, end)
+		isInvariant, _ := preprocessExprHelper(n.Expr, start, end)
 		if isInvariant {
 			n.Expr = newStepInvariantExpr(n.Expr)
 		}
@@ -3814,7 +3808,7 @@ func preprocessExprHelper(expr parser.Expr, start, end time.Time) bool {
 		case parser.END:
 			n.Timestamp = makeInt64Pointer(timestamp.FromTime(end))
 		}
-		return n.Timestamp != nil
+		return n.Timestamp != nil, n.Timestamp != nil
 
 	case *parser.ParenExpr:
 		return preprocessExprHelper(n.Expr, start, end)
@@ -3823,7 +3817,7 @@ func preprocessExprHelper(expr parser.Expr, start, end time.Time) bool {
 		return preprocessExprHelper(n.Expr, start, end)
 
 	case *parser.StringLiteral, *parser.NumberLiteral:
-		return true
+		return true, false
 	}
 
 	panic(fmt.Sprintf("found unexpected node %#v", expr))

promql/engine_test.go

@@ -2304,22 +2304,18 @@ func TestPreprocessAndWrapWithStepInvariantExpr(t *testing.T) {
 	}{
 		{
 			input: "123.4567",
-			expected: &parser.StepInvariantExpr{
+			expected: &parser.NumberLiteral{
-				Expr: &parser.NumberLiteral{
 				Val:      123.4567,
 				PosRange: posrange.PositionRange{Start: 0, End: 8},
 			},
 		},
-		},
 		{
 			input: `"foo"`,
-			expected: &parser.StepInvariantExpr{
+			expected: &parser.StringLiteral{
-				Expr: &parser.StringLiteral{
 				Val:      "foo",
 				PosRange: posrange.PositionRange{Start: 0, End: 5},
 			},
 		},
-		},
 		{
 			input: "foo * bar",
 			expected: &parser.BinaryExpr{
@@ -2427,8 +2423,7 @@ func TestPreprocessAndWrapWithStepInvariantExpr(t *testing.T) {
 		},
 		{
 			input: `test{a="b"}[5y] @ 1603774699`,
-			expected: &parser.StepInvariantExpr{
+			expected: &parser.MatrixSelector{
-				Expr: &parser.MatrixSelector{
 				VectorSelector: &parser.VectorSelector{
 					Name:      "test",
 					Timestamp: makeInt64Pointer(1603774699000),
@@ -2445,7 +2440,6 @@ func TestPreprocessAndWrapWithStepInvariantExpr(t *testing.T) {
 				EndPos: 28,
 			},
 		},
-		},
 		{
 			input: "sum by (foo)(some_metric)",
 			expected: &parser.AggregateExpr{
@@ -2942,9 +2936,16 @@ func TestPreprocessAndWrapWithStepInvariantExpr(t *testing.T) {
 			},
 		},
 		{
-			input: `test[5y] @ start()`,
+			input: `sum_over_time(test[5y] @ start())`,
 			expected: &parser.StepInvariantExpr{
-				Expr: &parser.MatrixSelector{
+				Expr: &parser.Call{
+					Func: &parser.Function{
+						Name:       "sum_over_time",
+						ArgTypes:   []parser.ValueType{parser.ValueTypeMatrix},
+						ReturnType: parser.ValueTypeVector,
+					},
+					Args: parser.Expressions{
+						&parser.MatrixSelector{
 							VectorSelector: &parser.VectorSelector{
 								Name:       "test",
 								Timestamp:  makeInt64Pointer(timestamp.FromTime(startTime)),
@@ -2953,19 +2954,21 @@ func TestPreprocessAndWrapWithStepInvariantExpr(t *testing.T) {
 									parser.MustLabelMatcher(labels.MatchEqual, "__name__", "test"),
 								},
 								PosRange: posrange.PositionRange{
-							Start: 0,
+									Start: 14,
-							End:   4,
+									End:   18,
 								},
 							},
 							Range:  5 * 365 * 24 * time.Hour,
-					EndPos: 18,
+							EndPos: 32,
+						},
+					},
+					PosRange: posrange.PositionRange{Start: 0, End: 33},
 				},
 			},
 		},
 		{
 			input: `test[5y] @ end()`,
-			expected: &parser.StepInvariantExpr{
+			expected: &parser.MatrixSelector{
-				Expr: &parser.MatrixSelector{
 				VectorSelector: &parser.VectorSelector{
 					Name:       "test",
 					Timestamp:  makeInt64Pointer(timestamp.FromTime(endTime)),
@@ -2982,7 +2985,6 @@ func TestPreprocessAndWrapWithStepInvariantExpr(t *testing.T) {
 				EndPos: 16,
 			},
 		},
-		},
 		{
 			input: `some_metric[10m:5s] @ start()`,
 			expected: &parser.StepInvariantExpr{

promql/functions.go

@@ -1946,9 +1946,7 @@ func createLabelsForAbsentFunction(expr parser.Expr) labels.Labels {
 }
 
 func stringFromArg(e parser.Expr) string {
-	tmp := unwrapStepInvariantExpr(e) // Unwrap StepInvariant
+	return e.(*parser.StringLiteral).Val
-	unwrapParenExpr(&tmp)             // Optionally unwrap ParenExpr
-	return tmp.(*parser.StringLiteral).Val
 }
 
 func stringSliceFromArgs(args parser.Expressions) []string {