diff --git a/compiler/kernel/vexpr.go b/compiler/kernel/vexpr.go index 391d077ae7..024c38e41a 100644 --- a/compiler/kernel/vexpr.go +++ b/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 +} diff --git a/runtime/sam/expr/values.go b/runtime/sam/expr/values.go index f63656af77..f027bd2513 100644 --- a/runtime/sam/expr/values.go +++ b/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()) @@ -238,7 +238,7 @@ func (a *SetExpr) 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()) diff --git a/runtime/sam/expr/ztests/array-spread.yaml b/runtime/sam/expr/ztests/array-spread.yaml deleted file mode 100644 index 5f14969b6c..0000000000 --- a/runtime/sam/expr/ztests/array-spread.yaml +++ /dev/null @@ -1,6 +0,0 @@ -zed: yield [...this] - -input: &input | - [{a:1},{b:2}] - -output: *input diff --git a/runtime/sam/expr/ztests/vector-spread.yaml b/runtime/sam/expr/ztests/vector-spread.yaml index 0894758339..150ea4a05f 100644 --- a/runtime/sam/expr/ztests/vector-spread.yaml +++ b/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"]| diff --git a/runtime/vam/expr/arrayexpr.go b/runtime/vam/expr/arrayexpr.go new file mode 100644 index 0000000000..ed32dc596c --- /dev/null +++ b/runtime/vam/expr/arrayexpr.go @@ -0,0 +1,104 @@ +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 +} diff --git a/runtime/ztests/expr/array-spread.yaml b/runtime/ztests/expr/array-spread.yaml new file mode 100644 index 0000000000..f35bc87bf7 --- /dev/null +++ b/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"] diff --git a/vector/union.go b/vector/union.go index 093a18d5c1..04e93ff7db 100644 --- a/vector/union.go +++ b/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() }