Repository: mars-project/mars

Issue:
[BUG] Dataframe groupby.agg raised `IndexError`
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
Executing dataframe groupby raised `IndexError` while the same sentences works well with pandas.

**To Reproduce**
``` Python
In [12]: mdf = md.DataFrame({'A' : ['foo', 'bar', 'foo', 'bar', 
    ...:                           'foo', 'bar'], 
    ...:                    'B' : ['one', 'one', 'two', 'three', 
    ...:                           'two', 'two'], 
    ...:                    'C' : [1, 5, 5, 2, None, 5], 
    ...:                    'D' : [2.0, 5., 8., 1., 2., None]})                                     

In [13]: mdf.groupby(['A'])[['B','C']].agg({'C': 'max'}).execute()                                  
  0%|                                                                       | 0/100 [00:00<?, ?it/s]--------------------------------------------
['C']           B        C
A                   
foo  object  float64
bar  object  float64
--------------------------------------------
Failed to run subtask jSWU1Zizz02M15ykbH3pNl8y on band numa-0
Traceback (most recent call last):
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 316, in internal_run_subtask
    subtask, band_name, subtask_api, batch_quota_req
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 404, in _retry_run_subtask
    return await _retry_run(subtask, subtask_info, _run_subtask_once)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 91, in _retry_run
    raise ex
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 66, in _retry_run
    return await target_async_func(*args)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 357, in _run_subtask_once
    return await asyncio.shield(aiotask)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/api.py", line 67, in run_subtask_in_slot
    subtask
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 186, in send
    return self._process_result_message(result)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 70, in _process_result_message
    raise message.error.with_traceback(message.traceback)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/pool.py", line 520, in send
    result = await self._run_coro(message.message_id, coro)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/pool.py", line 319, in _run_coro
    return await coro
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/api.py", line 115, in __on_receive__
    return await super().__on_receive__(message)
  File "mars/oscar/core.pyx", line 371, in __on_receive__
    debug_logger.exception('Got unhandled error when handling message %r'
  File "mars/oscar/core.pyx", line 365, in mars.oscar.core._BaseActor.__on_receive__
    raise ValueError(f'call_method {call_method} not valid')
  File "mars/oscar/core.pyx", line 250, in _handle_actor_result
    # asyncio.wait as it introduces much overhead
  File "mars/oscar/core.pyx", line 293, in _run_actor_async_generator
    res = None
  File "mars/oscar/core.pyx", line 295, in mars.oscar.core._BaseActor._run_actor_async_generator
    with debug_async_timeout('actor_lock_timeout',
  File "mars/oscar/core.pyx", line 299, in mars.oscar.core._BaseActor._run_actor_async_generator
    res = await gen.asend(res)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/runner.py", line 118, in run_subtask
    result = yield self._running_processor.run(subtask)
  File "mars/oscar/core.pyx", line 304, in mars.oscar.core._BaseActor._run_actor_async_generator
    message_trace = pop_message_trace()
  File "mars/oscar/core.pyx", line 224, in _handle_actor_result
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 186, in send
    return self._process_result_message(result)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 70, in _process_result_message
    raise message.error.with_traceback(message.traceback)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/pool.py", line 520, in send
    result = await self._run_coro(message.message_id, coro)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/pool.py", line 319, in _run_coro
    return await coro
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/api.py", line 115, in __on_receive__
    return await super().__on_receive__(message)
  File "mars/oscar/core.pyx", line 371, in __on_receive__
    debug_logger.exception('Got unhandled error when handling message %r'
  File "mars/oscar/core.pyx", line 365, in mars.oscar.core._BaseActor.__on_receive__
    raise ValueError(f'call_method {call_method} not valid')
  File "mars/oscar/core.pyx", line 250, in _handle_actor_result
    # asyncio.wait as it introduces much overhead
  File "mars/oscar/core.pyx", line 293, in _run_actor_async_generator
    res = None
  File "mars/oscar/core.pyx", line 295, in mars.oscar.core._BaseActor._run_actor_async_generator
    with debug_async_timeout('actor_lock_timeout',
  File "mars/oscar/core.pyx", line 299, in mars.oscar.core._BaseActor._run_actor_async_generator
    res = await gen.asend(res)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 596, in run
    result = yield self._running_aio_task
  File "mars/oscar/core.pyx", line 304, in mars.oscar.core._BaseActor._run_actor_async_generator
    message_trace = pop_message_trace()
  File "mars/oscar/core.pyx", line 224, in _handle_actor_result
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 457, in run
    await self._execute_graph(chunk_graph)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 209, in _execute_graph
    await to_wait
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/lib/aio/_threads.py", line 36, in to_thread
    return await loop.run_in_executor(None, func_call)
  File "/Users/hekaisheng/miniconda3/lib/python3.7/concurrent/futures/thread.py", line 57, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/core/mode.py", line 77, in _inner
    return func(*args, **kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 177, in _execute_operand
    return execute(ctx, op)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/core/operand/core.py", line 487, in execute
    result = executor(results, op)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/core/custom_log.py", line 94, in wrap
    return func(cls, ctx, op)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 1112, in wrapped
    result = func(cls, ctx, op)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/dataframe/groupby/aggregation.py", line 1028, in execute
    cls._execute_map(ctx, op)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/dataframe/groupby/aggregation.py", line 815, in _execute_map
    grouped if cols is None else grouped[cols]
  File "/Users/hekaisheng/miniconda3/lib/python3.7/site-packages/pandas/core/groupby/generic.py", line 1535, in __getitem__
    return super().__getitem__(key)
  File "/Users/hekaisheng/miniconda3/lib/python3.7/site-packages/pandas/core/base.py", line 255, in __getitem__
    raise IndexError(f"Column(s) {self._selection} already selected")
IndexError: Column(s) ['B', 'C'] already selected
```



Discussion:

Resolving patch:
diff --git a/mars/dataframe/groupby/aggregation.py b/mars/dataframe/groupby/aggregation.py
--- a/mars/dataframe/groupby/aggregation.py
+++ b/mars/dataframe/groupby/aggregation.py
@@ -825,9 +825,10 @@ def _execute_map(cls, ctx, op: "DataFrameGroupByAgg"):
 
         for input_key, output_key, cols, func in op.pre_funcs:
             if input_key == output_key:
-                ret_map_groupbys[output_key] = (
-                    grouped if cols is None else grouped[cols]
-                )
+                if cols is None or grouped._selection is not None:
+                    ret_map_groupbys[output_key] = grouped
+                else:
+                    ret_map_groupbys[output_key] = grouped[cols]
             else:
 
                 def _wrapped_func(col):


Test patch:
diff --git a/mars/dataframe/groupby/tests/test_groupby_execution.py b/mars/dataframe/groupby/tests/test_groupby_execution.py
--- a/mars/dataframe/groupby/tests/test_groupby_execution.py
+++ b/mars/dataframe/groupby/tests/test_groupby_execution.py
@@ -273,6 +273,32 @@ def test_groupby_getitem(setup):
         raw.groupby(0, as_index=False)[0].agg({"cnt": "count"}),
     )
 
+    # test groupby getitem then agg(#GH 2640)
+    rs = np.random.RandomState(0)
+    raw = pd.DataFrame(
+        {
+            "c1": np.arange(100).astype(np.int64),
+            "c2": rs.choice(["a", "b", "c"], (100,)),
+            "c3": rs.rand(100),
+            "c4": rs.rand(100),
+        }
+    )
+    mdf = md.DataFrame(raw, chunk_size=20)
+    r = mdf.groupby(["c2"])[["c1", "c3"]].agg({"c1": "max", "c3": "min"}, method="tree")
+    pd.testing.assert_frame_equal(
+        r.execute().fetch(),
+        raw.groupby(["c2"])[["c1", "c3"]].agg({"c1": "max", "c3": "min"}),
+    )
+
+    mdf = md.DataFrame(raw.copy(), chunk_size=30)
+    r = mdf.groupby(["c2"])[["c1", "c4"]].agg(
+        {"c1": "max", "c4": "mean"}, method="shuffle"
+    )
+    pd.testing.assert_frame_equal(
+        r.execute().fetch(),
+        raw.groupby(["c2"])[["c1", "c4"]].agg({"c1": "max", "c4": "mean"}),
+    )
+
 
 def test_dataframe_groupby_agg(setup):
     agg_funs = [