cython optimizations and cleanup

Author: jbrockmendel

pandas/_libs/algos.pyx

@@ -77,6 +77,8 @@ class NegInfinity(object):
     __ge__ = lambda self, other: isinstance(other, NegInfinity)
 
 
+@cython.wraparound(False)
+@cython.boundscheck(False)
 cpdef ndarray[int64_t, ndim=1] unique_deltas(int64_t[:] arr):
     """
     Efficiently find the unique first-differences of the given array.

pandas/_libs/groupby.pyx

@@ -98,7 +98,7 @@ cdef inline float64_t kth_smallest_c(float64_t* a,
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def group_median_float64(ndarray[float64_t, ndim=2] out,
+def group_median_float64(float64_t[:, :] out,
                          int64_t[:] counts,
                          ndarray[float64_t, ndim=2] values,
                          int64_t[:] labels,
@@ -109,7 +109,7 @@ def group_median_float64(ndarray[float64_t, ndim=2] out,
     cdef:
         Py_ssize_t i, j, N, K, ngroups, size
         int64_t[:] _counts
-        ndarray data
+        ndarray[float64_t, ndim=2] data
         float64_t* ptr
 
     assert min_count == -1, "'min_count' only used in add and prod"
@@ -291,16 +291,16 @@ def group_fillna_indexer(int64_t[:] out, ndarray[int64_t] labels,
     """
     cdef:
         Py_ssize_t i, N
-        ndarray[int64_t] sorted_labels
+        int64_t[:] sorted_labels
         int64_t idx, curr_fill_idx=-1, filled_vals=0
 
     N = len(out)
 
     # Make sure all arrays are the same size
     assert N == len(labels) == len(mask)
 
-    sorted_labels = np.argsort(labels, kind='mergesort').astype(
-        np.int64, copy=False)
+    sorted_labels = np.argsort(labels, kind='mergesort').astype(np.int64,
+                                                                copy=False)
     if direction == 'bfill':
         sorted_labels = sorted_labels[::-1]
 
@@ -327,7 +327,7 @@ def group_fillna_indexer(int64_t[:] out, ndarray[int64_t] labels,
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def group_any_all(ndarray[uint8_t] out,
+def group_any_all(uint8_t[:] out,
                   int64_t[:] labels,
                   uint8_t[:] values,
                   uint8_t[:] mask,
@@ -370,7 +370,7 @@ def group_any_all(ndarray[uint8_t] out,
     else:
         raise ValueError("'bool_func' must be either 'any' or 'all'!")
 
-    out.fill(1 - flag_val)
+    out[:] = 1 - flag_val
 
     with nogil:
         for i in range(N):

pandas/_libs/groupby_helper.pxi.in

@@ -220,7 +220,7 @@ def group_mean_{{name}}(ndarray[{{c_type}}, ndim=2] out,
 
 @cython.wraparound(False)
 @cython.boundscheck(False)
-def group_ohlc_{{name}}(ndarray[{{c_type}}, ndim=2] out,
+def group_ohlc_{{name}}({{c_type}}[:, :] out,
                   int64_t[:] counts,
                   ndarray[{{c_type}}, ndim=2] values,
                   int64_t[:] labels,
@@ -246,7 +246,8 @@ def group_ohlc_{{name}}(ndarray[{{c_type}}, ndim=2] out,
     if K > 1:
         raise NotImplementedError("Argument 'values' must have only "
                                   "one dimension")
-    out.fill(np.nan)
+
+    out[:] = np.nan
 
     with nogil:
         for i in range(N):
@@ -304,16 +305,16 @@ def group_last_{{name}}(ndarray[{{c_type}}, ndim=2] out,
     cdef:
         Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
         {{c_type}} val
-        ndarray[{{c_type}}, ndim=2] resx
-        ndarray[int64_t, ndim=2] nobs
+        {{c_type}}[:, :] resx
+        int64_t[:, :] nobs
 
     assert min_count == -1, "'min_count' only used in add and prod"
 
     if not len(values) == len(labels):
         raise AssertionError("len(index) != len(labels)")
 
     nobs = np.zeros((<object> out).shape, dtype=np.int64)
-    {{if name=='object'}}
+    {{if name == 'object'}}
     resx = np.empty((<object> out).shape, dtype=object)
     {{else}}
     resx = np.empty_like(out)
@@ -361,8 +362,8 @@ def group_nth_{{name}}(ndarray[{{c_type}}, ndim=2] out,
     cdef:
         Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
         {{c_type}} val
-        ndarray[{{c_type}}, ndim=2] resx
-        ndarray[int64_t, ndim=2] nobs
+        {{c_type}}[:, :] resx
+        int64_t[:, :] nobs
 
     assert min_count == -1, "'min_count' only used in add and prod"
 
@@ -411,7 +412,7 @@ def group_nth_{{name}}(ndarray[{{c_type}}, ndim=2] out,
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def group_rank_{{name}}(ndarray[float64_t, ndim=2] out,
+def group_rank_{{name}}(float64_t[:, :] out,
                         ndarray[{{c_type}}, ndim=2] values,
                         ndarray[int64_t] labels,
                         bint is_datetimelike, object ties_method,
@@ -453,8 +454,8 @@ def group_rank_{{name}}(ndarray[float64_t, ndim=2] out,
         TiebreakEnumType tiebreak
         Py_ssize_t i, j, N, K, grp_start=0, dups=0, sum_ranks=0
         Py_ssize_t grp_vals_seen=1, grp_na_count=0, grp_tie_count=0
-        ndarray[int64_t] _as
-        ndarray[float64_t, ndim=2] grp_sizes
+        int64_t[:] _as
+        float64_t[:, :] grp_sizes
         ndarray[{{c_type}}] masked_vals
         ndarray[uint8_t] mask
         bint keep_na
@@ -617,7 +618,7 @@ def group_max(ndarray[groupby_t, ndim=2] out,
     cdef:
         Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
         groupby_t val, count, nan_val
-        ndarray[groupby_t, ndim=2] maxx, nobs
+        groupby_t[:, :] maxx, nobs
 
     assert min_count == -1, "'min_count' only used in add and prod"
 
@@ -629,10 +630,10 @@ def group_max(ndarray[groupby_t, ndim=2] out,
     maxx = np.empty_like(out)
     if groupby_t is int64_t:
         # Note: evaluated at compile-time
-        maxx.fill(-_int64_max)
+        maxx[:] = -_int64_max
         nan_val = iNaT
     else:
-        maxx.fill(-np.inf)
+        maxx[:] = -np.inf
         nan_val = NAN
 
     N, K = (<object> values).shape
@@ -685,7 +686,7 @@ def group_min(ndarray[groupby_t, ndim=2] out,
     cdef:
         Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
         groupby_t val, count, nan_val
-        ndarray[groupby_t, ndim=2] minx, nobs
+        groupby_t[:, :] minx, nobs
 
     assert min_count == -1, "'min_count' only used in add and prod"
 
@@ -696,10 +697,10 @@ def group_min(ndarray[groupby_t, ndim=2] out,
 
     minx = np.empty_like(out)
     if groupby_t is int64_t:
-        minx.fill(_int64_max)
+        minx[:] = _int64_max
         nan_val = iNaT
     else:
-        minx.fill(np.inf)
+        minx[:] = np.inf
         nan_val = NAN
 
     N, K = (<object> values).shape
@@ -741,7 +742,7 @@ group_min_int64 = group_min["int64_t"]
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def group_cummin(ndarray[groupby_t, ndim=2] out,
+def group_cummin(groupby_t[:, :] out,
                  ndarray[groupby_t, ndim=2] values,
                  ndarray[int64_t] labels,
                  bint is_datetimelike):
@@ -751,15 +752,15 @@ def group_cummin(ndarray[groupby_t, ndim=2] out,
     cdef:
         Py_ssize_t i, j, N, K, size
         groupby_t val, mval
-        ndarray[groupby_t, ndim=2] accum
+        groupby_t[:, :] accum
         int64_t lab
 
     N, K = (<object> values).shape
     accum = np.empty_like(values)
     if groupby_t is int64_t:
-        accum.fill(_int64_max)
+        accum[:] = _int64_max
     else:
-        accum.fill(np.inf)
+        accum[:] = np.inf
 
     with nogil:
         for i in range(N):
@@ -794,7 +795,7 @@ group_cummin_int64 = group_cummin["int64_t"]
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def group_cummax(ndarray[groupby_t, ndim=2] out,
+def group_cummax(groupby_t[:, :] out,
                  ndarray[groupby_t, ndim=2] values,
                  ndarray[int64_t] labels,
                  bint is_datetimelike):
@@ -804,15 +805,15 @@ def group_cummax(ndarray[groupby_t, ndim=2] out,
     cdef:
         Py_ssize_t i, j, N, K, size
         groupby_t val, mval
-        ndarray[groupby_t, ndim=2] accum
+        groupby_t[:, :] accum
         int64_t lab
 
     N, K = (<object> values).shape
     accum = np.empty_like(values)
     if groupby_t is int64_t:
-        accum.fill(-_int64_max)
+        accum[:] = -_int64_max
     else:
-        accum.fill(-np.inf)
+        accum[:] = -np.inf
 
     with nogil:
         for i in range(N):

pandas/_libs/hashtable.pyx

@@ -2,10 +2,8 @@
 
 cimport cython
 
-from cpython cimport (PyObject, Py_INCREF, PyList_Check, PyTuple_Check,
-                      PyMem_Malloc, PyMem_Realloc, PyMem_Free,
-                      PyString_Check, PyBytes_Check,
-                      PyUnicode_Check)
+from cpython cimport (PyObject, Py_INCREF,
+                      PyMem_Malloc, PyMem_Realloc, PyMem_Free)
 
 from libc.stdlib cimport malloc, free
 
@@ -153,7 +151,7 @@ def unique_label_indices(ndarray[int64_t, ndim=1] labels):
     cdef:
         int ret = 0
         Py_ssize_t i, n = len(labels)
-        kh_int64_t * table = kh_init_int64()
+        kh_int64_t *table = kh_init_int64()
         Int64Vector idx = Int64Vector()
         ndarray[int64_t, ndim=1] arr
         Int64VectorData *ud = idx.data