Skip to content

Commit cd18cae

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

File tree

5 files changed

+115
-100
lines changed

5 files changed

+115
-100
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-60
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,6 +104,7 @@ 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 = 1
108108

109109
nlevels = len(list_of_arrays)
110110
n = len(list_of_arrays[0])
@@ -115,18 +115,20 @@ def is_lexsorted(list list_of_arrays):
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 from 1 <= i < n:
120+
for k from 0 <= k < 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 = 0
129+
break
128130
free(vecs)
129-
return True
131+
return result
130132

131133

132134
@cython.boundscheck(False)
@@ -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):
@@ -238,17 +241,18 @@ def max_subseq(ndarray[double_t] arr):
238241
S = m
239242
T = 0
240243

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

253257
return (s, e, m)
254258

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

269273
@cython.boundscheck(False)
270274
@cython.wraparound(False)
271-
def nancorr(ndarray[float64_t, ndim=2] mat, cov=False, minp=None):
275+
def nancorr(ndarray[float64_t, ndim=2] mat, bint cov=0, minp=None):
272276
cdef:
273277
Py_ssize_t i, j, xi, yi, N, K
278+
bint minpv
274279
ndarray[float64_t, ndim=2] result
275280
ndarray[uint8_t, ndim=2] mask
276281
int64_t nobs = 0
@@ -279,46 +284,49 @@ def nancorr(ndarray[float64_t, ndim=2] mat, cov=False, minp=None):
279284
N, K = (<object> mat).shape
280285

281286
if minp is None:
282-
minp = 1
287+
minpv = 1
288+
else:
289+
minpv = <int>minp
283290

284291
result = np.empty((K, K), dtype=np.float64)
285292
mask = np.isfinite(mat).view(np.uint8)
286293

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

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

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

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

323331
return result
324332

@@ -351,7 +359,7 @@ def nancorr_spearman(ndarray[float64_t, ndim=2] mat, Py_ssize_t minp=1):
351359
nobs += 1
352360

353361
if nobs < minp:
354-
result[xi, yi] = result[yi, xi] = np.NaN
362+
result[xi, yi] = result[yi, xi] = NaN
355363
else:
356364
maskedx = np.empty(nobs, dtype=np.float64)
357365
maskedy = np.empty(nobs, dtype=np.float64)
@@ -382,7 +390,7 @@ def nancorr_spearman(ndarray[float64_t, ndim=2] mat, Py_ssize_t minp=1):
382390
if divisor != 0:
383391
result[xi, yi] = result[yi, xi] = sumx / divisor
384392
else:
385-
result[xi, yi] = result[yi, xi] = np.NaN
393+
result[xi, yi] = result[yi, xi] = NaN
386394

387395
return result
388396

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)

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)