Skip to content

Commit 686993f

Browse files
committed
release GIL on median
release GIL on is_lexsorted / fix memory leak release GIL on nancorr
1 parent 07f2c1b commit 686993f

File tree

5 files changed

+113
-99
lines changed

5 files changed

+113
-99
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

+68-59
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ def is_lexsorted(list list_of_arrays):
105105
Py_ssize_t n, nlevels
106106
int64_t k, cur, pre
107107
ndarray arr
108+
bint result = 1
108109

109110
nlevels = len(list_of_arrays)
110111
n = len(list_of_arrays[0])
@@ -115,18 +116,20 @@ def is_lexsorted(list list_of_arrays):
115116
vecs[i] = <int64_t*> arr.data
116117

117118
# 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
119+
with nogil:
120+
for i from 1 <= i < n:
121+
for k from 0 <= k < nlevels:
122+
cur = vecs[k][i]
123+
pre = vecs[k][i -1]
124+
if cur == pre:
125+
continue
126+
elif cur > pre:
127+
break
128+
else:
129+
result = 0
130+
break
128131
free(vecs)
129-
return True
132+
return result
130133

131134

132135
@cython.boundscheck(False)
@@ -177,10 +180,11 @@ def groupsort_indexer(ndarray[int64_t] index, Py_ssize_t ngroups):
177180

178181
@cython.boundscheck(False)
179182
@cython.wraparound(False)
180-
cpdef numeric kth_smallest(numeric[:] a, Py_ssize_t k):
183+
cpdef numeric kth_smallest(numeric[:] a, Py_ssize_t k) nogil:
181184
cdef:
182-
Py_ssize_t i, j, l, m, n = a.size
185+
Py_ssize_t i, j, l, m, n = a.shape[0]
183186
numeric x
187+
184188
with nogil:
185189
l = 0
186190
m = n - 1
@@ -201,7 +205,7 @@ cpdef numeric kth_smallest(numeric[:] a, Py_ssize_t k):
201205

202206
if j < k: l = i
203207
if k < i: m = j
204-
return a[k]
208+
return a[k]
205209

206210

207211
cpdef numeric median(numeric[:] arr):
@@ -238,17 +242,18 @@ def max_subseq(ndarray[double_t] arr):
238242
S = m
239243
T = 0
240244

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
245+
with nogil:
246+
for i in range(1, n):
247+
# S = max { S + A[i], A[i] )
248+
if (S > 0):
249+
S = S + arr[i]
250+
else:
251+
S = arr[i]
252+
T = i
253+
if S > m:
254+
s = T
255+
e = i
256+
m = S
252257

253258
return (s, e, m)
254259

@@ -268,9 +273,10 @@ def min_subseq(ndarray[double_t] arr):
268273

269274
@cython.boundscheck(False)
270275
@cython.wraparound(False)
271-
def nancorr(ndarray[float64_t, ndim=2] mat, cov=False, minp=None):
276+
def nancorr(ndarray[float64_t, ndim=2] mat, bint cov=0, minp=None):
272277
cdef:
273278
Py_ssize_t i, j, xi, yi, N, K
279+
bint minpv
274280
ndarray[float64_t, ndim=2] result
275281
ndarray[uint8_t, ndim=2] mask
276282
int64_t nobs = 0
@@ -279,46 +285,49 @@ def nancorr(ndarray[float64_t, ndim=2] mat, cov=False, minp=None):
279285
N, K = (<object> mat).shape
280286

281287
if minp is None:
282-
minp = 1
288+
minpv = 1
289+
else:
290+
minpv = <int>minp
283291

284292
result = np.empty((K, K), dtype=np.float64)
285293
mask = np.isfinite(mat).view(np.uint8)
286294

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-
295+
with nogil:
296+
for xi in range(K):
297+
for yi in range(xi + 1):
298+
nobs = sumxx = sumyy = sumx = sumy = 0
307299
for i in range(N):
308300
if mask[i, xi] and mask[i, yi]:
309-
vx = mat[i, xi] - meanx
310-
vy = mat[i, yi] - meany
301+
vx = mat[i, xi]
302+
vy = mat[i, yi]
303+
nobs += 1
304+
sumx += vx
305+
sumy += vy
306+
307+
if nobs < minpv:
308+
result[xi, yi] = result[yi, xi] = NaN
309+
else:
310+
meanx = sumx / nobs
311+
meany = sumy / nobs
311312

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

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

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

323332
return result
324333

@@ -351,7 +360,7 @@ def nancorr_spearman(ndarray[float64_t, ndim=2] mat, Py_ssize_t minp=1):
351360
nobs += 1
352361

353362
if nobs < minp:
354-
result[xi, yi] = result[yi, xi] = np.NaN
363+
result[xi, yi] = result[yi, xi] = NaN
355364
else:
356365
maskedx = np.empty(nobs, dtype=np.float64)
357366
maskedy = np.empty(nobs, dtype=np.float64)
@@ -382,7 +391,7 @@ def nancorr_spearman(ndarray[float64_t, ndim=2] mat, Py_ssize_t minp=1):
382391
if divisor != 0:
383392
result[xi, yi] = result[yi, xi] = sumx / divisor
384393
else:
385-
result[xi, yi] = result[yi, xi] = np.NaN
394+
result[xi, yi] = result[yi, xi] = NaN
386395

387396
return result
388397

pandas/_libs/groupby.pyx

+28-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,31 @@ 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, Py_ssize_t k, Py_ssize_t n) nogil:
263+
cdef:
264+
Py_ssize_t i, j, l, m
265+
double_t x, t
266+
267+
l = 0
268+
m = n -1
269+
while (l<m):
270+
x = a[k]
271+
i = l
272+
j = m
273+
274+
while 1:
275+
while a[i] < x: i += 1
276+
while x < a[j]: j -= 1
277+
if i <= j:
278+
swap(&a[i], &a[j])
279+
i += 1; j -= 1
280+
281+
if i > j: break
282+
283+
if j < k: l = i
284+
if k < i: m = j
285+
return a[k]
286+
287+
262288
# generated from template
263289
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)

pandas/tests/test_algos.py

+4-3
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,8 @@
1010
import pandas as pd
1111

1212
from pandas import compat
13-
from pandas._libs import algos as libalgos, hashtable
13+
from pandas._libs import (groupby as libgroupby, algos as libalgos,
14+
hashtable)
1415
from pandas._libs.hashtable import unique_label_indices
1516
from pandas.compat import lrange
1617
import pandas.core.algorithms as algos
@@ -889,7 +890,7 @@ def test_group_var_constant(self):
889890
class TestGroupVarFloat64(tm.TestCase, GroupVarTestMixin):
890891
__test__ = True
891892

892-
algo = algos.algos.group_var_float64
893+
algo = libgroupby.group_var_float64
893894
dtype = np.float64
894895
rtol = 1e-5
895896

@@ -912,7 +913,7 @@ def test_group_var_large_inputs(self):
912913
class TestGroupVarFloat32(tm.TestCase, GroupVarTestMixin):
913914
__test__ = True
914915

915-
algo = algos.algos.group_var_float32
916+
algo = libgroupby.group_var_float32
916917
dtype = np.float32
917918
rtol = 1e-2
918919

0 commit comments

Comments
 (0)