Repository: mars-project/mars

Issue:
Provide tree reduction implementation for DataFrame groupby aggregations
<!--
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.**
Now the groupby aggregations will shuffle the data after the map stage, this implementation works well no matter how large the result is. However, aggregations will produce a 
collapsed data in general, we can provide tree reduction (no need to shuffle) method to achieve  a better performance in that situation.

**Describe the solution you'd like**
Provide a `method` parameter to specify the aggregation method.



Discussion:

Resolving patch:
diff --git a/mars/dataframe/datasource/from_tensor.py b/mars/dataframe/datasource/from_tensor.py
--- a/mars/dataframe/datasource/from_tensor.py
+++ b/mars/dataframe/datasource/from_tensor.py
@@ -105,11 +105,13 @@ def tile(cls, op):
                 i, = in_chunk.index
                 column_stop = 1
                 index = (in_chunk.index[0], 0)
+                dtypes = out_df.dtypes
                 columns_value = parse_index(out_df.columns_value.to_pandas()[0:1], store_data=True)
             else:
                 i, j = in_chunk.index
                 column_stop = cum_size[1][j]
                 index = in_chunk.index
+                dtypes = out_df.dtypes[column_stop - in_chunk.shape[1]:column_stop]
                 columns_value = parse_index(
                     out_df.columns_value.to_pandas()[column_stop - in_chunk.shape[1]:column_stop], store_data=True)
 
@@ -122,7 +124,7 @@ def tile(cls, op):
 
             out_op.extra_params['index_stop'] = index_stop
             out_op.extra_params['column_stop'] = column_stop
-            out_chunk = out_op.new_chunk([in_chunk], shape=in_chunk.shape, index=index,
+            out_chunk = out_op.new_chunk([in_chunk], shape=in_chunk.shape, index=index, dtypes=dtypes,
                                          index_value=index_value, columns_value=columns_value)
             out_chunks.append(out_chunk)
 
@@ -152,7 +154,8 @@ def dataframe_from_tensor(tensor, index=None, columns=None, gpu=None, sparse=Fal
     except IndexError:
         col_num = 1
     gpu = tensor.op.gpu if gpu is None else gpu
-    op = DataFrameFromTensor(dtypes=pd.Series([tensor.dtype] * col_num), gpu=gpu, sparse=sparse)
+    op = DataFrameFromTensor(dtypes=pd.Series([tensor.dtype] * col_num, index=columns),
+                             gpu=gpu, sparse=sparse)
     return op(tensor, index, columns)
 
 
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
@@ -24,10 +24,10 @@ def _install():
         setattr(cls, 'agg', agg)
         setattr(cls, 'aggregate', agg)
 
-        setattr(cls, 'sum', lambda groupby: agg(groupby, 'sum'))
-        setattr(cls, 'prod', lambda groupby: agg(groupby, 'prod'))
-        setattr(cls, 'max', lambda groupby: agg(groupby, 'max'))
-        setattr(cls, 'min', lambda groupby: agg(groupby, 'min'))
+        setattr(cls, 'sum', lambda groupby, **kw: agg(groupby, 'sum', **kw))
+        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))
 
 
 _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
@@ -17,8 +17,10 @@
 import numpy as np
 
 from ... import opcodes as OperandDef
+from ...config import options
 from ...serialize import BoolField, AnyField, StringField
 from ...compat import enum, six
+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
@@ -26,7 +28,7 @@
 
 
 class Stage(enum.Enum):
-    agg = 'agg'
+    map = 'map'
     combine = 'combine'
 
 
@@ -36,12 +38,13 @@ class DataFrameGroupByAgg(DataFrameOperand, DataFrameOperandMixin):
     _func = AnyField('func')
     _by = AnyField('by')
     _as_index = BoolField('as_index')
+    _method = StringField('method')
     _stage = StringField('stage', on_serialize=lambda x: getattr(x, 'value', None),
                          on_deserialize=Stage)
 
-    def __init__(self, func=None, by=None, as_index=None, stage=None, **kw):
-        super(DataFrameGroupByAgg, self).__init__(_func=func, _by=by, _as_index=as_index, _stage=stage,
-                                                  _object_type=ObjectType.dataframe, **kw)
+    def __init__(self, func=None, by=None, as_index=None, method=None, stage=None, **kw):
+        super(DataFrameGroupByAgg, self).__init__(_func=func, _by=by, _as_index=as_index, _method=method,
+                                                  _stage=stage, _object_type=ObjectType.dataframe, **kw)
 
     @property
     def func(self):
@@ -55,6 +58,10 @@ def by(self):
     def as_index(self):
         return self._as_index
 
+    @property
+    def method(self):
+        return self._method
+
     @property
     def stage(self):
         return self._stage
@@ -95,19 +102,24 @@ def _gen_shuffle_chunks(cls, op, in_df, chunks):
         return reduce_chunks
 
     @classmethod
-    def tile(cls, op):
-        in_df = build_concated_rows_frame(op.inputs[0])
-        out_df = op.outputs[0]
-
-        # First, perform groupby and aggregation on each chunk.
+    def _gen_map_chunks(cls, op, in_df, out_df):
         agg_chunks = []
         for chunk in in_df.chunks:
             agg_op = op.copy().reset_key()
-            agg_op._stage = Stage.agg
+            agg_op._stage = Stage.map
             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)
+        return agg_chunks
+
+    @classmethod
+    def _tile_with_shuffle(cls, op):
+        in_df = build_concated_rows_frame(op.inputs[0])
+        out_df = op.outputs[0]
+
+        # First, perform groupby and aggregation on each chunk.
+        agg_chunks = cls._gen_map_chunks(op, in_df, out_df)
 
         # Shuffle the aggregation chunk.
         reduce_chunks = cls._gen_shuffle_chunks(op, in_df, agg_chunks)
@@ -127,17 +139,66 @@ def tile(cls, op):
                                      columns_value=out_df.columns_value, chunks=combine_chunks,
                                      nsplits=((np.nan,) * len(combine_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]
+
+        combine_size = options.combine_size
+        chunks = cls._gen_map_chunks(op, in_df, out_df)
+        while len(chunks) > combine_size:
+            new_chunks = []
+            for idx, i in enumerate(range(0, len(chunks), combine_size)):
+                chks = chunks[i: i + combine_size]
+                if len(chks) == 1:
+                    chk = chks[0]
+                else:
+                    concat_op = DataFrameConcat(object_type=ObjectType.dataframe)
+                    # Change index for concatenate
+                    for j, c in enumerate(chks):
+                        c._index = (j, 0)
+                    chk = concat_op.new_chunk(chks, dtypes=chks[0].dtypes)
+                chunk_op = op.copy().reset_key()
+                chunk_op._stage = Stage.combine
+                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))
+            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 = Stage.combine
+        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()
+        nsplits = ((out_df.shape[0],), (out_df.shape[1],))
+        return new_op.new_tileables(op.inputs, chunks=[chunk], nsplits=nsplits,
+                                    dtypes=out_df.dtypes, shape=out_df.shape,
+                                    index_value=out_df.index_value,
+                                    columns_value=out_df.columns_value)
+
+    @classmethod
+    def tile(cls, op):
+        if op.method == 'shuffle':
+            return cls._tile_with_shuffle(op)
+        elif op.method == 'tree':
+            return cls._tile_with_tree(op)
+        else:
+            raise NotImplementedError
+
     @classmethod
     def execute(cls, ctx, op):
         df = ctx[op.inputs[0].key]
-        if op.stage == Stage.agg:
-            ret = cls._execute_agg(df, op)
+        if op.stage == Stage.map:
+            ret = cls._execute_map(df, op)
         else:
             ret = cls._execute_combine(df, op)
         ctx[op.outputs[0].key] = ret
 
     @classmethod
-    def _execute_agg(cls, df, op):
+    def _execute_map(cls, df, op):
         if isinstance(op.func, (six.string_types, dict)):
             return df.groupby(op.by, as_index=op.as_index).agg(op.func)
         else:
@@ -151,7 +212,16 @@ def _execute_combine(cls, df, op):
             raise NotImplementedError
 
 
-def agg(groupby, func):
+def agg(groupby, func, method='tree'):
+    """
+    Aggregate using one or more operations on grouped data.
+    :param groupby: Groupby data.
+    :param func: Aggregation functions.
+    :param method: 'shuffle' or 'tree', 'tree' method provide a better performance, 'shuffle' is recommended
+    if aggregated result is very large.
+    :return: Aggregated result.
+    """
+
     # When perform a computation on the grouped data, we won't shuffle
     # the data in the stage of groupby and do shuffle after aggregation.
     if not isinstance(groupby, GROUPBY_TYPE):
@@ -159,6 +229,9 @@ def agg(groupby, func):
     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)
+
     if isinstance(func, six.string_types):
         funcs = [func]
     elif isinstance(func, dict):
@@ -170,5 +243,5 @@ def agg(groupby, func):
             raise NotImplementedError('Aggregation function %s has not been supported' % f)
 
     in_df = groupby.inputs[0]
-    agg_op = DataFrameGroupByAgg(func=func, by=groupby.op.by, as_index=groupby.op.as_index)
+    agg_op = DataFrameGroupByAgg(func=func, by=groupby.op.by, method=method, as_index=groupby.op.as_index)
     return agg_op(in_df)


Test patch:
diff --git a/mars/dataframe/datasource/tests/test_datasource.py b/mars/dataframe/datasource/tests/test_datasource.py
--- a/mars/dataframe/datasource/tests/test_datasource.py
+++ b/mars/dataframe/datasource/tests/test_datasource.py
@@ -317,10 +317,15 @@ def testFromTensor(self):
         # from tensor with given columns
         df = dataframe_from_tensor(tensor, columns=list('abcdefghij'))
         df = df.tiles()
+        pd.testing.assert_index_equal(df.dtypes.index, pd.Index(list('abcdefghij')))
         pd.testing.assert_index_equal(df.chunks[0].columns_value.to_pandas(), pd.Index(['a', 'b', 'c', 'd', 'e']))
+        pd.testing.assert_index_equal(df.chunks[0].dtypes.index, pd.Index(['a', 'b', 'c', 'd', 'e']))
         pd.testing.assert_index_equal(df.chunks[1].columns_value.to_pandas(), pd.Index(['f', 'g', 'h', 'i', 'j']))
+        pd.testing.assert_index_equal(df.chunks[1].dtypes.index, pd.Index(['f', 'g', 'h', 'i', 'j']))
         pd.testing.assert_index_equal(df.chunks[2].columns_value.to_pandas(), pd.Index(['a', 'b', 'c', 'd', 'e']))
+        pd.testing.assert_index_equal(df.chunks[2].dtypes.index, pd.Index(['a', 'b', 'c', 'd', 'e']))
         pd.testing.assert_index_equal(df.chunks[3].columns_value.to_pandas(), pd.Index(['f', 'g', 'h', 'i', 'j']))
+        pd.testing.assert_index_equal(df.chunks[3].dtypes.index, pd.Index(['f', 'g', 'h', 'i', 'j']))
 
         # test series from tensor
         tensor = mt.random.rand(10, chunk_size=4)
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
@@ -23,7 +23,6 @@
 from mars.dataframe.groupby.aggregation import DataFrameGroupByAgg
 
 
-
 class Test(TestBase):
     def testGroupBy(self):
         df = pd.DataFrame({'a': [3, 4, 5, 3, 5, 4, 1, 2, 3],
@@ -40,11 +39,24 @@ def testGroupBy(self):
             self.assertIsInstance(chunk.op, DataFrameGroupByOperand)
 
     def testGroupByAgg(self):
+        df = pd.DataFrame({'a': np.random.choice([2, 3, 4], size=(20,)),
+                           'b': np.random.choice([2, 3, 4], size=(20,))})
+        mdf = md.DataFrame(df, chunk_size=3)
+        r = mdf.groupby('a').agg('sum')
+        self.assertIsInstance(r.op, DataFrameGroupByAgg)
+        self.assertIsInstance(r, DataFrame)
+        self.assertEqual(r.op.method, 'tree')
+        r = r.tiles()
+        self.assertEqual(len(r.chunks), 1)
+        self.assertEqual(r.chunks[0].op.stage.value, 'combine')
+        self.assertEqual(len(r.chunks[0].inputs), 1)
+        self.assertEqual(len(r.chunks[0].inputs[0].inputs), 2)
+
         df = pd.DataFrame({'c1': range(10),
                            'c2': np.random.choice(['a', 'b', 'c'], (10,)),
                            'c3': np.random.rand(10)})
         mdf = md.DataFrame(df, chunk_size=2)
-        r = mdf.groupby('c2').sum()
+        r = mdf.groupby('c2').sum(method='shuffle')
 
         self.assertIsInstance(r.op, DataFrameGroupByAgg)
         self.assertIsInstance(r, DataFrame)
@@ -59,4 +71,8 @@ def testGroupByAgg(self):
             self.assertIsInstance(chunk.inputs[0].inputs[0].inputs[0].op, DataFrameGroupByMap)
 
             agg_chunk = chunk.inputs[0].inputs[0].inputs[0].inputs[0]
-            self.assertEqual(agg_chunk.op.stage.value, 'agg')
+            self.assertEqual(agg_chunk.op.stage.value, 'map')
+
+        # test unknown method
+        with self.assertRaises(NotImplementedError):
+            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
@@ -43,8 +43,8 @@ def testGroupBy(self):
             pd.testing.assert_frame_equal(group, expected.get_group(key))
 
     def testGroupByAgg(self):
-        df1 = pd.DataFrame({'a': [3, 4, 5, 3, 5, 4, 1, 2, 3],
-                            'b': [1, 3, 4, 5, 6, 5, 4, 4, 4]})
+        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')
         pd.testing.assert_frame_equal(r1.execute(), df1.groupby('a').agg('sum'))
@@ -77,3 +77,27 @@ def testGroupByAgg(self):
 
         r7 = mdf2.groupby('c2').max()
         pd.testing.assert_frame_equal(r7.execute(), 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(r1.execute(), df1.groupby('a').agg('sum'))
+        r2 = mdf.groupby('b').agg('min', method='shuffle')
+        pd.testing.assert_frame_equal(r2.execute(), 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(r1.execute(), df2.groupby('c2').agg('prod'))
+        r2 = mdf2.groupby('c2').agg('max', method='shuffle')
+        pd.testing.assert_frame_equal(r2.execute(), df2.groupby('c2').agg('max'))
+
+        r3 = mdf2.groupby('c2').agg({'c1': 'min', 'c3': 'sum'}, method='shuffle')
+        pd.testing.assert_frame_equal(r3.execute(), df2.groupby('c2').agg({'c1': 'min', 'c3': 'sum'}))
+
+        r3 = mdf2.groupby('c2').agg({'c1': 'min'}, method='shuffle')
+        pd.testing.assert_frame_equal(r3.execute(), df2.groupby('c2').agg({'c1': 'min'}))