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/functions_test.go

@@ -180,10 +180,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,92 @@
+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/semconv"
+
+	"github.com/grafana/loki/v3/pkg/engine/internal/planner/physical"
+	"github.com/grafana/loki/v3/pkg/engine/internal/types"
+)
+
+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, arrow.Field{
+				Name: fmt.Sprintf("float64.generated.input_%d", i),
+				Type: types.Arrow.Float,
+			})
+		}
+
+		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()
+
+		tsColumnExpr := &physical.ColumnExpr{
+			Ref: types.ColumnRef{
+				Column: types.ColumnNameBuiltinTimestamp,
+				Type:   types.ColumnTypeBuiltin,
+			},
+		}
+		tsVec, err := evaluator.eval(tsColumnExpr, batch)
+		if err != nil {
+			return nil, err
+		}
+		tsCol := tsVec.ToArray().(*array.Timestamp)
+		defer tsCol.Release()
+
+		outputSchema := arrow.NewSchema([]arrow.Field{
+			semconv.FieldFromIdent(semconv.NewIdentifier(types.ColumnNameBuiltinTimestamp, types.ColumnTypeBuiltin, types.Loki.Timestamp), true),
+			semconv.FieldFromIdent(semconv.NewIdentifier(types.ColumnNameGeneratedValue, types.ColumnTypeGenerated, types.Loki.Float), true),
+		}, nil)
+		evaluatedRecord := array.NewRecord(outputSchema, []arrow.Array{tsCol, valCol}, 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/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,