frontend: add new 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. To ease migrating to the correct version, I'm not removing the old prune code yet, just adding the new method. 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) which become: subquery (vector(-1)==1) Signed-off-by: György Krajcsovits <[email protected]>
grafana · Nov 26, 2024 · 97a767d · 97a767d
1 parent aaae0dd
commit 97a767d
Show file tree

Hide file tree

Showing 10 changed files with 728 additions and 261 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -69,6 +69,7 @@
 * [ENHANCEMENT] Distributor: Initialize ha_tracker cache before ha_tracker and distributor reach running state and begin serving writes. #9826 #9976
 * [ENHANCEMENT] Ingester: `-ingest-storage.kafka.max-buffered-bytes` to limit the memory for buffered records when using concurrent fetching. #9892
 * [ENHANCEMENT] Querier: improve performance and memory consumption of queries that select many series. #9914
+* [ENHANCEMENT] Query-Frontend: prune `<subquery> and on() (vector(x)==y)` style queries. Triggered by https://github.com/prometheus/prometheus/pull/15245. #10026
 * [BUGFIX] Fix issue where functions such as `rate()` over native histograms could return incorrect values if a float stale marker was present in the selected range. #9508
 * [BUGFIX] Fix issue where negation of native histograms (eg. `-some_native_histogram_series`) did nothing. #9508
 * [BUGFIX] Fix issue where `metric might not be a counter, name does not end in _total/_sum/_count/_bucket` annotation would be emitted even if `rate` or `increase` did not have enough samples to compute a result. #9508

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 || 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
 }