@@ -88,7 +88,8 @@ def test_alignment_error():
88
88
89
89
@pytest .mark .parametrize ("dtype" , (float , int ))
90
90
@pytest .mark .parametrize ("chunk" , [False , True ])
91
- @pytest .mark .parametrize ("expected_groups" , [None , [0 , 1 , 2 ], np .array ([0 , 1 , 2 ])])
91
+ # TODO: make this intp when python 3.8 is dropped
92
+ @pytest .mark .parametrize ("expected_groups" , [None , [0 , 1 , 2 ], np .array ([0 , 1 , 2 ], dtype = np .int64 )])
92
93
@pytest .mark .parametrize (
93
94
"func, array, by, expected" ,
94
95
[
@@ -148,7 +149,12 @@ def test_groupby_reduce(
148
149
)
149
150
# we use pd.Index(expected_groups).to_numpy() which is always int64
150
151
# for the values in this tests
151
- g_dtype = by .dtype if expected_groups is None else np .intp
152
+ if expected_groups is None :
153
+ g_dtype = by .dtype
154
+ elif isinstance (expected_groups , np .ndarray ):
155
+ g_dtype = expected_groups .dtype
156
+ else :
157
+ g_dtype = np .int64
152
158
153
159
assert_equal (groups , np .array ([0 , 1 , 2 ], g_dtype ))
154
160
assert_equal (expected_result , result )
@@ -389,12 +395,12 @@ def test_groupby_agg_dask(func, shape, array_chunks, group_chunks, add_nan, dtyp
389
395
kwargs ["expected_groups" ] = [0 , 2 , 1 ]
390
396
with raise_if_dask_computes ():
391
397
actual , groups = groupby_reduce (array , by , engine = engine , ** kwargs , sort = False )
392
- assert_equal (groups , np .array ([0 , 2 , 1 ], dtype = np .intp ))
398
+ assert_equal (groups , np .array ([0 , 2 , 1 ], dtype = np .int64 ))
393
399
assert_equal (expected , actual [..., [0 , 2 , 1 ]])
394
400
395
401
with raise_if_dask_computes ():
396
402
actual , groups = groupby_reduce (array , by , engine = engine , ** kwargs , sort = True )
397
- assert_equal (groups , np .array ([0 , 1 , 2 ], np .intp ))
403
+ assert_equal (groups , np .array ([0 , 1 , 2 ], np .int64 ))
398
404
assert_equal (expected , actual )
399
405
400
406
0 commit comments