PRF: Optimize 2d take operations.

Use memory views

Author: dalejung

pandas/core/common.py

@@ -675,6 +675,17 @@ def take_nd(arr, indexer, axis=0, out=None, fill_value=np.nan,
                     # to crash when trying to cast it to dtype)
                     dtype, fill_value = arr.dtype, arr.dtype.type()
 
+    flip_order = False
+    if arr.ndim == 2:
+        if arr.flags.f_contiguous:
+            flip_order = True
+
+    if flip_order:
+        arr = arr.T
+        axis = arr.ndim - axis - 1
+        if out is not None:
+            out = out.T
+
     # at this point, it's guaranteed that dtype can hold both the arr values
     # and the fill_value
     if out is None:
@@ -692,7 +703,11 @@ def take_nd(arr, indexer, axis=0, out=None, fill_value=np.nan,
 
     func = _get_take_nd_function(arr.ndim, arr.dtype, out.dtype,
                                  axis=axis, mask_info=mask_info)
+
     func(arr, indexer, out, fill_value)
+
+    if flip_order:
+        out = out.T
     return out
 
 

pandas/src/generate_code.py

@@ -80,9 +80,9 @@ def take_1d_%(name)s_%(dest)s(ndarray[%(c_type_in)s] values,
 
 take_2d_axis0_template = """@cython.wraparound(False)
 @cython.boundscheck(False)
-def take_2d_axis0_%(name)s_%(dest)s(ndarray[%(c_type_in)s, ndim=2] values,
+def take_2d_axis0_%(name)s_%(dest)s(%(c_type_in)s[:, :] values,
                                     ndarray[int64_t] indexer,
-                                    ndarray[%(c_type_out)s, ndim=2] out,
+                                    %(c_type_out)s[:, :] out,
                                     fill_value=np.nan):
     cdef:
         Py_ssize_t i, j, k, n, idx
@@ -127,9 +127,9 @@ def take_2d_axis0_%(name)s_%(dest)s(ndarray[%(c_type_in)s, ndim=2] values,
 
 take_2d_axis1_template = """@cython.wraparound(False)
 @cython.boundscheck(False)
-def take_2d_axis1_%(name)s_%(dest)s(ndarray[%(c_type_in)s, ndim=2] values,
+def take_2d_axis1_%(name)s_%(dest)s(%(c_type_in)s[:, :] values,
                                     ndarray[int64_t] indexer,
-                                    ndarray[%(c_type_out)s, ndim=2] out,
+                                    %(c_type_out)s[:, :] out,
                                     fill_value=np.nan):
     cdef:
         Py_ssize_t i, j, k, n, idx
@@ -143,34 +143,12 @@ def take_2d_axis1_%(name)s_%(dest)s(ndarray[%(c_type_in)s, ndim=2] values,
 
     fv = fill_value
 
-    IF %(can_copy)s:
-        cdef:
-            %(c_type_out)s *v
-            %(c_type_out)s *o
-
-        #GH3130
-        if (values.strides[0] == out.strides[0] and
-            values.strides[0] == sizeof(%(c_type_out)s) and
-            sizeof(%(c_type_out)s) * n >= 256):
-
-            for j from 0 <= j < k:
-                idx = indexer[j]
-                if idx == -1:
-                    for i from 0 <= i < n:
-                        out[i, j] = fv
-                else:
-                    v = &values[0, idx]
-                    o = &out[0, j]
-                    memmove(o, v, <size_t>(sizeof(%(c_type_out)s) * n))
-            return
-
-    for j from 0 <= j < k:
-        idx = indexer[j]
-        if idx == -1:
-            for i from 0 <= i < n:
+    for i from 0 <= i < n:
+        for j from 0 <= j < k:
+            idx = indexer[j]
+            if idx == -1:
                 out[i, j] = fv
-        else:
-            for i from 0 <= i < n:
+            else:
                 out[i, j] = %(preval)svalues[i, idx]%(postval)s
 
 """