all: make coalesce a logical node

engine/user_defined_test.go

@@ -86,6 +86,9 @@ func (c logicalVectorSelector) MakeExecutionOperator(_ context.Context, vectors
 		maxt:        opts.End.UnixMilli(),
 		step:        opts.Step.Milliseconds(),
 		currentStep: opts.Start.UnixMilli(),
+
+		shard:     c.VectorSelector.Shard,
+		numShards: c.VectorSelector.NumShards,
 	}
 
 	return oper, nil
@@ -103,6 +106,9 @@ type vectorSelectorOperator struct {
 	maxt        int64
 	step        int64
 	currentStep int64
+
+	shard     int
+	numShards int
 }
 
 func (c *vectorSelectorOperator) Next(ctx context.Context) ([]model.StepVector, error) {
@@ -112,8 +118,12 @@ func (c *vectorSelectorOperator) Next(ctx context.Context) ([]model.StepVector,
 	vectors := c.vectors.GetVectorBatch()
 	for i := 0; i < c.stepsBatch && c.currentStep <= c.maxt; i++ {
 		vector := c.vectors.GetStepVector(c.currentStep)
-		vector.AppendSample(c.vectors, 1, 7)
-		vector.AppendSample(c.vectors, 2, 7)
+		if c.shard == 0 {
+			vector.AppendSample(c.vectors, 1, 7)
+		}
+		if c.shard == 1 {
+			vector.AppendSample(c.vectors, 2, 7)
+		}
 		vectors = append(vectors, vector)
 		c.currentStep += c.step
 	}

execution/execution.go

@@ -55,6 +55,8 @@ func New(ctx context.Context, expr logicalplan.Node, storage storage.Scanners, o
 
 func newOperator(ctx context.Context, expr logicalplan.Node, storage storage.Scanners, opts *query.Options, hints promstorage.SelectHints) (model.VectorOperator, error) {
 	switch e := expr.(type) {
+	case *logicalplan.Coalesce:
+		return newCoalesce(ctx, e, storage, opts, hints)
 	case *logicalplan.NumberLiteral:
 		return scan.NewNumberLiteralSelector(model.NewVectorPool(opts.StepsBatch), opts, e.Val), nil
 	case *logicalplan.VectorSelector:
@@ -86,6 +88,18 @@ func newOperator(ctx context.Context, expr logicalplan.Node, storage storage.Sca
 	}
 }
 
+func newCoalesce(ctx context.Context, e *logicalplan.Coalesce, scanners storage.Scanners, opts *query.Options, hints promstorage.SelectHints) (model.VectorOperator, error) {
+	operators := make([]model.VectorOperator, len(e.Exprs))
+	for i, expr := range e.Exprs {
+		operator, err := newOperator(ctx, expr, scanners, opts, hints)
+		if err != nil {
+			return nil, err
+		}
+		operators[i] = operator
+	}
+	return exchange.NewCoalesce(model.NewVectorPool(opts.StepsBatch), opts, 0, operators...), nil
+}
+
 func newVectorSelector(ctx context.Context, e *logicalplan.VectorSelector, scanners storage.Scanners, opts *query.Options, hints promstorage.SelectHints) (model.VectorOperator, error) {
 	start, end := getTimeRangesForVectorSelector(e, opts, 0)
 	hints.Start = start

execution/function/absent.go

@@ -57,13 +57,19 @@ func (o *absentOperator) loadSeries() {
 
 		// https://github.com/prometheus/prometheus/blob/df1b4da348a7c2f8c0b294ffa1f05db5f6641278/promql/functions.go#L1857
 		var lm []*labels.Matcher
-		switch n := o.funcExpr.Args[0].(type) {
-		case *logicalplan.VectorSelector:
-			lm = append(n.LabelMatchers, n.Filters...)
-		case *logicalplan.MatrixSelector:
-			v := n.VectorSelector
-			lm = append(v.LabelMatchers, v.Filters...)
-		default:
+		logicalplan.TraverseBottomUp(nil, &o.funcExpr.Args[0], func(parent, node *logicalplan.Node) bool {
+			switch n := (*node).(type) {
+			case *logicalplan.VectorSelector:
+				lm = append(n.LabelMatchers, n.Filters...)
+			case *logicalplan.CheckDuplicateLabels, *logicalplan.Coalesce, *logicalplan.MatrixSelector:
+			// ignore logical nodes between the absent function and the first vectorselector
+			default:
+				lm = lm[:0]
+			}
+			return false
+		})
+
+		if len(lm) == 0 {
 			o.series = []labels.Labels{labels.EmptyLabels()}
 			return
 		}

logicalplan/coalesce.go

@@ -0,0 +1,69 @@
+// Copyright (c) The Thanos Community Authors.
+// Licensed under the Apache License 2.0.
+
+package logicalplan
+
+import (
+	"github.com/prometheus/prometheus/util/annotations"
+
+	"github.com/thanos-io/promql-engine/query"
+)
+
+type CoalesceOptimizer struct{}
+
+func (c CoalesceOptimizer) Optimize(expr Node, opts *query.Options) (Node, annotations.Annotations) {
+	numShards := opts.NumShards()
+
+	TraverseBottomUp(nil, &expr, func(parent, e *Node) bool {
+		switch t := (*e).(type) {
+		case *VectorSelector:
+			if parent != nil {
+				// we coalesce matrix selectors in a different branch
+				if _, ok := (*parent).(*MatrixSelector); ok {
+					return false
+				}
+			}
+			exprs := make([]Node, numShards)
+			for i := range numShards {
+				vs := t.Clone().(*VectorSelector)
+				vs.Shard = i
+				vs.NumShards = numShards
+				exprs[i] = vs
+			}
+			*e = &Coalesce{Exprs: exprs}
+		case *MatrixSelector:
+			// handled in *parser.Call branch
+			return false
+		case *FunctionCall:
+			// non-recursively handled in execution.go
+			if t.Func.Name == "absent_over_time" {
+				return true
+			}
+			var (
+				ms   *MatrixSelector
+				marg int
+			)
+			for i := range t.Args {
+				if arg, ok := t.Args[i].(*MatrixSelector); ok {
+					ms = arg
+					marg = i
+				}
+			}
+			if ms == nil {
+				return false
+			}
+			exprs := make([]Node, numShards)
+			for i := range numShards {
+				aux := ms.Clone().(*MatrixSelector)
+				aux.VectorSelector.Shard = i
+				aux.VectorSelector.NumShards = numShards
+				f := t.Clone().(*FunctionCall)
+				f.Args[marg] = aux
+				exprs[i] = f
+			}
+			*e = &Coalesce{Exprs: exprs}
+		}
+		return true
+	})
+	return expr, nil
+}

logicalplan/codec.go

@@ -219,6 +219,19 @@ func unmarshalNode(data []byte) (Node, error) {
 			return nil, err
 		}
 		return u, nil
+	case CoalesceNode:
+		n := &Coalesce{}
+		if err := json.Unmarshal(t.Data, n); err != nil {
+			return nil, err
+		}
+		for _, c := range t.Children {
+			child, err := unmarshalNode(c)
+			if err != nil {
+				return nil, err
+			}
+			n.Exprs = append(n.Exprs, child)
+		}
+		return n, nil
 	}
 	return nil, nil
 }

logicalplan/logical_nodes.go

@@ -25,14 +25,16 @@ const (
 	NumberLiteralNode  = "number_literal"
 	StringLiteralNode  = "string_literal"
 	SubqueryNode       = "subquery"
-	CheckDuplicateNode = "check_duplicate"
 	StepInvariantNode  = "step_invariant"
 	ParensNode         = "parens"
 	UnaryNode          = "unary"
 
 	RemoteExecutionNode = "remote_exec"
 	DeduplicateNode     = "dedup"
 	NoopNode            = "noop"
+
+	CheckDuplicateNode = "check_duplicate"
+	CoalesceNode       = "coalesce"
 )
 
 type Cloneable interface {
@@ -78,6 +80,9 @@ type VectorSelector struct {
 	// CounterResetHint, Count and Sum values populated. Histogram buckets and spans
 	// will not be used during query evaluation.
 	DecodeNativeHistogramStats bool
+
+	Shard     int
+	NumShards int
 }
 
 func (f *VectorSelector) Clone() Node {
@@ -94,6 +99,11 @@ func (f *VectorSelector) Clone() Node {
 		clone.Timestamp = &ts
 	}
 
+	clone.Shard = f.Shard
+	clone.NumShards = f.NumShards
+	clone.DecodeNativeHistogramStats = f.DecodeNativeHistogramStats
+	clone.BatchSize = f.BatchSize
+
 	return &clone
 }
 
@@ -570,3 +580,33 @@ func isAvgAggregation(expr *Node) bool {
 	}
 	return false
 }
+
+type Coalesce struct {
+	// We assume to always have at least one expression
+	Exprs []Node `json:"-"`
+}
+
+func (c *Coalesce) ReturnType() parser.ValueType { return parser.ValueTypeVector }
+
+func (c *Coalesce) Type() NodeType { return CoalesceNode }
+
+func (c *Coalesce) Clone() Node {
+	clone := *c
+	clone.Exprs = make([]Node, 0, len(c.Exprs))
+	for _, arg := range c.Exprs {
+		clone.Exprs = append(clone.Exprs, arg.Clone())
+	}
+	return &clone
+}
+
+func (c *Coalesce) Children() []*Node {
+	children := make([]*Node, 0, len(c.Exprs))
+	for i := range c.Exprs {
+		children = append(children, &c.Exprs[i])
+	}
+	return children
+}
+
+func (c *Coalesce) String() string {
+	return c.Exprs[0].String()
+}

logicalplan/plan.go

@@ -25,6 +25,7 @@ var DefaultOptimizers = []Optimizer{
 	SortMatchers{},
 	MergeSelectsOptimizer{},
 	DetectHistogramStatsOptimizer{},
+	CoalesceOptimizer{},
 }
 
 type Plan interface {
@@ -255,20 +256,22 @@ func replacePrometheusNodes(plan parser.Expr) Node {
 		return &MatrixSelector{
 			VectorSelector: &VectorSelector{
 				VectorSelector: t.VectorSelector.(*parser.VectorSelector),
+				Shard:          0,
+				NumShards:      1,
 			},
 			Range:          t.Range,
 			OriginalString: t.String(),
 		}
 	case *parser.VectorSelector:
-		return &VectorSelector{VectorSelector: t}
+		return &VectorSelector{VectorSelector: t, Shard: 0, NumShards: 1}
 
 	// TODO: we dont yet have logical nodes for these, keep traversing here but set fields in-place
 	case *parser.Call:
 		if t.Func.Name == "timestamp" {
 			// pushed-down timestamp function
 			switch v := UnwrapParens(t.Args[0]).(type) {
 			case *parser.VectorSelector:
-				return &VectorSelector{VectorSelector: v, SelectTimestamp: true}
+				return &VectorSelector{VectorSelector: v, SelectTimestamp: true, Shard: 0, NumShards: 1}
 			case *parser.StepInvariantExpr:
 				vs, ok := UnwrapParens(v.Expr).(*parser.VectorSelector)
 				if ok {
@@ -277,7 +280,7 @@ func replacePrometheusNodes(plan parser.Expr) Node {
 						vs.OriginalOffset = 0
 					}
 					return &StepInvariantExpr{
-						Expr: &VectorSelector{VectorSelector: vs, SelectTimestamp: true},
+						Expr: &VectorSelector{VectorSelector: vs, SelectTimestamp: true, Shard: 0, NumShards: 1},
 					}
 				}
 			}

query/options.go

@@ -20,6 +20,10 @@ type Options struct {
 	DecodingConcurrency      int
 }
 
+func (o *Options) NumShards() int {
+	return max(o.DecodingConcurrency, 1)
+}
+
 func (o *Options) NumSteps() int {
 	// Instant evaluation is executed as a range evaluation with one step.
 	if o.Step.Milliseconds() == 0 {

storage/prometheus/scanners.go

@@ -58,23 +58,17 @@ func (p Scanners) NewVectorSelector(
 		selector = newHistogramStatsSelector(selector)
 	}
 
-	operators := make([]model.VectorOperator, 0, opts.DecodingConcurrency)
-	for i := range opts.DecodingConcurrency {
-		operator := exchange.NewConcurrent(
-			NewVectorSelector(
-				model.NewVectorPool(opts.StepsBatch),
-				selector,
-				opts,
-				logicalNode.Offset,
-				logicalNode.BatchSize,
-				logicalNode.SelectTimestamp,
-				i,
-				opts.DecodingConcurrency,
-			), 2, opts)
-		operators = append(operators, operator)
-	}
-
-	return exchange.NewCoalesce(model.NewVectorPool(opts.StepsBatch), opts, logicalNode.BatchSize*int64(opts.DecodingConcurrency), operators...), nil
+	return exchange.NewConcurrent(
+		NewVectorSelector(
+			model.NewVectorPool(opts.StepsBatch),
+			selector,
+			opts,
+			logicalNode.Offset,
+			logicalNode.BatchSize,
+			logicalNode.SelectTimestamp,
+			logicalNode.Shard,
+			logicalNode.NumShards,
+		), 2, opts), nil
 }
 
 func (p Scanners) NewMatrixSelector(
@@ -126,28 +120,23 @@ func (p Scanners) NewMatrixSelector(
 		selector = newHistogramStatsSelector(selector)
 	}
 
-	operators := make([]model.VectorOperator, 0, opts.DecodingConcurrency)
-	for i := range opts.DecodingConcurrency {
-		operator, err := NewMatrixSelector(
-			model.NewVectorPool(opts.StepsBatch),
-			selector,
-			call.Func.Name,
-			arg,
-			arg2,
-			opts,
-			logicalNode.Range,
-			vs.Offset,
-			vs.BatchSize,
-			i,
-			opts.DecodingConcurrency,
-		)
-		if err != nil {
-			return nil, err
-		}
-		operators = append(operators, exchange.NewConcurrent(operator, 2, opts))
+	mat, err := NewMatrixSelector(
+		model.NewVectorPool(opts.StepsBatch),
+		selector,
+		call.Func.Name,
+		arg,
+		arg2,
+		opts,
+		logicalNode.Range,
+		vs.Offset,
+		vs.BatchSize,
+		vs.Shard,
+		vs.NumShards,
+	)
+	if err != nil {
+		return nil, err
 	}
-
-	return exchange.NewCoalesce(model.NewVectorPool(opts.StepsBatch), opts, vs.BatchSize*int64(opts.DecodingConcurrency), operators...), nil
+	return exchange.NewConcurrent(mat, 2, opts), nil
 }
 
 type histogramStatsSelector struct {