Skip to content

Commit

Permalink
Implement vector put operator
Browse files Browse the repository at this point in the history
Unlike the sequence put operator, this does not yet support dynamic
field names (e.g., "put this[a] := 1") because it is implemented in
terms of record expressions.
  • Loading branch information
nwt committed Sep 16, 2024
1 parent 427a57b commit ff685a5
Show file tree
Hide file tree
Showing 19 changed files with 110 additions and 2 deletions.
33 changes: 33 additions & 0 deletions compiler/kernel/vop.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (

"github.com/brimdata/zed/compiler/ast/dag"
"github.com/brimdata/zed/compiler/optimizer"
"github.com/brimdata/zed/runtime/vam/expr"
vamop "github.com/brimdata/zed/runtime/vam/op"
"github.com/brimdata/zed/vector"
"github.com/brimdata/zed/zbuf"
Expand Down Expand Up @@ -81,6 +82,8 @@ func (b *Builder) compileVamLeaf(o dag.Op, parent vector.Puller) (vector.Puller,
} else {
return nil, fmt.Errorf("internal error: unhandled dag.Summarize: %#v", o)
}
case *dag.Put:
return b.compileVamPut(o, parent)
case *dag.Yield:
exprs, err := b.compileVamExprs(o.Exprs)
if err != nil {
Expand All @@ -95,6 +98,36 @@ func (b *Builder) compileVamLeaf(o dag.Op, parent vector.Puller) (vector.Puller,
}
}

func (b *Builder) compileVamPut(put *dag.Put, parent vector.Puller) (vector.Puller, error) {
elems := []dag.RecordElem{
&dag.Spread{Kind: "Spread", Expr: &dag.This{Kind: "This"}},
}
for _, a := range put.Args {
lhs, ok := a.LHS.(*dag.This)
if !ok {
return nil, fmt.Errorf("internal error: dynamic field name not yet supported in vector runtime: %#v", a.LHS)
}
elems = append(elems, newDagRecordExprForPath(lhs.Path, a.RHS).Elems...)
}
e, err := b.compileVamRecordExpr(&dag.RecordExpr{Kind: "RecordExpr", Elems: elems})
if err != nil {
return nil, err
}
return vamop.NewYield(b.rctx.Zctx, parent, []expr.Evaluator{expr.NewPutter(b.zctx(), e)}), nil
}

func newDagRecordExprForPath(path []string, expr dag.Expr) *dag.RecordExpr {
if len(path) > 1 {
expr = newDagRecordExprForPath(path[1:], expr)
}
return &dag.RecordExpr{
Kind: "RecordExpr",
Elems: []dag.RecordElem{
&dag.Field{Kind: "Field", Name: path[0], Value: expr},
},
}
}

func (b *Builder) compileVamSeq(seq dag.Seq, parents []vector.Puller) ([]vector.Puller, error) {
for _, o := range seq {
var err error
Expand Down
10 changes: 8 additions & 2 deletions runtime/vam/expr/coerce.go
Original file line number Diff line number Diff line change
Expand Up @@ -91,8 +91,6 @@ func coerceVals(zctx *zed.Context, a, b vector.Any) (vector.Any, vector.Any, vec
//return id, ok
}

//XXX need to handle const here

func promoteWider(id int, val vector.Any) vector.Any {
typ, err := zed.LookupPrimitiveByID(id)
if err != nil {
Expand All @@ -103,6 +101,14 @@ func promoteWider(id int, val vector.Any) vector.Any {
return val.Promote(typ)
case *vector.Uint:
return val.Promote(typ)
case *vector.Const:
var zedVal zed.Value
if zed.IsSigned(id) {
zedVal = zed.NewInt(typ, val.Value().Int())
} else {
zedVal = zed.NewUint(typ, val.Value().Uint())
}
return vector.NewConst(zedVal, val.Len(), val.Nulls)
case *vector.Dict:
promoted := val.Any.(vector.Promotable).Promote(typ)
return vector.NewDict(promoted, val.Index, val.Counts, val.Nulls)
Expand Down
30 changes: 30 additions & 0 deletions runtime/vam/expr/putter.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
package expr

import (
"github.com/brimdata/zed"
"github.com/brimdata/zed/vector"
)

// Putter adapts the behavior of recordExpr (obtained from NewRecordExpr) to
// match that of the put operator, which emits an error when an input value is
// not a record.
type Putter struct {
zctx *zed.Context
recordExpr Evaluator
}

func NewPutter(zctx *zed.Context, recordExpr Evaluator) *Putter {
return &Putter{zctx, recordExpr}
}

func (p *Putter) Eval(vec vector.Any) vector.Any {
return vector.Apply(false, p.eval, vec)
}

func (p *Putter) eval(vecs ...vector.Any) vector.Any {
vec := vecs[0]
if vec.Type().Kind() != zed.RecordKind {
return vector.NewWrappedError(p.zctx, "put: not a record", vec)
}
return p.recordExpr.Eval(vec)
}
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Tests a simple expression written into a new field
zed: put y := x + 1

vector: true

input: |
{x:1(int32)}
{x:2(int32)}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Tests a simple expression written into an existing field
zed: put x := x + 1

vector: true

input: |
{x:1(int32)}
{x:2(int32)}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Tests writing a new record
zed: put r2 := r

vector: true

input: |
{r:{s:"hello"}}
{r:{s:"world"}}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Tests overwriting a primitive value with a record
zed: put x := r

vector: true

input: |
{x:1(int32),r:{s:"hello"}}
{x:2(int32),r:{s:"world"}}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Tests overwriting a record with a primitive value
zed: put r := x

vector: true

input: |
{x:1(int32),r:{s:"hello"}}
{x:2(int32),r:{s:"world"}}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
zed: y := x + 1

vector: true

input: |
{x:1(int32)}
{x:2(int32)}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
zed: 'x:=upper(x), y:=lower(x)'

vector: true

input: |
{x:"vAlUe1"}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Tests multiple output expressions
zed: put a:=1, b:=2

vector: true

input: |
{x:1(int32)}
{x:2(int32)}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Test overwriting multiple fields in-order
zed: put a:=b, b:=a

vector: true

input: |
{a:1(int32),b:2(int32)}
{a:3(int32),b:4(int32)}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Test overwriting multiple fields out of order
zed: put b:=a, a:=b

vector: true

input: |
{a:1(int32),b:2(int32)}
{a:3(int32),b:4(int32)}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
# creates a new field
zed: put new:=1, x:=x+1

vector: true

input: |
{x:1(int32)}
{x:2(int32)}
Expand Down
9 changes: 9 additions & 0 deletions runtime/ztests/op/put-non-record.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
zed: put a:=1

vector: true

input: |
0
output: |
error({message:"put: not a record",on:0})
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
zed: put b:=null

vector: true

input: |
{a:1(int32)}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
# and this test will need to be updated.
zed: put c:=1,a:=3,b:=2

vector: true

input: |
{a:1(int32),b:2(int32),c:3(int32)}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
zed: put a:=3,b:=2,c:=1

vector: true

input: |
{a:1(int32),b:2(int32)}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
zed: yield a

vector: true

input: |
{"a": {"a": "a", "b": "b"}, "b": "b"}
Expand Down

0 comments on commit ff685a5

Please sign in to comment.