@@ -89,10 +89,6 @@ def generate_numba_agg_func(
89
89
90
90
numba_func = jit_user_function (func , nopython , nogil , parallel )
91
91
numba = import_optional_dependency ("numba" )
92
- if parallel :
93
- loop_range = numba .prange
94
- else :
95
- loop_range = range
96
92
97
93
@numba .jit (nopython = nopython , nogil = nogil , parallel = parallel )
98
94
def group_agg (
@@ -104,9 +100,9 @@ def group_agg(
104
100
num_columns : int ,
105
101
) -> np .ndarray :
106
102
result = np .empty ((num_groups , num_columns ))
107
- for i in loop_range (num_groups ):
103
+ for i in numba . prange (num_groups ):
108
104
group_index = index [begin [i ] : end [i ]]
109
- for j in loop_range (num_columns ):
105
+ for j in numba . prange (num_columns ):
110
106
group = values [begin [i ] : end [i ], j ]
111
107
result [i , j ] = numba_func (group , group_index , * args )
112
108
return result
@@ -153,10 +149,6 @@ def generate_numba_transform_func(
153
149
154
150
numba_func = jit_user_function (func , nopython , nogil , parallel )
155
151
numba = import_optional_dependency ("numba" )
156
- if parallel :
157
- loop_range = numba .prange
158
- else :
159
- loop_range = range
160
152
161
153
@numba .jit (nopython = nopython , nogil = nogil , parallel = parallel )
162
154
def group_transform (
@@ -168,9 +160,9 @@ def group_transform(
168
160
num_columns : int ,
169
161
) -> np .ndarray :
170
162
result = np .empty ((len (values ), num_columns ))
171
- for i in loop_range (num_groups ):
163
+ for i in numba . prange (num_groups ):
172
164
group_index = index [begin [i ] : end [i ]]
173
- for j in loop_range (num_columns ):
165
+ for j in numba . prange (num_columns ):
174
166
group = values [begin [i ] : end [i ], j ]
175
167
result [begin [i ] : end [i ], j ] = numba_func (group , group_index , * args )
176
168
return result
0 commit comments