REF: better use of fused_types for group_ohlc (#40668)

Author: jbrockmendel

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,

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

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)

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):