Repository: mars-project/mars

Issue:
Supports `DataFrame.mean()` and `Series.mean()`
<!--
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.**

So far, aggregation functions of Mars DataFrame include `max`, `min`, `prod`, `sum`, I wanna add `mean` support.



Discussion:

Resolving patch:
diff --git a/mars/dataframe/reduction/__init__.py b/mars/dataframe/reduction/__init__.py
--- a/mars/dataframe/reduction/__init__.py
+++ b/mars/dataframe/reduction/__init__.py
@@ -17,26 +17,26 @@
 from .max import DataFrameMax
 from .min import DataFrameMin
 from .count import DataFrameCount
+from .mean import DataFrameMean
 
 
 def _install():
-    from ..core import DataFrame, Series
+    from ..core import DATAFRAME_TYPE, SERIES_TYPE
     from .sum import sum_series, sum_dataframe
     from .prod import prod_series, prod_dataframe
     from .max import max_series, max_dataframe
     from .min import min_series, min_dataframe
     from .count import count_series, count_dataframe
+    from .mean import mean_series, mean_dataframe
 
-    setattr(DataFrame, 'sum', sum_dataframe)
-    setattr(Series, 'sum', sum_series)
-    setattr(DataFrame, 'prod', prod_dataframe)
-    setattr(Series, 'prod', prod_series)
-    setattr(DataFrame, 'max', max_dataframe)
-    setattr(Series, 'max', max_series)
-    setattr(DataFrame, 'min', min_dataframe)
-    setattr(Series, 'min', min_series)
-    setattr(DataFrame, 'count', count_dataframe)
-    setattr(Series, 'count', count_series)
+    func_names = ['sum', 'prod', 'max', 'min', 'count', 'mean']
+    series_funcs = [sum_series, prod_series, max_series, min_series, count_series, mean_series]
+    df_funcs = [sum_dataframe, prod_dataframe, max_dataframe, min_dataframe, count_dataframe, mean_dataframe]
+    for func_name, series_func, df_func in zip(func_names, series_funcs, df_funcs):
+        for t in DATAFRAME_TYPE:
+            setattr(t, func_name, df_func)
+        for t in SERIES_TYPE:
+            setattr(t, func_name, series_func)
 
 
 _install()
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
@@ -167,7 +167,7 @@ def _tile_dataframe(cls, op):
         reduction_chunks = np.empty(op.inputs[0].chunk_shape, dtype=np.object)
         for c in op.inputs[0].chunks:
             new_chunk_op = op.copy().reset_key()
-            op._stage = Stage.map
+            new_chunk_op._stage = Stage.map
             new_chunk_op._object_type = ObjectType.dataframe
             if op.axis == 0:
                 if op.numeric_only:
@@ -206,7 +206,7 @@ def _tile_series(cls, op):
         chunks = np.empty(op.inputs[0].chunk_shape, dtype=np.object)
         for c in op.inputs[0].chunks:
             new_chunk_op = op.copy().reset_key()
-            op._stage = Stage.map
+            new_chunk_op._stage = Stage.map
             chunks[c.index] = new_chunk_op.new_chunk([c], shape=(), dtype=df.dtype, index_value=df.index_value)
 
         while len(chunks) > combine_size:
@@ -249,7 +249,7 @@ def tile(cls, op):
             return cls._tile_series(op)
 
     @classmethod
-    def _execute_reduction(cls, in_data, op, min_count=None):
+    def _execute_reduction(cls, in_data, op, min_count=None, reduction_func=None):
         kwargs = dict()
         if op.axis is not None:
             kwargs['axis'] = op.axis
@@ -259,16 +259,20 @@ def _execute_reduction(cls, in_data, op, min_count=None):
             kwargs['numeric_only'] = op.numeric_only
         if min_count is not None:
             kwargs['min_count'] = op.min_count
-        return getattr(in_data, getattr(cls, '_func_name'))(**kwargs)
+        reduction_func = reduction_func or getattr(cls, '_func_name')
+        return getattr(in_data, reduction_func)(**kwargs)
 
     @classmethod
-    def _execute_map_with_count(cls, ctx, op):
+    def _execute_map_with_count(cls, ctx, op, reduction_func=None):
         # Execution with specified `min_count` in the map stage
 
         xdf = cudf if op.gpu else pd
         in_data = ctx[op.inputs[0].key]
-        count = in_data.notnull().sum(axis=op.axis)
-        r = cls._execute_reduction(in_data, op)
+        if isinstance(in_data, pd.Series):
+            count = in_data.count()
+        else:
+            count = in_data.count(axis=op.axis, numeric_only=op.numeric_only)
+        r = cls._execute_reduction(in_data, op, reduction_func=reduction_func)
         if isinstance(in_data, xdf.Series):
             ctx[op.outputs[0].key] = (r, count)
         else:
@@ -277,13 +281,13 @@ def _execute_map_with_count(cls, ctx, op):
                 else (xdf.DataFrame(r).transpose(), xdf.DataFrame(count).transpose())
 
     @classmethod
-    def _execute_combine_with_count(cls, ctx, op):
+    def _execute_combine_with_count(cls, ctx, op, reduction_func=None):
         # Execution with specified `min_count` in the combine stage
 
         xdf = cudf if op.gpu else pd
         in_data, concat_count = ctx[op.inputs[0].key]
         count = concat_count.sum(axis=op.axis)
-        r = cls._execute_reduction(in_data, op)
+        r = cls._execute_reduction(in_data, op, reduction_func=reduction_func)
         if isinstance(in_data, xdf.Series):
             ctx[op.outputs[0].key] = (r, count)
         else:
@@ -292,14 +296,14 @@ def _execute_combine_with_count(cls, ctx, op):
                 else (xdf.DataFrame(r).transpose(), xdf.DataFrame(count).transpose())
 
     @classmethod
-    def _execute_agg_with_count(cls, ctx, op):
+    def _execute_agg_with_count(cls, ctx, op, reduction_func=None):
         # Execution with specified `min_count` in the aggregate stage
 
         # When specify `min_count`, the terminal chunk has two inputs one of which is for count.
         # The output should be determined by comparing `count` with `min_count`.
         in_data, concat_count = ctx[op.inputs[0].key]
         count = concat_count.sum(axis=op.axis)
-        r = cls._execute_reduction(in_data, op)
+        r = cls._execute_reduction(in_data, op, reduction_func=reduction_func)
         if np.isscalar(r):
             ctx[op.outputs[0].key] = np.nan if count < op.min_count else r
         else:
@@ -307,13 +311,13 @@ def _execute_agg_with_count(cls, ctx, op):
             ctx[op.outputs[0].key] = r
 
     @classmethod
-    def _execute_without_count(cls, ctx, op):
+    def _execute_without_count(cls, ctx, op, reduction_func=None):
         # Execution for normal reduction operands.
 
         # For dataframe, will keep dimensions for intermediate results.
         xdf = cudf if op.gpu else pd
         in_data = ctx[op.inputs[0].key]
-        r = cls._execute_reduction(in_data, op, min_count=op.min_count)
+        r = cls._execute_reduction(in_data, op, min_count=op.min_count, reduction_func=reduction_func)
         if isinstance(in_data, xdf.Series) or op.object_type == ObjectType.series:
             ctx[op.outputs[0].key] = r
         else:
@@ -346,8 +350,12 @@ def execute(cls, ctx, op):
             cls._execute_combine(ctx, op)
         elif op.stage == Stage.agg:
             cls._execute_agg(ctx, op)
-        else:
+        elif op.stage == Stage.map:
             cls._execute_map(ctx, op)
+        else:
+            in_data = ctx[op.inputs[0].key]
+            min_count = getattr(op, 'min_count', None)
+            ctx[op.outputs[0].key] = cls._execute_reduction(in_data, op, min_count)
 
     def _call_dataframe(self, df):
         axis = getattr(self, 'axis', None)
diff --git a/mars/dataframe/reduction/mean.py b/mars/dataframe/reduction/mean.py
new file mode 100644
--- /dev/null
+++ b/mars/dataframe/reduction/mean.py
@@ -0,0 +1,51 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ... import opcodes as OperandDef
+from ...utils import lazy_import
+from .core import DataFrameReductionOperand, DataFrameReductionMixin, ObjectType
+
+cudf = lazy_import('cudf', globals=globals())
+
+
+class DataFrameMean(DataFrameReductionOperand, DataFrameReductionMixin):
+    _op_type_ = OperandDef.MEAN
+    _func_name = 'mean'
+
+    @classmethod
+    def _execute_map(cls, ctx, op):
+        cls._execute_map_with_count(ctx, op, reduction_func='sum')
+
+    @classmethod
+    def _execute_combine(cls, ctx, op):
+        cls._execute_combine_with_count(ctx, op, reduction_func='sum')
+
+    @classmethod
+    def _execute_agg(cls, ctx, op):
+        in_data, concat_count = ctx[op.inputs[0].key]
+        count = concat_count.sum(axis=op.axis)
+        r = cls._execute_reduction(in_data, op, reduction_func='sum')
+        ctx[op.outputs[0].key] = r / count
+
+
+def mean_series(df, axis=None, skipna=None, level=None, combine_size=None):
+    op = DataFrameMean(axis=axis, skipna=skipna, level=level, combine_size=combine_size,
+                       object_type=ObjectType.series)
+    return op(df)
+
+
+def mean_dataframe(df, axis=None, skipna=None, level=None, numeric_only=None, combine_size=None):
+    op = DataFrameMean(axis=axis, skipna=skipna, level=level, numeric_only=numeric_only,
+                       combine_size=combine_size, object_type=ObjectType.series)
+    return op(df)


Test patch:
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
@@ -19,7 +19,8 @@
 from mars import opcodes as OperandDef
 from mars.tests.core import TestBase, parameterized
 from mars.dataframe.core import IndexValue, Series
-from mars.dataframe.reduction import DataFrameSum, DataFrameProd, DataFrameMin, DataFrameMax, DataFrameCount
+from mars.dataframe.reduction import DataFrameSum, DataFrameProd, DataFrameMin, DataFrameMax, \
+    DataFrameCount, DataFrameMean
 from mars.dataframe.merge import DataFrameConcat
 from mars.dataframe.datasource.series import from_pandas as from_pandas_series
 from mars.dataframe.datasource.dataframe import from_pandas as from_pandas_df
@@ -31,6 +32,7 @@
     min=dict(func_name='min', op=DataFrameMin, has_skipna=True),
     max=dict(func_name='max', op=DataFrameMax, has_skipna=True),
     count=dict(func_name='count', op=DataFrameCount, has_skipna=False),
+    mean=dict(func_name='mean', op=DataFrameMean, has_skipna=True)
 )
 
 
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
@@ -27,6 +27,7 @@
     prod=dict(func_name='prod', has_min_count=True),
     min=dict(func_name='min', has_min_count=False),
     max=dict(func_name='max', has_min_count=False),
+    mean=dict(func_name='mean', has_min_count=False)
 )