move the tests to test_numba.py

Author: auderson

pandas/tests/window/test_apply.py

@@ -316,18 +316,3 @@ def test_center_reindex_frame(raw):
     )
     frame_rs = frame.rolling(window=25, min_periods=minp, center=True).apply(f, raw=raw)
     tm.assert_frame_equal(frame_xp, frame_rs)
-
-
-def test_apply_numba_with_kwargs():
-    # 58995
-    def func(sr, a=0):
-        return sr.sum() + a
-
-    data = DataFrame(range(10))
-
-    result = data.rolling(5).apply(func, engine="numba", raw=True, kwargs={"a": 1})
-    expected = data.rolling(5).sum() + 1
-    tm.assert_frame_equal(result, expected)
-
-    result = data.rolling(5).apply(func, engine="numba", raw=True, args=(1,))
-    tm.assert_frame_equal(result, expected)

pandas/tests/window/test_expanding.py

@@ -691,18 +691,3 @@ def test_numeric_only_corr_cov_series(kernel, use_arg, numeric_only, dtype):
         op2 = getattr(expanding2, kernel)
         expected = op2(*arg2, numeric_only=numeric_only)
         tm.assert_series_equal(result, expected)
-
-
-def test_apply_numba_with_kwargs():
-    # 58995
-    def func(sr, a=0):
-        return sr.sum() + a
-
-    data = DataFrame(range(10))
-
-    result = data.expanding().apply(func, engine="numba", raw=True, kwargs={"a": 1})
-    expected = data.expanding().sum() + 1
-    tm.assert_frame_equal(result, expected)
-
-    result = data.expanding().apply(func, engine="numba", raw=True, args=(1,))
-    tm.assert_frame_equal(result, expected)

pandas/tests/window/test_groupby.py

@@ -1024,26 +1024,6 @@ def test_datelike_on_not_monotonic_within_each_group(self):
         with pytest.raises(ValueError, match="Each group within B must be monotonic."):
             df.groupby("A").rolling("365D", on="B")
 
-    def test_groupby_rolling_apply_numba_with_kwargs(self, roll_frame):
-        def func(sr, a=0):
-            return sr.sum() + a
-
-        # 58995
-        result = (
-            roll_frame.groupby("A")
-            .rolling(5)
-            .apply(func, engine="numba", raw=True, kwargs={"a": 1})
-        )
-        expected = roll_frame.groupby("A").rolling(5).sum() + 1
-        tm.assert_frame_equal(result, expected)
-
-        result = (
-            roll_frame.groupby("A")
-            .rolling(5)
-            .apply(func, engine="numba", raw=True, args=(1,))
-        )
-        tm.assert_frame_equal(result, expected)
-
 
 class TestExpanding:
     @pytest.fixture
@@ -1154,26 +1134,6 @@ def test_expanding_apply(self, raw, frame):
         expected.index = expected_index
         tm.assert_frame_equal(result, expected)
 
-    def test_groupby_expanding_apply_numba_with_kwargs(self, roll_frame):
-        # 58995
-        def func(sr, a=0):
-            return sr.sum() + a
-
-        result = (
-            roll_frame.groupby("A")
-            .expanding()
-            .apply(func, engine="numba", raw=True, kwargs={"a": 1})
-        )
-        expected = roll_frame.groupby("A").expanding().sum() + 1
-        tm.assert_frame_equal(result, expected)
-
-        result = (
-            roll_frame.groupby("A")
-            .expanding()
-            .apply(func, engine="numba", raw=True, args=(1,))
-        )
-        tm.assert_frame_equal(result, expected)
-
 
 class TestEWM:
     @pytest.mark.parametrize(

pandas/tests/window/test_numba.py

@@ -38,6 +38,11 @@ def arithmetic_numba_supported_operators(request):
     return request.param
 
 
+@pytest.fixture
+def roll_frame():
+    return DataFrame({"A": [1] * 20 + [2] * 12 + [3] * 8, "B": np.arange(40)})
+
+
 @td.skip_if_no("numba")
 @pytest.mark.filterwarnings("ignore")
 # Filter warnings when parallel=True and the function can't be parallelized by Numba
@@ -67,6 +72,62 @@ def f(x, *args):
         )
         tm.assert_series_equal(result, expected)
 
+    def test_apply_numba_with_kwargs(self, roll_frame):
+        # GH 58995
+        # rolling apply
+        def func(sr, a=0):
+            return sr.sum() + a
+
+        data = DataFrame(range(10))
+
+        result = data.rolling(5).apply(func, engine="numba", raw=True, kwargs={"a": 1})
+        expected = data.rolling(5).sum() + 1
+        tm.assert_frame_equal(result, expected)
+
+        result = data.rolling(5).apply(func, engine="numba", raw=True, args=(1,))
+        tm.assert_frame_equal(result, expected)
+
+        # expanding apply
+
+        result = data.expanding().apply(func, engine="numba", raw=True, kwargs={"a": 1})
+        expected = data.expanding().sum() + 1
+        tm.assert_frame_equal(result, expected)
+
+        result = data.expanding().apply(func, engine="numba", raw=True, args=(1,))
+        tm.assert_frame_equal(result, expected)
+
+        # groupby rolling
+        result = (
+            roll_frame.groupby("A")
+            .rolling(5)
+            .apply(func, engine="numba", raw=True, kwargs={"a": 1})
+        )
+        expected = roll_frame.groupby("A").rolling(5).sum() + 1
+        tm.assert_frame_equal(result, expected)
+
+        result = (
+            roll_frame.groupby("A")
+            .rolling(5)
+            .apply(func, engine="numba", raw=True, args=(1,))
+        )
+        tm.assert_frame_equal(result, expected)
+        # groupby expanding
+
+        result = (
+            roll_frame.groupby("A")
+            .expanding()
+            .apply(func, engine="numba", raw=True, kwargs={"a": 1})
+        )
+        expected = roll_frame.groupby("A").expanding().sum() + 1
+        tm.assert_frame_equal(result, expected)
+
+        result = (
+            roll_frame.groupby("A")
+            .expanding()
+            .apply(func, engine="numba", raw=True, args=(1,))
+        )
+        tm.assert_frame_equal(result, expected)
+
     def test_numba_min_periods(self):
         # GH 58868
         def last_row(x):