diff --git a/pandas/_libs/groupby.pyx b/pandas/_libs/groupby.pyx index 7ddc087df9b11..f0f96fdfcc0a1 100644 --- a/pandas/_libs/groupby.pyx +++ b/pandas/_libs/groupby.pyx @@ -681,18 +681,17 @@ group_mean_float64 = _group_mean['double'] @cython.wraparound(False) @cython.boundscheck(False) -def _group_ohlc(floating[:, ::1] out, - int64_t[::1] counts, - ndarray[floating, ndim=2] values, - const intp_t[:] labels, - Py_ssize_t min_count=-1): +def group_ohlc(floating[:, ::1] out, + int64_t[::1] counts, + ndarray[floating, ndim=2] values, + const intp_t[:] labels, + Py_ssize_t min_count=-1): """ Only aggregates on axis=0 """ cdef: Py_ssize_t i, j, N, K, lab - floating val, count - Py_ssize_t ngroups = len(counts) + floating val assert min_count == -1, "'min_count' only used in add and prod" @@ -727,10 +726,6 @@ def _group_ohlc(floating[:, ::1] out, out[lab, 3] = val -group_ohlc_float32 = _group_ohlc['float'] -group_ohlc_float64 = _group_ohlc['double'] - - @cython.boundscheck(False) @cython.wraparound(False) def group_quantile(ndarray[float64_t] out, diff --git a/pandas/core/groupby/ops.py b/pandas/core/groupby/ops.py index 1350848741ad1..99b9aea4f82df 100644 --- a/pandas/core/groupby/ops.py +++ b/pandas/core/groupby/ops.py @@ -486,6 +486,12 @@ def _get_cython_func_and_vals( func = _get_cython_function(kind, how, values.dtype, is_numeric) else: raise + else: + if values.dtype.kind in ["i", "u"]: + if how in ["ohlc"]: + # The output may still include nans, so we have to cast + values = ensure_float64(values) + return func, values @final diff --git a/pandas/core/missing.py b/pandas/core/missing.py index 41d7fed66469d..feaecec382704 100644 --- a/pandas/core/missing.py +++ b/pandas/core/missing.py @@ -861,7 +861,4 @@ def _rolling_window(a: np.ndarray, window: int): # https://stackoverflow.com/a/6811241 shape = a.shape[:-1] + (a.shape[-1] - window + 1, window) strides = a.strides + (a.strides[-1],) - # error: Module has no attribute "stride_tricks" - return np.lib.stride_tricks.as_strided( # type: ignore[attr-defined] - a, shape=shape, strides=strides - ) + return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides) diff --git a/pandas/tests/groupby/test_libgroupby.py b/pandas/tests/groupby/test_libgroupby.py index febc12edf0b32..d776c34f5b5ec 100644 --- a/pandas/tests/groupby/test_libgroupby.py +++ b/pandas/tests/groupby/test_libgroupby.py @@ -138,7 +138,7 @@ def _check(dtype): counts = np.zeros(len(out), dtype=np.int64) labels = ensure_platform_int(np.repeat(np.arange(3), np.diff(np.r_[0, bins]))) - func = getattr(libgroupby, f"group_ohlc_{dtype}") + func = libgroupby.group_ohlc func(out, counts, obj[:, None], labels) def _ohlc(group):