diff --git a/pandas/core/groupby/numba_.py b/pandas/core/groupby/numba_.py index 76f50f1387196..5c983985628ad 100644 --- a/pandas/core/groupby/numba_.py +++ b/pandas/core/groupby/numba_.py @@ -89,10 +89,6 @@ def generate_numba_agg_func( numba_func = jit_user_function(func, nopython, nogil, parallel) numba = import_optional_dependency("numba") - if parallel: - loop_range = numba.prange - else: - loop_range = range @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel) def group_agg( @@ -104,9 +100,9 @@ def group_agg( num_columns: int, ) -> np.ndarray: result = np.empty((num_groups, num_columns)) - for i in loop_range(num_groups): + for i in numba.prange(num_groups): group_index = index[begin[i] : end[i]] - for j in loop_range(num_columns): + for j in numba.prange(num_columns): group = values[begin[i] : end[i], j] result[i, j] = numba_func(group, group_index, *args) return result @@ -153,10 +149,6 @@ def generate_numba_transform_func( numba_func = jit_user_function(func, nopython, nogil, parallel) numba = import_optional_dependency("numba") - if parallel: - loop_range = numba.prange - else: - loop_range = range @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel) def group_transform( @@ -168,9 +160,9 @@ def group_transform( num_columns: int, ) -> np.ndarray: result = np.empty((len(values), num_columns)) - for i in loop_range(num_groups): + for i in numba.prange(num_groups): group_index = index[begin[i] : end[i]] - for j in loop_range(num_columns): + for j in numba.prange(num_columns): group = values[begin[i] : end[i], j] result[begin[i] : end[i], j] = numba_func(group, group_index, *args) return result diff --git a/pandas/core/window/numba_.py b/pandas/core/window/numba_.py index 274586e1745b5..2bc701a16f452 100644 --- a/pandas/core/window/numba_.py +++ b/pandas/core/window/numba_.py @@ -50,17 +50,13 @@ def generate_numba_apply_func( numba_func = jit_user_function(func, nopython, nogil, parallel) numba = import_optional_dependency("numba") - if parallel: - loop_range = numba.prange - else: - loop_range = range @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel) def roll_apply( values: np.ndarray, begin: np.ndarray, end: np.ndarray, minimum_periods: int ) -> np.ndarray: result = np.empty(len(begin)) - for i in loop_range(len(result)): + for i in numba.prange(len(result)): start = begin[i] stop = end[i] window = values[start:stop] @@ -103,10 +99,6 @@ def generate_numba_groupby_ewma_func( return NUMBA_FUNC_CACHE[cache_key] numba = import_optional_dependency("numba") - if parallel: - loop_range = numba.prange - else: - loop_range = range @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel) def groupby_ewma( @@ -117,7 +109,7 @@ def groupby_ewma( ) -> np.ndarray: result = np.empty(len(values)) alpha = 1.0 / (1.0 + com) - for i in loop_range(len(begin)): + for i in numba.prange(len(begin)): start = begin[i] stop = end[i] window = values[start:stop]