wip

Author: jreback

pandas/src/generate_code.py

@@ -22,6 +22,7 @@
                       PyTuple_New)
 from cpython cimport PyFloat_Check
 cimport cpython
+cdef double NAN = <double> np.nan
 
 import numpy as np
 isnan = np.isnan
@@ -881,7 +882,7 @@ def group_add_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
     Only aggregates on axis=0
     '''
     cdef:
-        Py_ssize_t i, j, N, K, lab
+        Py_ssize_t i, j, N, K, lab, lcounts = len(counts)
         %(dest_type2)s val, count
         ndarray[%(dest_type2)s, ndim=2] sumx, nobs
 
@@ -894,39 +895,45 @@ def group_add_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
     N, K = (<object> values).shape
 
     if K > 1:
-        for i in range(N):
-            lab = labels[i]
-            if lab < 0:
-                continue
 
-            counts[lab] += 1
-            for j in range(K):
-                val = values[i, j]
+        with nogil:
+            for i in range(N):
+                lab = labels[i]
+                if lab < 0:
+                    continue
+
+                counts[lab] += 1
+                for j in range(K):
+                    val = values[i, j]
+
+                    # not nan
+                    if val == val:
+                        nobs[lab, j] += 1
+                        sumx[lab, j] += val
 
-                # not nan
-                if val == val:
-                    nobs[lab, j] += 1
-                    sumx[lab, j] += val
     else:
-        for i in range(N):
-            lab = labels[i]
-            if lab < 0:
-                continue
 
-            counts[lab] += 1
-            val = values[i, 0]
+        with nogil:
+            for i in range(N):
+                lab = labels[i]
+                if lab < 0:
+                    continue
 
-            # not nan
-            if val == val:
-                nobs[lab, 0] += 1
-                sumx[lab, 0] += val
+                counts[lab] += 1
+                val = values[i, 0]
 
-    for i in range(len(counts)):
-        for j in range(K):
-            if nobs[i, j] == 0:
-                out[i, j] = nan
-            else:
-                out[i, j] = sumx[i, j]
+                # not nan
+                if val == val:
+                    nobs[lab, 0] += 1
+                    sumx[lab, 0] += val
+
+    with nogil:
+        for i in range(lcounts):
+            for j in range(K):
+                if nobs[i, j] == 0:
+                    out[i, j] = NAN
+                else:
+                    out[i, j] = sumx[i, j]
 """
 
 group_add_bin_template = """@cython.boundscheck(False)
@@ -982,7 +989,7 @@ def group_add_bin_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
     for i in range(ngroups):
         for j in range(K):
             if nobs[i, j] == 0:
-                out[i, j] = nan
+                out[i, j] = NAN
             else:
                 out[i, j] = sumx[i, j]
 """
@@ -1040,7 +1047,7 @@ def group_prod_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
     for i in range(len(counts)):
         for j in range(K):
             if nobs[i, j] == 0:
-                out[i, j] = nan
+                out[i, j] = NAN
             else:
                 out[i, j] = prodx[i, j]
 """
@@ -1098,7 +1105,7 @@ def group_prod_bin_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
     for i in range(ngroups):
         for j in range(K):
             if nobs[i, j] == 0:
-                out[i, j] = nan
+                out[i, j] = NAN
             else:
                 out[i, j] = prodx[i, j]
 """
@@ -1160,7 +1167,7 @@ def group_var_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
         for j in range(K):
             ct = nobs[i, j]
             if ct < 2:
-                out[i, j] = nan
+                out[i, j] = NAN
             else:
                 out[i, j] = ((ct * sumxx[i, j] - sumx[i, j] * sumx[i, j]) /
                              (ct * ct - ct))
@@ -1223,7 +1230,7 @@ def group_var_bin_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
         for j in range(K):
             ct = nobs[i, j]
             if ct < 2:
-                out[i, j] = nan
+                out[i, j] = NAN
             else:
                 out[i, j] = ((ct * sumxx[i, j] - sumx[i, j] * sumx[i, j]) /
                              (ct * ct - ct))
@@ -1608,7 +1615,7 @@ def group_mean_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
         for j in range(K):
             count = nobs[i, j]
             if nobs[i, j] == 0:
-                out[i, j] = nan
+                out[i, j] = NAN
             else:
                 out[i, j] = sumx[i, j] / count
 """
@@ -1663,7 +1670,7 @@ def group_mean_bin_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
         for j in range(K):
             count = nobs[i, j]
             if count == 0:
-                out[i, j] = nan
+                out[i, j] = NAN
             else:
                 out[i, j] = sumx[i, j] / count
 """
@@ -1680,7 +1687,7 @@ def group_ohlc_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
     cdef:
         Py_ssize_t i, j, N, K, ngroups, b
         %(dest_type2)s val, count
-        %(dest_type2)s vopen, vhigh, vlow, vclose, NA
+        %(dest_type2)s vopen, vhigh, vlow, vclose
         bint got_first = 0
 
     if bins[len(bins) - 1] == len(values):
@@ -1693,8 +1700,6 @@ def group_ohlc_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
     if out.shape[1] != 4:
         raise ValueError('Output array must have 4 columns')
 
-    NA = np.nan
-
     b = 0
     if K > 1:
         raise NotImplementedError("Argument 'values' must have only "
@@ -1703,10 +1708,10 @@ def group_ohlc_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
         for i in range(N):
             while b < ngroups - 1 and i >= bins[b]:
                 if not got_first:
-                    out[b, 0] = NA
-                    out[b, 1] = NA
-                    out[b, 2] = NA
-                    out[b, 3] = NA
+                    out[b, 0] = NAN
+                    out[b, 1] = NAN
+                    out[b, 2] = NAN
+                    out[b, 3] = NAN
                 else:
                     out[b, 0] = vopen
                     out[b, 1] = vhigh
@@ -1733,10 +1738,10 @@ def group_ohlc_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
                 vclose = val
 
         if not got_first:
-            out[b, 0] = NA
-            out[b, 1] = NA
-            out[b, 2] = NA
-            out[b, 3] = NA
+            out[b, 0] = NAN
+            out[b, 1] = NAN
+            out[b, 2] = NAN
+            out[b, 3] = NAN
         else:
             out[b, 0] = vopen
             out[b, 1] = vhigh
@@ -2337,16 +2342,16 @@ def generate_put_template(template, use_ints=True, use_floats=True,
 def generate_put_min_max_template(template, use_ints=True, use_floats=True,
                                   use_objects=False, use_datelikes=False):
     floats_list = [
-        ('float64', 'float64_t', 'nan', 'np.inf'),
-        ('float32', 'float32_t', 'nan', 'np.inf'),
+        ('float64', 'float64_t', 'NAN', 'np.inf'),
+        ('float32', 'float32_t', 'NAN', 'np.inf'),
     ]
     ints_list = [
         ('int64', 'int64_t', 'iNaT', _int64_max),
     ]
     date_like_list = [
         ('int64', 'int64_t', 'iNaT', _int64_max),
     ]
-    object_list = [('object', 'object', 'nan', 'np.inf')]
+    object_list = [('object', 'object', 'np.nan', 'np.inf')]
     function_list = []
     if use_floats:
         function_list.extend(floats_list)
@@ -2369,16 +2374,16 @@ def generate_put_min_max_template(template, use_ints=True, use_floats=True,
 def generate_put_selection_template(template, use_ints=True, use_floats=True,
                                     use_objects=False, use_datelikes=False):
     floats_list = [
-        ('float64', 'float64_t', 'float64_t', 'nan'),
-        ('float32', 'float32_t', 'float32_t', 'nan'),
+        ('float64', 'float64_t', 'float64_t', 'NAN'),
+        ('float32', 'float32_t', 'float32_t', 'NAN'),
     ]
     ints_list = [
         ('int64', 'int64_t', 'int64_t', 'iNaT'),
     ]
     date_like_list = [
         ('int64', 'int64_t', 'int64_t', 'iNaT'),
     ]
-    object_list = [('object', 'object', 'object', 'nan')]
+    object_list = [('object', 'object', 'object', 'np.nan')]
     function_list = []
     if use_floats:
         function_list.extend(floats_list)