Skip to content

Commit

Permalink
Add query pruning middleware (#9086)
Browse files Browse the repository at this point in the history
* Write basic pruner

* Change behaviour

* Clean up

* Add pruning middleware

* Simplify into 1 no-data expression

* Support OR

* Update docs

* Fix lint

* Revise according to code review

* Revise test according to code review and add comments

* Refine logic to better handle parenthesis

* Revise according to code review

* Update changelog
  • Loading branch information
zenador authored Sep 12, 2024
1 parent 0a06baa commit 1e98df0
Show file tree
Hide file tree
Showing 11 changed files with 644 additions and 4 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
* `-ingester.partition-ring.*`: configures partitions ring backend.
* [FEATURE] Querier: added support for `limitk()` and `limit_ratio()` experimental PromQL functions. Experimental functions are disabled by default, but can be enabled setting `-querier.promql-experimental-functions-enabled=true` in the query-frontend and querier. #8632
* [FEATURE] Querier: experimental support for `X-Mimir-Chunk-Info-Logger` header that triggers logging information about TSDB chunks loaded from ingesters and store-gateways in the querier. The header should contain the comma separated list of labels for which their value will be included in the logs. #8599
* [FEATURE] Query frontend: added new query pruning middleware to enable pruning dead code (eg. expressions that cannot produce any results) and simplifying expressions (eg. expressions that can be evaluated immediately) in queries. #9086
* [FEATURE] Ruler: added experimental configuration, `-ruler.rule-evaluation-write-enabled`, to disable writing the result of rule evaluation to ingesters. This feature can be used for testing purposes. #9060
* [FEATURE] Ingester: added experimental configuration `ingester.ignore-ooo-exemplars`. When set to `true` out of order exemplars are no longer reported to the remote write client. #9151
* [ENHANCEMENT] Compactor: Add `cortex_compactor_compaction_job_duration_seconds` and `cortex_compactor_compaction_job_blocks` histogram metrics to track duration of individual compaction jobs and number of blocks per job. #8371
Expand Down
11 changes: 11 additions & 0 deletions cmd/mimir/config-descriptor.json
Original file line number Diff line number Diff line change
Expand Up @@ -6367,6 +6367,17 @@
"fieldFlag": "query-frontend.parallelize-shardable-queries",
"fieldType": "boolean"
},
{
"kind": "field",
"name": "prune_queries",
"required": false,
"desc": "True to enable pruning dead code (eg. expressions that cannot produce any results) and simplifying expressions (eg. expressions that can be evaluated immediately) in queries.",
"fieldValue": null,
"fieldDefaultValue": false,
"fieldFlag": "query-frontend.prune-queries",
"fieldType": "boolean",
"fieldCategory": "experimental"
},
{
"kind": "field",
"name": "query_sharding_target_series_per_shard",
Expand Down
2 changes: 2 additions & 0 deletions cmd/mimir/help-all.txt.tmpl
Original file line number Diff line number Diff line change
Expand Up @@ -2085,6 +2085,8 @@ Usage of ./cmd/mimir/mimir:
Maximum time to wait for the query-frontend to become ready before rejecting requests received before the frontend was ready. 0 to disable (i.e. fail immediately if a request is received while the frontend is still starting up) (default 2s)
-query-frontend.parallelize-shardable-queries
True to enable query sharding.
-query-frontend.prune-queries
[experimental] True to enable pruning dead code (eg. expressions that cannot produce any results) and simplifying expressions (eg. expressions that can be evaluated immediately) in queries.
-query-frontend.querier-forget-delay duration
[experimental] If a querier disconnects without sending notification about graceful shutdown, the query-frontend will keep the querier in the tenant's shard until the forget delay has passed. This feature is useful to reduce the blast radius when shuffle-sharding is enabled.
-query-frontend.query-result-response-format string
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1674,6 +1674,12 @@ results_cache:
# CLI flag: -query-frontend.parallelize-shardable-queries
[parallelize_shardable_queries: <boolean> | default = false]

# (experimental) True to enable pruning dead code (eg. expressions that cannot
# produce any results) and simplifying expressions (eg. expressions that can be
# evaluated immediately) in queries.
# CLI flag: -query-frontend.prune-queries
[prune_queries: <boolean> | default = false]

# (advanced) How many series a single sharded partial query should load at most.
# This is not a strict requirement guaranteed to be honoured by query sharding,
# but a hint given to the query sharding when the query execution is initially
Expand Down
217 changes: 217 additions & 0 deletions pkg/frontend/querymiddleware/astmapper/pruning.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,217 @@
// SPDX-License-Identifier: AGPL-3.0-only

package astmapper

import (
"context"
"math"
"strconv"

"github.com/go-kit/log"
"github.com/prometheus/prometheus/promql/parser"
)

func NewQueryPruner(ctx context.Context, logger log.Logger) ASTMapper {
pruner := newQueryPruner(ctx, logger)
return NewASTExprMapper(pruner)
}

type queryPruner struct {
ctx context.Context
logger log.Logger
}

func newQueryPruner(ctx context.Context, logger log.Logger) *queryPruner {
return &queryPruner{
ctx: ctx,
logger: logger,
}
}

func (pruner *queryPruner) MapExpr(expr parser.Expr) (mapped parser.Expr, finished bool, err error) {
if err := pruner.ctx.Err(); err != nil {
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
}
}

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

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
}
}
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
default:
return expr
}

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

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

return expr
}

// 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
}
switch e := expr.(type) {
case *parser.ParenExpr:
return pruner.isInfinite(e.Expr)
case *parser.NumberLiteral:
if math.IsInf(e.Val, 1) {
return true, true
}
if math.IsInf(e.Val, -1) {
return true, false
}
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 expr
}

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
}
switch e := expr.(type) {
case *parser.ParenExpr:
return pruner.isEmpty(e.Expr)
default:
if e.String() == `vector(0) < -Inf` {
return true
}
return false
}
}
Loading

0 comments on commit 1e98df0

Please sign in to comment.