frontend: update query pruning

The current method of excluding/including sub query results in PromQL by comparing to -Inf or +Inf is no longer valid after prometheus/prometheus#15245 due to comparison of native histograms to a float with < or > result in Jeanette's warning, not empty set. The new method uses logical AND operation to intersect the sub query with either a const vector() or an empty vector(). E.g. subquery and on() (vector(1)==1) subquery and on() (vector(-1)==1) Signed-off-by: György Krajcsovits <[email protected]>
grafana · Nov 26, 2024 · 9dcc588 · 9dcc588
1 parent aaae0dd
commit 9dcc588
Show file tree

Hide file tree

Showing 3 changed files with 120 additions and 262 deletions.
diff --git a/pkg/frontend/querymiddleware/astmapper/pruning.go b/pkg/frontend/querymiddleware/astmapper/pruning.go
@@ -4,8 +4,6 @@ package astmapper
 
 import (
 	"context"
-	"math"
-	"strconv"
 
 	"github.com/go-kit/log"
 	"github.com/prometheus/prometheus/promql/parser"
@@ -33,185 +31,89 @@ func (pruner *queryPruner) MapExpr(expr parser.Expr) (mapped parser.Expr, finish
 		return nil, false, err
 	}
 
-	switch e := expr.(type) {
-	case *parser.ParenExpr:
-		mapped, finished, err = pruner.MapExpr(e.Expr)
-		if err != nil {
-			return e, false, err
-		}
-		return &parser.ParenExpr{Expr: mapped, PosRange: e.PosRange}, finished, nil
-	case *parser.BinaryExpr:
-		return pruner.pruneBinOp(e)
-	default:
-		return e, false, nil
+	e, ok := expr.(*parser.BinaryExpr)
+	if !ok {
+		return expr, false, nil
 	}
-}
 
-func (pruner *queryPruner) pruneBinOp(expr *parser.BinaryExpr) (mapped parser.Expr, finished bool, err error) {
-	switch expr.Op {
-	case parser.MUL:
-		return pruner.handleMultiplyOp(expr), false, nil
-	case parser.GTR, parser.LSS:
-		return pruner.handleCompOp(expr), false, nil
-	case parser.LOR:
-		return pruner.handleOrOp(expr), false, nil
-	case parser.LAND:
-		return pruner.handleAndOp(expr), false, nil
-	case parser.LUNLESS:
-		return pruner.handleUnlessOp(expr), false, nil
-	default:
+	if e.Op != parser.LAND || e.VectorMatching == nil ||
+		e.VectorMatching.On == false || len(e.VectorMatching.MatchingLabels) != 0 {
+		// Return if not "<lhs> and on() <rhs>"
 		return expr, false, nil
 	}
-}
 
-// The bool signifies if the number evaluates to infinity, and if it does
-// we return the infinity of the correct sign.
-func calcInf(isPositive bool, num string) (*parser.NumberLiteral, bool) {
-	coeff, err := strconv.Atoi(num)
-	if err != nil || coeff == 0 {
-		return nil, false
+	isConst, isEmpty := pruner.isConst(e.RHS)
+	if !isConst {
+		return expr, false, nil
 	}
-	switch {
-	case isPositive && coeff > 0:
-		return &parser.NumberLiteral{Val: math.Inf(1)}, true
-	case isPositive && coeff < 0:
-		return &parser.NumberLiteral{Val: math.Inf(-1)}, true
-	case !isPositive && coeff > 0:
-		return &parser.NumberLiteral{Val: math.Inf(-1)}, true
-	case !isPositive && coeff < 0:
-		return &parser.NumberLiteral{Val: math.Inf(1)}, true
-	default:
-		return nil, false
+	if isEmpty {
+		// The right hand side is empty, so the whole expression is empty due to
+		// "and on()", return the right hand side.
+		return e.RHS, false, nil
 	}
+	// The right hand side is const no empty, so the whole expression is just the
+	// left side.
+	return e.LHS, false, nil
 }
 
-func (pruner *queryPruner) handleMultiplyOp(expr *parser.BinaryExpr) parser.Expr {
-	isInfR, signR := pruner.isInfinite(expr.RHS)
-	if isInfR {
-		newExpr, ok := calcInf(signR, expr.LHS.String())
-		if ok {
-			return newExpr
-		}
-	}
-	isInfL, signL := pruner.isInfinite(expr.LHS)
-	if isInfL {
-		newExpr, ok := calcInf(signL, expr.RHS.String())
-		if ok {
-			return newExpr
+func (pruner *queryPruner) isConst(expr parser.Expr) (isConst, isEmpty bool) {
+	var lhs, rhs parser.Expr
+	switch e := expr.(type) {
+	case *parser.ParenExpr:
+		return pruner.isConst(e.Expr)
+	case *parser.BinaryExpr:
+		if e.Op != parser.EQLC || e.ReturnBool {
+			return false, false
 		}
-	}
-	return expr
-}
-
-func (pruner *queryPruner) handleCompOp(expr *parser.BinaryExpr) parser.Expr {
-	var refNeg, refPos parser.Expr
-	switch expr.Op {
-	case parser.LSS:
-		refNeg = expr.RHS
-		refPos = expr.LHS
-	case parser.GTR:
-		refNeg = expr.LHS
-		refPos = expr.RHS
+		lhs = e.LHS
+		rhs = e.RHS
 	default:
-		return expr
+		return false, false
 	}
 
-	// foo < -Inf or -Inf > foo => vector(0) < -Inf
-	isInf, sign := pruner.isInfinite(refNeg)
-	if isInf && !sign {
-		return &parser.BinaryExpr{
-			LHS: &parser.Call{
-				Func: parser.Functions["vector"],
-				Args: []parser.Expr{&parser.NumberLiteral{Val: 0}},
-			},
-			Op:         parser.LSS,
-			RHS:        &parser.NumberLiteral{Val: math.Inf(-1)},
-			ReturnBool: false,
+	lIsVector, lValue := pruner.isConstVector(lhs)
+	if lIsVector {
+		rIsConst, rValue := pruner.isNumber(rhs)
+		if rIsConst {
+			return true, rValue != lValue
 		}
+		return false, false
 	}
-
-	// foo > +Inf or +Inf < foo => vector(0) > +Inf => vector(0) < -Inf
-	isInf, sign = pruner.isInfinite(refPos)
-	if isInf && sign {
-		return &parser.BinaryExpr{
-			LHS: &parser.Call{
-				Func: parser.Functions["vector"],
-				Args: []parser.Expr{&parser.NumberLiteral{Val: 0}},
-			},
-			Op:         parser.LSS,
-			RHS:        &parser.NumberLiteral{Val: math.Inf(-1)},
-			ReturnBool: false,
-		}
+	var lIsConst bool
+	lIsConst, lValue = pruner.isNumber(lhs)
+	if !lIsConst {
+		return false, false
 	}
-
-	return expr
+	rIsVector, rValue := pruner.isConstVector(rhs)
+	if !rIsVector {
+		return false, false
+	}
+	return true, lValue != rValue
 }
 
-// 1st bool is true if the number is infinite.
-// 2nd bool is true if the number is positive infinity.
-func (pruner *queryPruner) isInfinite(expr parser.Expr) (bool, bool) {
-	mapped, _, err := pruner.MapExpr(expr)
-	if err == nil {
-		expr = mapped
-	}
+func (pruner *queryPruner) isConstVector(expr parser.Expr) (isVector bool, value float64) {
 	switch e := expr.(type) {
 	case *parser.ParenExpr:
-		return pruner.isInfinite(e.Expr)
-	case *parser.NumberLiteral:
-		if math.IsInf(e.Val, 1) {
-			return true, true
+		return pruner.isConstVector(e.Expr)
+	case *parser.Call:
+		if e.Func.Name != "vector" || len(e.Args) != 1 {
+			return false, 0
 		}
-		if math.IsInf(e.Val, -1) {
-			return true, false
+		lit, ok := e.Args[0].(*parser.NumberLiteral)
+		if !ok {
+			return false, 0
 		}
-		return false, false
-	default:
-		return false, false
-	}
-}
-
-func (pruner *queryPruner) handleOrOp(expr *parser.BinaryExpr) parser.Expr {
-	switch {
-	case pruner.isEmpty(expr.LHS):
-		return expr.RHS
-	case pruner.isEmpty(expr.RHS):
-		return expr.LHS
-	}
-	return expr
-}
-
-func (pruner *queryPruner) handleAndOp(expr *parser.BinaryExpr) parser.Expr {
-	switch {
-	case pruner.isEmpty(expr.LHS):
-		return expr.LHS
-	case pruner.isEmpty(expr.RHS):
-		return expr.RHS
+		return true, lit.Val
 	}
-	return expr
+	return false, 0
 }
 
-func (pruner *queryPruner) handleUnlessOp(expr *parser.BinaryExpr) parser.Expr {
-	switch {
-	case pruner.isEmpty(expr.LHS):
-		return expr.LHS
-	case pruner.isEmpty(expr.RHS):
-		return expr.LHS
-	}
-	return expr
-}
-
-func (pruner *queryPruner) isEmpty(expr parser.Expr) bool {
-	mapped, _, err := pruner.MapExpr(expr)
-	if err == nil {
-		expr = mapped
-	}
+func (pruner *queryPruner) isNumber(expr parser.Expr) (isNumber bool, value float64) {
 	switch e := expr.(type) {
 	case *parser.ParenExpr:
-		return pruner.isEmpty(e.Expr)
-	default:
-		if e.String() == `vector(0) < -Inf` {
-			return true
-		}
-		return false
+		return pruner.isNumber(e.Expr)
+	case *parser.NumberLiteral:
+		return true, e.Val
 	}
+	return false, 0
 }