use fused types for reshape (#22454)

Author: jbrockmendel

pandas/_libs/reshape.pyx

@@ -1,15 +1,95 @@
 # -*- coding: utf-8 -*-
 
-cimport cython
-from cython cimport Py_ssize_t
+import cython
+from cython import Py_ssize_t
 
-import numpy as np
-from numpy cimport (ndarray,
-                    int8_t, int16_t, int32_t, int64_t, uint8_t, uint16_t,
+from numpy cimport (int8_t, int16_t, int32_t, int64_t, uint8_t, uint16_t,
                     uint32_t, uint64_t, float32_t, float64_t)
 
 
-cdef double NaN = <double> np.NaN
-cdef double nan = NaN
+ctypedef fused reshape_t:
+    uint8_t
+    uint16_t
+    uint32_t
+    uint64_t
+    int8_t
+    int16_t
+    int32_t
+    int64_t
+    float32_t
+    float64_t
+    object
 
-include "reshape_helper.pxi"
+
+@cython.wraparound(False)
+@cython.boundscheck(False)
+def unstack(reshape_t[:, :] values, uint8_t[:] mask,
+            Py_ssize_t stride, Py_ssize_t length, Py_ssize_t width,
+            reshape_t[:, :] new_values, uint8_t[:, :] new_mask):
+    """
+    transform long sorted_values to wide new_values
+
+    Parameters
+    ----------
+    values : typed ndarray
+    mask : boolean ndarray
+    stride : int
+    length : int
+    width : int
+    new_values : typed ndarray
+        result array
+    new_mask : boolean ndarray
+        result mask
+    """
+    cdef:
+        Py_ssize_t i, j, w, nulls, s, offset
+
+    if reshape_t is not object:
+        # evaluated at compile-time
+        with nogil:
+            for i in range(stride):
+
+                nulls = 0
+                for j in range(length):
+
+                    for w in range(width):
+
+                        offset = j * width + w
+
+                        if mask[offset]:
+                            s = i * width + w
+                            new_values[j, s] = values[offset - nulls, i]
+                            new_mask[j, s] = 1
+                        else:
+                            nulls += 1
+
+    else:
+        # object-dtype, identical to above but we cannot use nogil
+        for i in range(stride):
+
+            nulls = 0
+            for j in range(length):
+
+                for w in range(width):
+
+                    offset = j * width + w
+
+                    if mask[offset]:
+                        s = i * width + w
+                        new_values[j, s] = values[offset - nulls, i]
+                        new_mask[j, s] = 1
+                    else:
+                        nulls += 1
+
+
+unstack_uint8 = unstack["uint8_t"]
+unstack_uint16 = unstack["uint16_t"]
+unstack_uint32 = unstack["uint32_t"]
+unstack_uint64 = unstack["uint64_t"]
+unstack_int8 = unstack["int8_t"]
+unstack_int16 = unstack["int16_t"]
+unstack_int32 = unstack["int32_t"]
+unstack_int64 = unstack["int64_t"]
+unstack_float32 = unstack["float32_t"]
+unstack_float64 = unstack["float64_t"]
+unstack_object = unstack["object"]

pandas/_libs/reshape_helper.pxi.in

@@ -77,7 +77,6 @@ def is_platform_windows():
               '_libs/algos_rank_helper.pxi.in'],
     'groupby': ['_libs/groupby_helper.pxi.in'],
     'join': ['_libs/join_helper.pxi.in', '_libs/join_func_helper.pxi.in'],
-    'reshape': ['_libs/reshape_helper.pxi.in'],
     'hashtable': ['_libs/hashtable_class_helper.pxi.in',
                   '_libs/hashtable_func_helper.pxi.in'],
     'index': ['_libs/index_class_helper.pxi.in'],
@@ -558,7 +557,7 @@ def srcpath(name=None, suffix='.pyx', subdir='src'):
         'include': []},
     '_libs.reshape': {
         'pyxfile': '_libs/reshape',
-        'depends': _pxi_dep['reshape']},
+        'depends': []},
     '_libs.skiplist': {
         'pyxfile': '_libs/skiplist',
         'depends': ['pandas/_libs/src/skiplist.h']},