Data: (Breaking) Change enums to contain EnumItem (#632)

This makes FieldTypeFor and data.NewField work with enums

This is a breaking change for the newer Enum Field type since before this field would be passed unit16 values and is now passed a data.EnumItemIndex valuies instead.

Author: kylebrandt

data/arrow.go

@@ -607,17 +607,31 @@ func parseColumn(col array.Interface, i int, nullable []bool, frame *Frame) erro
 	case arrow.UINT16:
 		v := array.NewUint16Data(col.Data())
 		for rIdx := 0; rIdx < col.Len(); rIdx++ {
-			if nullable[i] {
-				if v.IsNull(rIdx) {
-					var ns *uint16
-					frame.Fields[i].vector.Append(ns)
+			if frame.Fields[i].Type().NullableType() == FieldTypeNullableEnum {
+				if nullable[i] {
+					if v.IsNull(rIdx) {
+						var ns *EnumItemIndex
+						frame.Fields[i].vector.Append(ns)
+						continue
+					}
+					rv := EnumItemIndex(v.Value(rIdx))
+					frame.Fields[i].vector.Append(&rv)
 					continue
 				}
-				rv := v.Value(rIdx)
-				frame.Fields[i].vector.Append(&rv)
-				continue
+				frame.Fields[i].vector.Append(EnumItemIndex(v.Value(rIdx)))
+			} else {
+				if nullable[i] {
+					if v.IsNull(rIdx) {
+						var ns *uint16
+						frame.Fields[i].vector.Append(ns)
+						continue
+					}
+					rv := v.Value(rIdx)
+					frame.Fields[i].vector.Append(&rv)
+					continue
+				}
+				frame.Fields[i].vector.Append(v.Value(rIdx))
 			}
-			frame.Fields[i].vector.Append(v.Value(rIdx))
 		}
 	case arrow.FLOAT32:
 		v := array.NewFloat32Data(col.Data())

data/arrow_column_builders.go

@@ -467,7 +467,7 @@ func buildNullableEnumColumn(pool memory.Allocator, field arrow.Field, vec *null
 			builder.AppendNull()
 			continue
 		}
-		builder.Append(*v)
+		builder.Append((uint16)(*v))
 	}
 
 	chunked := array.NewChunked(field.Type, []array.Interface{builder.NewArray()})
@@ -481,7 +481,7 @@ func buildEnumColumn(pool memory.Allocator, field arrow.Field, vec *enumVector)
 	defer builder.Release()
 
 	for _, v := range *vec {
-		builder.Append(v)
+		builder.Append(uint16(v))
 	}
 
 	chunked := array.NewChunked(field.Type, []array.Interface{builder.NewArray()})

data/arrow_test.go

@@ -271,7 +271,7 @@ func goldenDF() *data.Frame {
 			jsonRawMessagePtr(json.RawMessage("[{\"c\":3},{\"d\":4}]")),
 			jsonRawMessagePtr(json.RawMessage("{\"e\":{\"f\":5}}")),
 		}),
-		newField("enum", data.FieldTypeEnum, []uint16{
+		data.NewField("enum", nil, []data.EnumItemIndex{
 			1, 2, 2, 1, 1,
 		}).SetConfig(&data.FieldConfig{
 			TypeConfig: &data.FieldTypeConfig{
@@ -282,12 +282,12 @@ func goldenDF() *data.Frame {
 				},
 			},
 		}),
-		newField("nullable_enum", data.FieldTypeNullableEnum, []*uint16{
-			uint16Ptr(1),
-			uint16Ptr(2),
+		data.NewField("nullable_enum", nil, []*data.EnumItemIndex{
+			(*data.EnumItemIndex)(uint16Ptr(1)),
+			(*data.EnumItemIndex)(uint16Ptr(2)),
 			nil,
-			uint16Ptr(3),
-			uint16Ptr(0),
+			(*data.EnumItemIndex)(uint16Ptr(3)),
+			(*data.EnumItemIndex)(uint16Ptr(0)),
 		}),
 	).SetMeta(&data.FrameMeta{
 		Custom:              map[string]interface{}{"Hi": "there"},

data/field.go

@@ -58,6 +58,10 @@ type Fields []*Field
 //
 //	[]json.RawMessage, []*json.RawMessage
 //
+// Enum:
+//
+//	[]data.EnumItemIndex, []*data.EnumItemIndex
+//
 // If an unsupported values type is passed, NewField will panic.
 // nolint:gocyclo
 func NewField(name string, labels Labels, values interface{}) *Field {
@@ -119,6 +123,10 @@ func NewField(name string, labels Labels, values interface{}) *Field {
 		vec = newJsonRawMessageVectorWithValues(v)
 	case []*json.RawMessage:
 		vec = newNullableJsonRawMessageVectorWithValues(v)
+	case []EnumItemIndex:
+		vec = newEnumVectorWithValues(v)
+	case []*EnumItemIndex:
+		vec = newNullableEnumVectorWithValues(v)
 	default:
 		panic(fmt.Errorf("field '%s' specified with unsupported type %T", name, v))
 	}

data/field_type.go

@@ -84,9 +84,9 @@ const (
 	// FieldTypeNullableJSON indicates the underlying primitive is a []*json.RawMessage.
 	FieldTypeNullableJSON
 
-	// FieldTypeEnum indicates the underlying primitive is a []uint16, with field mapping metadata
+	// FieldTypeEnum indicates the underlying primitive is a []data.EnumItemIndex, with field mapping metadata
 	FieldTypeEnum
-	// FieldTypeNullableEnum indicates the underlying primitive is a []*uint16, with field mapping metadata
+	// FieldTypeNullableEnum indicates the underlying primitive is a []*data.EnumItemIndex, with field mapping metadata
 	FieldTypeNullableEnum
 )
 
@@ -116,10 +116,9 @@ func (p *FieldType) UnmarshalJSON(b []byte) error {
 
 // FieldTypeFor will return the FieldType that holds items of item's type.
 // If the FieldType is not recognized, FieldTypeUnknown is returned.
-// Enum fields are currently unrecognized.
 // For example, for an item of type *int8, FieldTypeNullableInt8 will be returned.
+// nolint:gocyclo
 func FieldTypeFor(item interface{}) FieldType {
-	// NOTE: enum does not have a native mapping ;(
 	switch item.(type) {
 	case int8:
 		return FieldTypeInt8
@@ -193,6 +192,11 @@ func FieldTypeFor(item interface{}) FieldType {
 		return FieldTypeJSON
 	case *json.RawMessage:
 		return FieldTypeNullableJSON
+
+	case EnumItemIndex:
+		return FieldTypeEnum
+	case *EnumItemIndex:
+		return FieldTypeNullableEnum
 	}
 
 	return FieldTypeUnknown

data/field_type_enum.go

@@ -1,27 +1,35 @@
 package data
 
 // this supports the enum type
-// it is diffent than the rest since it is backed by
-// a uint16, but has special semantics and interacts with the metadata
+// it is different than the rest since it is backed by
+// a uint16 but maps to the EnumItemIndex type, and has special semantics and interacts with the metadata
 // Unlike the other fields it can not be easily generated
 
-type enumVector []uint16
+type enumVector []EnumItemIndex
+
+type EnumItemIndex uint16
 
 func newEnumVector(n int) *enumVector {
-	v := enumVector(make([]uint16, n))
+	v := enumVector(make([]EnumItemIndex, n))
 	return &v
 }
 
+func newEnumVectorWithValues(s []EnumItemIndex) *enumVector {
+	v := make([]EnumItemIndex, len(s))
+	copy(v, s)
+	return (*enumVector)(&v)
+}
+
 func (v *enumVector) Set(idx int, i interface{}) {
-	(*v)[idx] = i.(uint16)
+	(*v)[idx] = i.(EnumItemIndex)
 }
 
 func (v *enumVector) SetConcrete(idx int, i interface{}) {
 	v.Set(idx, i)
 }
 
 func (v *enumVector) Append(i interface{}) {
-	*v = append(*v, i.(uint16))
+	*v = append(*v, i.(EnumItemIndex))
 }
 
 func (v *enumVector) At(i int) interface{} {
@@ -49,7 +57,7 @@ func (v *enumVector) Type() FieldType {
 }
 
 func (v *enumVector) Extend(i int) {
-	*v = append(*v, make([]uint16, i)...)
+	*v = append(*v, make([]EnumItemIndex, i)...)
 }
 
 func (v *enumVector) Insert(i int, val interface{}) {
@@ -69,23 +77,29 @@ func (v *enumVector) Delete(i int) {
 	*v = append((*v)[:i], (*v)[i+1:]...)
 }
 
-type nullableEnumVector []*uint16
+type nullableEnumVector []*EnumItemIndex
 
 func newNullableEnumVector(n int) *nullableEnumVector {
-	v := nullableEnumVector(make([]*uint16, n))
+	v := nullableEnumVector(make([]*EnumItemIndex, n))
 	return &v
 }
 
+func newNullableEnumVectorWithValues(s []*EnumItemIndex) *nullableEnumVector {
+	v := make([]*EnumItemIndex, len(s))
+	copy(v, s)
+	return (*nullableEnumVector)(&v)
+}
+
 func (v *nullableEnumVector) Set(idx int, i interface{}) {
 	if i == nil {
 		(*v)[idx] = nil
 		return
 	}
-	(*v)[idx] = i.(*uint16)
+	(*v)[idx] = i.(*EnumItemIndex)
 }
 
 func (v *nullableEnumVector) SetConcrete(idx int, i interface{}) {
-	val := i.(uint16)
+	val := i.(EnumItemIndex)
 	(*v)[idx] = &val
 }
 
@@ -94,7 +108,7 @@ func (v *nullableEnumVector) Append(i interface{}) {
 		*v = append(*v, nil)
 		return
 	}
-	*v = append(*v, i.(*uint16))
+	*v = append(*v, i.(*EnumItemIndex))
 }
 
 func (v *nullableEnumVector) At(i int) interface{} {
@@ -103,15 +117,15 @@ func (v *nullableEnumVector) At(i int) interface{} {
 
 func (v *nullableEnumVector) CopyAt(i int) interface{} {
 	if (*v)[i] == nil {
-		var g *uint16
+		var g *EnumItemIndex
 		return g
 	}
 	g := *(*v)[i]
 	return &g
 }
 
 func (v *nullableEnumVector) ConcreteAt(i int) (interface{}, bool) {
-	var g uint16
+	var g EnumItemIndex
 	val := (*v)[i]
 	if val == nil {
 		return g, false
@@ -133,7 +147,7 @@ func (v *nullableEnumVector) Type() FieldType {
 }
 
 func (v *nullableEnumVector) Extend(i int) {
-	*v = append(*v, make([]*uint16, i)...)
+	*v = append(*v, make([]*EnumItemIndex, i)...)
 }
 
 func (v *nullableEnumVector) Insert(i int, val interface{}) {

data/frame_json.gen.go

@@ -888,7 +888,7 @@ func readEnumVectorJSON(iter *jsoniter.Iterator, size int) (*enumVector, error)
 			iter.ReadNil()
 		} else {
 			v := iter.ReadUint16()
-			arr.Set(i, v)
+			arr.Set(i, EnumItemIndex(v))
 		}
 	}
 
@@ -911,7 +911,8 @@ func readNullableEnumVectorJSON(iter *jsoniter.Iterator, size int) (*nullableEnu
 			iter.ReadNil()
 		} else {
 			v := iter.ReadUint16()
-			arr.Set(i, &v)
+			eII := EnumItemIndex(v)
+			arr.Set(i, &eII)
 		}
 	}
 

data/frame_json.go

@@ -382,12 +382,18 @@ func jsonValuesToVector(ft FieldType, arr []interface{}) (vector, error) {
 			return uint8(iV), err
 		}
 
-	case FieldTypeUint16, FieldTypeEnum: // enums and uint16 share the same backings
+	case FieldTypeUint16: // enums and uint16 share the same backings
 		convert = func(v interface{}) (interface{}, error) {
 			iV, err := int64FromJSON(v)
 			return uint16(iV), err
 		}
 
+	case FieldTypeEnum: // enums and uint16 share the same backings
+		convert = func(v interface{}) (interface{}, error) {
+			iV, err := int64FromJSON(v)
+			return EnumItemIndex(iV), err
+		}
+
 	case FieldTypeUint32:
 		convert = func(v interface{}) (interface{}, error) {
 			iV, err := int64FromJSON(v)

data/frame_json_test.go

@@ -96,7 +96,7 @@ func TestFieldTypeToJSON(t *testing.T) {
 	assert.Equal(t, data.FieldTypeInt8, v.FType)
 	assert.Equal(t, data.FieldTypeTime, *v.FType2)
 
-	field := newField("enum", data.FieldTypeEnum, []uint16{
+	field := newField("enum", data.FieldTypeEnum, []data.EnumItemIndex{
 		1, 2, 2, 1, 1,
 	})