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

Author: eoincondron

pandas/core/nanops.py

@@ -48,6 +48,8 @@
     notna,
 )
 
+from pandas.core import nanops_numba
+
 if TYPE_CHECKING:
     from collections.abc import Callable
 
@@ -66,6 +68,18 @@ def set_use_bottleneck(v: bool = True) -> None:
 set_use_bottleneck(get_option("compute.use_bottleneck"))
 
 
+_USE_NUMBA = True
+
+
+def set_use_numba(v: bool = True) -> None:
+    # set/unset to use bottleneck
+    global _USE_NUMBA
+    _USE_NUMBA = v
+
+
+# set_use_numba(get_option("compute.use_numba"))
+
+
 class disallow:
     def __init__(self, *dtypes: Dtype) -> None:
         super().__init__()
@@ -97,6 +111,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
+            nb_func = None
+
+        @functools.wraps(alt)
+        def f(
+            values: np.ndarray,
+            *,
+            axis: AxisInt | None = None,
+            skipna: bool = True,
+            **kwds,
+        ):
+            disallowed = values.dtype == "O"
+            if _USE_NUMBA and 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 +639,7 @@ def nanall(
     return values.all(axis)  # type: ignore[return-value]
 
 
+@numba_switch()
 @disallow("M8")
 @_datetimelike_compat
 @maybe_operate_rowwise
@@ -660,7 +707,7 @@ def _mask_datetimelike_result(
     return result
 
 
-@bottleneck_switch()
+@numba_switch()
 @_datetimelike_compat
 def nanmean(
     values: np.ndarray,
@@ -910,7 +957,7 @@ def _get_counts_nanvar(
     return count, d
 
 
-@bottleneck_switch(ddof=1)
+@numba_switch(ddof=1)
 def nanstd(
     values,
     *,
@@ -957,7 +1004,7 @@ def nanstd(
 
 
 @disallow("M8", "m8")
-@bottleneck_switch(ddof=1)
+@numba_switch(ddof=1)
 def nanvar(
     values: np.ndarray,
     *,
@@ -1035,6 +1082,7 @@ def nanvar(
     return result
 
 
+@numba_switch()
 @disallow("M8", "m8")
 def nansem(
     values: np.ndarray,
@@ -1089,7 +1137,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,