Skip to content

Commit

Permalink
put: Assignment to array or set
Browse files Browse the repository at this point in the history
Update the put operator to allow assignments to elements in an array or
set.

Closes #4798
  • Loading branch information
mattnibs committed Dec 15, 2023
1 parent f530788 commit 692a0a3
Show file tree
Hide file tree
Showing 14 changed files with 606 additions and 318 deletions.
10 changes: 0 additions & 10 deletions compiler/semantic/op.go
Original file line number Diff line number Diff line change
Expand Up @@ -619,16 +619,6 @@ func (a *analyzer) semOp(o ast.Op, seq dag.Seq) (dag.Seq, error) {
if err != nil {
return nil, err
}
// We can do collision checking on static paths, so check what we can.
var fields field.List
for _, a := range assignments {
if this, ok := a.LHS.(*dag.This); ok {
fields = append(fields, this.Path)
}
}
if err := expr.CheckPutFields(fields); err != nil {
return nil, fmt.Errorf("put: %w", err)
}
return append(seq, &dag.Put{
Kind: "Put",
Args: assignments,
Expand Down
4 changes: 2 additions & 2 deletions runtime/expr/cutter.go
Original file line number Diff line number Diff line change
Expand Up @@ -94,11 +94,11 @@ func (c *Cutter) Eval(ectx Context, in *zed.Value) *zed.Value {
func (c *Cutter) lookupBuilder(ectx Context, in *zed.Value) (*recordBuilderCachedTypes, field.List, error) {
paths := c.fieldRefs[:0]
for _, p := range c.lvals {
path, err := p.Eval(ectx, in)
path, err := p.EvalAsRecordPath(ectx, in)
if err != nil {
return nil, nil, err
}
if path.IsEmpty() {
if len(path) == 0 {
return nil, nil, errors.New("'this' not allowed (use record literal)")
}
paths = append(paths, path)
Expand Down
64 changes: 64 additions & 0 deletions runtime/expr/dynfield/path.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
package dynfield

import (
"slices"

"github.com/brimdata/zed"
"github.com/brimdata/zed/pkg/field"
"github.com/brimdata/zed/zson"
)

type Path []zed.Value

func (p Path) Append(b []byte) []byte {
for i, v := range p {
if i > 0 {
b = append(b, 0)
}
b = append(b, v.Bytes()...)
}
return b
}

func (p Path) Equal(to Path) bool {
return slices.EqualFunc(p, to, func(a, b zed.Value) bool {
return a.Equal(b)
})
}

func (p Path) appendAsField(f field.Path) field.Path {
for _, v := range p {
if !v.IsString() {
return nil
}
f = append(f, v.String())
}
return f
}

func (p Path) HasPrefix(prefix Path) bool {
return len(p) >= len(prefix) && prefix.Equal(p[:len(prefix)])
}

func (p Path) String() string {
var b []byte
for i, v := range p {
if i > 0 {
b = append(b, '.')
}
b = append(b, zson.FormatValue(&v)...)
}
return string(b)
}

type List []Path

func (l List) Append(b []byte) []byte {
for i, path := range l {
if i > 0 {
b = append(b, ',')
}
b = path.Append(b)
}
return b
}
47 changes: 31 additions & 16 deletions runtime/expr/lval.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,14 @@ import (

"github.com/brimdata/zed"
"github.com/brimdata/zed/pkg/field"
"github.com/brimdata/zed/runtime/expr/dynfield"
"github.com/brimdata/zed/zson"
)

type Lval struct {
Elems []LvalElem
cache field.Path
Elems []LvalElem
cache []zed.Value
fieldCache field.Path
}

func NewLval(evals []LvalElem) *Lval {
Expand All @@ -19,18 +21,36 @@ func NewLval(evals []LvalElem) *Lval {

// Eval returns the path of the lval. If there's an error the returned *zed.Value
// will not be nill.
func (l *Lval) Eval(ectx Context, this *zed.Value) (field.Path, error) {
func (l *Lval) Eval(ectx Context, this *zed.Value) (dynfield.Path, error) {
l.cache = l.cache[:0]
for _, e := range l.Elems {
name, err := e.Eval(ectx, this)
val, err := e.Eval(ectx, this)
if err != nil {
return nil, err
}
l.cache = append(l.cache, name)
l.cache = append(l.cache, *val)
}
return l.cache, nil
}

func (l *Lval) EvalAsRecordPath(ectx Context, this *zed.Value) (field.Path, error) {
l.fieldCache = l.fieldCache[:0]
for _, e := range l.Elems {
val, err := e.Eval(ectx, this)
if err != nil {
return nil, err
}
if !val.IsString() {
// XXX Add context to error so we know what element is failing but
// let's wait until we can test this so we have a feel for what we
// want to see.
return nil, errors.New("field reference is not a string")
}
l.fieldCache = append(l.fieldCache, val.AsString())
}
return l.fieldCache, nil
}

// Path returns the receiver's path. Path returns false when the receiver
// contains a dynamic element.
func (l *Lval) Path() (field.Path, bool) {
Expand All @@ -46,15 +66,15 @@ func (l *Lval) Path() (field.Path, bool) {
}

type LvalElem interface {
Eval(ectx Context, this *zed.Value) (string, error)
Eval(ectx Context, this *zed.Value) (*zed.Value, error)
}

type StaticLvalElem struct {
Name string
}

func (l *StaticLvalElem) Eval(_ Context, _ *zed.Value) (string, error) {
return l.Name, nil
func (l *StaticLvalElem) Eval(_ Context, _ *zed.Value) (*zed.Value, error) {
return zed.NewString(l.Name), nil
}

type ExprLvalElem struct {
Expand All @@ -69,17 +89,12 @@ func NewExprLvalElem(zctx *zed.Context, e Evaluator) *ExprLvalElem {
}
}

func (l *ExprLvalElem) Eval(ectx Context, this *zed.Value) (string, error) {
func (l *ExprLvalElem) Eval(ectx Context, this *zed.Value) (*zed.Value, error) {
val := l.eval.Eval(ectx, this)
if val.IsError() {
return "", lvalErr(ectx, val)
}
if !val.IsString() {
if val = l.caster.Eval(ectx, val); val.IsError() {
return "", errors.New("field reference is not a string")
}
return nil, lvalErr(ectx, val)
}
return val.AsString(), nil
return val, nil
}

func lvalErr(ectx Context, errVal *zed.Value) error {
Expand Down
97 changes: 97 additions & 0 deletions runtime/expr/pathbuilder/builder.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
package pathbuilder

import (
"errors"

"github.com/brimdata/zed"
"github.com/brimdata/zed/runtime/expr/dynfield"
)

type builder struct {
inputCount int
base Step
}

func New(base zed.Type, paths []dynfield.Path, leafs []zed.Value) (Step, error) {
if len(paths) != len(leafs) {
return nil, errors.New("paths and leafs must be the same length")
}
b := &builder{base: newLeafStep(base, -1)}
for i, p := range paths {
if err := b.Put(p, leafs[i].Type); err != nil {
return nil, err
}
}
return b.base, nil
}

func (m *builder) Put(p dynfield.Path, leaf zed.Type) error {
defer func() { m.inputCount++ }()
return m.put(&m.base, p, leaf)
}

func (m *builder) put(parent *Step, p dynfield.Path, typ zed.Type) error {
// Actually let's do this differently. If current is a string then we are
// putting to a record. When we support maps we'll need to check for that.
if p[0].IsString() {
return m.putRecord(parent, p, typ)
}
// This could be for a map or a set but keep it simple for now.
if zed.IsInteger(p[0].Type.ID()) {
return m.putVector(parent, p, typ)
}
// if zed.TypeUnder(parent.typeof())
return errors.New("unsupported types")
}

func (m *builder) putRecord(s *Step, p dynfield.Path, typ zed.Type) error {
current, p := p[0], p[1:]
rstep, ok := (*s).(*recordStep)
if !ok {
// If this is a leafStep with a type of record than we need to
// initialize a recordStep with fields, otherwise just replace this will
// a recordStep.
var fields []zed.Field
if lstep, ok := (*s).(*leafStep); ok && zed.TypeRecordOf(lstep.typ) != nil {
fields = zed.TypeRecordOf(lstep.typ).Fields
}
rstep = newRecordStep(fields)
if *s == m.base {
rstep.isBase = true
}
*s = rstep
}
i := rstep.lookup(current.AsString())
field := &rstep.fields[i]
if len(p) == 0 {
field.step = newLeafStep(typ, m.inputCount)
return nil
}
return m.put(&field.step, p, typ)
}

func (m *builder) putVector(s *Step, p dynfield.Path, typ zed.Type) error {
current, p := p[0], p[1:]
vstep, ok := (*s).(*vectorStep)
if !ok {
// If this is a leafStep with a type of array than we need to
// initialize a arrayStep with fields, otherwise just replace this with
// an arrayStep.
vstep = &vectorStep{}
if lstep, ok := (*s).(*leafStep); ok && zed.InnerType(lstep.typ) != nil {
vstep.inner = zed.InnerType(lstep.typ)
_, vstep.isSet = zed.TypeUnder(lstep.typ).(*zed.TypeSet)
}
if *s == m.base {
vstep.isBase = true
}
*s = vstep
}
at := vstep.lookup(int(current.AsInt()))
elem := &vstep.elems[at]
if len(p) == 0 {
elem.step = newLeafStep(typ, m.inputCount)
return nil
}
return m.put(&elem.step, p, typ)
}
108 changes: 108 additions & 0 deletions runtime/expr/pathbuilder/builder_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,108 @@
package pathbuilder

import (
"testing"

"github.com/brimdata/zed"
"github.com/brimdata/zed/runtime/expr/dynfield"
"github.com/brimdata/zed/zcode"
"github.com/brimdata/zed/zson"
"github.com/stretchr/testify/require"
)

func parsePath(zctx *zed.Context, ss ...string) dynfield.Path {
var path dynfield.Path
for _, s := range ss {
path = append(path, *zson.MustParseValue(zctx, s))
}
return path
}

type testCase struct {
describe string
base string
paths [][]string
values []string
expected string
}

func runTestCase(t *testing.T, c testCase) {
zctx := zed.NewContext()
var baseTyp zed.Type
var baseBytes []byte
if c.base != "" {
base := zson.MustParseValue(zctx, c.base)
baseTyp, baseBytes = base.Type, base.Bytes()
}
var paths []dynfield.Path
for _, ss := range c.paths {
paths = append(paths, parsePath(zctx, ss...))
}
var values []zed.Value
for _, s := range c.values {
values = append(values, *zson.MustParseValue(zctx, s))
}
step, err := New(baseTyp, paths, values)
require.NoError(t, err)
var b zcode.Builder
typ, err := step.Build(zctx, &b, baseBytes, values)
require.NoError(t, err)
val := zed.NewValue(typ, b.Bytes())
require.Equal(t, c.expected, zson.FormatValue(val))
}

func TestIt(t *testing.T) {
runTestCase(t, testCase{
base: `{"a": 1, "b": 2}`,
paths: [][]string{
{`"c"`, `"a"`, `"a"`},
{`"c"`, `"b"`},
{`"c"`, `"c"`},
},
values: []string{
`45`,
`"string"`,
"127.0.0.1",
},
expected: `{a:1,b:2,c:{a:{a:45},b:"string",c:127.0.0.1}}`,
})
runTestCase(t, testCase{
base: `{"a": [1,{foo:"bar"}]}`,
paths: [][]string{
{`"a"`, `0`},
{`"a"`, `1`, `"foo"`},
},
values: []string{
`"hi"`,
`"baz"`,
},
expected: `{a:["hi",{foo:"baz"}]}`,
})
runTestCase(t, testCase{
describe: "create from empty base",
paths: [][]string{
{`"a"`},
{`"b"`},
},
values: []string{
`"foo"`,
`"bar"`,
},
expected: `{a:"foo",b:"bar"}`,
})
runTestCase(t, testCase{
describe: "assign to base level array",
base: `["a", "b", "c"]`,
paths: [][]string{
{`0`},
{`1`},
{`2`},
},
values: []string{
`"foo"`,
`"bar"`,
`"baz"`,
},
expected: `["foo","bar","baz"]`,
})
}
Loading

0 comments on commit 692a0a3

Please sign in to comment.