ENH: specialized cython diff functions

Author: wesm

RELEASE.rst

@@ -51,6 +51,7 @@ pandas 0.9.1
   - DataFrame.dot can accept ndarrays (#2042)
   - Support negative periods in Panel.shift (#2164)
   - Make .drop(...) work with non-unique indexes (#2101)
+  - Improve performance of Series/DataFrame.diff (re: #2087)
 
 **Bug fixes**
 

pandas/core/common.py

@@ -373,16 +373,38 @@ def mask_out_axis(arr, mask, axis, fill_value=np.nan):
 
     arr[tuple(indexer)] = fill_value
 
-def diff(arr, n, indexer, axis=0):
-    out_arr = arr - arr.take(indexer, axis=axis)
-    out_arr = _maybe_upcast(out_arr)
+_diff_special = {
+    'float64': lib.diff_2d_float64,
+    'int64': lib.diff_2d_int64,
+    'int32': lib.diff_2d_int32
+}
 
-    indexer = [slice(None)] * arr.ndim
-    if n > 0:
-        indexer[axis] = slice(None, n)
-    elif n < 0:
-        indexer[axis] = slice(None, n)
-    out_arr[tuple(indexer)] = np.nan
+def diff(arr, n, axis=0):
+    dtype = arr.dtype
+    if issubclass(dtype.type, np.integer):
+        dtype = np.float64
+    elif issubclass(dtype.type, np.bool_):
+        dtype = np.object_
+
+    out_arr = np.empty(arr.shape, dtype=dtype)
+
+    na_indexer = [slice(None)] * arr.ndim
+    na_indexer[axis] = slice(None, n)
+    out_arr[tuple(na_indexer)] = np.nan
+
+    if arr.ndim == 2 and arr.dtype.name in _diff_special:
+        f = _diff_special[arr.dtype.name]
+        f(arr, out_arr, n, axis)
+    else:
+        res_indexer = [slice(None)] * arr.ndim
+        res_indexer[axis] = slice(n, None)
+        res_indexer = tuple(res_indexer)
+
+        lag_indexer = [slice(None)] * arr.ndim
+        lag_indexer[axis] = slice(None, -n)
+        lag_indexer = tuple(lag_indexer)
+
+        out_arr[res_indexer] = arr[res_indexer] - arr[lag_indexer]
 
     return out_arr
 

pandas/core/frame.py

@@ -3680,8 +3680,7 @@ def diff(self, periods=1):
         -------
         diffed : DataFrame
         """
-        indexer = com._shift_indexer(len(self), periods)
-        new_blocks = [b.diff(periods, indexer) for b in self._data.blocks]
+        new_blocks = [b.diff(periods) for b in self._data.blocks]
         new_data = BlockManager(new_blocks, [self.columns, self.index])
         return self._constructor(new_data)
 

pandas/core/internals.py

@@ -294,10 +294,8 @@ def take(self, indexer, axis=1, fill_value=np.nan):
     def get_values(self, dtype):
         return self.values
 
-    def diff(self, n, indexer=None):
-        if indexer is None:
-            indexer = com._shift_indexer(self.shape[1], n)
-        new_values = com.diff(self.values, n, indexer, axis=1)
+    def diff(self, n):
+        new_values = com.diff(self.values, n, axis=1)
         return make_block(new_values, self.items, self.ref_items)
 
 

pandas/core/series.py

@@ -1511,9 +1511,8 @@ def diff(self, periods=1):
         -------
         diffed : Series
         """
-        indexer = com._shift_indexer(len(self), periods)
-        val = com.diff(self.values, periods, indexer)
-        return Series(val, self.index, name=self.name)
+        result = com.diff(self.values[:, np.newaxis], periods)
+        return Series(result.squeeze(), self.index, name=self.name)
 
     def autocorr(self):
         """

pandas/src/stats.pyx

@@ -565,3 +565,89 @@ def rank_2d_generic(object in_arr, axis=0, ties_method='average',
 #         for j in range(K):
 #             result[i, j] = values[i, indexer[i, j]]
 #     return result
+
+@cython.wraparound(False)
+@cython.boundscheck(False)
+def diff_2d_float64(ndarray[float64_t, ndim=2] arr,
+                    ndarray[float64_t, ndim=2] out,
+                    Py_ssize_t periods, int axis):
+    cdef:
+        Py_ssize_t i, j, sx, sy
+
+    sx, sy = (<object> arr).shape
+    if arr.flags.f_contiguous:
+        if axis == 0:
+            for j in range(sy):
+                for i in range(periods, sx):
+                    out[i, j] = arr[i, j] - arr[i - periods, j]
+        else:
+            for j in range(periods, sy):
+                for i in range(sx):
+                    out[i, j] = arr[i, j] - arr[i, j - periods]
+    else:
+        if axis == 0:
+            for i in range(periods, sx):
+                for j in range(sy):
+                    out[i, j] = arr[i, j] - arr[i - periods, j]
+        else:
+            for i in range(sx):
+                for j in range(periods, sy):
+                    out[i, j] = arr[i, j] - arr[i, j - periods]
+
+@cython.wraparound(False)
+@cython.boundscheck(False)
+def diff_2d_int64(ndarray[int64_t, ndim=2] arr,
+                  ndarray[float64_t, ndim=2] out,
+                  Py_ssize_t periods, int axis):
+    cdef:
+        Py_ssize_t i, j, sx, sy
+
+    sx, sy = (<object> arr).shape
+    if arr.flags.f_contiguous:
+        if axis == 0:
+            for j in range(sy):
+                for i in range(periods, sx):
+                    out[i, j] = arr[i, j] - arr[i - periods, j]
+        else:
+            for j in range(periods, sy):
+                for i in range(sx):
+                    out[i, j] = arr[i, j] - arr[i, j - periods]
+    else:
+        if axis == 0:
+            for i in range(periods, sx):
+                for j in range(sy):
+                    out[i, j] = arr[i, j] - arr[i - periods, j]
+        else:
+            for i in range(sx):
+                for j in range(periods, sy):
+                    out[i, j] = arr[i, j] - arr[i, j - periods]
+
+
+@cython.wraparound(False)
+@cython.boundscheck(False)
+def diff_2d_int32(ndarray[int64_t, ndim=2] arr,
+                  ndarray[float64_t, ndim=2] out,
+                  Py_ssize_t periods, int axis):
+    cdef:
+        Py_ssize_t i, j, sx, sy
+
+    sx, sy = (<object> arr).shape
+    if arr.flags.f_contiguous:
+        if axis == 0:
+            for j in range(sy):
+                for i in range(periods, sx):
+                    out[i, j] = arr[i, j] - arr[i - periods, j]
+        else:
+            for j in range(periods, sy):
+                for i in range(sx):
+                    out[i, j] = arr[i, j] - arr[i, j - periods]
+    else:
+        if axis == 0:
+            for i in range(periods, sx):
+                for j in range(sy):
+                    out[i, j] = arr[i, j] - arr[i - periods, j]
+        else:
+            for i in range(sx):
+                for j in range(periods, sy):
+                    out[i, j] = arr[i, j] - arr[i, j - periods]
+