feat(engine): basic math expressions

pkg/engine/internal/executor/executor.go

@@ -91,6 +91,8 @@ func (c *Context) execute(ctx context.Context, node physical.Node) Pipeline {
 		return tracePipeline("physical.VectorAggregation", c.executeVectorAggregation(ctx, n, inputs))
 	case *physical.ParseNode:
 		return tracePipeline("physical.ParseNode", c.executeParse(ctx, n, inputs))
+	case *physical.MathExpression:
+		return tracePipeline("physical.MathExpression", c.executeMathExpression(ctx, n, inputs))
 	case *physical.ColumnCompat:
 		return tracePipeline("physical.ColumnCompat", c.executeColumnCompat(ctx, n, inputs))
 	case *physical.Parallelize:
@@ -376,6 +378,18 @@ func (c *Context) executeParse(ctx context.Context, parse *physical.ParseNode, i
 	return NewParsePipeline(parse, inputs[0], allocator)
 }
 
+func (c *Context) executeMathExpression(ctx context.Context, expr *physical.MathExpression, inputs []Pipeline) Pipeline {
+	if len(inputs) == 0 {
+		return emptyPipeline()
+	}
+
+	if len(inputs) > 1 {
+		return errorPipeline(ctx, fmt.Errorf("math expression expects exactly one input, got %d", len(inputs)))
+	}
+
+	return NewMathExpressionPipeline(expr, inputs[0], c.evaluator)
+}
+
 func (c *Context) executeColumnCompat(ctx context.Context, compat *physical.ColumnCompat, inputs []Pipeline) Pipeline {
 	if len(inputs) == 0 {
 		return emptyPipeline()

pkg/engine/internal/executor/expressions.go

@@ -259,7 +259,7 @@ var _ ColumnVector = (*Array)(nil)
 
 // ToArray implements ColumnVector.
 func (a *Array) ToArray() arrow.Array {
-	a.array.Retain()
+	// this should already has 1 ref counter
 	return a.array
 }
 

pkg/engine/internal/executor/expressions_test.go

@@ -216,7 +216,6 @@ func TestEvaluateBinaryExpression(t *testing.T) {
 func collectBooleanColumnVector(vec ColumnVector) []bool {
 	res := make([]bool, 0, vec.Len())
 	arr := vec.ToArray().(*array.Boolean)
-	defer arr.Release()
 	for i := range int(vec.Len()) {
 		res = append(res, arr.Value(i))
 	}
@@ -327,10 +326,8 @@ null,null,null`
 		colVec, err := e.eval(colExpr, record)
 		require.NoError(t, err)
 		require.IsType(t, &CoalesceVector{}, colVec)
-		defer colVec.Release()
 
 		arr := colVec.ToArray()
-		defer arr.Release()
 		require.IsType(t, &array.String{}, arr)
 		stringArr := arr.(*array.String)
 

pkg/engine/internal/executor/filter.go

@@ -28,8 +28,6 @@ func NewFilterPipeline(filter *physical.Filter, input Pipeline, evaluator expres
 			if err != nil {
 				return nil, err
 			}
-			defer vec.Release()
-
 			arr := vec.ToArray()
 			// boolean filters are only used for filtering; they're not returned
 			// and must be released

pkg/engine/internal/executor/functions_test.go

@@ -120,7 +120,6 @@ func createFloat64Array(mem memory.Allocator, values []float64, nulls []bool) *A
 // Helper function to extract boolean values from result
 func extractBoolValues(result ColumnVector) []bool {
 	arr := result.ToArray().(*array.Boolean)
-	defer arr.Release()
 
 	values := make([]bool, arr.Len())
 	for i := 0; i < arr.Len(); i++ {
@@ -182,10 +181,34 @@ func TestBinaryFunctionRegistry_GetForSignature(t *testing.T) {
 			dataType:    arrow.BinaryTypes.String,
 			expectError: false,
 		},
+		{
+			name:        "valid div operation",
+			op:          types.BinaryOpDiv,
+			dataType:    arrow.PrimitiveTypes.Float64,
+			expectError: false,
+		},
+		{
+			name:        "valid add operation",
+			op:          types.BinaryOpAdd,
+			dataType:    arrow.PrimitiveTypes.Float64,
+			expectError: false,
+		},
+		{
+			name:        "valid Mul operation",
+			op:          types.BinaryOpMul,
+			dataType:    arrow.PrimitiveTypes.Float64,
+			expectError: false,
+		},
+		{
+			name:        "valid sub operation",
+			op:          types.BinaryOpSub,
+			dataType:    arrow.PrimitiveTypes.Float64,
+			expectError: false,
+		},
 		{
 			name:        "invalid operation",
-			op:          types.BinaryOpAdd, // Not registered
-			dataType:    arrow.PrimitiveTypes.Int64,
+			op:          types.BinaryOpAnd, // Not registered
+			dataType:    arrow.FixedWidthTypes.Boolean,
 			expectError: true,
 		},
 		{

pkg/engine/internal/executor/math_expression.go

@@ -0,0 +1,81 @@
+package executor
+
+import (
+	"context"
+	"fmt"
+	"slices"
+
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/apache/arrow-go/v18/arrow/array"
+
+	"github.com/grafana/loki/v3/pkg/engine/internal/types"
+
+	"github.com/grafana/loki/v3/pkg/engine/internal/semconv"
+
+	"github.com/grafana/loki/v3/pkg/engine/internal/planner/physical"
+)
+
+func NewMathExpressionPipeline(expr *physical.MathExpression, input Pipeline, evaluator expressionEvaluator) *GenericPipeline {
+	return newGenericPipeline(func(ctx context.Context, inputs []Pipeline) (arrow.Record, error) {
+		input := inputs[0]
+		batch, err := input.Read(ctx)
+		if err != nil {
+			return nil, err
+		}
+		defer batch.Release()
+
+		res, err := evaluator.eval(expr.Expression, batch)
+		if err != nil {
+			return nil, err
+		}
+		data := res.ToArray()
+		if data.DataType().ID() != arrow.FLOAT64 {
+			return nil, fmt.Errorf("expression returned non-float64 type %s", data.DataType())
+		}
+		valCol := data.(*array.Float64)
+		defer valCol.Release()
+
+		// Build output by moving all columns from the input except columns that were used for expression evaluation
+		// and add the result `value` column. In case of simple math expressions this will replace `value` with `value`.
+		// In case of math expressions that go after Joins this will replace `value_left` and `value_right` with `value`.
+		allColumnRefs := getAllColumnRefs(expr.Expression)
+		schema := batch.Schema()
+		outputFields := make([]arrow.Field, 0, schema.NumFields()+1)
+		outputCols := make([]arrow.Array, 0, schema.NumFields()+1)
+		for i := 0; i < schema.NumFields(); i++ {
+			if !slices.ContainsFunc(allColumnRefs, func(ref *types.ColumnRef) bool {
+				ident, err := semconv.ParseFQN(schema.Field(i).Name)
+				if err != nil {
+					return false
+				}
+				return ref.Column == ident.ShortName()
+			}) {
+				outputFields = append(outputFields, schema.Field(i))
+				outputCols = append(outputCols, batch.Column(i))
+			}
+		}
+
+		// Add `values` column
+		outputFields = append(outputFields, semconv.FieldFromIdent(semconv.ColumnIdentValue, false))
+		outputCols = append(outputCols, valCol)
+
+		outputSchema := arrow.NewSchema(outputFields, nil)
+		evaluatedRecord := array.NewRecord(outputSchema, outputCols, batch.NumRows())
+
+		return evaluatedRecord, nil
+	}, input)
+}
+
+// getAllColumnRefs finds all column ref from the expression. Should 1 or 2, but might be deep in the expression tree.
+func getAllColumnRefs(expr physical.Expression) []*types.ColumnRef {
+	switch expr := expr.(type) {
+	case *physical.BinaryExpr:
+		return append(getAllColumnRefs(expr.Left), getAllColumnRefs(expr.Right)...)
+	case *physical.UnaryExpr:
+		return getAllColumnRefs(expr.Left)
+	case *physical.ColumnExpr:
+		return []*types.ColumnRef{&expr.Ref}
+	}
+
+	return nil
+}