Add vector arithmetic

compiler/kernel/vexpr.go

@@ -98,8 +98,8 @@ func (b *Builder) compileVamBinary(e *dag.BinaryExpr) (vamexpr.Evaluator, error)
 	//	return vamexpr.NewIn(b.zctx(), lhs, rhs), nil
 	case "==", "!=", "<", "<=", ">", ">=":
 		return vamexpr.NewCompare(b.zctx(), lhs, rhs, op), nil
-	//case "+", "-", "*", "/", "%":
-	//	return vamexpr.NewArithmetic(b.zctx(), lhs, rhs, op)
+	case "+", "-", "*", "/", "%":
+		return vamexpr.NewArith(b.zctx(), lhs, rhs, op), nil
 	//case "[":
 	//	return vamexpr.NewIndexExpr(b.zctx(), lhs, rhs), nil
 	default:

runtime/vam/expr/arith.go

@@ -0,0 +1,52 @@
+package expr
+
+//go:generate go run genarithfuncs.go
+
+import (
+	"fmt"
+
+	"github.com/brimdata/zed"
+	"github.com/brimdata/zed/runtime/sam/expr/coerce"
+	"github.com/brimdata/zed/vector"
+)
+
+type Arith struct {
+	zctx   *zed.Context
+	opCode int
+	lhs    Evaluator
+	rhs    Evaluator
+}
+
+func NewArith(zctx *zed.Context, lhs, rhs Evaluator, op string) *Arith {
+	return &Arith{zctx, vector.ArithOpFromString(op), lhs, rhs}
+}
+
+func (a *Arith) Eval(val vector.Any) vector.Any {
+	return a.eval(a.lhs.Eval(val), a.rhs.Eval(val))
+}
+
+func (a *Arith) eval(lhs, rhs vector.Any) vector.Any {
+	lhs = vector.Under(lhs)
+	rhs = vector.Under(rhs)
+	lhs, rhs, errVal := coerceVals(a.zctx, lhs, rhs)
+	if errVal != nil {
+		return errVal
+	}
+	kind := vector.KindOf(lhs)
+	if kind != vector.KindOf(rhs) {
+		panic(fmt.Sprintf("vector kind mismatch after coerce (%#v and %#v)", lhs, rhs))
+	}
+	lform, ok := vector.FormOf(lhs)
+	if !ok {
+		return vector.NewStringError(a.zctx, coerce.ErrIncompatibleTypes.Error(), lhs.Len())
+	}
+	rform, ok := vector.FormOf(rhs)
+	if !ok {
+		return vector.NewStringError(a.zctx, coerce.ErrIncompatibleTypes.Error(), lhs.Len())
+	}
+	f, ok := arithFuncs[vector.FuncCode(a.opCode, kind, lform, rform)]
+	if !ok {
+		return vector.NewStringError(a.zctx, coerce.ErrIncompatibleTypes.Error(), lhs.Len())
+	}
+	return f(lhs, rhs)
+}

runtime/vam/expr/arith_test.go

@@ -0,0 +1,58 @@
+package expr
+
+import (
+	"testing"
+
+	"github.com/brimdata/zed"
+	"github.com/brimdata/zed/vector"
+	"github.com/stretchr/testify/assert"
+)
+
+// Test that Arith.Eval handles all ops for all vector forms.
+func TestArithOpsAndForms(t *testing.T) {
+	// These are both [0, 1, 2].
+	lhsFlat := vector.NewInt(zed.TypeInt64, []int64{0, 1, 2}, nil)
+	lhsDict := vector.NewDict(lhsFlat, []byte{0, 1, 2}, nil, nil)
+
+	// These are all [1, 1, 1].
+	rhsFlat := vector.NewInt(zed.TypeInt64, []int64{1, 1, 1}, nil)
+	rhsDict := vector.NewDict(rhsFlat, []byte{0, 0, 0}, nil, nil)
+	Const := vector.NewConst(nil, zed.NewInt64(1), 3, nil)
+
+	cases := []struct {
+		op                            string
+		expected, expectedForConstLHS []int64
+	}{
+		{"+", []int64{1, 2, 3}, []int64{2, 2, 2}},
+		{"-", []int64{-1, 0, 1}, []int64{0, 0, 0}},
+		{"*", []int64{0, 1, 2}, []int64{1, 1, 1}},
+		{"/", []int64{0, 1, 2}, []int64{1, 1, 1}},
+		{"%", []int64{0, 0, 0}, []int64{0, 0, 0}},
+	}
+	for _, c := range cases {
+		f := func(expected []int64, lhs, rhs vector.Any) {
+			t.Helper()
+			cmp := NewArith(zed.NewContext(), &testEval{lhs}, &testEval{rhs}, c.op)
+			assert.Equal(t, expected, cmp.Eval(nil).(*vector.Int).Values, "op: %s", c.op)
+		}
+
+		f(c.expected, lhsFlat, rhsFlat)
+		f(c.expected, lhsFlat, rhsDict)
+		f(c.expected, lhsFlat, Const)
+
+		f(c.expected, lhsDict, rhsFlat)
+		f(c.expected, lhsDict, rhsDict)
+		f(c.expected, lhsDict, Const)
+
+		f(c.expectedForConstLHS, Const, rhsDict)
+		f(c.expectedForConstLHS, Const, rhsFlat)
+
+		// Arithmetic on two vector.Consts yields another vector.Const.
+		cmp := NewArith(zed.NewContext(), &testEval{Const}, &testEval{Const}, c.op)
+		val := cmp.Eval(nil).(*vector.Const)
+		assert.Equal(t, uint32(3), val.Len(), "op: %s", c.op)
+		expected := zed.NewInt64(c.expectedForConstLHS[0])
+		assert.Equal(t, expected, val.Value(), "op: %s", c.op)
+	}
+
+}