execution: cleanup scalar binops (#645)

Signed-off-by: Michael Hoffmann <[email protected]>

Author: MichaHoffmann

execution/binary/scalar.go

@@ -6,87 +6,61 @@ package binary
 import (
 	"context"
 	"fmt"
-	"math"
 	"sync"
 
 	"github.com/thanos-io/promql-engine/execution/model"
 	"github.com/thanos-io/promql-engine/execution/telemetry"
+	"github.com/thanos-io/promql-engine/execution/warnings"
 	"github.com/thanos-io/promql-engine/extlabels"
 	"github.com/thanos-io/promql-engine/query"
 
+	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql/parser"
 )
 
-type ScalarSide int
-
-const (
-	ScalarSideBoth ScalarSide = iota
-	ScalarSideLeft
-	ScalarSideRight
-)
-
 // scalarOperator evaluates expressions where one operand is a scalarOperator.
 type scalarOperator struct {
 	seriesOnce sync.Once
 	series     []labels.Labels
 
-	pool          *model.VectorPool
-	scalar        model.VectorOperator
-	next          model.VectorOperator
-	opType        parser.ItemType
-	getOperands   getOperandsFunc
-	operandValIdx int
-	floatOp       operation
-	histOp        histogramFloatOperation
+	pool   *model.VectorPool
+	lhs    model.VectorOperator
+	rhs    model.VectorOperator
+	opType parser.ItemType
 
 	// If true then return the comparison result as 0/1.
 	returnBool bool
 
-	// Keep the result if both sides are scalars.
-	bothScalars bool
+	lhsType parser.ValueType
+	rhsType parser.ValueType
 }
 
 func NewScalar(
 	pool *model.VectorPool,
-	next model.VectorOperator,
-	scalar model.VectorOperator,
-	op parser.ItemType,
-	scalarSide ScalarSide,
+	lhs model.VectorOperator,
+	rhs model.VectorOperator,
+	lhsType parser.ValueType,
+	rhsType parser.ValueType,
+	opType parser.ItemType,
 	returnBool bool,
 	opts *query.Options,
 ) (model.VectorOperator, error) {
-	binaryOperation, err := newOperation(op, scalarSide != ScalarSideBoth)
-	if err != nil {
-		return nil, err
-	}
-	// operandValIdx 0 means to get lhs as the return value
-	// while 1 means to get rhs as the return value.
-	operandValIdx := 0
-	getOperands := getOperandsScalarRight
-	if scalarSide == ScalarSideLeft {
-		getOperands = getOperandsScalarLeft
-		operandValIdx = 1
+	op := &scalarOperator{
+		pool:       pool,
+		lhs:        lhs,
+		rhs:        rhs,
+		lhsType:    lhsType,
+		rhsType:    rhsType,
+		opType:     opType,
+		returnBool: returnBool,
 	}
 
-	oper := &scalarOperator{
-		pool:          pool,
-		next:          next,
-		scalar:        scalar,
-		floatOp:       binaryOperation,
-		histOp:        getHistogramFloatOperation(op, scalarSide),
-		opType:        op,
-		getOperands:   getOperands,
-		operandValIdx: operandValIdx,
-		returnBool:    returnBool,
-		bothScalars:   scalarSide == ScalarSideBoth,
-	}
-
-	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts), oper), nil
+	return telemetry.NewOperator(telemetry.NewTelemetry(op, opts), op), nil
 }
 
 func (o *scalarOperator) Explain() (next []model.VectorOperator) {
-	return []model.VectorOperator{o.next, o.scalar}
+	return []model.VectorOperator{o.lhs, o.rhs}
 }
 
 func (o *scalarOperator) Series(ctx context.Context) ([]labels.Labels, error) {
@@ -109,77 +83,69 @@ func (o *scalarOperator) Next(ctx context.Context) ([]model.StepVector, error) {
 	default:
 	}
 
-	in, err := o.next.Next(ctx)
-	if err != nil {
-		return nil, err
-	}
-	if in == nil {
-		return nil, nil
-	}
+	var err error
 	o.seriesOnce.Do(func() { err = o.loadSeries(ctx) })
 	if err != nil {
 		return nil, err
 	}
 
-	scalarIn, err := o.scalar.Next(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	out := o.pool.GetVectorBatch()
-	for v, vector := range in {
-		step := o.pool.GetStepVector(vector.T)
-		scalarVal := math.NaN()
-		if len(scalarIn) > v && len(scalarIn[v].Samples) > 0 {
-			scalarVal = scalarIn[v].Samples[0]
+	var lhs []model.StepVector
+	var lerrChan = make(chan error, 1)
+	go func() {
+		var err error
+		lhs, err = o.lhs.Next(ctx)
+		if err != nil {
+			lerrChan <- err
 		}
+		close(lerrChan)
+	}()
 
-		for i := range vector.Samples {
-			operands := o.getOperands(vector, i, scalarVal)
-			val, keep := o.floatOp(operands, o.operandValIdx)
-			if o.returnBool {
-				if !o.bothScalars {
-					val = 0.0
-					if keep {
-						val = 1.0
-					}
-				}
-			} else if !keep {
-				continue
-			}
-			step.AppendSample(o.pool, vector.SampleIDs[i], val)
-		}
-
-		for i := range vector.HistogramIDs {
-			val := o.histOp(ctx, vector.Histograms[i], scalarVal)
-			if val != nil {
-				step.AppendHistogram(o.pool, vector.HistogramIDs[i], val)
-			}
-		}
+	rhs, rerr := o.rhs.Next(ctx)
+	lerr := <-lerrChan
+	if rerr != nil {
+		return nil, rerr
+	}
+	if lerr != nil {
+		return nil, lerr
+	}
 
-		out = append(out, step)
-		o.next.GetPool().PutStepVector(vector)
+	// TODO(fpetkovski): When one operator becomes empty,
+	// we might want to drain or close the other one.
+	// We don't have a concept of closing an operator yet.
+	if len(lhs) == 0 || len(rhs) == 0 {
+		return nil, nil
 	}
 
-	for i := range scalarIn {
-		o.scalar.GetPool().PutStepVector(scalarIn[i])
+	batch := o.pool.GetVectorBatch()
+	for i := range lhs {
+		if i < len(rhs) {
+			step := o.execBinaryOperation(ctx, lhs[i], rhs[i])
+			batch = append(batch, step)
+			o.rhs.GetPool().PutStepVector(rhs[i])
+		}
+		o.lhs.GetPool().PutStepVector(lhs[i])
 	}
+	o.lhs.GetPool().PutVectors(lhs)
+	o.rhs.GetPool().PutVectors(rhs)
 
-	o.next.GetPool().PutVectors(in)
-	o.scalar.GetPool().PutVectors(scalarIn)
+	return batch, nil
 
-	return out, nil
 }
 
 func (o *scalarOperator) GetPool() *model.VectorPool {
 	return o.pool
 }
 
 func (o *scalarOperator) loadSeries(ctx context.Context) error {
-	vectorSeries, err := o.next.Series(ctx)
+	vectorSide := o.lhs
+	if o.lhsType == parser.ValueTypeScalar {
+		vectorSide = o.rhs
+	}
+	vectorSeries, err := vectorSide.Series(ctx)
 	if err != nil {
 		return err
 	}
+
 	series := make([]labels.Labels, len(vectorSeries))
 	b := labels.ScratchBuilder{}
 	for i := range vectorSeries {
@@ -198,12 +164,61 @@ func (o *scalarOperator) loadSeries(ctx context.Context) error {
 	return nil
 }
 
-type getOperandsFunc func(v model.StepVector, i int, scalar float64) [2]float64
+func (o *scalarOperator) execBinaryOperation(ctx context.Context, lhs, rhs model.StepVector) model.StepVector {
+	ts := lhs.T
+	step := o.pool.GetStepVector(ts)
 
-func getOperandsScalarLeft(v model.StepVector, i int, scalar float64) [2]float64 {
-	return [2]float64{scalar, v.Samples[i]}
-}
+	scalar, other := lhs, rhs
+	if o.lhsType != parser.ValueTypeScalar {
+		scalar, other = rhs, lhs
+	}
+
+	var (
+		v    float64
+		h    *histogram.FloatHistogram
+		keep bool
+		err  error
+	)
+	for i, otherVal := range other.Samples {
+		scalarVal := scalar.Samples[0]
+
+		if o.lhsType == parser.ValueTypeScalar {
+			v, _, keep, err = binOp(o.opType, scalarVal, otherVal, nil, nil)
+		} else {
+			v, _, keep, err = binOp(o.opType, otherVal, scalarVal, nil, nil)
+		}
+		if err != nil {
+			warnings.AddToContext(err, ctx)
+			continue
+		}
+		// in comparison operations between scalars and vectors, the vectors are filtered, regardless if lhs or rhs
+		if keep && o.opType.IsComparisonOperator() && (o.lhsType == parser.ValueTypeVector || o.rhsType == parser.ValueTypeVector) {
+			v = otherVal
+		}
+		if o.returnBool {
+			v = 0.0
+			if keep {
+				v = 1.0
+			}
+		} else if !keep {
+			continue
+		}
+		step.AppendSample(o.pool, other.SampleIDs[i], v)
+	}
+	for i, otherVal := range other.Histograms {
+		scalarVal := scalar.Samples[0]
+
+		if o.lhsType == parser.ValueTypeScalar {
+			_, h, _, err = binOp(o.opType, scalarVal, 0., nil, otherVal)
+		} else {
+			_, h, _, err = binOp(o.opType, 0., scalarVal, otherVal, nil)
+		}
+		if err != nil {
+			warnings.AddToContext(err, ctx)
+			continue
+		}
+		step.AppendHistogram(o.pool, other.HistogramIDs[i], h)
+	}
 
-func getOperandsScalarRight(v model.StepVector, i int, scalar float64) [2]float64 {
-	return [2]float64{v.Samples[i], scalar}
+	return step
 }