BUG: groupby.transform/agg caching *args with numba engine (#41656)

Author: mroeschke

doc/source/whatsnew/v1.3.0.rst

@@ -993,6 +993,7 @@ Groupby/resample/rolling
 - Bug in :meth:`DataFrameGroupBy.__getitem__` with non-unique columns incorrectly returning a malformed :class:`SeriesGroupBy` instead of :class:`DataFrameGroupBy` (:issue:`41427`)
 - Bug in :meth:`DataFrameGroupBy.transform` with non-unique columns incorrectly raising ``AttributeError`` (:issue:`41427`)
 - Bug in :meth:`Resampler.apply` with non-unique columns incorrectly dropping duplicated columns (:issue:`41445`)
+- Bug in :meth:`DataFrameGroupBy.transform` and :meth:`DataFrameGroupBy.agg` with ``engine="numba"`` where ``*args`` were being cached with the user passed function (:issue:`41647`)
 
 Reshaping
 ^^^^^^^^^

pandas/core/groupby/groupby.py

@@ -1131,10 +1131,16 @@ def _transform_with_numba(self, data, func, *args, engine_kwargs=None, **kwargs)
         group_keys = self.grouper._get_group_keys()
 
         numba_transform_func = numba_.generate_numba_transform_func(
-            tuple(args), kwargs, func, engine_kwargs
+            kwargs, func, engine_kwargs
         )
         result = numba_transform_func(
-            sorted_data, sorted_index, starts, ends, len(group_keys), len(data.columns)
+            sorted_data,
+            sorted_index,
+            starts,
+            ends,
+            len(group_keys),
+            len(data.columns),
+            *args,
         )
 
         cache_key = (func, "groupby_transform")
@@ -1157,11 +1163,15 @@ def _aggregate_with_numba(self, data, func, *args, engine_kwargs=None, **kwargs)
         starts, ends, sorted_index, sorted_data = self._numba_prep(func, data)
         group_keys = self.grouper._get_group_keys()
 
-        numba_agg_func = numba_.generate_numba_agg_func(
-            tuple(args), kwargs, func, engine_kwargs
-        )
+        numba_agg_func = numba_.generate_numba_agg_func(kwargs, func, engine_kwargs)
         result = numba_agg_func(
-            sorted_data, sorted_index, starts, ends, len(group_keys), len(data.columns)
+            sorted_data,
+            sorted_index,
+            starts,
+            ends,
+            len(group_keys),
+            len(data.columns),
+            *args,
         )
 
         cache_key = (func, "groupby_agg")

pandas/core/groupby/numba_.py

@@ -56,11 +56,12 @@ def f(values, index, ...):
 
 
 def generate_numba_agg_func(
-    args: tuple,
     kwargs: dict[str, Any],
     func: Callable[..., Scalar],
     engine_kwargs: dict[str, bool] | None,
-) -> Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, int], np.ndarray]:
+) -> Callable[
+    [np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, int, Any], np.ndarray
+]:
     """
     Generate a numba jitted agg function specified by values from engine_kwargs.
 
@@ -72,8 +73,6 @@ def generate_numba_agg_func(
 
     Parameters
     ----------
-    args : tuple
-        *args to be passed into the function
     kwargs : dict
         **kwargs to be passed into the function
     func : function
@@ -103,6 +102,7 @@ def group_agg(
         end: np.ndarray,
         num_groups: int,
         num_columns: int,
+        *args: Any,
     ) -> np.ndarray:
         result = np.empty((num_groups, num_columns))
         for i in numba.prange(num_groups):
@@ -116,11 +116,12 @@ def group_agg(
 
 
 def generate_numba_transform_func(
-    args: tuple,
     kwargs: dict[str, Any],
     func: Callable[..., np.ndarray],
     engine_kwargs: dict[str, bool] | None,
-) -> Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, int], np.ndarray]:
+) -> Callable[
+    [np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, int, Any], np.ndarray
+]:
     """
     Generate a numba jitted transform function specified by values from engine_kwargs.
 
@@ -132,8 +133,6 @@ def generate_numba_transform_func(
 
     Parameters
     ----------
-    args : tuple
-        *args to be passed into the function
     kwargs : dict
         **kwargs to be passed into the function
     func : function
@@ -163,6 +162,7 @@ def group_transform(
         end: np.ndarray,
         num_groups: int,
         num_columns: int,
+        *args: Any,
     ) -> np.ndarray:
         result = np.empty((len(values), num_columns))
         for i in numba.prange(num_groups):

pandas/tests/groupby/aggregate/test_numba.py

@@ -6,7 +6,9 @@
 
 from pandas import (
     DataFrame,
+    Index,
     NamedAgg,
+    Series,
     option_context,
 )
 import pandas._testing as tm
@@ -154,3 +156,20 @@ def test_multifunc_notimplimented(agg_func):
 
     with pytest.raises(NotImplementedError, match="Numba engine can"):
         grouped[1].agg(agg_func, engine="numba")
+
+
+@td.skip_if_no("numba", "0.46.0")
+def test_args_not_cached():
+    # GH 41647
+    def sum_last(values, index, n):
+        return values[-n:].sum()
+
+    df = DataFrame({"id": [0, 0, 1, 1], "x": [1, 1, 1, 1]})
+    grouped_x = df.groupby("id")["x"]
+    result = grouped_x.agg(sum_last, 1, engine="numba")
+    expected = Series([1.0] * 2, name="x", index=Index([0, 1], name="id"))
+    tm.assert_series_equal(result, expected)
+
+    result = grouped_x.agg(sum_last, 2, engine="numba")
+    expected = Series([2.0] * 2, name="x", index=Index([0, 1], name="id"))
+    tm.assert_series_equal(result, expected)

pandas/tests/groupby/transform/test_numba.py

@@ -5,6 +5,7 @@
 
 from pandas import (
     DataFrame,
+    Series,
     option_context,
 )
 import pandas._testing as tm
@@ -146,3 +147,20 @@ def test_multifunc_notimplimented(agg_func):
 
     with pytest.raises(NotImplementedError, match="Numba engine can"):
         grouped[1].transform(agg_func, engine="numba")
+
+
+@td.skip_if_no("numba", "0.46.0")
+def test_args_not_cached():
+    # GH 41647
+    def sum_last(values, index, n):
+        return values[-n:].sum()
+
+    df = DataFrame({"id": [0, 0, 1, 1], "x": [1, 1, 1, 1]})
+    grouped_x = df.groupby("id")["x"]
+    result = grouped_x.transform(sum_last, 1, engine="numba")
+    expected = Series([1.0] * 4, name="x")
+    tm.assert_series_equal(result, expected)
+
+    result = grouped_x.transform(sum_last, 2, engine="numba")
+    expected = Series([2.0] * 4, name="x")
+    tm.assert_series_equal(result, expected)