Repository: mars-project/mars

Issue:
DataFrameGroupBy missing size() method.
<!--
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.**

I have processed a amount of data which read from 'md.read_csv()' method and got the DataFrameGroupBy object by invoking `groupby()` method. While it could not chain the `size()` method as `Pandas` do.

I wish repo-maintainers could add more feature on Mars to keep API be consistent with Pandas.

**Describe the solution you'd like**
A clear and concise description of what you want to happen.

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**

Mars should consider that support more sql-like functions on DataFrame in order to migrate SQL users could make use of Mars to empower their workflow.

https://medium.com/jbennetcodes/how-to-rewrite-your-sql-queries-in-pandas-and-more-149d341fc53e



Discussion:
We will add this method in the next release. btw, you can use ``df.groupby('key').apply(lambda df: len(df))`` to compute the result of ``df.groupby('key').size()``.

On the question of SQL-like operands, we will add them in next releases, but it takes time. Therefore if you need some operands please let us know, and we may implement them first.

Resolving patch:
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
@@ -40,6 +40,7 @@ def _install():
         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, 'size', lambda groupby, **kw: agg(groupby, 'size', **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))
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
@@ -34,7 +34,7 @@
 from .core import DataFrameGroupByOperand
 
 
-_available_aggregation_functions = {'sum', 'prod', 'min', 'max', 'count',
+_available_aggregation_functions = {'sum', 'prod', 'min', 'max', 'count', 'size',
                                     'mean', 'var', 'std'}
 
 _stage_infos = namedtuple('stage_infos', ('intermediate_cols', 'agg_cols',
@@ -188,15 +188,18 @@ def _call_series(self, groupby, in_series):
         else:
             self._object_type = ObjectType.series
             return self.new_series(inputs, shape=(np.nan,), dtype=agg_result.dtype,
-                                   name=groupby.name, index_value=index_value)
+                                   name=agg_result.name, index_value=index_value)
 
     def __call__(self, groupby):
         df = groupby
         while df.op.object_type not in (ObjectType.dataframe, ObjectType.series):
             df = df.inputs[0]
 
-        self._object_type = ObjectType.dataframe if groupby.op.object_type == ObjectType.dataframe_groupby \
-            else ObjectType.series
+        if self.func == 'size':
+            self._object_type = ObjectType.series
+        else:
+            self._object_type = ObjectType.dataframe \
+                if groupby.op.object_type == ObjectType.dataframe_groupby else ObjectType.series
 
         if self.object_type == ObjectType.dataframe:
             return self._call_dataframe(groupby, df)
@@ -250,7 +253,7 @@ def _add_column_to_functions(col, fun_name, mappers, combiners, aggregator):
             for f in functions:
                 if f in {'sum', 'prod', 'min', 'max'}:
                     _add_column_to_functions(col, f, [f], [f], f)
-                elif f == 'count':
+                elif f in {'count', 'size'}:
                     _add_column_to_functions(col, f, [f], ['sum'], 'sum')
                 elif f == 'mean':
                     def _mean(_, grouped, columns):
@@ -435,8 +438,11 @@ def _tile_with_tree(cls, op):
                 chunk_op._groupby_params['level'] = level
                 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]),
+                new_shape = (np.nan, out_df.shape[1]) if len(out_df.shape) == 2 else (np.nan,)
+
+                new_chunks.append(chunk_op.new_chunk([chk], index=(idx, 0), shape=new_shape,
                                                      index_value=chks[0].index_value,
                                                      columns_value=columns_value))
             chunks = new_chunks
@@ -514,7 +520,7 @@ def _is_raw_one_func(cls, op):
             func = raw_func[0]
         if func is None:
             return False
-        return func in {'min', 'max', 'prod', 'sum', 'count'}
+        return func in {'min', 'max', 'prod', 'sum', 'count', 'size'}
 
     @classmethod
     def execute(cls, ctx, op):
diff --git a/mars/dataframe/reduction/aggregation.py b/mars/dataframe/reduction/aggregation.py
--- a/mars/dataframe/reduction/aggregation.py
+++ b/mars/dataframe/reduction/aggregation.py
@@ -28,9 +28,10 @@
 from ..operands import DataFrameOperand, DataFrameOperandMixin, ObjectType
 from ..utils import build_empty_df, build_empty_series, parse_index, validate_axis
 
+cp = lazy_import('cupy', globals=globals(), rename='cp')
 cudf = lazy_import('cudf', globals=globals())
 
-_builtin_aggregation_functions = {'sum', 'prod', 'min', 'max', 'count', 'mean', 'var', 'std'}
+_builtin_aggregation_functions = {'sum', 'prod', 'min', 'max', 'count', 'size', 'mean', 'var', 'std'}
 _stage_info = namedtuple('_stage_info', ('map_groups', 'map_sources', 'combine_groups', 'combine_sources',
                                          'agg_sources', 'agg_columns', 'agg_funcs', 'key_to_funcs',
                                          'valid_columns'))
@@ -41,6 +42,7 @@ class DataFrameAggregate(DataFrameOperand, DataFrameOperandMixin):
     _op_type_ = opcodes.AGGREGATE
 
     _func = AnyField('func')
+    _raw_func = AnyField('raw_func')
     _axis = AnyField('axis')
     _use_inf_as_na = BoolField('use_inf_as_na')
 
@@ -53,11 +55,11 @@ class DataFrameAggregate(DataFrameOperand, DataFrameOperandMixin):
     _agg_funcs = DictField('agg_funcs')
     _key_to_funcs = DictField('keys_to_funcs')
 
-    def __init__(self, func=None, axis=None, use_inf_as_na=None, map_groups=None, map_sources=None,
-                 combine_groups=None, combine_sources=None, agg_sources=None, agg_columns=None,
-                 agg_funcs=None, key_to_funcs=None, object_type=None, stage=None, **kw):
-        super().__init__(_func=func, _axis=axis, _use_inf_as_na=use_inf_as_na, _map_groups=map_groups,
-                         _map_sources=map_sources, _combine_groups=combine_groups,
+    def __init__(self, func=None, raw_func=None, axis=None, use_inf_as_na=None, map_groups=None,
+                 map_sources=None, combine_groups=None, combine_sources=None, agg_sources=None,
+                 agg_columns=None, agg_funcs=None, key_to_funcs=None, object_type=None, stage=None, **kw):
+        super().__init__(_func=func, _raw_func=raw_func, _axis=axis, _use_inf_as_na=use_inf_as_na,
+                         _map_groups=map_groups, _map_sources=map_sources, _combine_groups=combine_groups,
                          _combine_sources=combine_sources, _agg_sources=agg_sources,
                          _agg_columns=agg_columns, _agg_funcs=agg_funcs, _key_to_funcs=key_to_funcs,
                          _stage=stage, _object_type=object_type, **kw)
@@ -66,6 +68,10 @@ def __init__(self, func=None, axis=None, use_inf_as_na=None, map_groups=None, ma
     def func(self):
         return self._func
 
+    @property
+    def raw_func(self):
+        return self._raw_func
+
     @property
     def axis(self) -> int:
         return self._axis
@@ -125,6 +131,8 @@ def _calc_result_shape(self, df):
             return np.array(result_df).dtype, None
 
     def _normalize_funcs(self):
+        self._raw_func = self._func
+
         if isinstance(self._func, dict):
             new_func = OrderedDict()
             for k, v in self._func.items():
@@ -218,7 +226,7 @@ def _add_column_to_functions(col, fun_name, mappers, combiners, aggregator):
             for func in funcs:
                 if func in {'sum', 'prod', 'min', 'max'}:
                     _add_column_to_functions(col, func, [func], [func], func)
-                elif func == 'count':
+                elif func in {'count', 'size'}:
                     _add_column_to_functions(col, func, [func], ['sum'], 'sum')
                 elif func == 'mean':
                     def _mean(sum_data, count_data, axis=0):
@@ -306,15 +314,34 @@ def _tile_single_chunk(cls, op: "DataFrameAggregate"):
             chunk = chunk_op.new_chunk(in_df.chunks, index=(0, 0), shape=out_df.shape,
                                        index_value=out_df.index_value, columns_value=out_df.columns_value,
                                        dtypes=out_df.dtypes)
-        else:
-            chunk = chunk_op.new_chunk(in_df.chunks, index=(0,), shape=out_df.shape,
+        elif op.object_type == ObjectType.series:
+            chunk = chunk_op.new_chunk(in_df.chunks, index=(0,), shape=out_df.shape, dtype=out_df.dtype,
                                        index_value=out_df.index_value, name=out_df.name)
+        else:
+            chunk = chunk_op.new_chunk(in_df.chunks, dtype=out_df.dtype, shape=())
 
         tileable_op = op.copy().reset_key()
         kw = out_df.params.copy()
         kw.update(dict(chunks=[chunk], nsplits=((x,) for x in out_df.shape)))
         return tileable_op.new_tileables([in_df], **kw)
 
+    @classmethod
+    def _tile_size(cls, op: "DataFrameAggregate"):
+        in_df = op.inputs[0]
+        out_df = op.outputs[0]
+
+        chunks = []
+        for c in in_df.chunks:
+            chunk_op = op.copy().reset_key()
+            chunks.append(chunk_op.new_chunk([c], index=c.index, shape=(1,) * len(in_df.shape),
+                                             dtype=out_df.dtype))
+
+        tileable_op = op.copy().reset_key()
+        nsplits = tuple((1,) * s for s in in_df.chunk_shape)
+        tileable = tileable_op.new_tileable(out_df.inputs, chunks=chunks, nsplits=nsplits,
+                                            shape=in_df.chunk_shape, dtype=out_df.dtype)
+        return [tileable.sum()._inplace_tile()]
+
     @classmethod
     def _tile_tree(cls, op: "DataFrameAggregate"):
         in_df = op.inputs[0]
@@ -433,7 +460,7 @@ def _tile_tree(cls, op: "DataFrameAggregate"):
                     kw.update(dict(name=out_df.name, dtype=out_df.dtype, index=(idx,),
                                    index_value=index_value, shape=shape))
                 else:
-                    kw.update(dict(index=(0,), shape=(), dtype=out_df.dtype))
+                    kw.update(dict(index=(), shape=(), dtype=out_df.dtype))
             agg_chunks.append(chunk_op.new_chunk([chk], **kw))
 
         new_op = op.copy()
@@ -455,14 +482,18 @@ def _tile_tree(cls, op: "DataFrameAggregate"):
 
     @classmethod
     def tile(cls, op: "DataFrameAggregate"):
-        if len(op.inputs[0].chunks) == 1:
+        in_df = op.inputs[0]
+
+        if len(in_df.chunks) == 1:
             return cls._tile_single_chunk(op)
+        elif in_df.ndim == 2 and op.raw_func == 'size':
+            return cls._tile_size(op)
         else:
             return cls._tile_tree(op)
 
     @staticmethod
     def _wrap_df(xdf, value, columns=None, transform=False):
-        if isinstance(value, np.generic):
+        if isinstance(value, (np.generic, int, float, complex)):
             value = xdf.DataFrame([value], columns=columns)
         else:
             value = xdf.DataFrame(value, columns=columns)
@@ -482,7 +513,8 @@ def _execute_map(cls, ctx, op: "DataFrameAggregate"):
                 src_df = in_data
             else:
                 src_df = in_data[cols]
-            ret_map_dfs[map_func_str] = cls._wrap_df(xdf, src_df.agg(map_func_str, axis=axis),
+            map_func = map_func_str if map_func_str != 'size' else (lambda x: x.size)
+            ret_map_dfs[map_func_str] = cls._wrap_df(xdf, src_df.agg(map_func, axis=axis),
                                                      columns=[axis_index], transform=axis == 0)
 
         ret_combine_dfs = OrderedDict()
@@ -541,7 +573,7 @@ def _execute_agg(cls, ctx, op: "DataFrameAggregate"):
 
         concat_df = xdf.concat(aggs, axis=axis)
         if op.object_type == ObjectType.series:
-            if concat_df.shape[0] == 1:
+            if concat_df.shape[1] > 1:
                 concat_df = concat_df.iloc[0, :]
             else:
                 concat_df = concat_df.iloc[:, 0]
@@ -569,8 +601,12 @@ def execute(cls, ctx, op: "DataFrameAggregate"):
                 cls._execute_combine(ctx, op)
             elif op.stage == OperandStage.agg:
                 cls._execute_agg(ctx, op)
+            elif op.raw_func == 'size':
+                xp = cp if op.gpu else np
+                ctx[op.outputs[0].key] = xp.array(ctx[op.inputs[0].key].agg(op.raw_func, axis=op.axis)) \
+                    .reshape(op.outputs[0].shape)
             else:
-                ctx[op.outputs[0].key] = ctx[op.inputs[0].key].agg(op.func, axis=op.axis)
+                ctx[op.outputs[0].key] = ctx[op.inputs[0].key].agg(op.raw_func, axis=op.axis)
         finally:
             pd.reset_option('mode.use_inf_as_na')
 


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
@@ -168,6 +168,10 @@ def testDataFrameGroupByAgg(self):
         mdf2 = md.DataFrame(df2, chunk_size=2)
 
         for method in ['tree', 'shuffle']:
+            r0 = mdf2.groupby('c2').agg('size', method=method)
+            pd.testing.assert_series_equal(self.executor.execute_dataframe(r0, concat=True)[0],
+                                           df2.groupby('c2').agg('size'))
+
             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'))
@@ -194,7 +198,7 @@ def testDataFrameGroupByAgg(self):
             pd.testing.assert_frame_equal(self.executor.execute_dataframe(r6, concat=True)[0],
                                           df2.groupby('c2').agg('std'))
 
-            agg = ['std', 'mean', 'var', 'max', 'count']
+            agg = ['std', 'mean', 'var', 'max', 'count', 'size']
             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))
@@ -218,6 +222,10 @@ def testDataFrameGroupByAgg(self):
             pd.testing.assert_frame_equal(self.executor.execute_dataframe(r3, concat=True)[0],
                                           df2.groupby(df2['c2']).sum())
 
+        r8 = mdf2.groupby('c2').size()
+        pd.testing.assert_series_equal(self.executor.execute_dataframe(r8, concat=True)[0],
+                                       df2.groupby('c2').size())
+
         r4 = mdf2.groupby('c2').sum()
         pd.testing.assert_frame_equal(self.executor.execute_dataframe(r4, concat=True)[0],
                                       df2.groupby('c2').sum())
@@ -272,6 +280,10 @@ def testSeriesGroupByAgg(self):
         ms1 = md.Series(series1, chunk_size=3)
 
         for method in ['tree', 'shuffle']:
+            r0 = ms1.groupby(lambda x: x % 2).agg('size', method=method)
+            pd.testing.assert_series_equal(self.executor.execute_dataframe(r0, concat=True)[0],
+                                           series1.groupby(lambda x: x % 2).agg('size'))
+
             r1 = ms1.groupby(lambda x: x % 2).agg('sum', method=method)
             pd.testing.assert_series_equal(self.executor.execute_dataframe(r1, concat=True)[0],
                                            series1.groupby(lambda x: x % 2).agg('sum'))
@@ -298,7 +310,7 @@ def testSeriesGroupByAgg(self):
             pd.testing.assert_series_equal(self.executor.execute_dataframe(r6, concat=True)[0],
                                            series1.groupby(lambda x: x % 2).agg('std'))
 
-            agg = ['std', 'mean', 'var', 'max', 'count']
+            agg = ['std', 'mean', 'var', 'max', 'count', 'size']
             r3 = ms1.groupby(lambda x: x % 2).agg(agg, method=method)
             pd.testing.assert_frame_equal(self.executor.execute_dataframe(r3, concat=True)[0],
                                           series1.groupby(lambda x: x % 2).agg(agg))
@@ -308,6 +320,10 @@ def testSeriesGroupByAgg(self):
             pd.testing.assert_series_equal(self.executor.execute_dataframe(r3, concat=True)[0],
                                            series1.groupby(series1).sum())
 
+        r4 = ms1.groupby(lambda x: x % 2).size()
+        pd.testing.assert_series_equal(self.executor.execute_dataframe(r4, concat=True)[0],
+                                       series1.groupby(lambda x: x % 2).size())
+
         r4 = ms1.groupby(lambda x: x % 2).sum()
         pd.testing.assert_series_equal(self.executor.execute_dataframe(r4, concat=True)[0],
                                        series1.groupby(lambda x: x % 2).sum())
diff --git a/mars/dataframe/reduction/tests/test_reduction_execute.py b/mars/dataframe/reduction/tests/test_reduction_execute.py
--- a/mars/dataframe/reduction/tests/test_reduction_execute.py
+++ b/mars/dataframe/reduction/tests/test_reduction_execute.py
@@ -421,7 +421,7 @@ def setUp(self):
         self.executor = ExecutorForTest()
 
     def testDataFrameAggregate(self):
-        all_aggs = ['sum', 'prod', 'min', 'max', 'count', 'mean', 'var', 'std']
+        all_aggs = ['sum', 'prod', 'min', 'max', 'count', 'size', 'mean', 'var', 'std']
         data = pd.DataFrame(np.random.rand(20, 20))
 
         df = from_pandas_df(data)
@@ -429,6 +429,18 @@ def testDataFrameAggregate(self):
         pd.testing.assert_frame_equal(self.executor.execute_dataframe(result, concat=True)[0],
                                       data.agg(all_aggs))
 
+        result = df.agg('size')
+        self.assertEqual(self.executor.execute_dataframe(result)[0], data.agg('size'))
+
+        for func in (a for a in all_aggs if a != 'size'):
+            result = df.agg(func)
+            pd.testing.assert_series_equal(self.executor.execute_dataframe(result, concat=True)[0],
+                                           data.agg(func))
+
+            result = df.agg(func, axis=1)
+            pd.testing.assert_series_equal(self.executor.execute_dataframe(result, concat=True)[0],
+                                           data.agg(func, axis=1))
+
         df = from_pandas_df(data, chunk_size=3)
 
         # will redirect to transform
@@ -436,7 +448,10 @@ def testDataFrameAggregate(self):
         pd.testing.assert_frame_equal(self.executor.execute_dataframe(result, concat=True)[0],
                                       data.agg(['cumsum', 'cummax']))
 
-        for func in all_aggs:
+        result = df.agg('size')
+        self.assertEqual(self.executor.execute_dataframe(result)[0], data.agg('size'))
+
+        for func in (a for a in all_aggs if a != 'size'):
             result = df.agg(func)
             pd.testing.assert_series_equal(self.executor.execute_dataframe(result, concat=True)[0],
                                            data.agg(func))
@@ -462,7 +477,7 @@ def testDataFrameAggregate(self):
                                       data.agg({0: ['sum', 'min', 'var'], 9: ['mean', 'var', 'std']}))
 
     def testSeriesAggregate(self):
-        all_aggs = ['sum', 'prod', 'min', 'max', 'count', 'mean', 'var', 'std']
+        all_aggs = ['sum', 'prod', 'min', 'max', 'count', 'size', 'mean', 'var', 'std']
         data = pd.Series(np.random.rand(20), index=[str(i) for i in range(20)], name='a')
 
         series = from_pandas_series(data)
@@ -470,6 +485,11 @@ def testSeriesAggregate(self):
         pd.testing.assert_series_equal(self.executor.execute_dataframe(result, concat=True)[0],
                                        data.agg(all_aggs))
 
+        for func in all_aggs:
+            result = series.agg(func)
+            self.assertAlmostEqual(self.executor.execute_dataframe(result, concat=True)[0],
+                                   data.agg(func))
+
         series = from_pandas_series(data, chunk_size=3)
 
         for func in all_aggs: