feat(engine): basic math expressions

pkg/engine/internal/executor/executor.go

@@ -95,6 +95,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))
 	default:
@@ -429,6 +431,14 @@ 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) > 1 {
+		return errorPipeline(ctx, fmt.Errorf("only one input is currently supported in mathExpression, got %d", len(inputs)))
+	}
+
+	return NewMathExpressionPipeline(expr, inputs, 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,88 @@
+package executor
+
+import (
+	"context"
+	"fmt"
+
+	"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, inputs []Pipeline, evaluator expressionEvaluator) *GenericPipeline {
+	return newGenericPipeline(Local, func(ctx context.Context, inputs []Pipeline) (arrow.Record, error) {
+		// Works only with a single input for now.
+		input := inputs[0]
+		batch, err := input.Read(ctx)
+		if err != nil {
+			return nil, err
+		}
+		defer batch.Release()
+
+		// TODO make sure all inputs matches on timestamps
+
+		fields := make([]arrow.Field, 0, len(inputs))
+		for i := range inputs {
+			fields = append(fields, semconv.FieldFromFQN(fmt.Sprintf("float64.generated.input_%d", i), false))
+		}
+
+		valColumnExpr := &physical.ColumnExpr{
+			Ref: types.ColumnRef{
+				Column: types.ColumnNameGeneratedValue,
+				Type:   types.ColumnTypeGenerated,
+			},
+		}
+		cols := make([]arrow.Array, 0, len(inputs))
+		inputVec, err := evaluator.eval(valColumnExpr, batch) // TODO read input[i] instead of batch
+		if err != nil {
+			return nil, err
+		}
+		inputData := inputVec.ToArray()
+		if inputData.DataType().ID() != arrow.FLOAT64 {
+			return nil, fmt.Errorf("expression returned non-float64 type %s", inputData.DataType())
+		}
+		inputCol := inputData.(*array.Float64)
+		defer inputCol.Release()
+
+		cols = append(cols, inputCol)
+
+		inputSchema := arrow.NewSchema(fields, nil)
+		evalInput := array.NewRecord(inputSchema, cols, batch.NumRows())
+		defer evalInput.Release()
+
+		res, err := evaluator.eval(expr.Expression, evalInput)
+		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()
+
+		schema := batch.Schema()
+		valueCol := semconv.NewIdentifier(types.ColumnNameGeneratedValue, types.ColumnTypeGenerated, types.Loki.Float)
+		outputFields := make([]arrow.Field, 0, schema.NumFields())
+		outputCols := make([]arrow.Array, 0, schema.NumFields())
+		for i := 0; i < schema.NumFields(); i++ {
+			field := schema.Field(i)
+			if field.Name != valueCol.FQN() {
+				outputFields = append(outputFields, field)
+				outputCols = append(outputCols, batch.Column(i))
+			}
+		}
+		outputFields = append(outputFields, semconv.FieldFromIdent(valueCol, false))
+		outputCols = append(outputCols, valCol)
+		outputSchema := arrow.NewSchema(outputFields, nil)
+
+		evaluatedRecord := array.NewRecord(outputSchema, outputCols, batch.NumRows())
+
+		return evaluatedRecord, nil
+	}, inputs...)
+}

pkg/engine/internal/executor/math_expression_test.go

@@ -0,0 +1,75 @@
+package executor
+
+import (
+	"testing"
+	"time"
+
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/apache/arrow-go/v18/arrow/memory"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/loki/v3/pkg/engine/internal/semconv"
+
+	"github.com/grafana/loki/v3/pkg/util/arrowtest"
+
+	"github.com/grafana/loki/v3/pkg/engine/internal/planner/physical"
+	"github.com/grafana/loki/v3/pkg/engine/internal/types"
+)
+
+func TestNewMathExpressionPipeline(t *testing.T) {
+	colTs := "timestamp_ns.builtin.timestamp"
+	colVal := "float64.generated.value"
+
+	t.Run("calculates input_1 / 10", func(t *testing.T) {
+		alloc := memory.NewCheckedAllocator(memory.DefaultAllocator)
+		defer alloc.AssertSize(t, 0)
+
+		schema := arrow.NewSchema([]arrow.Field{
+			semconv.FieldFromFQN(colTs, false),
+			semconv.FieldFromFQN(colVal, false),
+		}, nil)
+
+		rowsPipeline1 := []arrowtest.Rows{
+			{
+				{colTs: time.Unix(20, 0).UTC(), colVal: float64(230)},
+				{colTs: time.Unix(15, 0).UTC(), colVal: float64(120)},
+				{colTs: time.Unix(10, 0).UTC(), colVal: float64(260)},
+				{colTs: time.Unix(12, 0).UTC(), colVal: float64(250)},
+			},
+		}
+		input1 := NewArrowtestPipeline(alloc, schema, rowsPipeline1...)
+
+		mathExpr := &physical.MathExpression{
+			Expression: &physical.BinaryExpr{
+				Left: &physical.ColumnExpr{
+					Ref: types.ColumnRef{
+						Column: "input_0",
+						Type:   types.ColumnTypeGenerated,
+					},
+				},
+				Right: physical.NewLiteral(float64(10)),
+				Op:    types.BinaryOpDiv,
+			},
+		}
+
+		pipeline := NewMathExpressionPipeline(mathExpr, []Pipeline{input1}, expressionEvaluator{})
+		defer pipeline.Close()
+
+		// Read the pipeline output
+		record, err := pipeline.Read(t.Context())
+		require.NoError(t, err)
+		defer record.Release()
+
+		expect := arrowtest.Rows{
+			{colTs: time.Unix(20, 0).UTC(), colVal: float64(23)},
+			{colTs: time.Unix(15, 0).UTC(), colVal: float64(12)},
+			{colTs: time.Unix(10, 0).UTC(), colVal: float64(26)},
+			{colTs: time.Unix(12, 0).UTC(), colVal: float64(25)},
+		}
+
+		rows, err := arrowtest.RecordRows(record)
+		require.NoError(t, err, "should be able to convert record back to rows")
+		require.Equal(t, len(expect), len(rows), "number of rows should match")
+		require.ElementsMatch(t, expect, rows)
+	})
+}

pkg/engine/internal/executor/project.go

@@ -40,7 +40,6 @@ func NewProjectPipeline(input Pipeline, columns []physical.ColumnExpression, eva
 			if err != nil {
 				return nil, err
 			}
-			defer vec.Release()
 			ident := semconv.NewIdentifier(columnNames[i], vec.ColumnType(), vec.Type())
 			fields = append(fields, semconv.FieldFromIdent(ident, true))
 			arr := vec.ToArray()

pkg/engine/internal/executor/range_aggregation.go

@@ -163,7 +163,6 @@ func (r *rangeAggregationPipeline) read(ctx context.Context) (arrow.Record, erro
 				if err != nil {
 					return nil, err
 				}
-				defer vec.Release()
 
 				if vec.Type() != types.Loki.String {
 					return nil, fmt.Errorf("unsupported datatype for partitioning %s", vec.Type())
@@ -180,7 +179,6 @@ func (r *rangeAggregationPipeline) read(ctx context.Context) (arrow.Record, erro
 			if err != nil {
 				return nil, err
 			}
-			defer tsVec.Release()
 			tsCol := tsVec.ToArray().(*array.Timestamp)
 			defer tsCol.Release()
 

pkg/engine/internal/executor/sortmerge.go

@@ -154,14 +154,12 @@ loop:
 		if err != nil {
 			return nil, err
 		}
-		defer col.Release()
 
 		tsCol, ok := col.ToArray().(*array.Timestamp)
 		if !ok {
 			return nil, errors.New("column is not a timestamp column")
 		}
 		ts := tsCol.Value(int(p.offsets[i]))
-		tsCol.Release()
 
 		// Populate slices for sorting
 		inputIndexes = append(inputIndexes, i)
@@ -202,13 +200,11 @@ loop:
 	if err != nil {
 		return nil, err
 	}
-	defer col.Release()
 	// We assume the column is a Uint64 array
 	tsCol, ok := col.ToArray().(*array.Timestamp)
 	if !ok {
 		return nil, errors.New("column is not a timestamp column")
 	}
-	defer tsCol.Release()
 
 	// Calculate start/end of the sub-slice of the record
 	start := p.offsets[j]

pkg/engine/internal/executor/sortmerge_test.go

@@ -128,10 +128,8 @@ func TestSortMerge(t *testing.T) {
 			}
 
 			tsCol, err := c.evaluator.eval(merge.Column, batch)
-			defer tsCol.Release()
 			require.NoError(t, err)
 			arr := tsCol.ToArray().(*array.Timestamp)
-			defer arr.Release()
 
 			timestamps = append(timestamps, arr.Values()...)
 			batches++

pkg/engine/internal/executor/vector_aggregate.go

@@ -102,7 +102,6 @@ func (v *vectorAggregationPipeline) read(ctx context.Context) (arrow.Record, err
 			if err != nil {
 				return nil, err
 			}
-			defer tsVec.Release()
 			tsCol := tsVec.ToArray().(*array.Timestamp)
 			defer tsCol.Release()
 
@@ -111,7 +110,6 @@ func (v *vectorAggregationPipeline) read(ctx context.Context) (arrow.Record, err
 			if err != nil {
 				return nil, err
 			}
-			defer valueVec.Release()
 			valueArr := valueVec.ToArray().(*array.Float64)
 			defer valueArr.Release()
 
@@ -123,7 +121,6 @@ func (v *vectorAggregationPipeline) read(ctx context.Context) (arrow.Record, err
 				if err != nil {
 					return nil, err
 				}
-				defer vec.Release()
 
 				if vec.Type() != types.Loki.String {
 					return nil, fmt.Errorf("unsupported datatype for grouping %s", vec.Type())

pkg/engine/internal/planner/logical/builder.go

@@ -62,6 +62,28 @@ func (b *Builder) Sort(column ColumnRef, ascending, nullsFirst bool) *Builder {
 	}
 }
 
+// BinOpRight adds a binary arithmetic operation with a given right value
+func (b *Builder) BinOpRight(op types.BinaryOp, right Value) *Builder {
+	return &Builder{
+		val: &BinOp{
+			Left:  b.val,
+			Right: right,
+			Op:    op,
+		},
+	}
+}
+
+// BinOpLeft adds a binary arithmetic operation with a given left value
+func (b *Builder) BinOpLeft(op types.BinaryOp, left Value) *Builder {
+	return &Builder{
+		val: &BinOp{
+			Left:  left,
+			Right: b.val,
+			Op:    op,
+		},
+	}
+}
+
 // RangeAggregation applies a [RangeAggregation] operation to the Builder.
 func (b *Builder) RangeAggregation(
 	partitionBy []ColumnRef,