diff --git a/pandas/tests/groupby/test_grouping.py b/pandas/tests/groupby/test_grouping.py index 6c74e1521eeeb..39edc2b8d2527 100644 --- a/pandas/tests/groupby/test_grouping.py +++ b/pandas/tests/groupby/test_grouping.py @@ -124,6 +124,20 @@ def test_getitem_single_column(self): tm.assert_series_equal(result, expected) + def test_indices_grouped_by_tuple_with_lambda(self): + # GH 36158 + df = DataFrame( + {"Tuples": ((x, y) for x in [0, 1] for y in np.random.randint(3, 5, 5))} + ) + + gb = df.groupby("Tuples") + gb_lambda = df.groupby(lambda x: df.iloc[x, 0]) + + expected = gb.indices + result = gb_lambda.indices + + tm.assert_dict_equal(result, expected) + # grouping # -------------------------------- @@ -775,6 +789,20 @@ def test_get_group_grouped_by_tuple(self): expected = DataFrame({"ids": [(dt[0],), (dt[0],)]}, index=[0, 2]) tm.assert_frame_equal(result, expected) + def test_get_group_grouped_by_tuple_with_lambda(self): + # GH 36158 + df = DataFrame( + {"Tuples": ((x, y) for x in [0, 1] for y in np.random.randint(3, 5, 5))} + ) + + gb = df.groupby("Tuples") + gb_lambda = df.groupby(lambda x: df.iloc[x, 0]) + + expected = gb.get_group(list(gb.groups.keys())[0]) + result = gb_lambda.get_group(list(gb_lambda.groups.keys())[0]) + + tm.assert_frame_equal(result, expected) + def test_groupby_with_empty(self): index = pd.DatetimeIndex(()) data = ()