def test_groupby_iterate_groups():
np.random.seed(0)
df = DataFrame()
nelem = 20
df['key1'] = np.random.randint(0, 3, nelem)
df['key2'] = np.random.randint(0, 2, nelem)
df['val1'] = np.random.random(nelem)
df['val2'] = np.random.random(nelem)
def assert_values_equal(arr):
np.testing.assert_array_equal(arr[0], arr)
for grp in df.groupby(['key1', 'key2']):
pddf = grp.to_pandas()
for k in 'key1,key2'.split(','):
assert_values_equal(pddf[k].values)
def test_groupby_cats():
df = DataFrame()
df['cats'] = pd.Categorical(list('aabaacaab'))
df['vals'] = np.random.random(len(df))
cats = np.asarray(list(df['cats']))
vals = df['vals'].to_array()
grouped = df.groupby(['cats']).mean()
got_vals = grouped['vals']
got_cats = grouped['cats']
for c, v in zip(got_cats, got_vals):
print(c, v)
expect = vals[cats == c].mean()
np.testing.assert_almost_equal(v, expect)
def test_groupby_as_df():
np.random.seed(0)
df = DataFrame()
nelem = 20
df['key1'] = np.random.randint(0, 3, nelem)
df['key2'] = np.random.randint(0, 2, nelem)
df['val1'] = np.random.random(nelem)
df['val2'] = np.random.random(nelem)
def assert_values_equal(arr):
np.testing.assert_array_equal(arr[0], arr)
df, segs = df.groupby(['key1', 'key2']).as_df()
for s, e in zip(segs, list(segs[1:]) + [None]):
grp = df[s:e]
pddf = grp.to_pandas()
for k in 'key1,key2'.split(','):
assert_values_equal(pddf[k].values)
def test_groupby_apply():
np.random.seed(0)
df = DataFrame()
nelem = 20
df['key1'] = np.random.randint(0, 3, nelem)
df['key2'] = np.random.randint(0, 2, nelem)
df['val1'] = np.random.random(nelem)
df['val2'] = np.random.random(nelem)
expect_grpby = df.to_pandas().groupby(['key1', 'key2'], as_index=False)
got_grpby = df.groupby(['key1', 'key2'])
def foo(df):
df['out'] = df['val1'] + df['val2']
return df
expect = expect_grpby.apply(foo)
expect = expect.sort_values(['key1', 'key2']).reset_index(drop=True)
got = got_grpby.apply(foo).to_pandas()
pd.util.testing.assert_frame_equal(expect, got)
def test_groupby_apply_grouped():
from numba import cuda
np.random.seed(0)
df = DataFrame()
nelem = 20
df['key1'] = np.random.randint(0, 3, nelem)
df['key2'] = np.random.randint(0, 2, nelem)
df['val1'] = np.random.random(nelem)
df['val2'] = np.random.random(nelem)
expect_grpby = df.to_pandas().groupby(['key1', 'key2'], as_index=False)
got_grpby = df.groupby(['key1', 'key2'])
def foo(key1, val1, com1, com2):
for i in range(cuda.threadIdx.x, len(key1), cuda.blockDim.x):
com1[i] = key1[i] * 10000 + val1[i]
com2[i] = i
got = got_grpby.apply_grouped(foo,
incols=['key1', 'val1'],
outcols={
'com1': np.float64,
'com2': np.int32
},
tpb=8)
got = got.to_pandas()
# Get expected result by emulating the operation in pandas
def emulate(df):
df['com1'] = df.key1 * 10000 + df.val1
df['com2'] = np.arange(len(df), dtype=np.int32)
return df
expect = expect_grpby.apply(emulate)
expect = expect.sort_values(['key1', 'key2']).reset_index(drop=True)
pd.util.testing.assert_frame_equal(expect, got)
