github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/querier/astmapper/parallel.go (about)

     1  package astmapper
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/go-kit/log/level"
     7  	"github.com/prometheus/prometheus/promql/parser"
     8  
     9  	util_log "github.com/cortexproject/cortex/pkg/util/log"
    10  )
    11  
    12  var summableAggregates = map[parser.ItemType]struct{}{
    13  	parser.SUM:     {},
    14  	parser.MIN:     {},
    15  	parser.MAX:     {},
    16  	parser.TOPK:    {},
    17  	parser.BOTTOMK: {},
    18  	parser.COUNT:   {},
    19  }
    20  
    21  var nonParallelFuncs = []string{
    22  	"histogram_quantile",
    23  	"quantile_over_time",
    24  	"absent",
    25  }
    26  
    27  // CanParallelize tests if a subtree is parallelizable.
    28  // A subtree is parallelizable if all of its components are parallelizable.
    29  func CanParallelize(node parser.Node) bool {
    30  	switch n := node.(type) {
    31  	case nil:
    32  		// nil handles cases where we check optional fields that are not set
    33  		return true
    34  
    35  	case parser.Expressions:
    36  		for _, e := range n {
    37  			if !CanParallelize(e) {
    38  				return false
    39  			}
    40  		}
    41  		return true
    42  
    43  	case *parser.AggregateExpr:
    44  		_, ok := summableAggregates[n.Op]
    45  		if !ok {
    46  			return false
    47  		}
    48  
    49  		// Ensure there are no nested aggregations
    50  		nestedAggs, err := Predicate(n.Expr, func(node parser.Node) (bool, error) {
    51  			_, ok := node.(*parser.AggregateExpr)
    52  			return ok, nil
    53  		})
    54  
    55  		return err == nil && !nestedAggs && CanParallelize(n.Expr)
    56  
    57  	case *parser.BinaryExpr:
    58  		// since binary exprs use each side for merging, they cannot be parallelized
    59  		return false
    60  
    61  	case *parser.Call:
    62  		if n.Func == nil {
    63  			return false
    64  		}
    65  		if !ParallelizableFunc(*n.Func) {
    66  			return false
    67  		}
    68  
    69  		for _, e := range n.Args {
    70  			if !CanParallelize(e) {
    71  				return false
    72  			}
    73  		}
    74  		return true
    75  
    76  	case *parser.SubqueryExpr:
    77  		return CanParallelize(n.Expr)
    78  
    79  	case *parser.ParenExpr:
    80  		return CanParallelize(n.Expr)
    81  
    82  	case *parser.UnaryExpr:
    83  		// Since these are only currently supported for Scalars, should be parallel-compatible
    84  		return true
    85  
    86  	case *parser.EvalStmt:
    87  		return CanParallelize(n.Expr)
    88  
    89  	case *parser.MatrixSelector, *parser.NumberLiteral, *parser.StringLiteral, *parser.VectorSelector:
    90  		return true
    91  
    92  	default:
    93  		level.Error(util_log.Logger).Log("err", fmt.Sprintf("CanParallel: unhandled node type %T", node)) //lint:ignore faillint allow global logger for now
    94  		return false
    95  	}
    96  
    97  }
    98  
    99  // ParallelizableFunc ensures that a promql function can be part of a parallel query.
   100  func ParallelizableFunc(f parser.Function) bool {
   101  
   102  	for _, v := range nonParallelFuncs {
   103  		if v == f.Name {
   104  			return false
   105  		}
   106  	}
   107  	return true
   108  }