Skip to content

Commit 4e2bfec

Browse files
committed
release GIL on median
release GIL on is_lexsorted / fix memory leak release GIL on nancorr
1 parent ce28bb5 commit 4e2bfec

File tree

5 files changed

+123
-104
lines changed

5 files changed

+123
-104
lines changed

pandas/_libs/algos.pxd

+2-28
Original file line numberDiff line numberDiff line change
@@ -1,39 +1,13 @@
11
from util cimport numeric
22
from numpy cimport float64_t, double_t
33

4-
cpdef numeric kth_smallest(numeric[:] a, Py_ssize_t k)
4+
cpdef numeric kth_smallest(numeric[:] a, Py_ssize_t k) nogil
55

6-
cdef inline Py_ssize_t swap(numeric *a, numeric *b) nogil except -1:
6+
cdef inline Py_ssize_t swap(numeric *a, numeric *b) nogil:
77
cdef numeric t
88

99
# cython doesn't allow pointer dereference so use array syntax
1010
t = a[0]
1111
a[0] = b[0]
1212
b[0] = t
1313
return 0
14-
15-
16-
cdef inline kth_smallest_c(float64_t* a, Py_ssize_t k, Py_ssize_t n):
17-
cdef:
18-
Py_ssize_t i, j, l, m
19-
double_t x, t
20-
21-
l = 0
22-
m = n -1
23-
while (l<m):
24-
x = a[k]
25-
i = l
26-
j = m
27-
28-
while 1:
29-
while a[i] < x: i += 1
30-
while x < a[j]: j -= 1
31-
if i <= j:
32-
swap(&a[i], &a[j])
33-
i += 1; j -= 1
34-
35-
if i > j: break
36-
37-
if j < k: l = i
38-
if k < i: m = j
39-
return a[k]

pandas/_libs/algos.pyx

+76-64
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,6 @@ class NegInfinity(object):
9696
__ge__ = lambda self, other: self is other
9797

9898

99-
10099
@cython.wraparound(False)
101100
@cython.boundscheck(False)
102101
def is_lexsorted(list list_of_arrays):
@@ -105,28 +104,31 @@ def is_lexsorted(list list_of_arrays):
105104
Py_ssize_t n, nlevels
106105
int64_t k, cur, pre
107106
ndarray arr
107+
bint result = True
108108

109109
nlevels = len(list_of_arrays)
110110
n = len(list_of_arrays[0])
111111

112112
cdef int64_t **vecs = <int64_t**> malloc(nlevels * sizeof(int64_t*))
113-
for i from 0 <= i < nlevels:
113+
for i in range(nlevels):
114114
arr = list_of_arrays[i]
115115
vecs[i] = <int64_t*> arr.data
116116

117117
# Assume uniqueness??
118-
for i from 1 <= i < n:
119-
for k from 0 <= k < nlevels:
120-
cur = vecs[k][i]
121-
pre = vecs[k][i -1]
122-
if cur == pre:
123-
continue
124-
elif cur > pre:
125-
break
126-
else:
127-
return False
118+
with nogil:
119+
for i in range(n):
120+
for k in range(nlevels):
121+
cur = vecs[k][i]
122+
pre = vecs[k][i -1]
123+
if cur == pre:
124+
continue
125+
elif cur > pre:
126+
break
127+
else:
128+
result = False
129+
break
128130
free(vecs)
129-
return True
131+
return result
130132

131133

132134
@cython.boundscheck(False)
@@ -159,15 +161,15 @@ def groupsort_indexer(ndarray[int64_t] index, Py_ssize_t ngroups):
159161
with nogil:
160162

161163
# count group sizes, location 0 for NA
162-
for i from 0 <= i < n:
164+
for i in range(n):
163165
counts[index[i] + 1] += 1
164166

165167
# mark the start of each contiguous group of like-indexed data
166-
for i from 1 <= i < ngroups + 1:
168+
for i in range(1, ngroups + 1):
167169
where[i] = where[i - 1] + counts[i - 1]
168170

169171
# this is our indexer
170-
for i from 0 <= i < n:
172+
for i in range(n):
171173
label = index[i] + 1
172174
result[where[label]] = i
173175
where[label] += 1
@@ -177,10 +179,11 @@ def groupsort_indexer(ndarray[int64_t] index, Py_ssize_t ngroups):
177179

178180
@cython.boundscheck(False)
179181
@cython.wraparound(False)
180-
cpdef numeric kth_smallest(numeric[:] a, Py_ssize_t k):
182+
cpdef numeric kth_smallest(numeric[:] a, Py_ssize_t k) nogil:
181183
cdef:
182-
Py_ssize_t i, j, l, m, n = a.size
184+
Py_ssize_t i, j, l, m, n = a.shape[0]
183185
numeric x
186+
184187
with nogil:
185188
l = 0
186189
m = n - 1
@@ -201,7 +204,7 @@ cpdef numeric kth_smallest(numeric[:] a, Py_ssize_t k):
201204

202205
if j < k: l = i
203206
if k < i: m = j
204-
return a[k]
207+
return a[k]
205208

206209

207210
cpdef numeric median(numeric[:] arr):
@@ -224,6 +227,8 @@ cpdef numeric median(numeric[:] arr):
224227

225228
# -------------- Min, Max subsequence
226229

230+
@cython.boundscheck(False)
231+
@cython.wraparound(False)
227232
def max_subseq(ndarray[double_t] arr):
228233
cdef:
229234
Py_ssize_t i=0, s=0, e=0, T, n
@@ -238,21 +243,24 @@ def max_subseq(ndarray[double_t] arr):
238243
S = m
239244
T = 0
240245

241-
for i in range(1, n):
242-
# S = max { S + A[i], A[i] )
243-
if (S > 0):
244-
S = S + arr[i]
245-
else:
246-
S = arr[i]
247-
T = i
248-
if S > m:
249-
s = T
250-
e = i
251-
m = S
246+
with nogil:
247+
for i in range(1, n):
248+
# S = max { S + A[i], A[i] )
249+
if (S > 0):
250+
S = S + arr[i]
251+
else:
252+
S = arr[i]
253+
T = i
254+
if S > m:
255+
s = T
256+
e = i
257+
m = S
252258

253259
return (s, e, m)
254260

255261

262+
@cython.boundscheck(False)
263+
@cython.wraparound(False)
256264
def min_subseq(ndarray[double_t] arr):
257265
cdef:
258266
Py_ssize_t s, e
@@ -268,9 +276,10 @@ def min_subseq(ndarray[double_t] arr):
268276

269277
@cython.boundscheck(False)
270278
@cython.wraparound(False)
271-
def nancorr(ndarray[float64_t, ndim=2] mat, cov=False, minp=None):
279+
def nancorr(ndarray[float64_t, ndim=2] mat, bint cov=0, minp=None):
272280
cdef:
273281
Py_ssize_t i, j, xi, yi, N, K
282+
bint minpv
274283
ndarray[float64_t, ndim=2] result
275284
ndarray[uint8_t, ndim=2] mask
276285
int64_t nobs = 0
@@ -279,46 +288,49 @@ def nancorr(ndarray[float64_t, ndim=2] mat, cov=False, minp=None):
279288
N, K = (<object> mat).shape
280289

281290
if minp is None:
282-
minp = 1
291+
minpv = 1
292+
else:
293+
minpv = <int>minp
283294

284295
result = np.empty((K, K), dtype=np.float64)
285296
mask = np.isfinite(mat).view(np.uint8)
286297

287-
for xi in range(K):
288-
for yi in range(xi + 1):
289-
nobs = sumxx = sumyy = sumx = sumy = 0
290-
for i in range(N):
291-
if mask[i, xi] and mask[i, yi]:
292-
vx = mat[i, xi]
293-
vy = mat[i, yi]
294-
nobs += 1
295-
sumx += vx
296-
sumy += vy
297-
298-
if nobs < minp:
299-
result[xi, yi] = result[yi, xi] = np.NaN
300-
else:
301-
meanx = sumx / nobs
302-
meany = sumy / nobs
303-
304-
# now the cov numerator
305-
sumx = 0
306-
298+
with nogil:
299+
for xi in range(K):
300+
for yi in range(xi + 1):
301+
nobs = sumxx = sumyy = sumx = sumy = 0
307302
for i in range(N):
308303
if mask[i, xi] and mask[i, yi]:
309-
vx = mat[i, xi] - meanx
310-
vy = mat[i, yi] - meany
304+
vx = mat[i, xi]
305+
vy = mat[i, yi]
306+
nobs += 1
307+
sumx += vx
308+
sumy += vy
309+
310+
if nobs < minpv:
311+
result[xi, yi] = result[yi, xi] = NaN
312+
else:
313+
meanx = sumx / nobs
314+
meany = sumy / nobs
311315

312-
sumx += vx * vy
313-
sumxx += vx * vx
314-
sumyy += vy * vy
316+
# now the cov numerator
317+
sumx = 0
315318

316-
divisor = (nobs - 1.0) if cov else sqrt(sumxx * sumyy)
319+
for i in range(N):
320+
if mask[i, xi] and mask[i, yi]:
321+
vx = mat[i, xi] - meanx
322+
vy = mat[i, yi] - meany
317323

318-
if divisor != 0:
319-
result[xi, yi] = result[yi, xi] = sumx / divisor
320-
else:
321-
result[xi, yi] = result[yi, xi] = np.NaN
324+
sumx += vx * vy
325+
sumxx += vx * vx
326+
sumyy += vy * vy
327+
328+
divisor = (nobs - 1.0) if cov else sqrt(sumxx * sumyy)
329+
330+
if divisor != 0:
331+
result[xi, yi] = result[yi, xi] = sumx / divisor
332+
else:
333+
result[xi, yi] = result[yi, xi] = NaN
322334

323335
return result
324336

@@ -351,7 +363,7 @@ def nancorr_spearman(ndarray[float64_t, ndim=2] mat, Py_ssize_t minp=1):
351363
nobs += 1
352364

353365
if nobs < minp:
354-
result[xi, yi] = result[yi, xi] = np.NaN
366+
result[xi, yi] = result[yi, xi] = NaN
355367
else:
356368
maskedx = np.empty(nobs, dtype=np.float64)
357369
maskedy = np.empty(nobs, dtype=np.float64)
@@ -382,7 +394,7 @@ def nancorr_spearman(ndarray[float64_t, ndim=2] mat, Py_ssize_t minp=1):
382394
if divisor != 0:
383395
result[xi, yi] = result[yi, xi] = sumx / divisor
384396
else:
385-
result[xi, yi] = result[yi, xi] = np.NaN
397+
result[xi, yi] = result[yi, xi] = NaN
386398

387399
return result
388400

pandas/_libs/groupby.pyx

+30-2
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ from numpy cimport (int8_t, int16_t, int32_t, int64_t, uint8_t, uint16_t,
1616
from libc.stdlib cimport malloc, free
1717

1818
from util cimport numeric, get_nat
19-
from algos cimport kth_smallest, kth_smallest_c
19+
from algos cimport swap
2020
from algos import take_2d_axis1_float64_float64, groupsort_indexer
2121

2222
cdef int64_t iNaT = get_nat()
@@ -219,7 +219,7 @@ def group_last_bin_object(ndarray[object, ndim=2] out,
219219
out[i, j] = resx[i, j]
220220

221221

222-
cdef inline float64_t _median_linear(float64_t* a, int n):
222+
cdef inline float64_t _median_linear(float64_t* a, int n) nogil:
223223
cdef int i, j, na_count = 0
224224
cdef float64_t result
225225
cdef float64_t* tmp
@@ -259,5 +259,33 @@ cdef inline float64_t _median_linear(float64_t* a, int n):
259259
return result
260260

261261

262+
cdef inline float64_t kth_smallest_c(float64_t* a,
263+
Py_ssize_t k,
264+
Py_ssize_t n) nogil:
265+
cdef:
266+
Py_ssize_t i, j, l, m
267+
double_t x, t
268+
269+
l = 0
270+
m = n -1
271+
while (l<m):
272+
x = a[k]
273+
i = l
274+
j = m
275+
276+
while 1:
277+
while a[i] < x: i += 1
278+
while x < a[j]: j -= 1
279+
if i <= j:
280+
swap(&a[i], &a[j])
281+
i += 1; j -= 1
282+
283+
if i > j: break
284+
285+
if j < k: l = i
286+
if k < i: m = j
287+
return a[k]
288+
289+
262290
# generated from template
263291
include "groupby_helper.pxi"

pandas/_libs/groupby_helper.pxi.in

+11-7
Original file line numberDiff line numberDiff line change
@@ -681,6 +681,8 @@ def group_cummax_{{name}}(ndarray[{{dest_type2}}, ndim=2] out,
681681
#----------------------------------------------------------------------
682682

683683

684+
@cython.boundscheck(False)
685+
@cython.wraparound(False)
684686
def group_median_float64(ndarray[float64_t, ndim=2] out,
685687
ndarray[int64_t] counts,
686688
ndarray[float64_t, ndim=2] values,
@@ -704,13 +706,15 @@ def group_median_float64(ndarray[float64_t, ndim=2] out,
704706

705707
take_2d_axis1_float64_float64(values.T, indexer, out=data)
706708

707-
for i in range(K):
708-
# exclude NA group
709-
ptr += _counts[0]
710-
for j in range(ngroups):
711-
size = _counts[j + 1]
712-
out[j, i] = _median_linear(ptr, size)
713-
ptr += size
709+
with nogil:
710+
711+
for i in range(K):
712+
# exclude NA group
713+
ptr += _counts[0]
714+
for j in range(ngroups):
715+
size = _counts[j + 1]
716+
out[j, i] = _median_linear(ptr, size)
717+
ptr += size
714718

715719

716720
@cython.boundscheck(False)

0 commit comments

Comments
 (0)