vector: Support Enums

runtime/vam/expr/arith.go

@@ -27,8 +27,8 @@ func (a *Arith) Eval(val vector.Any) vector.Any {
 }
 
 func (a *Arith) eval(vecs ...vector.Any) (out vector.Any) {
-	lhs := vector.Under(vecs[0])
-	rhs := vector.Under(vecs[1])
+	lhs := enumToIndex(vector.Under(vecs[0]))
+	rhs := enumToIndex(vector.Under(vecs[1]))
 	lhs, rhs, errVal := coerceVals(a.zctx, lhs, rhs)
 	if errVal != nil {
 		return errVal
@@ -67,6 +67,18 @@ func (a *Arith) eval(vecs ...vector.Any) (out vector.Any) {
 	return vector.CopyAndSetNulls(out, vector.Or(vector.NullsOf(lhs), vector.NullsOf(rhs)))
 }
 
+func enumToIndex(vec vector.Any) vector.Any {
+	switch vec := vec.(type) {
+	case *vector.View:
+		if enum, ok := vec.Any.(*vector.Enum); ok {
+			return vector.NewView(enum.Uint, vec.Index)
+		}
+	case *vector.Enum:
+		return vec.Uint
+	}
+	return vec
+}
+
 func (a *Arith) evalDivideByZero(kind vector.Kind, lhs, rhs vector.Any) vector.Any {
 	var errs []uint32
 	var out vector.Any

runtime/vam/expr/cast/string.go

@@ -89,6 +89,18 @@ func castToString(vec vector.Any, index []uint32) (vector.Any, []uint32, bool) {
 			bytes = append(bytes, vec.Values[idx].String()...)
 			offs = append(offs, uint32(len(bytes)))
 		}
+	case *vector.Enum:
+		for i := range n {
+			idx := i
+			if index != nil {
+				idx = index[i]
+			}
+			if !nulls.Value(i) {
+				val := vec.Uint.Values[idx]
+				bytes = append(bytes, vec.Typ.Symbols[val]...)
+			}
+			offs = append(offs, uint32(len(bytes)))
+		}
 	default:
 		var b zcode.Builder
 		for i := range n {

runtime/vcache/loader.go

@@ -286,6 +286,14 @@ func (l *loader) loadVals(typ super.Type, s *primitive, nulls *vector.Bool) (vec
 		}
 		offs[length] = off
 		return vector.NewTypeValue(offs, bytes, nulls), nil
+	case *super.TypeEnum:
+		values := make([]uint64, length)
+		for slot := range length {
+			if !nulls.Value(slot) {
+				values[slot] = super.DecodeUint(it.Next())
+			}
+		}
+		return vector.NewEnum(typ, values, nulls), nil
 	case *super.TypeOfNull:
 		return vector.NewConst(super.Null, s.length(), nil), nil
 	}

ztests/enum.yaml → runtime/ztests/expr/enum.yaml

@@ -1,5 +1,7 @@
 zed: "put s:=string(e), v:=e+1"
 
+vector: true
+
 input: |
   {e:%foo(enum(foo,bar,baz))}
   {e:%bar(enum(foo,bar,baz))}

vector/bool.go

@@ -188,6 +188,8 @@ func NullsOf(v Any) *Bool {
 		return v.Nulls
 	case *Dict:
 		return v.Nulls
+	case *Enum:
+		return v.Nulls
 	case *Error:
 		return Or(v.Nulls, NullsOf(v.Vals))
 	case *Float:
@@ -254,6 +256,11 @@ func CopyAndSetNulls(v Any, nulls *Bool) Any {
 			Counts: v.Counts,
 			Nulls:  nulls,
 		}
+	case *Enum:
+		return &Enum{
+			Typ:  v.Typ,
+			Uint: CopyAndSetNulls(v.Uint, nulls).(*Uint),
+		}
 	case *Error:
 		return &Error{
 			Typ:   v.Typ,

vector/enum.go

@@ -0,0 +1,19 @@
+package vector
+
+import (
+	"github.com/brimdata/super"
+)
+
+type Enum struct {
+	*Uint
+	Typ *super.TypeEnum
+}
+
+func NewEnum(typ *super.TypeEnum, vals []uint64, nulls *Bool) *Enum {
+	return &Enum{
+		Typ:  typ,
+		Uint: NewUint(super.TypeUint64, vals, nulls),
+	}
+}
+
+func (e *Enum) Type() super.Type { return e.Typ }

vng/encoder.go

@@ -46,6 +46,8 @@ func NewEncoder(typ super.Type) Encoder {
 		return NewNullsEncoder(NewMapEncoder(typ))
 	case *super.TypeUnion:
 		return NewNullsEncoder(NewUnionEncoder(typ))
+	case *super.TypeEnum:
+		return NewNullsEncoder(NewPrimitiveEncoder(typ, false))
 	default:
 		if !super.IsPrimitiveType(typ) {
 			panic(fmt.Sprintf("unsupported type in VNG file: %T", typ))