Repository: mars-project/mars

Issue:
[BUG] Wrong index produced when grouping by multiple columns with `as_index=False`
**Describe the bug**
Wrong index produced when grouping by multiple columns with `as_index=False`.

**To Reproduce**

```
id1,id2,id3,id4,id5,id6,v1,v2,v3
id016,id016,id0000042202,15,24,5971,5,11,37.211254
id039,id045,id0000029558,40,49,39457,5,4,48.951141
id047,id023,id0000071286,68,20,74463,2,14,60.469241
id043,id057,id0000015141,32,43,63743,1,15,7.692145
id054,id040,id0000011083,9,25,16920,2,9,22.863525
id029,id020,id0000017974,40,43,14435,3,13,87.521355
id047,id023,id0000084849,90,96,35790,2,9,93.348148
id091,id022,id0000031441,50,44,71525,3,11,81.013682
id090,id048,id0000067778,24,2,51862,4,9,30.718821
id070,id008,id0000091167,78,4,23333,5,15,70.95464
id039,id084,id0000013708,94,81,44406,1,3,54.368009
id023,id061,id0000011331,36,67,86498,5,2,13.847979
id070,id054,id0000099110,24,15,47054,4,2,92.057305
id022,id008,id0000038862,38,92,63088,3,10,33.517765
id020,id070,id0000028952,17,57,50831,4,15,48.060814
id078,id022,id0000082008,69,44,15891,1,4,95.75571
```

Code
```
md_x = md.read_csv('testdata.csv', dtype={'id1': 'category', 'id2': 'category', 'id3': 'category', 'id4': 'Int32', 'id5': 'Int32', 'id6': 'Int32', 'v1': 'Int32', 'v2': 'Int32', 'v3': 'float64'})
md_x.groupby(['id1','id2'], as_index=False, sort=False).agg({'v1':'sum'}).execute()
```

Wrong result:
```
      id1  id2  v1
0   id016  NaN   5
1   id039  NaN   6
2   id047  NaN   4
3   id043  NaN   1
4   id054  NaN   2
5   id029  NaN   3
6   id091  NaN   3
7   id090  NaN   4
8   id070  NaN   9
9   id023  NaN   5
10  id022  NaN   3
11  id020  NaN   4
12  id078  NaN   1
```

Expected result:
```
      id1    id2  v1
0   id016  id016   5
1   id039  id045   5
2   id047  id023   4
3   id043  id057   1
4   id054  id040   2
5   id029  id020   3
6   id091  id022   3
7   id090  id048   4
8   id070  id008   5
9   id039  id084   1
10  id023  id061   5
11  id070  id054   4
12  id022  id008   3
13  id020  id070   4
14  id078  id022   1
```


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
@@ -98,15 +98,17 @@ class DataFrameGroupByAgg(DataFrameOperand, DataFrameOperandMixin):
     _pre_funcs = ListField('pre_funcs')
     _agg_funcs = ListField('agg_funcs')
     _post_funcs = ListField('post_funcs')
+    _index_levels = Int32Field('index_levels')
 
     def __init__(self, raw_func=None, raw_func_kw=None, func=None, func_rename=None,
                  method=None, groupby_params=None, use_inf_as_na=None, combine_size=None,
-                 pre_funcs=None, agg_funcs=None, post_funcs=None, output_types=None, **kw):
+                 pre_funcs=None, agg_funcs=None, post_funcs=None, index_levels=None,
+                 output_types=None, **kw):
         super().__init__(_raw_func=raw_func, _raw_func_kw=raw_func_kw, _func=func,
                          _func_rename=func_rename, _method=method, _groupby_params=groupby_params,
                          _combine_size=combine_size, _use_inf_as_na=use_inf_as_na,
                          _pre_funcs=pre_funcs, _agg_funcs=agg_funcs, _post_funcs=post_funcs,
-                         _output_types=output_types, **kw)
+                         _index_levels=index_levels, _output_types=output_types, **kw)
 
     @property
     def raw_func(self):
@@ -152,6 +154,10 @@ def agg_funcs(self) -> List[ReductionAggStep]:
     def post_funcs(self) -> List[ReductionPostStep]:
         return self._post_funcs
 
+    @property
+    def index_levels(self) -> int:
+        return self._index_levels
+
     def _set_inputs(self, inputs):
         super()._set_inputs(inputs)
         inputs_iter = iter(self._inputs[1:])
@@ -180,14 +186,22 @@ def _call_dataframe(self, groupby, input_df):
 
         as_index = self.groupby_params.get('as_index')
         # make sure if as_index=False takes effect
+        if isinstance(agg_df.index, pd.MultiIndex):
+            # if MultiIndex, as_index=False definitely takes no effect
+            self.groupby_params['as_index'] = as_index = True
+        elif agg_df.index.name is not None:
+            # if not MultiIndex and agg_df.index has a name
+            # means as_index=False takes no effect
+            self.groupby_params['as_index'] = as_index = True
+
+        # determine num of indices to group in intermediate steps
         if not as_index:
-            if isinstance(agg_df.index, pd.MultiIndex):
-                # if MultiIndex, as_index=False definitely takes no effect
-                self.groupby_params['as_index'] = True
-            elif agg_df.index.name is not None:
-                # if not MultiIndex and agg_df.index has a name
-                # means as_index=False takes no effect
-                self.groupby_params['as_index'] = True
+            as_index_agg_df = groupby.op.build_mock_groupby(as_index=True) \
+                .aggregate(self.raw_func, **self.raw_func_kw)
+            pd_index = as_index_agg_df.index
+        else:
+            pd_index = agg_df.index
+        self._index_levels = 1 if not isinstance(pd_index, pd.MultiIndex) else len(pd_index.levels)
 
         inputs = self._get_inputs([input_df])
         return self.new_dataframe(inputs, shape=shape, dtypes=agg_df.dtypes,
@@ -201,6 +215,11 @@ def _call_series(self, groupby, in_series):
         index_value.value.should_be_monotonic = True
 
         inputs = self._get_inputs([in_series])
+
+        # determine num of indices to group in intermediate steps
+        pd_index = agg_result.index
+        self._index_levels = 1 if not isinstance(pd_index, pd.MultiIndex) else len(pd_index.levels)
+
         # update value type
         if isinstance(agg_result, pd.DataFrame):
             return self.new_dataframe(inputs, shape=(np.nan, len(agg_result.columns)),
@@ -316,9 +335,6 @@ def _tile_with_shuffle(cls, op: "DataFrameGroupByAgg"):
             in_df = build_concatenated_rows_frame(in_df)
         out_df = op.outputs[0]
 
-        index = out_df.index_value.to_pandas()
-        level = 0 if not isinstance(index, pd.MultiIndex) else list(range(len(index.levels)))
-
         func_infos = cls._compile_funcs(op, in_df)
 
         # First, perform groupby and aggregation on each chunk.
@@ -336,7 +352,7 @@ def _tile_with_shuffle(cls, op: "DataFrameGroupByAgg"):
             agg_op._groupby_params.pop('selection', None)
             # use levels instead of by for reducer
             agg_op._groupby_params.pop('by', None)
-            agg_op._groupby_params['level'] = level
+            agg_op._groupby_params['level'] = list(range(op.index_levels))
             agg_op.stage = OperandStage.agg
             agg_op._agg_funcs = func_infos.agg_funcs
             agg_op._post_funcs = func_infos.post_funcs
@@ -366,9 +382,6 @@ def _tile_with_tree(cls, op: "DataFrameGroupByAgg"):
             in_df = build_concatenated_rows_frame(in_df)
         out_df = op.outputs[0]
 
-        index = out_df.index_value.to_pandas()
-        level = 0 if not isinstance(index, pd.MultiIndex) else list(range(len(index.levels)))
-
         func_infos = cls._compile_funcs(op, in_df)
         combine_size = op.combine_size
         chunks = cls._gen_map_chunks(op, in_df, out_df, func_infos)
@@ -392,7 +405,7 @@ def _tile_with_tree(cls, op: "DataFrameGroupByAgg"):
                 chunk_op._groupby_params.pop('selection', None)
                 # use levels instead of by for agg
                 chunk_op._groupby_params.pop('by', None)
-                chunk_op._groupby_params['level'] = level
+                chunk_op._groupby_params['level'] = list(range(op.index_levels))
                 chunk_op._agg_funcs = func_infos.agg_funcs
 
                 new_shape = (np.nan, out_df.shape[1]) if len(out_df.shape) == 2 else (np.nan,)
@@ -411,7 +424,7 @@ def _tile_with_tree(cls, op: "DataFrameGroupByAgg"):
         chunk_op._groupby_params.pop('selection', None)
         # use levels instead of by for agg
         chunk_op._groupby_params.pop('by', None)
-        chunk_op._groupby_params['level'] = level
+        chunk_op._groupby_params['level'] = list(range(op.index_levels))
         chunk_op._agg_funcs = func_infos.agg_funcs
         chunk_op._post_funcs = func_infos.post_funcs
         kw = out_df.params.copy()
diff --git a/mars/dataframe/reduction/core.py b/mars/dataframe/reduction/core.py
--- a/mars/dataframe/reduction/core.py
+++ b/mars/dataframe/reduction/core.py
@@ -595,6 +595,17 @@ def _check_function_valid(cls, func):
                 raise ValueError(f'Variable {var_name} used by {func.__name__} '
                                  'cannot be a Mars object')
 
+    @staticmethod
+    def _update_col_dict(col_dict: Dict, key: str, cols: List):
+        if key in col_dict:
+            existing_cols = col_dict[key]
+            if existing_cols is not None:
+                existing_col_set = set(existing_cols)
+                col_dict[key].extend(
+                    [c for c in cols if c not in existing_col_set])
+        else:
+            col_dict[key] = list(cols) if cols is not None else None
+
     def add_function(self, func, ndim, cols=None, func_name=None):
         cols = cols if cols is not None and self._axis == 0 else None
 
@@ -611,21 +622,14 @@ def add_function(self, func, ndim, cols=None, func_name=None):
 
         for step in compile_result.pre_funcs:
             self._output_key_to_pre_steps[step.output_key] = step
-            if step.output_key in self._output_key_to_pre_cols:
-                existing_cols = self._output_key_to_pre_cols[step.output_key]
-                if existing_cols is not None:
-                    existing_col_set = set(existing_cols)
-                    self._output_key_to_pre_cols[step.output_key].extend(
-                        [c for c in cols if c not in existing_col_set])
-            else:
-                self._output_key_to_pre_cols[step.output_key] = list(cols) if cols is not None else None
+            self._update_col_dict(self._output_key_to_pre_cols, step.output_key, cols)
 
         for step in compile_result.agg_funcs:
             self._output_key_to_agg_steps[step.output_key] = step
 
         for step in compile_result.post_funcs:
             self._output_key_to_post_steps[step.output_key] = step
-            self._output_key_to_post_cols[step.output_key] = cols
+            self._update_col_dict(self._output_key_to_post_cols, step.output_key, cols)
 
     @functools.lru_cache(100)
     def _compile_expr_function(self, py_src):


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
@@ -250,6 +250,10 @@ def test_dataframe_groupby_agg(setup):
         pd.testing.assert_frame_equal(r.execute().fetch().sort_index(),
                                       raw.groupby('c2').agg(agg).sort_index())
 
+        r = mdf.groupby('c2').agg({'c1': 'min', 'c3': 'min'}, method=method)
+        pd.testing.assert_frame_equal(r.execute().fetch().sort_index(),
+                                      raw.groupby('c2').agg({'c1': 'min', 'c3': 'min'}).sort_index())
+
         r = mdf.groupby('c2').agg({'c1': 'min'}, method=method)
         pd.testing.assert_frame_equal(r.execute().fetch().sort_index(),
                                       raw.groupby('c2').agg({'c1': 'min'}).sort_index())
@@ -275,14 +279,20 @@ def test_dataframe_groupby_agg(setup):
         pd.testing.assert_frame_equal(r.execute().fetch(),
                                       getattr(raw.groupby('c2'), agg_fun)())
 
+    # test as_index=False
     for method in ['tree', 'shuffle']:
-        # test as_index=False
         r = mdf.groupby('c2', as_index=False).agg('mean', method=method)
         pd.testing.assert_frame_equal(
             r.execute().fetch().sort_values('c2', ignore_index=True),
             raw.groupby('c2', as_index=False).agg('mean').sort_values('c2', ignore_index=True))
         assert r.op.groupby_params['as_index'] is False
 
+        r = mdf.groupby(['c1', 'c2'], as_index=False).agg('mean', method=method)
+        pd.testing.assert_frame_equal(
+            r.execute().fetch().sort_values(['c1', 'c2'], ignore_index=True),
+            raw.groupby(['c1', 'c2'], as_index=False).agg('mean').sort_values(['c1', 'c2'], ignore_index=True))
+        assert r.op.groupby_params['as_index'] is False
+
     # test as_index=False takes no effect
     r = mdf.groupby(['c1', 'c2'], as_index=False).agg(['mean', 'count'], method='tree')
     pd.testing.assert_frame_equal(r.execute().fetch(),
diff --git a/mars/dataframe/reduction/tests/test_reduction.py b/mars/dataframe/reduction/tests/test_reduction.py
--- a/mars/dataframe/reduction/tests/test_reduction.py
+++ b/mars/dataframe/reduction/tests/test_reduction.py
@@ -480,7 +480,24 @@ def test_compile_function():
     # check post_funcs
     assert len(result.post_funcs) == 2
     assert set(''.join(sorted(result.post_funcs[i].columns)) for i in range(2)) == {'ab', 'bc'}
-    
+
+    # test agg for multiple columns
+    compiler = ReductionCompiler(store_source=True)
+    compiler.add_function(lambda x: x.sum(), ndim=2, cols=['a'])
+    compiler.add_function(lambda x: x.sum(), ndim=2, cols=['b'])
+    compiler.add_function(lambda x: x.min(), ndim=2, cols=['c'])
+    result = compiler.compile()
+    # check pre_funcs
+    assert len(result.pre_funcs) == 1
+    assert set(result.pre_funcs[0].columns) == set('abc')
+    # check agg_funcs
+    assert len(result.agg_funcs) == 2
+    assert result.agg_funcs[0].map_func_name == 'sum'
+    assert result.agg_funcs[0].agg_func_name == 'sum'
+    # check post_funcs
+    assert len(result.post_funcs) == 2
+    assert set(result.post_funcs[0].columns) == set('ab')
+
     
 def test_custom_aggregation():
     class MockReduction1(CustomReduction):