vam: Array Expressions

compiler/kernel/vexpr.go

@@ -17,6 +17,8 @@ func (b *Builder) compileVamExpr(e dag.Expr) (vamexpr.Evaluator, error) {
 		return nil, errors.New("null expression not allowed")
 	}
 	switch e := e.(type) {
+	case *dag.ArrayExpr:
+		return b.compileVamArrayExpr(e)
 	case *dag.Literal:
 		val, err := zson.ParseValue(b.zctx(), e.Value)
 		if err != nil {
@@ -47,8 +49,6 @@ func (b *Builder) compileVamExpr(e dag.Expr) (vamexpr.Evaluator, error) {
 	//	return b.compileVamRegexpSearch(e)
 	case *dag.RecordExpr:
 		return b.compileVamRecordExpr(e)
-	//case *dag.ArrayExpr:
-	//	return b.compileVamArrayExpr(e)
 	//case *dag.SetExpr:
 	//	return b.compileVamSetExpr(e)
 	//case *dag.MapCall:
@@ -218,3 +218,32 @@ func (b *Builder) compileVamRecordExpr(e *dag.RecordExpr) (vamexpr.Evaluator, er
 	}
 	return vamexpr.NewRecordExpr(b.zctx(), elems), nil
 }
+
+func (b *Builder) compileVamArrayExpr(e *dag.ArrayExpr) (vamexpr.Evaluator, error) {
+	elems, err := b.compileVamVectorElems(e.Elems)
+	if err != nil {
+		return nil, err
+	}
+	return vamexpr.NewArrayExpr(b.zctx(), elems), nil
+}
+
+func (b *Builder) compileVamVectorElems(elems []dag.VectorElem) ([]vamexpr.VectorElem, error) {
+	var out []vamexpr.VectorElem
+	for _, elem := range elems {
+		switch elem := elem.(type) {
+		case *dag.Spread:
+			e, err := b.compileVamExpr(elem.Expr)
+			if err != nil {
+				return nil, err
+			}
+			out = append(out, vamexpr.VectorElem{Spread: e})
+		case *dag.VectorValue:
+			e, err := b.compileVamExpr(elem.Expr)
+			if err != nil {
+				return nil, err
+			}
+			out = append(out, vamexpr.VectorElem{Value: e})
+		}
+	}
+	return out, nil
+}

runtime/sam/expr/values.go

@@ -198,7 +198,7 @@ func (a *ArrayExpr) Eval(ectx Context, this super.Value) super.Value {
 		val := e.Spread.Eval(ectx, this)
 		inner := super.InnerType(val.Type())
 		if inner == nil {
-			// Treat non-list spread values values like missing.
+			a.collection.append(val)
 			continue
 		}
 		a.collection.appendSpread(inner, val.Bytes())

runtime/sam/expr/ztests/vector-spread.yaml

@@ -7,5 +7,5 @@ input: |
 output: |
   [2,3,0,1]
   |[0,1,2,3]|
-  [1,2]
-  |[1,2]|
+  [1,2,"hi"]
+  |[1,2,"hi"]|

runtime/vam/expr/arrayexpr.go

@@ -0,0 +1,112 @@
+package expr
+
+import (
+	"github.com/brimdata/super"
+	"github.com/brimdata/super/vector"
+)
+
+type VectorElem struct {
+	Value  Evaluator
+	Spread Evaluator
+}
+
+type ArrayExpr struct {
+	elems []VectorElem
+	zctx  *super.Context
+}
+
+func NewArrayExpr(zctx *super.Context, elems []VectorElem) *ArrayExpr {
+	return &ArrayExpr{
+		elems: elems,
+		zctx:  zctx,
+	}
+}
+
+func (a *ArrayExpr) Eval(this vector.Any) vector.Any {
+	var vecs []vector.Any
+	for _, e := range a.elems {
+		if e.Spread != nil {
+			vecs = append(vecs, e.Spread.Eval(this))
+		} else {
+			vecs = append(vecs, e.Value.Eval(this))
+		}
+	}
+	return vector.Apply(false, a.eval, vecs...)
+}
+
+func (a *ArrayExpr) eval(vecs ...vector.Any) vector.Any {
+	n := vecs[0].Len()
+	if n == 0 {
+		return vector.NewConst(super.Null, 0, nil)
+	}
+	spreadOffs := make([][]uint32, len(a.elems))
+	viewIndexes := make([][]uint32, len(a.elems))
+	for i, elem := range a.elems {
+		if elem.Spread != nil {
+			vecs[i], spreadOffs[i], viewIndexes[i] = a.unwrapSpread(vecs[i])
+		}
+	}
+	offsets := []uint32{0}
+	var tags []uint32
+	for i := range n {
+		var size uint32
+		for tag, spreadOff := range spreadOffs {
+			if len(spreadOff) == 0 {
+				tags = append(tags, uint32(tag))
+				size++
+				continue
+			} else {
+				if index := viewIndexes[tag]; index != nil {
+					i = index[i]
+				}
+				off := spreadOff[i]
+				for end := spreadOff[i+1]; off < end; off++ {
+					tags = append(tags, uint32(tag))
+					size++
+				}
+			}
+		}
+		offsets = append(offsets, offsets[i]+size)
+	}
+	var typ super.Type
+	var innerVals vector.Any
+	if len(vecs) == 1 {
+		typ = vecs[0].Type()
+		innerVals = vecs[0]
+	} else {
+		var all []super.Type
+		for _, vec := range vecs {
+			all = append(all, vec.Type())
+		}
+		types := super.UniqueTypes(all)
+		if len(types) == 1 {
+			typ = types[0]
+			innerVals = vector.NewDynamic(tags, vecs)
+		} else {
+			typ = a.zctx.LookupTypeUnion(types)
+			innerVals = vector.NewUnion(typ.(*super.TypeUnion), tags, vecs, nil)
+		}
+	}
+	return vector.NewArray(a.zctx.LookupTypeArray(typ), offsets, innerVals, nil)
+}
+
+func (a *ArrayExpr) unwrapSpread(vec vector.Any) (vector.Any, []uint32, []uint32) {
+	switch vec := vec.(type) {
+	case *vector.Array:
+		return vec.Values, vec.Offsets, nil
+	case *vector.Set:
+		return vec.Values, vec.Offsets, nil
+	case *vector.View:
+		vals, offsets, _ := a.unwrapSpread(vec.Any)
+		return vals, offsets, vec.Index
+	}
+	return vec, nil, nil
+}
+
+func viewOfOffsets(index []uint32, offsets []uint32) []uint32 {
+	var out []uint32
+	for _, i := range index {
+		out = append(out, offsets[i], offsets[i+1])
+	}
+	return out
+}

runtime/ztests/expr/array-spread.yaml

@@ -0,0 +1,11 @@
+zed: yield [...a,...b]
+
+vector: true
+
+input: |
+  {a:|[1,2]|,b:[0,1],c:"hi"}
+  {a:[1,2],b:"hi"}
+
+output: |
+  [1,2,0,1]
+  [1,2,"hi"]

vector/union.go

@@ -25,7 +25,8 @@ func (u *Union) Type() super.Type {
 
 func (u *Union) Serialize(b *zcode.Builder, slot uint32) {
 	b.BeginContainer()
-	b.Append(super.EncodeInt(int64(u.Tags[slot])))
+	tag := u.Typ.TagOf(u.TypeOf(slot))
+	b.Append(super.EncodeInt(int64(tag)))
 	u.Dynamic.Serialize(b, slot)
 	b.EndContainer()
 }