github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/querier/astmapper/subtree_folder.go (about)

     1  package astmapper
     2  
     3  import (
     4  	"github.com/prometheus/prometheus/promql/parser"
     5  )
     6  
     7  /*
     8  subtreeFolder is a NodeMapper which embeds an entire parser.Node in an embedded query
     9  if it does not contain any previously embedded queries. This allows the frontend to "zip up" entire
    10  subtrees of an AST that have not already been parallelized.
    11  
    12  */
    13  type subtreeFolder struct{}
    14  
    15  // NewSubtreeFolder creates a subtreeFolder which can reduce an AST
    16  // to one embedded query if it contains no embedded queries yet
    17  func NewSubtreeFolder() ASTMapper {
    18  	return NewASTNodeMapper(&subtreeFolder{})
    19  }
    20  
    21  // MapNode implements NodeMapper
    22  func (f *subtreeFolder) MapNode(node parser.Node) (parser.Node, bool, error) {
    23  	switch n := node.(type) {
    24  	// do not attempt to fold number or string leaf nodes
    25  	case *parser.NumberLiteral, *parser.StringLiteral:
    26  		return n, true, nil
    27  	}
    28  
    29  	containsEmbedded, err := Predicate(node, isEmbedded)
    30  	if err != nil {
    31  		return nil, true, err
    32  	}
    33  
    34  	if containsEmbedded {
    35  		return node, false, nil
    36  	}
    37  
    38  	expr, err := VectorSquasher(node)
    39  	return expr, true, err
    40  }
    41  
    42  func isEmbedded(node parser.Node) (bool, error) {
    43  	switch n := node.(type) {
    44  	case *parser.VectorSelector:
    45  		if n.Name == EmbeddedQueriesMetricName {
    46  			return true, nil
    47  		}
    48  
    49  	case *parser.MatrixSelector:
    50  		return isEmbedded(n.VectorSelector)
    51  	}
    52  	return false, nil
    53  }
    54  
    55  type predicate = func(parser.Node) (bool, error)
    56  
    57  // Predicate is a helper which uses parser.Walk under the hood determine if any node in a subtree
    58  // returns true for a specified function
    59  func Predicate(node parser.Node, fn predicate) (bool, error) {
    60  	v := &visitor{
    61  		fn: fn,
    62  	}
    63  
    64  	if err := parser.Walk(v, node, nil); err != nil {
    65  		return false, err
    66  	}
    67  	return v.result, nil
    68  }
    69  
    70  type visitor struct {
    71  	fn     predicate
    72  	result bool
    73  }
    74  
    75  // Visit implements parser.Visitor
    76  func (v *visitor) Visit(node parser.Node, path []parser.Node) (parser.Visitor, error) {
    77  	// if the visitor has already seen a predicate success, don't overwrite
    78  	if v.result {
    79  		return nil, nil
    80  	}
    81  
    82  	var err error
    83  
    84  	v.result, err = v.fn(node)
    85  	if err != nil {
    86  		return nil, err
    87  	}
    88  	if v.result {
    89  		return nil, nil
    90  	}
    91  	return v, nil
    92  }