use nogil in aggregations and groupby

Author: CloseChoice

pandas/_libs/groupby.pyx

@@ -51,7 +51,14 @@ from pandas._libs.missing cimport checknull
 cdef int64_t NPY_NAT = get_nat()
 _int64_max = np.iinfo(np.int64).max
 
-cdef float64_t NaN = <float64_t>np.NaN
+cdef:
+    float32_t MINfloat32 = np.NINF
+    float64_t MINfloat64 = np.NINF
+
+    float32_t MAXfloat32 = np.inf
+    float64_t MAXfloat64 = np.inf
+
+    float64_t NaN = <float64_t>np.NaN
 
 cdef enum InterpolationEnumType:
     INTERPOLATION_LINEAR,
@@ -240,39 +247,58 @@ def group_cumsum(numeric_t[:, ::1] out,
     accum = np.zeros((ngroups, K), dtype=np.asarray(values).dtype)
     compensation = np.zeros((ngroups, K), dtype=np.asarray(values).dtype)
 
-    for i in range(N):
-        lab = labels[i]
+    with nogil:
+        for i in range(N):
+            lab = labels[i]
 
-        if lab < 0:
-            continue
-        for j in range(K):
-            val = values[i, j]
+            if lab < 0:
+                continue
+            for j in range(K):
+                val = values[i, j]
 
-            # For floats, use Kahan summation to reduce floating-point
-            # error (https://en.wikipedia.org/wiki/Kahan_summation_algorithm)
-            if numeric_t == float32_t or numeric_t == float64_t:
-                if np.isinf(val):
-                    out[i, j] = val
-                    accum[lab, j] = val
-                    break
-                elif val == val:
-                    y = val - compensation[lab, j]
-                    t = accum[lab, j] + y
-                    compensation[lab, j] = t - accum[lab, j] - y
-                    accum[lab, j] = t
-                    out[i, j] = t
-                    if np.isinf(t):
-                        compensation[lab, j] = 0
+                # For floats, use Kahan summation to reduce floating-point
+                # error (https://en.wikipedia.org/wiki/Kahan_summation_algorithm)
+                if numeric_t == float32_t:
+                    if (val == MAXfloat32) or (val == MINfloat32):
+                        if (t == MAXfloat32) or (t == MINfloat32):
+                            val = t
+                        out[i, j] = val
+                        accum[lab, j] = val
                         break
-                else:
-                    out[i, j] = NaN
-                    if not skipna:
-                        accum[lab, j] = NaN
+                    elif val == val:
+                        y = val - compensation[lab, j]
+                        t = accum[lab, j] + y
+                        compensation[lab, j] = t - accum[lab, j] - y
+                        accum[lab, j] = t
+                        out[i, j] = t
+                    else:
+                        out[i, j] = NaN
+                        if not skipna:
+                            accum[lab, j] = NaN
+                            break
+                elif numeric_t == float64_t:
+                    if (val == MAXfloat64) or (val == MINfloat64):
+                        out[i, j] = val
+                        accum[lab, j] = val
                         break
-            else:
-                t = val + accum[lab, j]
-                accum[lab, j] = t
-                out[i, j] = t
+                    elif val == val:
+                        y = val - compensation[lab, j]
+                        t = accum[lab, j] + y
+                        compensation[lab, j] = t - accum[lab, j] - y
+                        accum[lab, j] = t
+                        out[i, j] = t
+                        if (t == MAXfloat64) or (t == MINfloat64):
+                            compensation[lab, j] = 0
+                            break
+                    else:
+                        out[i, j] = NaN
+                        if not skipna:
+                            accum[lab, j] = NaN
+                            break
+                else:
+                    t = val + accum[lab, j]
+                    accum[lab, j] = t
+                    out[i, j] = t
 
 
 @cython.boundscheck(False)

pandas/_libs/window/aggregations.pyx

@@ -87,7 +87,7 @@ cdef inline float64_t calc_sum(int64_t minp, int64_t nobs, float64_t sum_x) nogi
 
 
 cdef inline void add_sum(float64_t val, int64_t *nobs, float64_t *sum_x,
-                         float64_t *compensation):
+                         float64_t *compensation) nogil:
     """ add a value from the sum calc using Kahan summation """
 
     cdef:
@@ -100,14 +100,14 @@ cdef inline void add_sum(float64_t val, int64_t *nobs, float64_t *sum_x,
         t = sum_x[0] + y
         compensation[0] = t - sum_x[0] - y
         sum_x[0] = t
-    if np.isinf(val):
+    if (val == MINfloat64) or (val == MAXfloat64):
         sum_x[0] = val
         nobs[0] = nobs[0] + 1
         compensation[0] = 0
 
 
 cdef inline void remove_sum(float64_t val, int64_t *nobs, float64_t *sum_x,
-                            float64_t *compensation):
+                            float64_t *compensation) nogil:
     """ remove a value from the sum calc using Kahan summation """
 
     cdef:
@@ -120,7 +120,7 @@ cdef inline void remove_sum(float64_t val, int64_t *nobs, float64_t *sum_x,
         t = sum_x[0] + y
         compensation[0] = t - sum_x[0] - y
         sum_x[0] = t
-    if np.isinf(val):
+    if (val == MINfloat64) or (val == MAXfloat64):
         sum_x[0] = val
         nobs[0] = nobs[0] - 1
         compensation[0] = 0
@@ -141,31 +141,35 @@ def roll_sum(const float64_t[:] values, ndarray[int64_t] start,
     )
     output = np.empty(N, dtype=np.float64)
 
-    for i in range(0, N):
-        s = start[i]
-        e = end[i]
+    with nogil:
 
-        if i == 0 or not is_monotonic_increasing_bounds:
+        for i in range(0, N):
+            s = start[i]
+            e = end[i]
 
-            # setup
-            for j in range(s, e):
-                add_sum(values[j], &nobs, &sum_x, &compensation_add)
+            if i == 0 or not is_monotonic_increasing_bounds:
 
-        else:
-            # calculate deletes
-            for j in range(start[i - 1], s):
-                remove_sum(values[j], &nobs, &sum_x, &compensation_remove)
+                # setup
+
+                for j in range(s, e):
+                    add_sum(values[j], &nobs, &sum_x, &compensation_add)
 
-            # calculate adds
-            for j in range(end[i - 1], e):
-                add_sum(values[j], &nobs, &sum_x, &compensation_add)
+            else:
+
+                # calculate deletes
+                for j in range(start[i - 1], s):
+                    remove_sum(values[j], &nobs, &sum_x, &compensation_remove)
 
-        output[i] = calc_sum(minp, nobs, sum_x)
+                # calculate adds
+                for j in range(end[i - 1], e):
+                    add_sum(values[j], &nobs, &sum_x, &compensation_add)
 
-        if not is_monotonic_increasing_bounds:
-            nobs = 0
-            sum_x = 0.0
-            compensation_remove = 0.0
+            output[i] = calc_sum(minp, nobs, sum_x)
+
+            if not is_monotonic_increasing_bounds:
+                nobs = 0
+                sum_x = 0.0
+                compensation_remove = 0.0
 
     return output