ENH: implement Cython OHLC function for groupby #152

Author: wesm

pandas/core/groupby.py

@@ -853,11 +853,18 @@ def agg_series(self, obj, func):
         'add' : lib.group_add_bin,
         'mean' : lib.group_mean_bin,
         'var' : lib.group_var_bin,
-        'std' : lib.group_var_bin
+        'std' : lib.group_var_bin,
+        'ohlc' : lib.group_ohlc
+    }
+
+    _cython_arity = {
+        'ohlc' : 4, # OHLC
     }
 
     def aggregate(self, values, how):
         agg_func = self._cython_functions[how]
+        arity = self._cython_arity.get(how, 1)
+
         if values.ndim == 1:
             squeeze = True
             values = values[:, None]

pandas/src/engines.pyx

@@ -27,9 +27,15 @@ PyDateTime_IMPORT
 cdef extern from "Python.h":
     int PySlice_Check(object)
     int PyList_Check(object)
-
+    int PyTuple_Check(object)
 
 cdef inline is_definitely_invalid_key(object val):
+    if PyTuple_Check(val):
+        try:
+            hash(val)
+        except TypeError:
+            return True
+
     return (PySlice_Check(val) or cnp.PyArray_Check(val)
             or PyList_Check(val))
 

pandas/src/groupby.pyx

@@ -532,6 +532,78 @@ def group_add_bin(ndarray[float64_t, ndim=2] out,
             else:
                 out[i, j] = sumx[i, j]
 
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def group_ohlc(ndarray[float64_t, ndim=2] out,
+                  ndarray[int32_t] counts,
+                  ndarray[float64_t, ndim=2] values,
+                  ndarray[int32_t] bins):
+    '''
+    Only aggregates on axis=0
+    '''
+    cdef:
+        Py_ssize_t i, j, N, K, ngroups, b
+        float64_t val, count
+        float64_t vopen, vhigh, vlow, vclose, NA
+        bint got_first = 0
+
+    ngroups = len(bins) + 1
+    N, K = (<object> values).shape
+
+    if out.shape[1] != 4:
+        raise ValueError('Output array must have 4 columns')
+
+    NA = np.nan
+
+    b = 0
+    if K > 1:
+        raise NotImplementedError
+    else:
+        for i in range(N):
+            if b < ngroups - 1 and i >= bins[b]:
+                if not got_first:
+                    out[b, 0] = NA
+                    out[b, 1] = NA
+                    out[b, 2] = NA
+                    out[b, 3] = NA
+                else:
+                    out[b, 0] = vopen
+                    out[b, 1] = vlow
+                    out[b, 2] = vhigh
+                    out[b, 3] = vclose
+                b += 1
+                got_first = 0
+
+            counts[b] += 1
+            val = values[i, 0]
+
+            # not nan
+            if val == val:
+                if not got_first:
+                    got_first = 1
+                    vopen = val
+                    vlow = val
+                    vhigh = val
+                else:
+                    if val < vlow:
+                        vlow = val
+                    if val > vhigh:
+                        vhigh = val
+                vclose = val
+
+        if not got_first:
+            out[b, 0] = NA
+            out[b, 1] = NA
+            out[b, 2] = NA
+            out[b, 3] = NA
+        else:
+            out[b, 0] = vopen
+            out[b, 1] = vlow
+            out[b, 2] = vhigh
+            out[b, 3] = vclose
+
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def group_mean_bin(ndarray[float64_t, ndim=2] out,

pandas/tests/test_tseries.py

@@ -1,7 +1,7 @@
 import unittest
 
 import numpy as np
-from pandas import Index
+from pandas import Index, isnull
 from pandas.util.testing import assert_almost_equal
 import pandas.util.testing as common
 import pandas._tseries as lib
@@ -317,7 +317,7 @@ def test_group_add_bin():
     # bin-based group_add
     bins = np.array([3, 6], dtype=np.int32)
     out  = np.zeros((3, 1), np.float64)
-    counts = np.empty(len(out), dtype=np.int32)
+    counts = np.zeros(len(out), dtype=np.int32)
     lib.group_add_bin(out, counts, obj, bins)
 
     assert_almost_equal(out, exp)
@@ -334,7 +334,7 @@ def test_group_mean_bin():
     # bin-based group_mean
     bins = np.array([3, 6], dtype=np.int32)
     out  = np.zeros((3, 1), np.float64)
-    counts = np.empty(len(out), dtype=np.int32)
+    counts = np.zeros(len(out), dtype=np.int32)
     lib.group_mean_bin(out, counts, obj, bins)
 
     assert_almost_equal(out, exp)
@@ -351,12 +351,37 @@ def test_group_var_bin():
     # bin-based group_var
     bins = np.array([3, 6], dtype=np.int32)
     out  = np.zeros((3, 1), np.float64)
-    counts = np.empty(len(out), dtype=np.int32)
+    counts = np.zeros(len(out), dtype=np.int32)
 
     lib.group_var_bin(out, counts, obj, bins)
 
     assert_almost_equal(out, exp)
 
+def test_group_ohlc():
+    obj = np.random.randn(20)
+
+    bins = np.array([6, 12], dtype=np.int32)
+    out  = np.zeros((3, 4), np.float64)
+    counts = np.zeros(len(out), dtype=np.int32)
+
+    lib.group_ohlc(out, counts, obj[:, None], bins)
+
+    def _ohlc(group):
+        if isnull(group).all():
+            return np.repeat(np.nan, 4)
+        return [group[0], group.min(), group.max(), group[-1]]
+
+    expected = np.array([_ohlc(obj[:6]), _ohlc(obj[6:12]),
+                         _ohlc(obj[12:])])
+
+    assert_almost_equal(out, expected)
+    assert_almost_equal(counts, [6, 6, 8])
+
+    obj[:6] = np.nan
+    lib.group_ohlc(out, counts, obj[:, None], bins)
+    expected[0] = np.nan
+    assert_almost_equal(out, expected)
+
 class TestTypeInference(unittest.TestCase):
 
     def test_length_zero(self):

ts_todo.txt

@@ -15,3 +15,6 @@ gp- get rid of Ts class, simplify timestamp creation
 - attach tz in DatetimeIndex.asobject
 - failing duplicate timestamp test
 - _tseries.pyd depends on datetime.pyx
+
+
+- BUG: time_rule DateRange tests