From 2e77cef7dadd9607871f486371bb395751e293f3 Mon Sep 17 00:00:00 2001 From: VirosaLi Date: Mon, 8 Jun 2020 10:08:49 -0500 Subject: [PATCH] TST: groupby apply with indexing and colunm aggregation returns the column (#7002) --- pandas/tests/groupby/test_apply.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/pandas/tests/groupby/test_apply.py b/pandas/tests/groupby/test_apply.py index bc8067212d60e..ec454f1b3e93f 100644 --- a/pandas/tests/groupby/test_apply.py +++ b/pandas/tests/groupby/test_apply.py @@ -961,3 +961,13 @@ def fn(x): name="col2", ) tm.assert_series_equal(result, expected) + + +def test_apply_function_with_indexing_return_column(): + # GH: 7002 + df = DataFrame({'foo1': ['one', 'two', 'two', 'three', 'one', 'two'], + 'foo2': np.random.randn(6)}) + result = df.groupby('foo1', as_index=False).apply(lambda x: x.mean()) + expected = df.groupby('foo1', as_index=False).mean() + tm.assert_frame_equal(result, expected) + assert 'foo1' in result.columns