cython methods for group bins #1809

Author: changhiskhan

pandas/core/groupby.py

@@ -754,12 +754,8 @@ def _aggregate(self, result, counts, values, how, is_numeric):
             raise NotImplementedError
         elif values.ndim > 2:
             for i, chunk in enumerate(values.transpose(2, 0, 1)):
-                if not is_numeric:
-                    print 'getting results for %d' % i
                 agg_func(result[:, :, i], counts, chunk.squeeze(),
                          comp_ids)
-                if not is_numeric:
-                    print 'got results for %d' % i
         else:
             agg_func(result, counts, values, comp_ids)
 
@@ -937,14 +933,22 @@ def names(self):
         'last': lib.group_last_bin
     }
 
+    _cython_object_functions = {
+        'first' : lambda a, b, c, d: lib.group_nth_bin_object(a, b, c, d, 1),
+        'last' : lib.group_last_bin_object
+    }
+
     _name_functions = {
         'ohlc' : lambda *args: ['open', 'high', 'low', 'close']
     }
 
     _filter_empty_groups = True
 
-    def _aggregate(self, result, counts, values, how):
-        agg_func = self._cython_functions[how]
+    def _aggregate(self, result, counts, values, how, is_numeric=True):
+        fdict = self._cython_functions
+        if not is_numeric:
+            fdict = self._cython_object_functions
+        agg_func = fdict[how]
         trans_func = self._cython_transforms.get(how, lambda x: x)
 
         if values.ndim > 3:

pandas/src/groupby.pyx

@@ -466,6 +466,54 @@ def group_nth_bin(ndarray[float64_t, ndim=2] out,
             else:
                 out[i, j] = resx[i, j]
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def group_nth_bin_object(ndarray[object, ndim=2] out,
+                         ndarray[int64_t] counts,
+                         ndarray[object, ndim=2] values,
+                         ndarray[int64_t] bins, int64_t rank):
+    '''
+    Only aggregates on axis=0
+    '''
+    cdef:
+        Py_ssize_t i, j, N, K, ngroups, b
+        object val
+        float64_t count
+        ndarray[object, ndim=2] resx
+        ndarray[float64_t, ndim=2] nobs
+
+    nobs = np.zeros((<object> out).shape, dtype=np.float64)
+    resx = np.empty((<object> out).shape, dtype=object)
+
+    if bins[len(bins) - 1] == len(values):
+        ngroups = len(bins)
+    else:
+        ngroups = len(bins) + 1
+
+    N, K = (<object> values).shape
+
+    b = 0
+    for i in range(N):
+        while b < ngroups - 1 and i >= bins[b]:
+            b += 1
+
+        counts[b] += 1
+        for j in range(K):
+            val = values[i, j]
+
+            # not nan
+            if val == val:
+                nobs[b, j] += 1
+                if nobs[b, j] == rank:
+                    resx[b, j] = val
+
+    for i in range(ngroups):
+        for j in range(K):
+            if nobs[i, j] == 0:
+                out[i, j] = nan
+            else:
+                out[i, j] = resx[i, j]
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def group_last(ndarray[float64_t, ndim=2] out,
@@ -595,6 +643,53 @@ def group_last_bin(ndarray[float64_t, ndim=2] out,
             else:
                 out[i, j] = resx[i, j]
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def group_last_bin_object(ndarray[object, ndim=2] out,
+                          ndarray[int64_t] counts,
+                          ndarray[object, ndim=2] values,
+                          ndarray[int64_t] bins):
+    '''
+    Only aggregates on axis=0
+    '''
+    cdef:
+        Py_ssize_t i, j, N, K, ngroups, b
+        object val
+        float64_t count
+        ndarray[object, ndim=2] resx
+        ndarray[float64_t, ndim=2] nobs
+
+    nobs = np.zeros((<object> out).shape, dtype=np.float64)
+    resx = np.empty((<object> out).shape, dtype=object)
+
+    if bins[len(bins) - 1] == len(values):
+        ngroups = len(bins)
+    else:
+        ngroups = len(bins) + 1
+
+    N, K = (<object> values).shape
+
+    b = 0
+    for i in range(N):
+        while b < ngroups - 1 and i >= bins[b]:
+            b += 1
+
+        counts[b] += 1
+        for j in range(K):
+            val = values[i, j]
+
+            # not nan
+            if val == val:
+                nobs[b, j] += 1
+                resx[b, j] = val
+
+    for i in range(ngroups):
+        for j in range(K):
+            if nobs[i, j] == 0:
+                out[i, j] = nan
+            else:
+                out[i, j] = resx[i, j]
+
 #----------------------------------------------------------------------
 # group_min, group_max
 

pandas/tseries/tests/test_timeseries.py

@@ -1080,7 +1080,7 @@ def test_frame_datetime64_handling_groupby(self):
                         (3,np.datetime64('2012-07-04'))],
                        columns = ['a', 'date'])
         result = df.groupby('a').first()
-        self.assertEqual(result['date'][3].year, 2012)
+        self.assertEqual(result['date'][3], np.datetime64('2012-07-03'))
 
     def test_series_interpolate_intraday(self):
         # #1698
@@ -2190,4 +2190,3 @@ def test_hash_equivalent(self):
 if __name__ == '__main__':
     nose.runmodule(argv=[__file__,'-vvs','-x','--pdb', '--pdb-failure'],
                    exit=False)
-