Add a numba_switch decorator to nanops and replace most of the bottleneck switches

Author: eoincondron

pandas/core/nanops.py

@@ -48,6 +48,12 @@
     notna,
 )
 
+from pandas.core.util.numba_ import GLOBAL_USE_NUMBA
+
+if GLOBAL_USE_NUMBA:
+    from pandas.core import nanops_numba
+
+
 if TYPE_CHECKING:
     from collections.abc import Callable
 
@@ -97,6 +103,38 @@ def _f(*args, **kwargs):
         return cast(F, _f)
 
 
+class numba_switch:
+    def __init__(self, name=None, **kwargs) -> None:
+        self.name = name
+        self.kwargs = kwargs
+
+    def __call__(self, alt: F) -> F:
+        nb_name = self.name or alt.__name__
+
+        try:
+            nb_func = getattr(nanops_numba, nb_name)
+        except (AttributeError, NameError):  # pragma: no cover
+            return alt
+
+        @functools.wraps(alt)
+        def f(
+            values: np.ndarray,
+            *,
+            axis: AxisInt | None = None,
+            skipna: bool = True,
+            **kwds,
+        ):
+            disallowed = values.dtype == "O"
+            if not disallowed:
+                result = nb_func(values, skipna=skipna, axis=axis, **kwds)
+            else:
+                result = alt(values, axis=axis, skipna=skipna, **kwds)
+
+            return result
+
+        return cast(F, f)
+
+
 class bottleneck_switch:
     def __init__(self, name=None, **kwargs) -> None:
         self.name = name
@@ -593,6 +631,7 @@ def nanall(
     return values.all(axis)  # type: ignore[return-value]
 
 
+@numba_switch()
 @disallow("M8")
 @_datetimelike_compat
 @maybe_operate_rowwise
@@ -658,7 +697,7 @@ def _mask_datetimelike_result(
     return result
 
 
-@bottleneck_switch()
+@numba_switch()
 @_datetimelike_compat
 def nanmean(
     values: np.ndarray,
@@ -908,7 +947,7 @@ def _get_counts_nanvar(
     return count, d
 
 
-@bottleneck_switch(ddof=1)
+@numba_switch(ddof=1)
 def nanstd(
     values,
     *,
@@ -955,7 +994,7 @@ def nanstd(
 
 
 @disallow("M8", "m8")
-@bottleneck_switch(ddof=1)
+@numba_switch(ddof=1)
 def nanvar(
     values: np.ndarray,
     *,
@@ -1033,6 +1072,7 @@ def nanvar(
     return result
 
 
+@numba_switch()
 @disallow("M8", "m8")
 def nansem(
     values: np.ndarray,
@@ -1087,7 +1127,7 @@ def nansem(
 
 
 def _nanminmax(meth, fill_value_typ):
-    @bottleneck_switch(name=f"nan{meth}")
+    @numba_switch(name=f"nan{meth}")
     @_datetimelike_compat
     def reduction(
         values: np.ndarray,