Skip to content

Commit e44744e

Browse files
authored
BUG: handle multi-dimensional grouping better (#36605)
1 parent 86ee235 commit e44744e

File tree

1 file changed

+28
-0
lines changed

1 file changed

+28
-0
lines changed

pandas/tests/groupby/test_grouping.py

+28
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,20 @@ def test_getitem_single_column(self):
124124

125125
tm.assert_series_equal(result, expected)
126126

127+
def test_indices_grouped_by_tuple_with_lambda(self):
128+
# GH 36158
129+
df = DataFrame(
130+
{"Tuples": ((x, y) for x in [0, 1] for y in np.random.randint(3, 5, 5))}
131+
)
132+
133+
gb = df.groupby("Tuples")
134+
gb_lambda = df.groupby(lambda x: df.iloc[x, 0])
135+
136+
expected = gb.indices
137+
result = gb_lambda.indices
138+
139+
tm.assert_dict_equal(result, expected)
140+
127141

128142
# grouping
129143
# --------------------------------
@@ -773,6 +787,20 @@ def test_get_group_grouped_by_tuple(self):
773787
expected = DataFrame({"ids": [(dt[0],), (dt[0],)]}, index=[0, 2])
774788
tm.assert_frame_equal(result, expected)
775789

790+
def test_get_group_grouped_by_tuple_with_lambda(self):
791+
# GH 36158
792+
df = DataFrame(
793+
{"Tuples": ((x, y) for x in [0, 1] for y in np.random.randint(3, 5, 5))}
794+
)
795+
796+
gb = df.groupby("Tuples")
797+
gb_lambda = df.groupby(lambda x: df.iloc[x, 0])
798+
799+
expected = gb.get_group(list(gb.groups.keys())[0])
800+
result = gb_lambda.get_group(list(gb_lambda.groups.keys())[0])
801+
802+
tm.assert_frame_equal(result, expected)
803+
776804
def test_groupby_with_empty(self):
777805
index = pd.DatetimeIndex(())
778806
data = ()

0 commit comments

Comments
 (0)