From 20f335b14188044eb23124e77259b8b928ac78a4 Mon Sep 17 00:00:00 2001
From: Matthew Nibecker <hello@mattnibecker.com>
Date: Tue, 7 Jan 2025 11:35:30 -0800
Subject: [PATCH] vector: Support Enums

---
 runtime/vam/expr/arith.go                 | 16 ++++++++++++++--
 runtime/vam/expr/cast/string.go           | 12 ++++++++++++
 runtime/vcache/loader.go                  |  8 ++++++++
 {ztests => runtime/ztests/expr}/enum.yaml |  2 ++
 vector/bool.go                            |  7 +++++++
 vector/enum.go                            | 19 +++++++++++++++++++
 vng/encoder.go                            |  2 ++
 7 files changed, 64 insertions(+), 2 deletions(-)
 rename {ztests => runtime/ztests/expr}/enum.yaml (95%)
 create mode 100644 vector/enum.go

diff --git a/runtime/vam/expr/arith.go b/runtime/vam/expr/arith.go
index df65859ccd..7cc28b77a8 100644
--- a/runtime/vam/expr/arith.go
+++ b/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
diff --git a/runtime/vam/expr/cast/string.go b/runtime/vam/expr/cast/string.go
index b516f2c74b..230d6d9c04 100644
--- a/runtime/vam/expr/cast/string.go
+++ b/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 {
diff --git a/runtime/vcache/loader.go b/runtime/vcache/loader.go
index 660f53a512..f2a9e06f8f 100644
--- a/runtime/vcache/loader.go
+++ b/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
 	}
diff --git a/ztests/enum.yaml b/runtime/ztests/expr/enum.yaml
similarity index 95%
rename from ztests/enum.yaml
rename to runtime/ztests/expr/enum.yaml
index a4d49518b0..d996114d3c 100644
--- a/ztests/enum.yaml
+++ b/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))}
diff --git a/vector/bool.go b/vector/bool.go
index ed3cac6671..93e4f76f4e 100644
--- a/vector/bool.go
+++ b/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,
diff --git a/vector/enum.go b/vector/enum.go
new file mode 100644
index 0000000000..435d70023d
--- /dev/null
+++ b/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 }
diff --git a/vng/encoder.go b/vng/encoder.go
index dd378d438f..7d4f661c00 100644
--- a/vng/encoder.go
+++ b/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))