Skip to content

Commit 5bbeb56

Browse files
jrebackjreback
committed
fix up group_count,group_bin_count
1 parent 0ab0389 commit 5bbeb56

File tree

2 files changed

+160
-142
lines changed

2 files changed

+160
-142
lines changed

‎pandas/src/generate_code.py

+58-50
Original file line numberDiff line numberDiff line change
@@ -1262,21 +1262,22 @@ def group_count_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
12621262
raise AssertionError("len(index) != len(labels)")
12631263
12641264
1265-
for i in range(N):
1266-
lab = labels[i]
1267-
if lab < 0:
1268-
continue
1269-
1270-
counts[lab] += 1
1271-
for j in range(K):
1272-
val = values[i, j]
1273-
1274-
# not nan
1275-
nobs[lab, j] += val == val and val != iNaT
1276-
1277-
for i in range(ncounts):
1278-
for j in range(K):
1279-
out[i, j] = nobs[i, j]
1265+
%(nogil)s
1266+
%(tab)sfor i in range(N):
1267+
%(tab)s lab = labels[i]
1268+
%(tab)s if lab < 0:
1269+
%(tab)s continue
1270+
1271+
%(tab)s counts[lab] += 1
1272+
%(tab)s for j in range(K):
1273+
%(tab)s val = values[i, j]
1274+
1275+
%(tab)s # not nan
1276+
%(tab)s nobs[lab, j] += val == val and val != iNaT
1277+
1278+
%(tab)sfor i in range(ncounts):
1279+
%(tab)s for j in range(K):
1280+
%(tab)s out[i, j] = nobs[i, j]
12801281
"""
12811282

12821283
group_count_bin_template = """@cython.wraparound(False)
@@ -1297,20 +1298,21 @@ def group_count_bin_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
12971298
12981299
ngroups = len(bins) + (bins[len(bins) - 1] != N)
12991300
1300-
for i in range(N):
1301-
while b < ngroups - 1 and i >= bins[b]:
1302-
b += 1
1301+
%(nogil)s
1302+
%(tab)sfor i in range(N):
1303+
%(tab)s while b < ngroups - 1 and i >= bins[b]:
1304+
%(tab)s b += 1
13031305
1304-
counts[b] += 1
1305-
for j in range(K):
1306-
val = values[i, j]
1306+
%(tab)s counts[b] += 1
1307+
%(tab)s for j in range(K):
1308+
%(tab)s val = values[i, j]
13071309
1308-
# not nan
1309-
nobs[b, j] += val == val and val != iNaT
1310+
%(tab)s # not nan
1311+
%(tab)s nobs[b, j] += val == val and val != iNaT
13101312
1311-
for i in range(ngroups):
1312-
for j in range(K):
1313-
out[i, j] = nobs[i, j]
1313+
%(tab)sfor i in range(ngroups):
1314+
%(tab)s for j in range(K):
1315+
%(tab)s out[i, j] = nobs[i, j]
13141316
"""
13151317

13161318
# add passing bin edges, instead of labels
@@ -2312,19 +2314,19 @@ def put2d_%(name)s_%(dest_type)s(ndarray[%(c_type)s, ndim=2, cast=True] values,
23122314
def generate_put_template(template, use_ints=True, use_floats=True,
23132315
use_objects=False, use_datelikes=False):
23142316
floats_list = [
2315-
('float64', 'float64_t', 'float64_t', 'np.float64'),
2316-
('float32', 'float32_t', 'float32_t', 'np.float32'),
2317+
('float64', 'float64_t', 'float64_t', 'np.float64', True),
2318+
('float32', 'float32_t', 'float32_t', 'np.float32', True),
23172319
]
23182320
ints_list = [
2319-
('int8', 'int8_t', 'float32_t', 'np.float32'),
2320-
('int16', 'int16_t', 'float32_t', 'np.float32'),
2321-
('int32', 'int32_t', 'float64_t', 'np.float64'),
2322-
('int64', 'int64_t', 'float64_t', 'np.float64'),
2321+
('int8', 'int8_t', 'float32_t', 'np.float32', True),
2322+
('int16', 'int16_t', 'float32_t', 'np.float32', True),
2323+
('int32', 'int32_t', 'float64_t', 'np.float64', True),
2324+
('int64', 'int64_t', 'float64_t', 'np.float64', True),
23232325
]
23242326
date_like_list = [
2325-
('int64', 'int64_t', 'float64_t', 'np.float64'),
2327+
('int64', 'int64_t', 'float64_t', 'np.float64', True),
23262328
]
2327-
object_list = [('object', 'object', 'object', 'np.object_')]
2329+
object_list = [('object', 'object', 'object', 'np.object_', False)]
23282330
function_list = []
23292331
if use_floats:
23302332
function_list.extend(floats_list)
@@ -2336,29 +2338,31 @@ def generate_put_template(template, use_ints=True, use_floats=True,
23362338
function_list.extend(date_like_list)
23372339

23382340
output = StringIO()
2339-
for name, c_type, dest_type, dest_dtype in function_list:
2341+
for name, c_type, dest_type, dest_dtype, nogil in function_list:
23402342
func = template % {'name': name,
23412343
'c_type': c_type,
23422344
'dest_type': dest_type.replace('_t', ''),
23432345
'dest_type2': dest_type,
2344-
'dest_dtype': dest_dtype}
2346+
'dest_dtype': dest_dtype,
2347+
'nogil' : 'with nogil:' if nogil else '',
2348+
'tab' : ' ' if nogil else '' }
23452349
output.write(func)
23462350
output.write("\n")
23472351
return output.getvalue()
23482352

23492353
def generate_put_min_max_template(template, use_ints=True, use_floats=True,
23502354
use_objects=False, use_datelikes=False):
23512355
floats_list = [
2352-
('float64', 'float64_t', 'NAN', 'np.inf'),
2353-
('float32', 'float32_t', 'NAN', 'np.inf'),
2356+
('float64', 'float64_t', 'NAN', 'np.inf', True),
2357+
('float32', 'float32_t', 'NAN', 'np.inf', True),
23542358
]
23552359
ints_list = [
2356-
('int64', 'int64_t', 'iNaT', _int64_max),
2360+
('int64', 'int64_t', 'iNaT', _int64_max, True),
23572361
]
23582362
date_like_list = [
2359-
('int64', 'int64_t', 'iNaT', _int64_max),
2363+
('int64', 'int64_t', 'iNaT', _int64_max, True),
23602364
]
2361-
object_list = [('object', 'object', 'np.nan', 'np.inf')]
2365+
object_list = [('object', 'object', 'np.nan', 'np.inf', False)]
23622366
function_list = []
23632367
if use_floats:
23642368
function_list.extend(floats_list)
@@ -2370,28 +2374,30 @@ def generate_put_min_max_template(template, use_ints=True, use_floats=True,
23702374
function_list.extend(date_like_list)
23712375

23722376
output = StringIO()
2373-
for name, dest_type, nan_val, inf_val in function_list:
2377+
for name, dest_type, nan_val, inf_val, nogil in function_list:
23742378
func = template % {'name': name,
23752379
'dest_type2': dest_type,
23762380
'nan_val': nan_val,
2377-
'inf_val': inf_val}
2381+
'inf_val': inf_val,
2382+
'nogil' : "with nogil:" if nogil else '',
2383+
'tab' : ' ' if nogil else '' }
23782384
output.write(func)
23792385
output.write("\n")
23802386
return output.getvalue()
23812387

23822388
def generate_put_selection_template(template, use_ints=True, use_floats=True,
23832389
use_objects=False, use_datelikes=False):
23842390
floats_list = [
2385-
('float64', 'float64_t', 'float64_t', 'NAN'),
2386-
('float32', 'float32_t', 'float32_t', 'NAN'),
2391+
('float64', 'float64_t', 'float64_t', 'NAN', True),
2392+
('float32', 'float32_t', 'float32_t', 'NAN', True),
23872393
]
23882394
ints_list = [
2389-
('int64', 'int64_t', 'int64_t', 'iNaT'),
2395+
('int64', 'int64_t', 'int64_t', 'iNaT', True),
23902396
]
23912397
date_like_list = [
2392-
('int64', 'int64_t', 'int64_t', 'iNaT'),
2398+
('int64', 'int64_t', 'int64_t', 'iNaT', True),
23932399
]
2394-
object_list = [('object', 'object', 'object', 'np.nan')]
2400+
object_list = [('object', 'object', 'object', 'np.nan', False)]
23952401
function_list = []
23962402
if use_floats:
23972403
function_list.extend(floats_list)
@@ -2403,11 +2409,13 @@ def generate_put_selection_template(template, use_ints=True, use_floats=True,
24032409
function_list.extend(date_like_list)
24042410

24052411
output = StringIO()
2406-
for name, c_type, dest_type, nan_val in function_list:
2412+
for name, c_type, dest_type, nan_val, nogil in function_list:
24072413
func = template % {'name': name,
24082414
'c_type': c_type,
24092415
'dest_type2': dest_type,
2410-
'nan_val': nan_val}
2416+
'nan_val': nan_val,
2417+
'nogil' : "with nogil:" if nogil else "",
2418+
'tab' : ' ' if nogil else '' }
24112419
output.write(func)
24122420
output.write("\n")
24132421
return output.getvalue()

0 commit comments

Comments
 (0)