Backport PR #55586 on branch 2.1.x (REGR: Groupby methods not supporting numba raising TypeError when the…) (#55686)

Backport PR #55586: REGR: Groupby methods not supporting numba raising TypeError when the…

Co-authored-by: Thomas Li <[email protected]>

Author: meeseeksmachine

doc/source/whatsnew/v2.1.2.rst

@@ -16,6 +16,7 @@ Fixed regressions
 - Fixed regression in :meth:`DataFrame.join` where result has missing values and dtype is arrow backed string (:issue:`55348`)
 - Fixed regression in :meth:`DataFrame.resample` which was extrapolating back to ``origin`` when ``origin`` was outside its bounds (:issue:`55064`)
 - Fixed regression in :meth:`DataFrame.sort_index` which was not sorting correctly when the index was a sliced :class:`MultiIndex` (:issue:`55379`)
+- Fixed regression in :meth:`DataFrameGroupBy.agg` and :meth:`SeriesGroupBy.agg` where if the option ``compute.use_numba`` was set to True, groupby methods not supported by the numba engine would raise a ``TypeError`` (:issue:`55520`)
 - Fixed performance regression with wide DataFrames, typically involving methods where all columns were accessed individually (:issue:`55256`, :issue:`55245`)
 - Fixed regression in :func:`merge_asof` raising ``TypeError`` for ``by`` with datetime and timedelta dtypes (:issue:`55453`)
 

pandas/core/groupby/generic.py

@@ -236,10 +236,13 @@ def aggregate(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs)
             kwargs = {}
 
         if isinstance(func, str):
-            if maybe_use_numba(engine):
+            if maybe_use_numba(engine) and engine is not None:
                 # Not all agg functions support numba, only propagate numba kwargs
-                # if user asks for numba
+                # if user asks for numba, and engine is not None
+                # (if engine is None, the called function will handle the case where
+                # numba is requested via the global option)
                 kwargs["engine"] = engine
+            if engine_kwargs is not None:
                 kwargs["engine_kwargs"] = engine_kwargs
             return getattr(self, func)(*args, **kwargs)
 

pandas/tests/groupby/test_numba.py

@@ -3,6 +3,7 @@
 from pandas import (
     DataFrame,
     Series,
+    option_context,
 )
 import pandas._testing as tm
 
@@ -66,3 +67,14 @@ def test_axis_1_unsupported(self, numba_supported_reductions):
         gb = df.groupby("a", axis=1)
         with pytest.raises(NotImplementedError, match="axis=1"):
             getattr(gb, func)(engine="numba", **kwargs)
+
+    def test_no_engine_doesnt_raise(self):
+        # GH55520
+        df = DataFrame({"a": [3, 2, 3, 2], "b": range(4), "c": range(1, 5)})
+        gb = df.groupby("a")
+        # Make sure behavior of functions w/out engine argument don't raise
+        # when the global use_numba option is set
+        with option_context("compute.use_numba", True):
+            res = gb.agg({"b": "first"})
+        expected = gb.agg({"b": "first"})
+        tm.assert_frame_equal(res, expected)