-
Notifications
You must be signed in to change notification settings - Fork 0
/
filter.go
165 lines (130 loc) · 4.65 KB
/
filter.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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
package dataframe
import "context"
// FilterAction is the return value of FilterSeriesFn and FilterDataFrameFn.
type FilterAction int
const (
// DROP is used to signify that a row must be dropped.
DROP FilterAction = 0
// KEEP is used to signify that a row must be kept.
KEEP FilterAction = 1
// CHOOSE is used to signify that a row must be kept.
CHOOSE FilterAction = 1
)
// FilterSeriesFn is used by the Filter function to determine which rows are selected.
// val contains the value of the current row.
// If the function returns DROP, then the row is removed. If KEEP or CHOOSE is chosen, the row is kept.
type FilterSeriesFn[T any] func(val T, row, nRows int) (FilterAction, error)
// FilterDataFrameFn is used by the Filter function to determine which rows are selected.
// vals contains the values for the current row. The keys contain ints (index of Series) and strings (name of Series).
// If the function returns DROP, then the row is removed. If KEEP or CHOOSE is chosen, the row is kept.
type FilterDataFrameFn func(vals map[string]any, row, nRows int) (FilterAction, error)
// FilterSeries applies filter function to series. If FilterOptions are set as `FilterOptions { InPlace: true }`
// then series is modified, otherwise new series is returned.
func FilterSeries[T any](ctx context.Context, s *Series[T], fn FilterSeriesFn[T], options ...FilterOptions) (*Series[T], error) {
if fn == nil {
panic("fn is required")
}
opts := DefaultOptions(options...)
if !opts.DontLock {
s.Lock()
defer s.Unlock()
}
transfer := []int{}
iterator := s.Iterator(IteratorOptions { InitialRow: 0, Step: 1, DontLock: true })
for iterator.Next() {
if err := ctx.Err(); err != nil {
return nil, err
}
fa, err := fn(iterator.Value, iterator.Index, iterator.Total)
if err != nil {
return nil, err
}
if fa == DROP {
if opts.InPlace {
transfer = append(transfer, iterator.Index)
}
} else if fa == KEEP || fa == CHOOSE {
if !opts.InPlace {
transfer = append(transfer, iterator.Index)
}
} else {
panic("unrecognized FilterAction returned by fn")
}
}
if !opts.InPlace {
ns := NewSeries[T](s.Name(dontLock), &SeriesInit{ Capacity: len(transfer) })
for _, rowToTransfer := range transfer {
val := s.Value(rowToTransfer, dontLock)
ns.Append([] T { val }, dontLock)
}
return ns, nil
}
// Remove rows that need to be removed
for idx := len(transfer) - 1; idx >= 0; idx-- {
s.Remove(transfer[idx], dontLock)
}
return s, nil
}
// Filter applies filter function to series. If FilterOptions are set as `FilterOptions { InPlace: true }`
// then series is modified, otherwise new series is returned.
func (s *Series[T]) Filter(ctx context.Context, fn FilterSeriesFn[T], options ...FilterOptions) (*Series[T], error) {
return FilterSeries(ctx, s, fn, options...)
}
// FilterDataFrame applies filter function to DataFrame. If FilterOptions are set as `FilterOptions { InPlace: true }`
// then dataframe is modified, otherwise new dataframe is returned.
func FilterDataFrame(ctx context.Context, df *DataFrame, fn FilterDataFrameFn, options ...FilterOptions) (*DataFrame, error) {
if fn == nil {
panic("fn is required")
}
opts := DefaultOptions(options...)
if !opts.DontLock {
df.Lock()
defer df.Unlock()
}
transfer := []int{}
iterator := df.Iterator(IteratorOptions { InitialRow: 0, Step: 1, DontLock: true })
for iterator.Next() {
if err := ctx.Err(); err != nil {
return nil, err
}
fa, err := fn(iterator.Value, iterator.Index, iterator.Total)
if err != nil {
return nil, err
}
if fa == DROP {
if opts.InPlace {
transfer = append(transfer, iterator.Index)
}
} else if fa == KEEP || fa == CHOOSE {
if !opts.InPlace {
transfer = append(transfer, iterator.Index)
}
} else {
panic("unrecognized FilterAction returned by fn")
}
}
if !opts.InPlace {
// Create all series
seriess := []SeriesAny{}
for _, s := range df.Series {
seriess = append(seriess, s.cloneAsEmpty(len(transfer), len(transfer)))
}
// Create a new dataframe
ndf := NewDataFrame(seriess...)
for _, rowToTransfer := range transfer {
vals := df.Row(rowToTransfer, dontLock)
ndf.Append(vals, dontLock)
}
return ndf, nil
}
// Remove rows that need to be removed
for idx := len(transfer) - 1; idx >= 0; idx-- {
df.Remove(transfer[idx], dontLock)
}
return df, nil
}
// Filter applies filter function to DataFrame. If FilterOptions are set as `FilterOptions { InPlace: true }`
// then dataframe is modified, otherwise new dataframe is returned.
func (df *DataFrame) Filter(ctx context.Context, fn FilterDataFrameFn, options ...FilterOptions) (*DataFrame, error) {
return FilterDataFrame(ctx, df, fn, options...)
}