vam: Add fields() func

runtime/sam/expr/function/fields.go

@@ -19,8 +19,7 @@ func NewFields(zctx *zed.Context) *Fields {
 	}
 }
 
-func buildPath(typ *zed.TypeRecord, b *zcode.Builder, prefix []string) []string {
-	var out []string
+func buildPath(typ *zed.TypeRecord, b *zcode.Builder, prefix []string) {
 	for _, f := range typ.Fields {
 		if typ, ok := zed.TypeUnder(f.Type).(*zed.TypeRecord); ok {
 			buildPath(typ, b, append(prefix, f.Name))
@@ -33,17 +32,15 @@ func buildPath(typ *zed.TypeRecord, b *zcode.Builder, prefix []string) []string
 			b.EndContainer()
 		}
 	}
-	return out
 }
 
 func (f *Fields) Call(ectx expr.Context, args []zed.Value) zed.Value {
 	arena := ectx.Arena()
-	subjectVal := args[0]
+	subjectVal := args[0].Under(arena)
 	typ := f.recordType(subjectVal)
 	if typ == nil {
 		return f.zctx.Missing(ectx.Arena())
 	}
-	//XXX should have a way to append into allocator
 	var b zcode.Builder
 	buildPath(typ, &b, nil)
 	return arena.New(f.typ, b.Bytes())

runtime/vam/expr/function/fields.go

@@ -0,0 +1,89 @@
+package function
+
+import (
+	"github.com/brimdata/zed"
+	"github.com/brimdata/zed/vector"
+)
+
+// https://github.com/brimdata/zed/blob/main/docs/language/functions.md#fields
+type Fields struct {
+	zctx     *zed.Context
+	innerTyp *zed.TypeArray
+	outerTyp *zed.TypeArray
+}
+
+func NewFields(zctx *zed.Context) *Fields {
+	inner := zctx.LookupTypeArray(zed.TypeString)
+	return &Fields{
+		zctx:     zctx,
+		innerTyp: inner,
+		outerTyp: zctx.LookupTypeArray(inner),
+	}
+}
+
+func (f *Fields) Call(args ...vector.Any) vector.Any {
+	val := vector.Under(args[0])
+	switch typ := val.Type().(type) {
+	case *zed.TypeRecord:
+		paths := buildPath(typ, nil)
+		s := vector.NewStringEmpty(val.Len(), nil)
+		inOffs, outOffs := []uint32{0}, []uint32{0}
+		for i := uint32(0); i < val.Len(); i++ {
+			inOffs, outOffs = appendPaths(paths, s, inOffs, outOffs)
+		}
+		inner := vector.NewArray(f.innerTyp, inOffs, s, nil)
+		return vector.NewArray(f.outerTyp, outOffs, inner, nil)
+	case *zed.TypeOfType:
+		var errs []uint32
+		s := vector.NewStringEmpty(val.Len(), nil)
+		inOffs, outOffs := []uint32{0}, []uint32{0}
+		for i := uint32(0); i < val.Len(); i++ {
+			b, _ := vector.TypeValueValue(val, i)
+			rtyp := f.recordType(b)
+			if rtyp == nil {
+				errs = append(errs, i)
+				continue
+			}
+			inOffs, outOffs = appendPaths(buildPath(rtyp, nil), s, inOffs, outOffs)
+		}
+		inner := vector.NewArray(f.innerTyp, inOffs, s, nil)
+		out := vector.NewArray(f.outerTyp, outOffs, inner, nil)
+		if len(errs) > 0 {
+			return mix(val.Len(), out, errs, vector.NewStringError(f.zctx, "missing", uint32(len(errs))))
+		}
+		return out
+	default:
+		return vector.NewStringError(f.zctx, "missing", val.Len())
+	}
+}
+
+func (f *Fields) recordType(b []byte) *zed.TypeRecord {
+	typ, err := f.zctx.LookupByValue(b)
+	if err != nil {
+		return nil
+	}
+	rtyp, _ := typ.(*zed.TypeRecord)
+	return rtyp
+}
+
+func buildPath(typ *zed.TypeRecord, prefix []string) [][]string {
+	var out [][]string
+	for _, f := range typ.Fields {
+		if typ, ok := zed.TypeUnder(f.Type).(*zed.TypeRecord); ok {
+			out = append(out, buildPath(typ, append(prefix, f.Name))...)
+		} else {
+			out = append(out, append(prefix, f.Name))
+		}
+	}
+	return out
+}
+
+func appendPaths(paths [][]string, s *vector.String, inner, outer []uint32) ([]uint32, []uint32) {
+	for _, path := range paths {
+		for _, f := range path {
+			s.Append(f)
+		}
+		inner = append(inner, inner[len(inner)-1]+uint32(len(path)))
+	}
+	return inner, append(outer, outer[len(outer)-1]+uint32(len(paths)))
+}

runtime/vam/expr/function/function.go

@@ -1,6 +1,8 @@
 package function
 
 import (
+	"slices"
+
 	"github.com/brimdata/zed"
 	"github.com/brimdata/zed/pkg/field"
 	"github.com/brimdata/zed/runtime/sam/expr/function"
@@ -16,6 +18,8 @@ func New(zctx *zed.Context, name string, narg int) (expr.Function, field.Path, e
 	switch name {
 	case "base64":
 		f = &Base64{zctx}
+	case "fields":
+		f = NewFields(zctx)
 	case "hex":
 		f = &Hex{zctx}
 	case "join":
@@ -55,3 +59,23 @@ func underAll(args []vector.Any) []vector.Any {
 	}
 	return args
 }
+
+func mix(desiredLen uint32, vec vector.Any, index []uint32, add vector.Any) vector.Any {
+	var tags []uint32
+	var vecs []vector.Any
+	if variant, ok := vec.(*vector.Variant); ok {
+		vecs = variant.Values
+		tags = slices.Clone(variant.Tags)
+		if len(tags) < int(desiredLen) {
+			tags = slices.Grow(tags, int(desiredLen))[:desiredLen]
+		}
+	} else {
+		vecs = []vector.Any{vec}
+		tags = make([]uint32, desiredLen)
+	}
+	n := uint32(len(vecs))
+	for _, k := range index {
+		tags[k] = n
+	}
+	return vector.NewVariant(tags, append(vecs, add))
+}

runtime/ztests/expr/function/fields.yaml

@@ -0,0 +1,19 @@
+zed: fields(this)
+
+vector: true
+
+input: |
+  {r:{a:1(int32)},s:123(int32)}((string,{r:{a:int32},s:int32}))
+  null({r:{a:int32},s:int32})
+  {r:{a:1(int8),b:2(int8)},s:"a"}
+  <{r:{a:int32},s:int32}>
+  <{r:{a:int8,b:int8},s:string}>((int8,type))
+  <int64>
+
+output: |
+  [["r","a"],["s"]]
+  [["r","a"],["s"]]
+  [["r","a"],["r","b"],["s"]]
+  [["r","a"],["s"]]
+  [["r","a"],["r","b"],["s"]]
+  error("missing")