dynamic cut: fix bug in keys

runtime/expr/cutter.go

@@ -9,15 +9,13 @@ import (
 )
 
 type Cutter struct {
-	zctx        *zed.Context
-	fieldRefs   field.List
-	fieldExprs  []Evaluator
-	lvals       []*Lval
-	outTypes    *zed.TypeVectorTable
-	recordTypes map[int]*zed.TypeRecord
-	typeCache   []zed.Type
+	zctx       *zed.Context
+	fieldRefs  field.List
+	fieldExprs []Evaluator
+	lvals      []*Lval
+	typeCache  []zed.Type
 
-	builders     map[string]*zed.RecordBuilder
+	builders     map[string]*recordBuilderCachedTypes
 	droppers     map[string]*Dropper
 	dropperCache []*Dropper
 	dirty        bool
@@ -32,12 +30,10 @@ func NewCutter(zctx *zed.Context, fieldRefs []*Lval, fieldExprs []Evaluator) *Cu
 	n := len(fieldRefs)
 	return &Cutter{
 		zctx:         zctx,
-		builders:     make(map[string]*zed.RecordBuilder),
+		builders:     make(map[string]*recordBuilderCachedTypes),
 		fieldRefs:    make(field.List, n),
 		fieldExprs:   fieldExprs,
 		lvals:        fieldRefs,
-		outTypes:     zed.NewTypeVectorTable(),
-		recordTypes:  make(map[int]*zed.TypeRecord),
 		typeCache:    make([]zed.Type, n),
 		droppers:     make(map[string]*Dropper),
 		dropperCache: make([]*Dropper, n),
@@ -79,13 +75,11 @@ func (c *Cutter) Eval(ectx Context, in *zed.Value) *zed.Value {
 		rb.Append(val.Bytes())
 		types[k] = val.Type
 	}
-	// check paths
 	bytes, err := rb.Encode()
 	if err != nil {
 		panic(err)
 	}
-	typ := c.lookupTypeRecord(types, rb)
-	rec := ectx.NewValue(typ, bytes)
+	rec := ectx.NewValue(rb.Type(types), bytes)
 	for _, d := range droppers {
 		rec = d.Eval(ectx, rec)
 	}
@@ -95,7 +89,7 @@ func (c *Cutter) Eval(ectx Context, in *zed.Value) *zed.Value {
 	return rec
 }
 
-func (c *Cutter) lookupBuilder(ectx Context, in *zed.Value) (*zed.RecordBuilder, field.List, error) {
+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)
@@ -110,20 +104,38 @@ func (c *Cutter) lookupBuilder(ectx Context, in *zed.Value) (*zed.RecordBuilder,
 	builder, ok := c.builders[paths.String()]
 	if !ok {
 		var err error
-		if builder, err = zed.NewRecordBuilder(c.zctx, paths); err != nil {
+		if builder, err = newRecordBuilderCachedTypes(c.zctx, paths); err != nil {
 			return nil, nil, err
 		}
 		c.builders[paths.String()] = builder
 	}
 	return builder, paths, nil
 }
 
-func (c *Cutter) lookupTypeRecord(types []zed.Type, builder *zed.RecordBuilder) *zed.TypeRecord {
-	id := c.outTypes.Lookup(types)
-	typ, ok := c.recordTypes[id]
+type recordBuilderCachedTypes struct {
+	*zed.RecordBuilder
+	outTypes    *zed.TypeVectorTable
+	recordTypes map[int]*zed.TypeRecord
+}
+
+func newRecordBuilderCachedTypes(zctx *zed.Context, paths field.List) (*recordBuilderCachedTypes, error) {
+	b, err := zed.NewRecordBuilder(zctx, paths)
+	if err != nil {
+		return nil, err
+	}
+	return &recordBuilderCachedTypes{
+		RecordBuilder: b,
+		outTypes:      zed.NewTypeVectorTable(),
+		recordTypes:   make(map[int]*zed.TypeRecord),
+	}, nil
+}
+
+func (r *recordBuilderCachedTypes) Type(types []zed.Type) *zed.TypeRecord {
+	id := r.outTypes.Lookup(types)
+	typ, ok := r.recordTypes[id]
 	if !ok {
-		typ = builder.Type(types)
-		c.recordTypes[id] = typ
+		typ = r.RecordBuilder.Type(types)
+		r.recordTypes[id] = typ
 	}
 	return typ
 }

runtime/op/ztests/cut-dynamic-field.yaml

@@ -9,6 +9,8 @@ script: |
   echo "// ==="
   echo '{a:"foo"}' | zq -z 'cut this[a]["bar"] := "baz"' -
   echo "// ==="
+  echo '{key:"foo",v1:1,v2:2} {key:"bar",v1:2,v2:3}' | zq -z 'cut this[key] := [v1,v2]' -
+  echo "// ==="
   # runtime error cases
   echo '{a:"hello",b:"hello"}' | zq -z 'cut this[a] := "world1", this[b] := "world2"' -
   echo "// ==="
@@ -31,6 +33,9 @@ outputs:
       // ===
       {foo:{bar:"baz"}}
       // ===
+      {foo:[1,2]}
+      {bar:[2,3]}
+      // ===
       error({message:"cut: duplicate field: \"hello\"",on:{a:"hello",b:"hello"}})
       // ===
       error({message:"cut: duplicate field: \"foo\"",on:{a:"foo",b:"bar"}})