forked from rocketlaunchr/dataframe-go
/
search.go
222 lines (182 loc) · 5.15 KB
/
search.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
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
// Copyright 2018-20 PJ Engineering and Business Solutions Pty. Ltd. All rights reserved.
package utils
import (
"context"
"golang.org/x/sync/errgroup"
"runtime"
"sync"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
dataframe "github.com/mmavka/dataframe-go"
)
// SearchOptions modifies the behavior of Search.
type SearchOptions struct {
// Max is used to stop the search after it finds Max number of results.
// When Max is set to nil (or 0), all results are sought.
Max *int
// NoConcurrency can be set to ensure that results are stable (& reproducible).
// This can be useful if the context is canceled or a Max is set.
NoConcurrency bool
// R is used to limit the range of the Series for search purposes.
R *dataframe.Range
// DontLock can be set to true if the Series should not be locked.
DontLock bool
}
// Search is used to find particular values in a given Series.
// It will find all values that are between lower and upper bounds (inclusive).
// It will return a slice containing the rows which contain values within the bounds.
// If Search is canceled, an incomplete list of the rows "found so far" is returned.
// s will be locked for the duration of the operation.
//
// Example:
//
// s1 := dataframe.NewSeriesInt64("", nil, 11, 10, 9, 8, 7, 6, 5, 23, 25, 2, 1, 5, 4)
//
// fmt.Println(utils.Search(ctx, s1, int64(4), int64(6)))
// // Output: [5 6 11 12]
func Search(ctx context.Context, s dataframe.Series, lower, upper interface{}, opts ...SearchOptions) ([]int, error) {
if len(opts) == 0 {
opts = append(opts, SearchOptions{R: &dataframe.Range{}})
} else if opts[0].R == nil {
opts[0].R = &dataframe.Range{}
}
if !opts[0].DontLock {
s.Lock()
defer s.Unlock()
}
fullRowCount := s.NRows(dataframe.DontLock)
if fullRowCount == 0 {
return []int{}, nil
}
var equalCheck bool
if cmp.Equal(lower, upper, cmpopts.IgnoreUnexported()) {
equalCheck = true
}
start, end, err := opts[0].R.Limits(fullRowCount)
if err != nil {
return nil, err
}
var nCores int
if opts[0].NoConcurrency {
nCores = 1
} else {
nCores = runtime.NumCPU()
}
// Group search range equally amongst each core
div := (end - start + 1) / nCores
subRanges := []dataframe.Range{}
for i := 0; i < nCores; i++ {
subStart := i * div
var subEnd int
if i != nCores-1 {
subEnd = (i+1)*div - 1
} else {
// last core
subEnd = end
}
subRanges = append(subRanges, dataframe.Range{
Start: &subStart,
End: &subEnd,
})
}
// Concurrently search each subRange for values in range
g, newCtx := errgroup.WithContext(ctx)
var mapProtect sync.Mutex
mapRows := map[int][]int{} // For each core store the rows we have found so far
var foundSoFarProtect sync.RWMutex
foundSoFar := 0
for i := 0; i < nCores; i++ {
i := i
g.Go(func() error {
rowsFound := []int{} // Store all rows that we have found
defer func() {
mapProtect.Lock()
mapRows[i] = rowsFound
mapProtect.Unlock()
}()
for row := *subRanges[i].Start; row < *subRanges[i].End+1; row++ {
if nCores != 1 {
// Exit early if max has been found already.
if opts[0].Max != nil && *opts[0].Max != 0 {
foundSoFarProtect.RLock()
if foundSoFar >= *opts[0].Max {
foundSoFarProtect.RUnlock()
return nil
}
foundSoFarProtect.RUnlock()
}
}
// Cancel for loop if context is canceled
if err := newCtx.Err(); err != nil {
return err
}
val := s.Value(row, dataframe.DontLock)
// Check if val is in range
if equalCheck {
if s.IsEqualFunc(val, lower) {
rowsFound = append(rowsFound, row)
if nCores == 1 {
foundSoFar++
if opts[0].Max != nil && *opts[0].Max != 0 {
if foundSoFar >= *opts[0].Max {
return nil
}
}
} else {
if opts[0].Max != nil && *opts[0].Max != 0 {
foundSoFarProtect.Lock()
foundSoFar++
if foundSoFar >= *opts[0].Max {
foundSoFarProtect.Unlock()
return nil
}
foundSoFarProtect.Unlock()
}
}
}
} else {
if !s.IsLessThanFunc(val, lower) && (s.IsLessThanFunc(val, upper) || s.IsEqualFunc(val, upper)) {
rowsFound = append(rowsFound, row)
if nCores == 1 {
foundSoFar++
if opts[0].Max != nil && *opts[0].Max != 0 {
if foundSoFar >= *opts[0].Max {
return nil
}
}
} else {
if opts[0].Max != nil && *opts[0].Max != 0 {
foundSoFarProtect.Lock()
foundSoFar++
if foundSoFar >= *opts[0].Max {
foundSoFarProtect.Unlock()
return nil
}
foundSoFarProtect.Unlock()
}
}
}
}
}
return nil
})
}
err = g.Wait()
// Convert rows found to Range slice
var rows []int
var count int
for i := 0; i < nCores; i++ {
count = count + len(mapRows[i])
}
rows = make([]int, 0, count)
// Store found rows into []int
for i := 0; i < nCores; i++ {
foundRows := mapRows[i]
rows = append(rows, foundRows...)
}
// Ensure only max number of results is returned
if opts[0].Max != nil && *opts[0].Max != 0 && len(rows) > *opts[0].Max {
return rows[:*opts[0].Max], err
}
return rows, err
}