-
Notifications
You must be signed in to change notification settings - Fork 67
/
dropper.go
114 lines (105 loc) · 2.71 KB
/
dropper.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
package expr
import (
"slices"
"github.com/brimdata/zed"
"github.com/brimdata/zed/pkg/field"
)
type dropper struct {
typ zed.Type
builder *zed.RecordBuilder
fieldRefs []Evaluator
}
func (d *dropper) drop(ectx Context, in *zed.Value) *zed.Value {
if d.typ == in.Type {
return in
}
b := d.builder
b.Reset()
for _, e := range d.fieldRefs {
val := e.Eval(ectx, in)
b.Append(val.Bytes())
}
val, err := b.Encode()
if err != nil {
panic(err)
}
return ectx.NewValue(d.typ, val)
}
type Dropper struct {
zctx *zed.Context
fields field.List
droppers map[int]*dropper
}
func NewDropper(zctx *zed.Context, fields field.List) *Dropper {
return &Dropper{
zctx: zctx,
fields: fields,
droppers: make(map[int]*dropper),
}
}
func (d *Dropper) newDropper(zctx *zed.Context, r *zed.Value) *dropper {
fields, fieldTypes, match := complementFields(d.fields, nil, zed.TypeRecordOf(r.Type))
if !match {
// r.Type contains no fields matching d.fields, so we set
// dropper.typ to r.Type to indicate that records of this type
// should not be modified.
return &dropper{typ: r.Type}
}
// If the set of dropped fields is equal to the all of record's
// fields, then there is no output for this input type.
// We return nil to block this input type.
if len(fieldTypes) == 0 {
return nil
}
var fieldRefs []Evaluator
for _, f := range fields {
fieldRefs = append(fieldRefs, NewDottedExpr(zctx, f))
}
builder, err := zed.NewRecordBuilder(d.zctx, fields)
if err != nil {
panic(err)
}
typ := builder.Type(fieldTypes)
return &dropper{typ, builder, fieldRefs}
}
// complementFields returns the slice of fields and associated types that make
// up the complement of the set of fields in drops along with a boolean that is
// true if typ contains any the fields in drops.
func complementFields(drops field.List, prefix field.Path, typ *zed.TypeRecord) (field.List, []zed.Type, bool) {
var fields field.List
var types []zed.Type
var match bool
for _, f := range typ.Fields {
fld := append(prefix, f.Name)
if drops.Has(fld) {
match = true
continue
}
if typ, ok := zed.TypeUnder(f.Type).(*zed.TypeRecord); ok {
if fs, ts, m := complementFields(drops, fld, typ); m {
fields = append(fields, fs...)
types = append(types, ts...)
match = true
continue
}
}
fields = append(fields, slices.Clone(fld))
types = append(types, f.Type)
}
return fields, types, match
}
func (d *Dropper) Eval(ectx Context, in *zed.Value) *zed.Value {
if !zed.IsRecordType(in.Type) {
return in
}
id := in.Type.ID()
dropper, ok := d.droppers[id]
if !ok {
dropper = d.newDropper(d.zctx, in)
d.droppers[id] = dropper
}
if dropper == nil {
return d.zctx.Quiet()
}
return dropper.drop(ectx, in)
}