Repository: mars-project/mars

Issue:
Support list argument for DataFrameGroupBy.aggregate
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**

Currently, DataFrameGroupBy.aggregate only accept a string or a dict, list is not supported, when dict is passed, the value could not be list as well.

What's more, `sum`, `prod`, `min`, `max` are available for now, more aggregation functions should be implemented as well.


Discussion:

Resolving patch:
diff --git a/mars/dataframe/core.py b/mars/dataframe/core.py
--- a/mars/dataframe/core.py
+++ b/mars/dataframe/core.py
@@ -186,7 +186,7 @@ def to_pandas(self):
             if data is None:
                 return pd.MultiIndex.from_arrays([[], []], sortorder=self._sortorder,
                                                  names=self._names)
-            return pd.MultiIndex.from_tuples(np.asarray(data), sortorder=self._sortorder,
+            return pd.MultiIndex.from_tuples([tuple(d) for d in data], sortorder=self._sortorder,
                                              names=self._names)
 
     _index_value = OneOfField('index_value', index=Index,
diff --git a/mars/dataframe/groupby/__init__.py b/mars/dataframe/groupby/__init__.py
--- a/mars/dataframe/groupby/__init__.py
+++ b/mars/dataframe/groupby/__init__.py
@@ -28,6 +28,10 @@ def _install():
         setattr(cls, 'prod', lambda groupby, **kw: agg(groupby, 'prod', **kw))
         setattr(cls, 'max', lambda groupby, **kw: agg(groupby, 'max', **kw))
         setattr(cls, 'min', lambda groupby, **kw: agg(groupby, 'min', **kw))
+        setattr(cls, 'count', lambda groupby, **kw: agg(groupby, 'count', **kw))
+        setattr(cls, 'mean', lambda groupby, **kw: agg(groupby, 'mean', **kw))
+        setattr(cls, 'var', lambda groupby, **kw: agg(groupby, 'var', **kw))
+        setattr(cls, 'std', lambda groupby, **kw: agg(groupby, 'std', **kw))
 
 
 _install()
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
@@ -13,37 +13,65 @@
 # limitations under the License.
 
 import itertools
+from collections import OrderedDict, namedtuple
+from collections.abc import Iterable
+from functools import partial
 
+import cloudpickle
 import numpy as np
+import pandas as pd
 
 from ... import opcodes as OperandDef
 from ...config import options
 from ...operands import OperandStage
-from ...serialize import BoolField, AnyField, StringField
+from ...serialize import ValueType, BoolField, AnyField, StringField, BytesField, ListField
 from ..merge import DataFrameConcat
 from ..operands import DataFrameOperand, DataFrameOperandMixin, DataFrameShuffleProxy, ObjectType
 from ..core import GROUPBY_TYPE
-from ..utils import build_empty_df, parse_index, build_concated_rows_frame
+from ..utils import build_empty_df, parse_index, build_concated_rows_frame, tokenize
 from .core import DataFrameGroupByOperand
 
 
+_available_aggregation_functions = {'sum', 'prod', 'min', 'max', 'count',
+                                    'mean', 'var', 'std'}
+
+_stage_infos = namedtuple('stage_infos', ('intermediate_cols', 'agg_cols',
+                                          'map_func', 'map_output_column_to_func',
+                                          'combine_func', 'combine_output_column_to_func',
+                                          'agg_func', 'agg_output_column_to_func'))
+
+
 class DataFrameGroupByAgg(DataFrameOperand, DataFrameOperandMixin):
     _op_type_ = OperandDef.GROUPBY_AGG
 
     _func = AnyField('func')
+    _raw_func = AnyField('raw_func')
     _by = AnyField('by')
     _as_index = BoolField('as_index')
     _sort = BoolField('sort')
     _method = StringField('method')
-
-    def __init__(self, func=None, by=None, as_index=None, sort=None, method=None, stage=None, **kw):
+    # for chunk
+    # store the intermediate aggregated columns for the result
+    _agg_columns = ListField('agg_columns', ValueType.string)
+    # store output columns -> function to apply on DataFrameGroupBy
+    _output_column_to_func = BytesField('output_column_to_funcs',
+                                        on_serialize=cloudpickle.dumps,
+                                        on_deserialize=cloudpickle.loads)
+
+    def __init__(self, func=None, by=None, as_index=None, sort=None, method=None,
+                 raw_func=None, agg_columns=None, output_column_to_func=None, stage=None, **kw):
         super().__init__(_func=func, _by=by, _as_index=as_index, _sort=sort, _method=method,
-                         _stage=stage, _object_type=ObjectType.dataframe, **kw)
+                         _agg_columns=agg_columns, _output_column_to_func=output_column_to_func,
+                         _raw_func=raw_func, _stage=stage, _object_type=ObjectType.dataframe, **kw)
 
     @property
     def func(self):
         return self._func
 
+    @property
+    def raw_func(self):
+        return self._raw_func
+
     @property
     def by(self):
         return self._by
@@ -60,19 +88,170 @@ def sort(self):
     def method(self):
         return self._method
 
+    @property
+    def agg_columns(self):
+        return self._agg_columns
+
+    @property
+    def output_column_to_func(self):
+        return self._output_column_to_func
+
     @property
     def stage(self):
         return self._stage
 
+    def _normalize_keyword_aggregations(self, empty_df):
+        raw_func = self._raw_func
+        if isinstance(raw_func, dict):
+            func = OrderedDict()
+            for col, function in raw_func.items():
+                if isinstance(function, Iterable) and not isinstance(function, str):
+                    func[col] = list(function)
+                else:
+                    func[col] = [function]
+            self._func = func
+        else:
+            # force as_index=True
+            agg_df = empty_df.groupby(self.by, as_index=True).aggregate(self.func)
+
+            if not isinstance(agg_df.columns, pd.MultiIndex):
+                # 1 func, the columns of agg_df is the columns to aggregate
+                self._func = OrderedDict((c, [raw_func]) for c in agg_df.columns)
+            else:
+                func = OrderedDict()
+                for c, f in agg_df.columns:
+                    self._safe_append(func, c, f)
+                self._func = func
+
     def __call__(self, df):
         empty_df = build_empty_df(df.dtypes)
-        agg_df = empty_df.groupby(self.by).agg(self.func)
+        agg_df = empty_df.groupby(self.by, as_index=self._as_index).aggregate(self.func)
+
         shape = (np.nan, agg_df.shape[1])
         index_value = parse_index(agg_df.index, df.key)
         index_value.value.should_be_monotonic = True
+
+        # convert func to dict always
+        self._normalize_keyword_aggregations(empty_df)
+
+        # make sure if as_index=False takes effect
+        if not self._as_index:
+            if isinstance(agg_df.index, pd.MultiIndex):
+                # if MultiIndex, as_index=False definitely takes no effect
+                self._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._as_index = True
+
         return self.new_dataframe([df], shape=shape, dtypes=agg_df.dtypes,
                                   index_value=index_value,
-                                  columns_value=parse_index(agg_df.columns, self))
+                                  columns_value=parse_index(agg_df.columns, store_data=True))
+
+    @staticmethod
+    def _safe_append(d, key, val):
+        if key not in d:
+            d[key] = []
+        d[key].append(val)
+
+    @classmethod
+    def _gen_stages_columns_and_funcs(cls, func):
+        intermediate_cols = []
+        intermediate_cols_set = set()
+        agg_cols = []
+        map_func = OrderedDict()
+        map_output_column_to_func = dict()
+        combine_func = OrderedDict()
+        combine_output_column_to_func = dict()
+        agg_func = OrderedDict()
+        agg_output_column_to_func = dict()
+
+        for col, functions in func.items():
+            for f in functions:
+                new_col = tokenize(col, f)
+                agg_cols.append(new_col)
+                if f in {'sum', 'prod', 'min', 'max'}:
+                    if new_col not in intermediate_cols_set:
+                        intermediate_cols.append(new_col)
+                        intermediate_cols_set.add(new_col)
+
+                        # function identical for all stages
+                        cls._safe_append(map_func, col, f)
+                        cls._safe_append(combine_func, new_col, f)
+
+                    cls._safe_append(agg_func, new_col, f)
+                elif f == 'count':
+                    if new_col not in intermediate_cols_set:
+                        intermediate_cols.append(new_col)
+                        intermediate_cols_set.add(new_col)
+
+                        # do count for map
+                        cls._safe_append(map_func, col, f)
+                        # do sum for combine and agg
+                        cls._safe_append(combine_func, new_col, 'sum')
+                    cls._safe_append(agg_func, new_col, 'sum')
+                elif f in {'mean', 'var', 'std'}:
+                    # handle special funcs that have intermediate results
+                    sum_col = tokenize(col, 'sum')
+                    if sum_col not in intermediate_cols_set:
+                        intermediate_cols.append(sum_col)
+                        intermediate_cols_set.add(sum_col)
+
+                        cls._safe_append(map_func, col, 'sum')
+                        cls._safe_append(combine_func, sum_col, 'sum')
+                    count_col = tokenize(col, 'count')
+                    if count_col not in intermediate_cols_set:
+                        intermediate_cols.append(count_col)
+                        intermediate_cols_set.add(count_col)
+
+                        cls._safe_append(map_func, col, 'count')
+                        cls._safe_append(combine_func, count_col, 'sum')
+                    if f == 'mean':
+                        def _mean(df, columns):
+                            return df[columns[0]].sum() / df[columns[1]].sum()
+
+                        cls._safe_append(agg_func, new_col, None)
+                        agg_output_column_to_func[new_col] = \
+                            partial(_mean, columns=[sum_col, count_col])
+                    else:  # var, std
+                        # calculate var for map
+                        var_col = tokenize(col, 'var')
+
+                        def _reduce_var(df, columns):
+                            cnt = df[columns[1]]
+                            reduced_cnt = cnt.sum()
+                            data = df[columns[0]]
+                            var_square = df[columns[2]] * (cnt - 1)
+                            avg = data.sum() / reduced_cnt
+                            avg_diff = data / cnt - avg
+                            var_square = (var_square.sum() + (cnt * avg_diff ** 2).sum())
+                            return var_square / (reduced_cnt - 1)
+
+                        def _reduce_std(df, columns):
+                            return np.sqrt(_reduce_var(df, columns))
+
+                        if var_col not in intermediate_cols_set:
+                            intermediate_cols.append(var_col)
+                            intermediate_cols_set.add(var_col)
+
+                            cls._safe_append(map_func, col, 'var')
+                            cls._safe_append(combine_func, var_col, None)
+                            combine_output_column_to_func[var_col] = \
+                                partial(_reduce_var, columns=[sum_col, count_col, var_col])
+                        cls._safe_append(agg_func, new_col, None)
+                        fun = _reduce_var if f == 'var' else _reduce_std
+                        agg_output_column_to_func[new_col] = \
+                           partial(fun, columns=[sum_col, count_col, var_col])
+                else:  # pragma: no cover
+                    raise NotImplementedError
+
+        return _stage_infos(intermediate_cols=intermediate_cols, agg_cols=agg_cols,
+                            map_func=map_func,
+                            map_output_column_to_func=map_output_column_to_func,
+                            combine_func=combine_func,
+                            combine_output_column_to_func=combine_output_column_to_func,
+                            agg_func=agg_func,
+                            agg_output_column_to_func=agg_output_column_to_func)
 
     @classmethod
     def _gen_shuffle_chunks(cls, op, in_df, chunks):
@@ -80,10 +259,9 @@ def _gen_shuffle_chunks(cls, op, in_df, chunks):
         map_chunks = []
         chunk_shape = (in_df.chunk_shape[0], 1)
         for chunk in chunks:
-            if op.as_index:
-                map_op = DataFrameGroupByOperand(stage=OperandStage.map, shuffle_size=chunk_shape[0])
-            else:
-                map_op = DataFrameGroupByOperand(stage=OperandStage.map, by=op.by, shuffle_size=chunk_shape[0])
+            # no longer consider as_index=False for the intermediate phases,
+            # will do reset_index at last if so
+            map_op = DataFrameGroupByOperand(stage=OperandStage.map, shuffle_size=chunk_shape[0])
             map_chunks.append(map_op.new_chunk([chunk], shape=(np.nan, np.nan), index=chunk.index,
                                                index_value=op.outputs[0].index_value))
 
@@ -101,14 +279,19 @@ def _gen_shuffle_chunks(cls, op, in_df, chunks):
         return reduce_chunks
 
     @classmethod
-    def _gen_map_chunks(cls, op, in_df, out_df):
+    def _gen_map_chunks(cls, op, in_df, out_df, stage_infos: _stage_infos):
         agg_chunks = []
         for chunk in in_df.chunks:
             agg_op = op.copy().reset_key()
+            # force as_index=True for map phase
+            agg_op._as_index = True
             agg_op._stage = OperandStage.map
+            agg_op._func = stage_infos.map_func
+            agg_op._output_column_to_func = stage_infos.map_output_column_to_func
+            columns_value = parse_index(pd.Index(stage_infos.intermediate_cols), store_data=True)
             agg_chunk = agg_op.new_chunk([chunk], shape=out_df.shape, index=chunk.index,
                                          index_value=out_df.index_value,
-                                         columns_value=out_df.columns_value)
+                                         columns_value=columns_value)
             agg_chunks.append(agg_chunk)
         return agg_chunks
 
@@ -117,34 +300,40 @@ def _tile_with_shuffle(cls, op):
         in_df = build_concated_rows_frame(op.inputs[0])
         out_df = op.outputs[0]
 
+        stage_infos = cls._gen_stages_columns_and_funcs(op.func)
+
         # First, perform groupby and aggregation on each chunk.
-        agg_chunks = cls._gen_map_chunks(op, in_df, out_df)
+        agg_chunks = cls._gen_map_chunks(op, in_df, out_df, stage_infos)
 
         # Shuffle the aggregation chunk.
         reduce_chunks = cls._gen_shuffle_chunks(op, in_df, agg_chunks)
 
         # Combine groups
-        combine_chunks = []
+        agg_chunks = []
         for chunk in reduce_chunks:
-            combine_op = op.copy().reset_key()
-            combine_op._stage = OperandStage.combine
-            combine_chunk = combine_op.new_chunk([chunk], shape=out_df.shape, index=chunk.index,
-                                                 index_value=out_df.index_value,
-                                                 columns_value=out_df.columns_value)
-            combine_chunks.append(combine_chunk)
+            agg_op = op.copy().reset_key()
+            agg_op._stage = OperandStage.agg
+            agg_op._func = stage_infos.agg_func
+            agg_op._output_column_to_func = stage_infos.agg_output_column_to_func
+            agg_op._agg_columns = stage_infos.agg_cols
+            agg_chunk = agg_op.new_chunk([chunk], shape=out_df.shape, index=chunk.index,
+                                         index_value=out_df.index_value,
+                                         columns_value=out_df.columns_value)
+            agg_chunks.append(agg_chunk)
 
         new_op = op.copy()
         return new_op.new_dataframes([in_df], shape=out_df.shape, index_value=out_df.index_value,
-                                     columns_value=out_df.columns_value, chunks=combine_chunks,
-                                     nsplits=((np.nan,) * len(combine_chunks), (out_df.shape[1],)))
+                                     columns_value=out_df.columns_value, chunks=agg_chunks,
+                                     nsplits=((np.nan,) * len(agg_chunks), (out_df.shape[1],)))
 
     @classmethod
     def _tile_with_tree(cls, op):
         in_df = build_concated_rows_frame(op.inputs[0])
         out_df = op.outputs[0]
 
+        stage_infos = cls._gen_stages_columns_and_funcs(op.func)
         combine_size = options.combine_size
-        chunks = cls._gen_map_chunks(op, in_df, out_df)
+        chunks = cls._gen_map_chunks(op, in_df, out_df, stage_infos)
         while len(chunks) > combine_size:
             new_chunks = []
             for idx, i in enumerate(range(0, len(chunks), combine_size)):
@@ -159,16 +348,21 @@ def _tile_with_tree(cls, op):
                     chk = concat_op.new_chunk(chks, dtypes=chks[0].dtypes)
                 chunk_op = op.copy().reset_key()
                 chunk_op._stage = OperandStage.combine
+                chunk_op._func = stage_infos.combine_func
+                chunk_op._output_column_to_func = stage_infos.combine_output_column_to_func
+                columns_value = parse_index(pd.Index(stage_infos.intermediate_cols), store_data=True)
                 new_chunks.append(chunk_op.new_chunk([chk], index=(idx, 0), shape=(np.nan, out_df.shape[1]),
                                                      index_value=chks[0].index_value,
-                                                     columns_value=chks[0].columns_value,
-                                                     dtypes=chks[0].dtypes))
+                                                     columns_value=columns_value))
             chunks = new_chunks
 
         concat_op = DataFrameConcat(object_type=ObjectType.dataframe)
         chk = concat_op.new_chunk(chunks, dtypes=chunks[0].dtypes)
         chunk_op = op.copy().reset_key()
-        chunk_op._stage = OperandStage.combine
+        chunk_op._stage = OperandStage.agg
+        chunk_op._func = stage_infos.agg_func
+        chunk_op._output_column_to_func = stage_infos.agg_output_column_to_func
+        chunk_op._agg_columns = stage_infos.agg_cols
         chunk = chunk_op.new_chunk([chk], index=(0, 0), shape=out_df.shape, index_value=out_df.index_value,
                                    columns_value=out_df.columns_value, dtypes=out_df.dtypes)
         new_op = op.copy()
@@ -184,31 +378,114 @@ def tile(cls, op):
             return cls._tile_with_shuffle(op)
         elif op.method == 'tree':
             return cls._tile_with_tree(op)
-        else:
+        else:  # pragma: no cover
             raise NotImplementedError
 
     @classmethod
-    def execute(cls, ctx, op):
-        df = ctx[op.inputs[0].key]
-        if op.stage == OperandStage.map:
-            ret = cls._execute_map(df, op)
+    def _get_grouped(cls, op, df, copy=False):
+        if copy:
+            df = df.copy()
+        if op.stage == OperandStage.agg:
+            return df.groupby(op.by, sort=op.sort)
         else:
-            ret = cls._execute_combine(df, op)
-        ctx[op.outputs[0].key] = ret
+            # for the intermediate phases, do not sort
+            return df.groupby(op.by, sort=False)
 
     @classmethod
-    def _execute_map(cls, df, op):
-        if isinstance(op.func, (str, dict)):
-            return df.groupby(op.by, as_index=op.as_index, sort=False).agg(op.func)
-        else:
-            raise NotImplementedError
+    def _is_raw_one_func(cls, op):
+        raw_func = op.raw_func
+        func = None
+        if isinstance(raw_func, str):
+            func = raw_func
+        elif isinstance(raw_func, list) and len(raw_func) == 1:
+            func = raw_func[0]
+        if func is None:
+            return False
+        return func in {'min', 'max', 'prod', 'sum', 'count'}
 
     @classmethod
-    def _execute_combine(cls, df, op):
-        if isinstance(op.func, (str, dict)):
-            return df.groupby(level=0, as_index=op.as_index, sort=op.sort).agg(op.func)
+    def execute(cls, ctx, op):
+        df = ctx[op.inputs[0].key]
+        out = op.outputs[0]
+
+        grouped = cls._get_grouped(op, df)
+        if op.stage == OperandStage.agg:
+            # use intermediate columns
+            columns = op.agg_columns
         else:
-            raise NotImplementedError
+            columns = op.outputs[0].columns_value.to_pandas().tolist()
+
+        func = OrderedDict()
+        processed_cols = []
+        col_iter = iter(columns)
+        for col, funcs in op.func.items():
+            for f in funcs:
+                c = next(col_iter)
+                if f is not None:
+                    cls._safe_append(func, col, f)
+                    processed_cols.append(c)
+
+        if cls._is_raw_one_func(op):
+            # do some optimization if the raw func is a str or list whose length is 1
+            func = next(iter(func.values()))[0]
+            result = grouped.agg(func)
+            if result.empty and not df.empty:
+                # fix for py35, due to buffer read-only
+                grouped = cls._get_grouped(op, df, copy=True)
+                result = grouped.agg(func)
+        else:
+            # agg the funcs that can be done
+            try:
+                result = grouped.agg(func)
+            except ValueError:  # pragma: no cover
+                # fail due to buffer read-only
+                # force to get grouped again by copy
+                grouped = cls._get_grouped(op, df, copy=True)
+                result = grouped.agg(func)
+        result.columns = processed_cols
+
+        if len(op.output_column_to_func) > 0:
+            # process the functions that require operating on the grouped data
+            for out_col, f in op.output_column_to_func.items():
+                if len(df) > 0:
+                    result[out_col] = grouped.apply(f)
+                else:
+                    result[out_col] = []
+
+            # sort columns as origin
+            result = result[columns]
+
+        if len(result) == 0:
+            # empty data, set index manually
+            result.index = out.index_value.to_pandas()
+
+        if op.stage == OperandStage.agg:
+            if not op.as_index:
+                result.reset_index(inplace=True)
+            result.columns = out.columns_value.to_pandas()
+
+        ctx[out.key] = result
+
+
+def _check_if_func_available(func):
+    to_check = []
+    if isinstance(func, list):
+        to_check.extend(func)
+    elif isinstance(func, dict):
+        for f in func.values():
+            if isinstance(f, Iterable) and not isinstance(f, str):
+                to_check.extend(f)
+            else:
+                to_check.append(f)
+    else:
+        to_check.append(func)
+
+    for f in to_check:
+        if f not in _available_aggregation_functions:
+            raise NotImplementedError(
+                'Aggregation function {} has not been implemented, '
+                'available functions include: {}'.format(
+                    f, _available_aggregation_functions))
 
 
 def agg(groupby, func, method='tree'):
@@ -225,23 +502,14 @@ def agg(groupby, func, method='tree'):
     # the data in the stage of groupby and do shuffle after aggregation.
     if not isinstance(groupby, GROUPBY_TYPE):
         raise TypeError('Input should be type of groupby, not %s' % type(groupby))
-    elif isinstance(func, list):
-        raise NotImplementedError('Function list is not supported now.')
 
     if method not in ['shuffle', 'tree']:
-        raise NotImplementedError('Method %s has not been implemented' % method)
+        raise ValueError("Method %s is not available, "
+                         "please specify 'tree' or 'shuffle" % method)
 
-    if isinstance(func, str):
-        funcs = [func]
-    elif isinstance(func, dict):
-        funcs = func.values()
-    else:
-        raise NotImplementedError('Type %s is not support' % type(func))
-    for f in funcs:
-        if f not in ['sum', 'prod', 'min', 'max']:
-            raise NotImplementedError('Aggregation function %s has not been supported' % f)
+    _check_if_func_available(func)
 
     in_df = groupby.inputs[0]
-    agg_op = DataFrameGroupByAgg(func=func, by=groupby.op.by, method=method,
+    agg_op = DataFrameGroupByAgg(func=func, by=groupby.op.by, method=method, raw_func=func,
                                  as_index=groupby.op.as_index, sort=groupby.op.sort)
     return agg_op(in_df)


Test patch:
diff --git a/mars/dataframe/groupby/tests/test_groupby.py b/mars/dataframe/groupby/tests/test_groupby.py
--- a/mars/dataframe/groupby/tests/test_groupby.py
+++ b/mars/dataframe/groupby/tests/test_groupby.py
@@ -48,7 +48,7 @@ def testGroupByAgg(self):
         self.assertEqual(r.op.method, 'tree')
         r = r.tiles()
         self.assertEqual(len(r.chunks), 1)
-        self.assertEqual(r.chunks[0].op.stage, OperandStage.combine)
+        self.assertEqual(r.chunks[0].op.stage, OperandStage.agg)
         self.assertEqual(len(r.chunks[0].inputs), 1)
         self.assertEqual(len(r.chunks[0].inputs[0].inputs), 2)
 
@@ -65,7 +65,7 @@ def testGroupByAgg(self):
         self.assertEqual(len(r.chunks), 5)
         for chunk in r.chunks:
             self.assertIsInstance(chunk.op, DataFrameGroupByAgg)
-            self.assertEqual(chunk.op.stage, OperandStage.combine)
+            self.assertEqual(chunk.op.stage, OperandStage.agg)
             self.assertIsInstance(chunk.inputs[0].op, DataFrameGroupByOperand)
             self.assertEqual(chunk.inputs[0].op.stage, OperandStage.reduce)
             self.assertIsInstance(chunk.inputs[0].inputs[0].op, DataFrameShuffleProxy)
@@ -76,5 +76,5 @@ def testGroupByAgg(self):
             self.assertEqual(agg_chunk.op.stage, OperandStage.map)
 
         # test unknown method
-        with self.assertRaises(NotImplementedError):
+        with self.assertRaises(ValueError):
             mdf.groupby('c2').sum(method='not_exist')
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
@@ -49,35 +49,61 @@ def testGroupBy(self):
             pd.testing.assert_frame_equal(group, expected.get_group(key))
 
     def testGroupByAgg(self):
-        df1 = pd.DataFrame({'a': np.random.choice([2, 3, 4], size=(100,)),
-                            'b': np.random.choice([2, 3, 4], size=(100,))})
+        rs = np.random.RandomState(0)
+        df1 = pd.DataFrame({'a': rs.choice([2, 3, 4], size=(100,)),
+                            'b': rs.choice([2, 3, 4], size=(100,))})
         mdf = md.DataFrame(df1, chunk_size=3)
-        r1 = mdf.groupby('a').agg('sum')
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r1, concat=True)[0],
-                                      df1.groupby('a').agg('sum'))
-        r2 = mdf.groupby('b').agg('min')
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r2, concat=True)[0],
-                                      df1.groupby('b').agg('min'))
 
         df2 = pd.DataFrame({'c1': range(10),
-                            'c2': np.random.choice(['a', 'b', 'c'], (10,)),
-                            'c3': np.random.rand(10)})
+                            'c2': rs.choice(['a', 'b', 'c'], (10,)),
+                            'c3': rs.rand(10)})
         mdf2 = md.DataFrame(df2, chunk_size=2)
-        r1 = mdf2.groupby('c2').agg('prod')
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r1, concat=True)[0],
-                                      df2.groupby('c2').agg('prod'))
-        r2 = mdf2.groupby('c2').agg('max')
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r2, concat=True)[0],
-                                      df2.groupby('c2').agg('max'))
-
-        agg = OrderedDict([('c1', 'min'), ('c3', 'sum')])
-        r3 = mdf2.groupby('c2').agg(agg)
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r3, concat=True)[0],
-                                      df2.groupby('c2').agg(agg))
-
-        r3 = mdf2.groupby('c2').agg({'c1': 'min'})
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r3, concat=True)[0],
-                                      df2.groupby('c2').agg({'c1': 'min'}))
+
+        for method in ['tree', 'shuffle']:
+            r1 = mdf.groupby('a').agg('sum', method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r1, concat=True)[0],
+                                          df1.groupby('a').agg('sum'))
+            r2 = mdf.groupby('b').agg('min', method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r2, concat=True)[0],
+                                          df1.groupby('b').agg('min'))
+
+            r1 = mdf2.groupby('c2').agg('prod', method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r1, concat=True)[0],
+                                          df2.groupby('c2').agg('prod'))
+            r2 = mdf2.groupby('c2').agg('max', method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r2, concat=True)[0],
+                                          df2.groupby('c2').agg('max'))
+            r3 = mdf2.groupby('c2').agg('count', method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r3, concat=True)[0],
+                                          df2.groupby('c2').agg('count'))
+            r4 = mdf2.groupby('c2').agg('mean', method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r4, concat=True)[0],
+                                          df2.groupby('c2').agg('mean'))
+            r5 = mdf2.groupby('c2').agg('var', method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r5, concat=True)[0],
+                                          df2.groupby('c2').agg('var'))
+            r6 = mdf2.groupby('c2').agg('std', method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r6, concat=True)[0],
+                                          df2.groupby('c2').agg('std'))
+
+            agg = ['std', 'mean', 'var', 'max', 'count']
+            r3 = mdf2.groupby('c2').agg(agg, method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r3, concat=True)[0],
+                                          df2.groupby('c2').agg(agg))
+
+            agg = OrderedDict([('c1', ['min', 'mean']), ('c3', 'std')])
+            r3 = mdf2.groupby('c2').agg(agg, method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r3, concat=True)[0],
+                                          df2.groupby('c2').agg(agg))
+
+            agg = OrderedDict([('c1', 'min'), ('c3', 'sum')])
+            r3 = mdf2.groupby('c2').agg(agg, method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r3, concat=True)[0],
+                                          df2.groupby('c2').agg(agg))
+
+            r3 = mdf2.groupby('c2').agg({'c1': 'min'}, method=method)
+            pd.testing.assert_frame_equal(self.executor.execute_dataframe(r3, concat=True)[0],
+                                          df2.groupby('c2').agg({'c1': 'min'}))
 
         r4 = mdf2.groupby('c2').sum()
         pd.testing.assert_frame_equal(self.executor.execute_dataframe(r4, concat=True)[0],
@@ -95,33 +121,30 @@ def testGroupByAgg(self):
         pd.testing.assert_frame_equal(self.executor.execute_dataframe(r7, concat=True)[0],
                                       df2.groupby('c2').max())
 
-        # test shuffle method
-        df1 = pd.DataFrame({'a': np.random.choice([2, 3, 4], size=(100,)),
-                            'b': np.random.choice([2, 3, 4], size=(100,))})
-        mdf = md.DataFrame(df1, chunk_size=3)
-        r1 = mdf.groupby('a').agg('sum', method='shuffle')
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r1, concat=True)[0],
-                                      df1.groupby('a').agg('sum'))
-        r2 = mdf.groupby('b').agg('min', method='shuffle')
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r2, concat=True)[0],
-                                      df1.groupby('b').agg('min'))
-
-        df2 = pd.DataFrame({'c1': range(10),
-                            'c2': np.random.choice(['a', 'b', 'c'], (10,)),
-                            'c3': np.random.rand(10)})
-        mdf2 = md.DataFrame(df2, chunk_size=2)
-        r1 = mdf2.groupby('c2').agg('prod', method='shuffle')
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r1, concat=True)[0],
-                                      df2.groupby('c2').agg('prod'))
-        r2 = mdf2.groupby('c2').agg('max', method='shuffle')
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r2, concat=True)[0],
-                                      df2.groupby('c2').agg('max'))
-
-        agg = OrderedDict([('c1', 'min'), ('c3', 'sum')])
-        r3 = mdf2.groupby('c2').agg(agg, method='shuffle')
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r3, concat=True)[0],
-                                      df2.groupby('c2').agg(agg))
-
-        r3 = mdf2.groupby('c2').agg({'c1': 'min'}, method='shuffle')
-        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r3, concat=True)[0],
-                                      df2.groupby('c2').agg({'c1': 'min'}))
+        r8 = mdf2.groupby('c2').count()
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r8, concat=True)[0],
+                                      df2.groupby('c2').count())
+
+        r9 = mdf2.groupby('c2').mean()
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r9, concat=True)[0],
+                                      df2.groupby('c2').mean())
+
+        r10 = mdf2.groupby('c2').var()
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r10, concat=True)[0],
+                                      df2.groupby('c2').var())
+
+        r11 = mdf2.groupby('c2').std()
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r11, concat=True)[0],
+                                      df2.groupby('c2').std())
+
+        # test as_index=False
+        r12 = mdf2.groupby('c2', as_index=False).agg('mean')
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r12, concat=True)[0],
+                                      df2.groupby('c2', as_index=False).agg('mean'))
+        self.assertFalse(r12.op.as_index)
+
+        # test as_index=False takes no effect
+        r13 = mdf2.groupby(['c1', 'c2'], as_index=False).agg(['mean', 'count'])
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r13, concat=True)[0],
+                                      df2.groupby(['c1', 'c2'], as_index=False).agg(['mean', 'count']))
+        self.assertTrue(r13.op.as_index)
diff --git a/mars/deploy/kubernetes/tests/test_kubernetes.py b/mars/deploy/kubernetes/tests/test_kubernetes.py
--- a/mars/deploy/kubernetes/tests/test_kubernetes.py
+++ b/mars/deploy/kubernetes/tests/test_kubernetes.py
@@ -102,7 +102,7 @@ def testRunInKubernetes(self):
                                   worker_spill_paths=[temp_spill_dir],
                                   extra_volumes=[extra_vol_config],
                                   pre_stop_command=['rm', '/tmp/stopping.tmp'],
-                                  timeout=120, log_when_fail=True)
+                                  timeout=600, log_when_fail=True)
             self.assertIsNotNone(cluster.endpoint)
 
             pod_items = kube_api.list_namespaced_pod(cluster.namespace).to_dict()
@@ -115,7 +115,7 @@ def testRunInKubernetes(self):
             a = mt.ones((100, 100), chunk_size=30) * 2 * 1 + 1
             b = mt.ones((100, 100), chunk_size=30) * 2 * 1 + 1
             c = (a * b * 2 + 1).sum()
-            r = cluster.session.run(c, timeout=120)
+            r = cluster.session.run(c, timeout=600)
 
             expected = (np.ones(a.shape) * 2 * 1 + 1) ** 2 * 2 + 1
             assert_array_equal(r, expected.sum())