def test_arrange():
df = diamonds.groupby('cut').apply(arrange_apply_helperfunc).reset_index(drop=True)
d = (diamonds >> groupby('cut') >> arrange('depth', ascending=False) >>
head(2) >> ungroup()).reset_index(drop=True)
print(df.head(5))
print(d.head(5))
assert df.equals(d)
def test_row_slice():
df = diamonds.iloc[[0,1],:]
assert df.equals(diamonds >> row_slice([0,1]))
df = diamonds.groupby('cut').apply(lambda df: df.iloc[0,:]).reset_index(drop=True)
d = diamonds >> groupby(X.cut) >> row_slice(0)
assert df.equals(d.reset_index(drop=True))
df = diamonds.loc[diamonds.table > 61, :]
assert df.equals(diamonds >> row_slice(X.table > 61))
def test_grouped_tail():
df = diamonds.groupby(['cut','color']).apply(lambda x: x.tail(2)).reset_index(drop=True)
d = diamonds >> groupby('cut','color') >> tail(2)
assert df.equals(d.reset_index(drop=True))
