Skip to content

Commit

Permalink
vam: nest_dotted (#5597)
Browse files Browse the repository at this point in the history
  • Loading branch information
mattnibs authored Jan 29, 2025
1 parent 55d4764 commit 8d1c21a
Show file tree
Hide file tree
Showing 7 changed files with 95 additions and 16 deletions.
14 changes: 12 additions & 2 deletions runtime/sam/expr/function/nestdotted.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,13 +57,23 @@ func (n *NestDotted) lookupBuilderAndType(in *super.TypeRecord) (*super.RecordBu

func (n *NestDotted) Call(_ super.Allocator, args []super.Value) super.Value {
val := args[len(args)-1]
b, typ, err := n.lookupBuilderAndType(super.TypeRecordOf(val.Type()))
if val.Type().ID() == super.IDNull {
return val
}
rtyp := super.TypeRecordOf(val.Type())
if rtyp == nil {
return n.zctx.WrapError("nest_dotted: non-record value", val)
}
b, typ, err := n.lookupBuilderAndType(rtyp)
if err != nil {
return n.zctx.WrapError("nest_dotted(): "+err.Error(), val)
return n.zctx.WrapError("nest_dotted: "+err.Error(), val)
}
if b == nil {
return val
}
if val.IsNull() {
return super.NewValue(typ, nil)
}
b.Reset()
for it := val.Bytes().Iter(); !it.Done(); {
b.Append(it.Next())
Expand Down
11 changes: 0 additions & 11 deletions runtime/sam/expr/ztests/nest_dotted.yaml

This file was deleted.

4 changes: 4 additions & 0 deletions runtime/vam/expr/function/function.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,10 @@ func New(zctx *super.Context, name string, narg int) (expr.Function, field.Path,
f = &Log{zctx}
case "lower":
f = &ToLower{zctx}
case "nest_dotted":
path = field.Path{}
argmin = 0
f = &NestDotted{zctx}
case "network_of":
argmax = 2
f = &NetworkOf{zctx}
Expand Down
54 changes: 54 additions & 0 deletions runtime/vam/expr/function/nestdotted.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
package function

import (
"github.com/brimdata/super"
"github.com/brimdata/super/pkg/field"
"github.com/brimdata/super/vector"
)

type NestDotted struct {
zctx *super.Context
}

func (n *NestDotted) Call(args ...vector.Any) vector.Any {
vec := vector.Under(args[len(args)-1])
if vec.Type().ID() == super.IDNull {
return vec
}
view, ok := vec.(*vector.View)
if ok {
vec = view.Any
}
record, ok := vec.(*vector.Record)
if !ok {
return vector.NewWrappedError(n.zctx, "nest_dotted: non-record value", args[len(args)-1])
}
b, err := n.getBuilder(record.Typ)
if err != nil {
return vector.NewWrappedError(n.zctx, "nest_dotted: "+err.Error(), args[len(args)-1])
}
if b == nil {
return args[len(args)-1]
}
out := vector.Any(b.New(record.Fields, record.Nulls))
if view != nil {
out = vector.NewView(out, view.Index)
}
return out
}

func (n *NestDotted) getBuilder(in *super.TypeRecord) (*vector.RecordBuilder, error) {
var foundDotted bool
var fields field.List
for _, f := range in.Fields {
dotted := field.Dotted(f.Name)
if len(dotted) > 1 {
foundDotted = true
}
fields = append(fields, dotted)
}
if !foundDotted {
return nil, nil
}
return vector.NewRecordBuilder(n.zctx, fields)
}
4 changes: 2 additions & 2 deletions runtime/vam/op/summarize/agg.go
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ func (s *superTable) materialize() vector.Any {
}
// Since aggs can return dynamic values need to do apply to create record.
return vector.Apply(false, func(vecs ...vector.Any) vector.Any {
return s.builder.New(vecs)
return s.builder.New(vecs, nil)
}, vecs...)
}

Expand Down Expand Up @@ -271,7 +271,7 @@ func (c *countByString) materialize() vector.Any {
}
keyVec := vector.NewString(offs, bytes, nulls)
countVec := vector.NewUint(super.TypeUint64, counts, nil)
return c.builder.New([]vector.Any{keyVec, countVec})
return c.builder.New([]vector.Any{keyVec, countVec}, nil)
}

type Sum struct {
Expand Down
21 changes: 21 additions & 0 deletions runtime/ztests/expr/function/nest_dotted.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
zed: nest_dotted(this)

vector: true

input: |
{a:1,"b.a":2,"b.b":3,"b.c.a":4,c:5}
{a:1,b:{a:2,b:3,c:{a:4}},c:5}
{a:1,"b.a":2}
null({a:int64,"b.a":int64,"b.b":int64})
{a:1,b:null({a:int64,b:int64})}
null
"foo"
output: |
{a:1,b:{a:2,b:3,c:{a:4}},c:5}
{a:1,b:{a:2,b:3,c:{a:4}},c:5}
{a:1,b:{a:2}}
null({a:int64,b:{a:int64,b:int64}})
{a:1,b:null({a:int64,b:int64})}
null
error({message:"nest_dotted: non-record value",on:"foo"})
3 changes: 2 additions & 1 deletion vector/recordbuilder.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,9 @@ func NewRecordBuilder(zctx *super.Context, fields field.List) (*RecordBuilder, e
return &RecordBuilder{zctx: zctx, base: base}, nil
}

func (r *RecordBuilder) New(vecs []Any) *Record {
func (r *RecordBuilder) New(vecs []Any, nulls *Bool) *Record {
rec, _ := r.base.build(r.zctx, vecs)
rec.Nulls = nulls
return rec
}

Expand Down

0 comments on commit 8d1c21a

Please sign in to comment.