Repository: mars-project/mars

Issue:
[BUG] df.groupby().size() returns NaN 
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
The result of df.groupby().size() is wrong when specify `as_index=False`.

``` Python
In [6]: a = md.DataFrame(np.random.rand(10, 2), columns=['a', 'b'], chunk_size=5)                   

In [7]: a['c'] = a['a'].map(lambda x: x+1)                                                          

In [8]: a.groupby('c', as_index=False).size().execute()                                             
100%|██████████████████████████████████████████████████████████| 100.0/100 [00:00<00:00, 321.49it/s]
Out[8]: 
          c  size
0  1.045709   NaN
1  1.185217   NaN
2  1.360065   NaN
3  1.386844   NaN
4  1.506352   NaN
5  1.676305   NaN
6  1.715007   NaN
7  1.833545   NaN
8  1.890987   NaN
9  1.943286   NaN
```


Discussion:
Caused by special output by `df.agg('size')` when all columns are ignored.

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
@@ -630,7 +630,7 @@ def _do_custom_agg(op, custom_reduction, *input_objs):
         return concat_result
 
     @staticmethod
-    def _do_predefined_agg(input_obj, agg_func, **kwds):
+    def _do_predefined_agg(input_obj, agg_func, single_func=False, **kwds):
         ndim = getattr(input_obj, 'ndim', None) or input_obj.obj.ndim
         if agg_func == 'str_concat':
             agg_func = lambda x: x.str.cat(**kwds)
@@ -640,8 +640,14 @@ def _do_predefined_agg(input_obj, agg_func, **kwds):
             agg_func.__name__ = func_name
 
         if ndim == 2:
-            result = input_obj.agg([agg_func])
-            result.columns = result.columns.droplevel(-1)
+            if single_func:
+                result = input_obj.agg(agg_func)
+                if result.ndim == 1:
+                    # when agg_func == size, agg only returns one single series.
+                    result = result.to_frame(agg_func)
+            else:
+                result = input_obj.agg([agg_func])
+                result.columns = result.columns.droplevel(-1)
             return result
         else:
             return input_obj.agg(agg_func)
@@ -704,7 +710,9 @@ def _wrapped_func(col):
             if map_func_name == 'custom_reduction':
                 agg_dfs.extend(cls._do_custom_agg(op, custom_reduction, input_obj))
             else:
-                agg_dfs.append(cls._do_predefined_agg(input_obj, map_func_name, **kwds))
+                single_func = map_func_name == op.raw_func
+                agg_dfs.append(cls._do_predefined_agg(
+                    input_obj, map_func_name, single_func, **kwds))
 
         if op._size_recorder_name is not None:
             # record_size


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
@@ -282,6 +282,11 @@ def test_dataframe_groupby_agg(setup):
 
     # test as_index=False
     for method in ['tree', 'shuffle']:
+        r = mdf.groupby('c2', as_index=False).agg('size', method=method)
+        pd.testing.assert_frame_equal(
+            r.execute().fetch().sort_values('c2', ignore_index=True),
+            raw.groupby('c2', as_index=False).agg('size').sort_values('c2', ignore_index=True))
+
         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),