Repository: polarsignals/frostdb
Pull request #669 (2024-01-03 03:49:48)

*: Add mathematical projections
After this, we can remove the average aggregation and instead have a logical optimizer that replaces an average expression with `sum(col) / count(col) as avg(col)`.

This is also some groundwork, so we can perform `count * period` calculations in Parca, which is required so we can compare profiling data with different periods.

@thorfour @asubiotto 


Patch:
diff --git a/query/logicalplan/expr.go b/query/logicalplan/expr.go
index 5b1887f..45e53ff 100644
--- a/query/logicalplan/expr.go
+++ b/query/logicalplan/expr.go
@@ -2,6 +2,7 @@ package logicalplan
 
 import (
 	"errors"
+	"fmt"
 	"regexp"
 	"strings"
 	"time"
@@ -27,6 +28,10 @@ const (
 	OpRegexNotMatch
 	OpAnd
 	OpOr
+	OpAdd
+	OpSub
+	OpMul
+	OpDiv
 )
 
 func (o Op) String() string {
@@ -51,6 +56,14 @@ func (o Op) String() string {
 		return "&&"
 	case OpOr:
 		return "||"
+	case OpAdd:
+		return "+"
+	case OpSub:
+		return "-"
+	case OpMul:
+		return "*"
+	case OpDiv:
+		return "/"
 	default:
 		panic("unknown operator")
 	}
@@ -94,8 +107,29 @@ func (e *BinaryExpr) Accept(visitor Visitor) bool {
 	return visitor.PostVisit(e)
 }
 
-func (e *BinaryExpr) DataType(_ *parquet.Schema) (arrow.DataType, error) {
-	return &arrow.BooleanType{}, nil
+func (e *BinaryExpr) DataType(s *parquet.Schema) (arrow.DataType, error) {
+	switch e.Op {
+	case OpEq, OpNotEq, OpLt, OpLtEq, OpGt, OpGtEq, OpAnd, OpOr:
+		return arrow.FixedWidthTypes.Boolean, nil
+	case OpAdd, OpSub, OpMul, OpDiv:
+		leftType, err := e.Left.DataType(s)
+		if err != nil {
+			return nil, fmt.Errorf("left operand: %w", err)
+		}
+
+		rightType, err := e.Right.DataType(s)
+		if err != nil {
+			return nil, fmt.Errorf("right operand: %w", err)
+		}
+
+		if arrow.TypeEqual(leftType, rightType, arrow.CheckMetadata()) {
+			return leftType, nil
+		}
+
+		return nil, fmt.Errorf("left and right operands must be of the same type, got %s and %s", leftType, rightType)
+	default:
+		return nil, errors.New("unknown operator")
+	}
 }
 
 func (e *BinaryExpr) Name() string {
@@ -287,6 +321,38 @@ func And(exprs ...Expr) Expr {
 	return and(exprs)
 }
 
+func Add(left, right Expr) *BinaryExpr {
+	return &BinaryExpr{
+		Left:  left,
+		Op:    OpAdd,
+		Right: right,
+	}
+}
+
+func Sub(left, right Expr) *BinaryExpr {
+	return &BinaryExpr{
+		Left:  left,
+		Op:    OpSub,
+		Right: right,
+	}
+}
+
+func Mul(left, right Expr) *BinaryExpr {
+	return &BinaryExpr{
+		Left:  left,
+		Op:    OpMul,
+		Right: right,
+	}
+}
+
+func Div(left, right Expr) *BinaryExpr {
+	return &BinaryExpr{
+		Left:  left,
+		Op:    OpDiv,
+		Right: right,
+	}
+}
+
 func and(exprs []Expr) Expr {
 	return computeBinaryExpr(exprs, OpAnd)
 }
diff --git a/query/physicalplan/project.go b/query/physicalplan/project.go
index f924dfb..18427f0 100644
--- a/query/physicalplan/project.go
+++ b/query/physicalplan/project.go
@@ -14,7 +14,9 @@ import (
 )
 
 type columnProjection interface {
+	// Name returns the name of the column that this projection will return.
 	Name() string
+	// Projects one or more columns. Each element in the field list corresponds to an element in the list of arrays.
 	Project(mem memory.Allocator, ar arrow.Record) ([]arrow.Field, []arrow.Array, error)
 }
 
@@ -34,7 +36,7 @@ func (a aliasProjection) Project(mem memory.Allocator, ar arrow.Record) ([]arrow
 		if err != nil {
 			return nil, nil, err
 		}
-		fields, array, err := binaryExprProjection{boolExpr: boolExpr}.Project(mem, ar)
+		fields, array, err := boolExprProjection{boolExpr: boolExpr}.Project(mem, ar)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -53,20 +55,252 @@ func (a aliasProjection) Project(mem memory.Allocator, ar arrow.Record) ([]arrow
 				return []arrow.Field{field}, []arrow.Array{ar.Column(i)}, nil
 			}
 		}
+	default:
+		return nil, nil, fmt.Errorf("unsupported alias expression type: %T", e)
 	}
 
 	return nil, nil, nil
 }
 
 type binaryExprProjection struct {
-	boolExpr BooleanExpression
+	expr *logicalplan.BinaryExpr
+
+	left  columnProjection
+	right columnProjection
 }
 
 func (b binaryExprProjection) Name() string {
-	return b.boolExpr.String()
+	return b.expr.String()
 }
 
 func (b binaryExprProjection) Project(mem memory.Allocator, ar arrow.Record) ([]arrow.Field, []arrow.Array, error) {
+	leftFields, leftArrays, err := b.left.Project(mem, ar)
+	if err != nil {
+		return nil, nil, fmt.Errorf("project left side of binary expression: %w", err)
+	}
+
+	rightFields, rightArrays, err := b.right.Project(mem, ar)
+	if err != nil {
+		return nil, nil, fmt.Errorf("project right side of binary expression: %w", err)
+	}
+
+	if len(leftFields) != 1 || len(rightFields) != 1 || len(leftArrays) != 1 || len(rightArrays) != 1 {
+		return nil, nil, fmt.Errorf("binary expression projection expected one field and one array for each side")
+	}
+
+	leftArray := leftArrays[0]
+	rightArray := rightArrays[0]
+
+	switch leftArray := leftArray.(type) {
+	case *array.Int64:
+		switch b.expr.Op {
+		case logicalplan.OpAdd:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{AddInt64s(mem, leftArray, rightArray.(*array.Int64))}, nil
+		case logicalplan.OpSub:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{SubInt64s(mem, leftArray, rightArray.(*array.Int64))}, nil
+		case logicalplan.OpMul:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{MulInt64s(mem, leftArray, rightArray.(*array.Int64))}, nil
+		case logicalplan.OpDiv:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{DivInt64s(mem, leftArray, rightArray.(*array.Int64))}, nil
+		default:
+			return nil, nil, fmt.Errorf("unsupported binary expression: %s", b.expr.String())
+		}
+	case *array.Float64:
+		switch b.expr.Op {
+		case logicalplan.OpAdd:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{AddFloat64s(mem, leftArray, rightArray.(*array.Float64))}, nil
+		case logicalplan.OpSub:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{SubFloat64s(mem, leftArray, rightArray.(*array.Float64))}, nil
+		case logicalplan.OpMul:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{MulFloat64s(mem, leftArray, rightArray.(*array.Float64))}, nil
+		case logicalplan.OpDiv:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{DivFloat64s(mem, leftArray, rightArray.(*array.Float64))}, nil
+		default:
+			return nil, nil, fmt.Errorf("unsupported binary expression: %s", b.expr.String())
+		}
+	case *array.Int32:
+		switch b.expr.Op {
+		case logicalplan.OpAdd:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{AddInt32s(mem, leftArray, rightArray.(*array.Int32))}, nil
+		case logicalplan.OpSub:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{SubInt32s(mem, leftArray, rightArray.(*array.Int32))}, nil
+		case logicalplan.OpMul:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{MulInt32s(mem, leftArray, rightArray.(*array.Int32))}, nil
+		case logicalplan.OpDiv:
+			return []arrow.Field{leftFields[0]}, []arrow.Array{DivInt32s(mem, leftArray, rightArray.(*array.Int32))}, nil
+		default:
+			return nil, nil, fmt.Errorf("unsupported binary expression: %s", b.expr.String())
+		}
+	default:
+		return nil, nil, fmt.Errorf("unsupported type: %T", leftArray)
+	}
+}
+
+func AddInt64s(mem memory.Allocator, left, right *array.Int64) *array.Int64 {
+	res := array.NewInt64Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) + right.Value(i))
+	}
+
+	return res.NewInt64Array()
+}
+
+func SubInt64s(mem memory.Allocator, left, right *array.Int64) *array.Int64 {
+	res := array.NewInt64Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) - right.Value(i))
+	}
+
+	return res.NewInt64Array()
+}
+
+func MulInt64s(mem memory.Allocator, left, right *array.Int64) *array.Int64 {
+	res := array.NewInt64Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) * right.Value(i))
+	}
+
+	return res.NewInt64Array()
+}
+
+func DivInt64s(mem memory.Allocator, left, right *array.Int64) *array.Int64 {
+	res := array.NewInt64Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) / right.Value(i))
+	}
+
+	return res.NewInt64Array()
+}
+
+func AddFloat64s(mem memory.Allocator, left, right *array.Float64) *array.Float64 {
+	res := array.NewFloat64Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) + right.Value(i))
+	}
+
+	return res.NewFloat64Array()
+}
+
+func SubFloat64s(mem memory.Allocator, left, right *array.Float64) *array.Float64 {
+	res := array.NewFloat64Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) - right.Value(i))
+	}
+
+	return res.NewFloat64Array()
+}
+
+func MulFloat64s(mem memory.Allocator, left, right *array.Float64) *array.Float64 {
+	res := array.NewFloat64Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) * right.Value(i))
+	}
+
+	return res.NewFloat64Array()
+}
+
+func DivFloat64s(mem memory.Allocator, left, right *array.Float64) *array.Float64 {
+	res := array.NewFloat64Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) / right.Value(i))
+	}
+
+	return res.NewFloat64Array()
+}
+
+func AddInt32s(mem memory.Allocator, left, right *array.Int32) *array.Int32 {
+	res := array.NewInt32Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) + right.Value(i))
+	}
+
+	return res.NewInt32Array()
+}
+
+func SubInt32s(mem memory.Allocator, left, right *array.Int32) *array.Int32 {
+	res := array.NewInt32Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) - right.Value(i))
+	}
+
+	return res.NewInt32Array()
+}
+
+func MulInt32s(mem memory.Allocator, left, right *array.Int32) *array.Int32 {
+	res := array.NewInt32Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) * right.Value(i))
+	}
+
+	return res.NewInt32Array()
+}
+
+func DivInt32s(mem memory.Allocator, left, right *array.Int32) *array.Int32 {
+	res := array.NewInt32Builder(mem)
+	defer res.Release()
+
+	res.Resize(left.Len())
+
+	for i := 0; i < left.Len(); i++ {
+		res.Append(left.Value(i) / right.Value(i))
+	}
+
+	return res.NewInt32Array()
+}
+
+type boolExprProjection struct {
+	boolExpr BooleanExpression
+}
+
+func (b boolExprProjection) Name() string {
+	return b.boolExpr.String()
+}
+
+func (b boolExprProjection) Project(mem memory.Allocator, ar arrow.Record) ([]arrow.Field, []arrow.Array, error) {
 	var partiallyComputedExprRes arrow.Array
 	if ar.Schema().HasField(b.boolExpr.String()) {
 		arr := ar.Column(ar.Schema().FieldIndices(b.boolExpr.String())[0])
@@ -187,13 +421,35 @@ func projectionFromExpr(expr logicalplan.Expr) (columnProjection, error) {
 			name: e.Name(),
 		}, nil
 	case *logicalplan.BinaryExpr:
-		boolExpr, err := binaryBooleanExpr(e)
-		if err != nil {
-			return nil, err
+		switch e.Op {
+		case logicalplan.OpEq, logicalplan.OpNotEq, logicalplan.OpGt, logicalplan.OpGtEq, logicalplan.OpLt, logicalplan.OpLtEq, logicalplan.OpRegexMatch, logicalplan.OpRegexNotMatch, logicalplan.OpAnd, logicalplan.OpOr:
+			boolExpr, err := binaryBooleanExpr(e)
+			if err != nil {
+				return nil, err
+			}
+			return boolExprProjection{
+				boolExpr: boolExpr,
+			}, nil
+		case logicalplan.OpAdd, logicalplan.OpSub, logicalplan.OpMul, logicalplan.OpDiv:
+			left, err := projectionFromExpr(e.Left)
+			if err != nil {
+				return nil, err
+			}
+
+			right, err := projectionFromExpr(e.Right)
+			if err != nil {
+				return nil, err
+			}
+
+			return binaryExprProjection{
+				expr: e,
+
+				left:  left,
+				right: right,
+			}, nil
+		default:
+			return nil, fmt.Errorf("unknown binary expression: %s", e.String())
 		}
-		return binaryExprProjection{
-			boolExpr: boolExpr,
-		}, nil
 	case *logicalplan.AverageExpr:
 		return &averageProjection{expr: e}, nil
 	default:
diff --git a/sqlparse/visitor.go b/sqlparse/visitor.go
index 462a138..52fc687 100644
--- a/sqlparse/visitor.go
+++ b/sqlparse/visitor.go
@@ -129,6 +129,14 @@ func (v *astVisitor) leaveImpl(n ast.Node) error {
 			frostDBOp = logicalplan.OpEq
 		case opcode.NE:
 			frostDBOp = logicalplan.OpNotEq
+		case opcode.Plus:
+			frostDBOp = logicalplan.OpAdd
+		case opcode.Minus:
+			frostDBOp = logicalplan.OpSub
+		case opcode.Mul:
+			frostDBOp = logicalplan.OpMul
+		case opcode.Div:
+			frostDBOp = logicalplan.OpDiv
 		case opcode.LogicAnd:
 			v.exprStack = append(v.exprStack, logicalplan.And(leftExpr, rightExpr))
 			return nil


Test patch:
diff --git a/logictest/testdata/exec/projection/math_projection b/logictest/testdata/exec/projection/math_projection
new file mode 100644
index 0000000..e6f3ad7
--- /dev/null
+++ b/logictest/testdata/exec/projection/math_projection
@@ -0,0 +1,15 @@
+createtable schema=default
+----
+
+insert cols=(labels.label1, timestamp, value)
+value1 1 2
+value1 3 4
+value1 5 6
+----
+
+exec
+select value * timestamp
+----
+2
+12
+30