Fix IndexError raise by aggregation of DataFrameGroupBy (#2641)

Author: hekaisheng

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):

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 = [