| 60 | class TestEngine: |
| 61 | @pytest.mark.parametrize("jit", [True, False]) |
| 62 | def test_numba_vs_cython_apply(self, jit, nogil, parallel, nopython, center, step): |
| 63 | def f(x, *args): |
| 64 | arg_sum = 0 |
| 65 | for arg in args: |
| 66 | arg_sum += arg |
| 67 | return np.mean(x) + arg_sum |
| 68 | |
| 69 | if jit: |
| 70 | import numba |
| 71 | |
| 72 | f = numba.jit(f) |
| 73 | |
| 74 | engine_kwargs = {"nogil": nogil, "parallel": parallel, "nopython": nopython} |
| 75 | args = (2,) |
| 76 | |
| 77 | s = Series(range(10)) |
| 78 | result = s.rolling(2, center=center, step=step).apply( |
| 79 | f, args=args, engine="numba", engine_kwargs=engine_kwargs, raw=True |
| 80 | ) |
| 81 | expected = s.rolling(2, center=center, step=step).apply( |
| 82 | f, engine="cython", args=args, raw=True |
| 83 | ) |
| 84 | tm.assert_series_equal(result, expected) |
| 85 | |
| 86 | def test_apply_numba_with_kwargs(self, roll_frame): |
| 87 | # GH 58995 |