| 210 | |
| 211 | @pytest.mark.parametrize("jit", [True, False]) |
| 212 | def test_cache_apply(self, jit, nogil, parallel, nopython, step): |
| 213 | # Test that the functions are cached correctly if we switch functions |
| 214 | def func_1(x): |
| 215 | return np.mean(x) + 4 |
| 216 | |
| 217 | def func_2(x): |
| 218 | return np.std(x) * 5 |
| 219 | |
| 220 | if jit: |
| 221 | import numba |
| 222 | |
| 223 | func_1 = numba.jit(func_1) |
| 224 | func_2 = numba.jit(func_2) |
| 225 | |
| 226 | engine_kwargs = {"nogil": nogil, "parallel": parallel, "nopython": nopython} |
| 227 | |
| 228 | roll = Series(range(10)).rolling(2, step=step) |
| 229 | result = roll.apply( |
| 230 | func_1, engine="numba", engine_kwargs=engine_kwargs, raw=True |
| 231 | ) |
| 232 | expected = roll.apply(func_1, engine="cython", raw=True) |
| 233 | tm.assert_series_equal(result, expected) |
| 234 | |
| 235 | result = roll.apply( |
| 236 | func_2, engine="numba", engine_kwargs=engine_kwargs, raw=True |
| 237 | ) |
| 238 | expected = roll.apply(func_2, engine="cython", raw=True) |
| 239 | tm.assert_series_equal(result, expected) |
| 240 | # This run should use the cached func_1 |
| 241 | result = roll.apply( |
| 242 | func_1, engine="numba", engine_kwargs=engine_kwargs, raw=True |
| 243 | ) |
| 244 | expected = roll.apply(func_1, engine="cython", raw=True) |
| 245 | tm.assert_series_equal(result, expected) |
| 246 | |
| 247 | @pytest.mark.parametrize( |
| 248 | "window,window_kwargs", |