feat(engine): basic math expressions (#19407)

Author: spiridonov

pkg/engine/internal/executor/expressions_test.go

@@ -240,10 +240,7 @@ func collectBooleanArray(arr *array.Boolean) []bool {
 var words = []string{"one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten"}
 
 func batch(n int, now time.Time) arrow.Record {
-	// 1. Create a memory allocator
-	mem := memory.NewGoAllocator()
-
-	// 2. Define the schema
+	// Define the schema
 	schema := arrow.NewSchema(
 		[]arrow.Field{
 			semconv.FieldFromIdent(semconv.ColumnIdentMessage, false),
@@ -252,12 +249,11 @@ func batch(n int, now time.Time) arrow.Record {
 		nil, // No metadata
 	)
 
-	// 3. Create builders for each column
-	logBuilder := array.NewStringBuilder(mem)
-
-	tsBuilder := array.NewTimestampBuilder(mem, &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"})
+	// Create builders for each column
+	logBuilder := array.NewStringBuilder(memory.DefaultAllocator)
+	tsBuilder := array.NewTimestampBuilder(memory.DefaultAllocator, &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: "UTC"})
 
-	// 4. Append data to the builders
+	// Append data to the builders
 	logs := make([]string, n)
 	ts := make([]arrow.Timestamp, n)
 
@@ -269,12 +265,12 @@ func batch(n int, now time.Time) arrow.Record {
 	tsBuilder.AppendValues(ts, nil)
 	logBuilder.AppendValues(logs, nil)
 
-	// 5. Build the arrays
+	// Build the arrays
 	logArray := logBuilder.NewArray()
 
 	tsArray := tsBuilder.NewArray()
 
-	// 6. Create the record
+	// Create the record
 	columns := []arrow.Array{logArray, tsArray}
 	record := array.NewRecord(schema, columns, int64(n))
 

pkg/engine/internal/executor/functions.go

@@ -151,10 +151,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/functions_test.go

@@ -21,19 +21,25 @@ func NewProjectPipeline(input Pipeline, proj *physical.Projection, evaluator *ex
 
 	// Get the column names from the projection expressions
 	colRefs := make([]types.ColumnRef, 0, len(proj.Expressions))
-	unaryExprs := make([]physical.UnaryExpression, 0, len(proj.Expressions))
+	mathExprs := make([]physical.Expression, 0, len(proj.Expressions))
 
 	for i, expr := range proj.Expressions {
 		switch expr := expr.(type) {
 		case *physical.ColumnExpr:
 			colRefs = append(colRefs, expr.Ref)
 		case *physical.UnaryExpr:
-			unaryExprs = append(unaryExprs, expr)
+			mathExprs = append(mathExprs, expr)
+		case *physical.BinaryExpr:
+			mathExprs = append(mathExprs, expr)
 		default:
 			return nil, fmt.Errorf("projection expression %d is unsupported", i)
 		}
 	}
 
+	if len(mathExprs) > 1 {
+		return nil, fmt.Errorf("there might be only one math expression for `value` column at a time")
+	}
+
 	// Create KEEP projection pipeline:
 	// Drop all columns except the ones referenced in proj.Expressions.
 	if !proj.All && !proj.Drop && !proj.Expand {
@@ -66,9 +72,9 @@ func NewProjectPipeline(input Pipeline, proj *physical.Projection, evaluator *ex
 
 	// Create EXPAND projection pipeline:
 	// Keep all columns and expand the ones referenced in proj.Expressions.
-	// TODO: as implemented, expanding and keeping/dropping cannot happen in the same projection. Is this desired?
-	if proj.All && proj.Expand {
-		return newExpandPipeline(unaryExprs, evaluator, input)
+	// TODO: as implemented, epanding and keeping/dropping cannot happen in the same projection. Is this desired?
+	if proj.All && proj.Expand && len(mathExprs) > 0 {
+		return newExpandPipeline(mathExprs[0], evaluator, input)
 	}
 
 	return nil, errNotImplemented
@@ -105,7 +111,7 @@ func newKeepPipeline(colRefs []types.ColumnRef, keepFunc func([]types.ColumnRef,
 	}, input), nil
 }
 
-func newExpandPipeline(expressions []physical.UnaryExpression, evaluator *expressionEvaluator, input Pipeline) (*GenericPipeline, error) {
+func newExpandPipeline(expr physical.Expression, evaluator *expressionEvaluator, input Pipeline) (*GenericPipeline, error) {
 	return newGenericPipeline(func(ctx context.Context, inputs []Pipeline) (arrow.Record, error) {
 		if len(inputs) != 1 {
 			return nil, fmt.Errorf("expected 1 input, got %d", len(inputs))
@@ -116,34 +122,46 @@ func newExpandPipeline(expressions []physical.UnaryExpression, evaluator *expres
 			return nil, err
 		}
 
-		columns := []arrow.Array{}
-		fields := []arrow.Field{}
+		outputFields := make([]arrow.Field, 0)
+		outputCols := make([]arrow.Array, 0)
+		schema := batch.Schema()
 
+		// move all columns into the output except `value`
 		for i, field := range batch.Schema().Fields() {
-			columns = append(columns, batch.Column(i))
-			fields = append(fields, field)
-		}
-
-		for _, expr := range expressions {
-			vec, err := evaluator.eval(expr, batch)
+			ident, err := semconv.ParseFQN(schema.Field(i).Name)
 			if err != nil {
 				return nil, err
 			}
-			if arrStruct, ok := vec.(*array.Struct); ok {
-				structSchema, ok := arrStruct.DataType().(*arrow.StructType)
-				if !ok {
-					return nil, fmt.Errorf("unexpected type returned from evaluation, expected *arrow.StructType, got %T", arrStruct.DataType())
-				}
+			if !ident.Equal(semconv.ColumnIdentValue) {
+				outputCols = append(outputCols, batch.Column(i))
+				outputFields = append(outputFields, field)
+			}
+		}
 
-				for i := range arrStruct.NumField() {
-					columns = append(columns, arrStruct.Field(i))
-					fields = append(fields, structSchema.Field(i))
-				}
+		vec, err := evaluator.eval(expr, batch)
+		if err != nil {
+			return nil, err
+		}
+
+		switch arrCasted := vec.(type) {
+		case *array.Struct:
+			structSchema, ok := arrCasted.DataType().(*arrow.StructType)
+			if !ok {
+				return nil, fmt.Errorf("unexpected type returned from evaluation, expected *arrow.StructType, got %T", arrCasted.DataType())
+			}
+			for i := range arrCasted.NumField() {
+				outputCols = append(outputCols, arrCasted.Field(i))
+				outputFields = append(outputFields, structSchema.Field(i))
 			}
+		case *array.Float64:
+			outputFields = append(outputFields, semconv.FieldFromIdent(semconv.ColumnIdentValue, false))
+			outputCols = append(outputCols, arrCasted)
+		default:
+			return nil, fmt.Errorf("unexpected type returned from evaluation %T", arrCasted.DataType())
 		}
 
-		metadata := batch.Schema().Metadata()
-		schema := arrow.NewSchema(fields, &metadata)
-		return array.NewRecord(schema, columns, batch.NumRows()), nil
+		metadata := schema.Metadata()
+		outputSchema := arrow.NewSchema(outputFields, &metadata)
+		return array.NewRecord(outputSchema, outputCols, batch.NumRows()), nil
 	}, input), nil
 }