Repository: mars-project/mars

Issue:
Implements fill methods in GroupBy
**Is your feature request related to a problem? Please describe.**
Mars now lack fill methods in GroupBy objects:

* [ ] [`GroupBy.bfill`](https://pandas.pydata.org/docs/reference/api/pandas.core.groupby.GroupBy.bfill.html#)
* [ ] [`GroupBy.ffill`](https://pandas.pydata.org/docs/reference/api/pandas.core.groupby.GroupBy.ffill.html#)
* [ ] [`DataFrameGroupBy.fillna`](https://pandas.pydata.org/docs/reference/api/pandas.core.groupby.DataFrameGroupBy.fillna.html#)

**Describe the solution you'd like**
Possible to implement with `groupby(xxx).apply(...)`. Maybe possible to keep original index to reduce shuffle in data alignment.


Discussion:

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
@@ -26,6 +26,7 @@ def _install():
     from .head import head
     from .transform import groupby_transform
     from .sample import groupby_sample
+    from .fill import ffill, bfill, fillna
 
     for cls in DATAFRAME_TYPE:
         setattr(cls, "groupby", groupby)
@@ -66,6 +67,11 @@ def _install():
 
         setattr(cls, "sample", groupby_sample)
 
+        setattr(cls, "ffill", ffill)
+        setattr(cls, "bfill", bfill)
+        setattr(cls, "backfill", bfill)
+        setattr(cls, "fillna", fillna)
+
     for cls in DATAFRAME_GROUPBY_TYPE:
         setattr(cls, "__getitem__", df_groupby_getitem)
 
diff --git a/mars/dataframe/groupby/fill.py b/mars/dataframe/groupby/fill.py
new file mode 100644
--- /dev/null
+++ b/mars/dataframe/groupby/fill.py
@@ -0,0 +1,212 @@
+# Copyright 1999-2021 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.
+
+import numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...core import OutputType
+from ...serialization.serializables import AnyField, DictField, StringField, Int64Field
+from ..utils import parse_index, build_empty_df, build_empty_series
+from ..operands import DataFrameOperandMixin, DataFrameOperand
+
+
+class GroupByFillOperand(DataFrameOperand, DataFrameOperandMixin):
+    _op_module_ = "dataframe.groupby"
+
+    value = AnyField("value", default=None)
+    method = StringField("method", default=None)
+    axis = AnyField("axis", default=0)
+    limit = Int64Field("limit", default=None)
+    downcast = DictField("downcast", default=None)
+
+    def _calc_out_dtypes(self, in_groupby):
+        mock_groupby = in_groupby.op.build_mock_groupby()
+        func_name = getattr(self, "_func_name")
+
+        if func_name == "fillna":
+            result_df = mock_groupby.fillna(
+                value=self.value,
+                method=self.method,
+                axis=self.axis,
+                limit=self.limit,
+                downcast=self.downcast,
+            )
+        else:
+            result_df = getattr(mock_groupby, func_name)(limit=self.limit)
+
+        if isinstance(result_df, pd.DataFrame):
+            self.output_types = [OutputType.dataframe]
+            return result_df.dtypes
+        else:
+            self.output_types = [OutputType.series]
+            return result_df.name, result_df.dtype
+
+    def __call__(self, groupby):
+        in_df = groupby
+        while in_df.op.output_types[0] not in (OutputType.dataframe, OutputType.series):
+            in_df = in_df.inputs[0]
+        out_dtypes = self._calc_out_dtypes(groupby)
+
+        kw = in_df.params.copy()
+        kw["index_value"] = parse_index(pd.RangeIndex(-1), groupby.key)
+        if self.output_types[0] == OutputType.dataframe:
+            kw.update(
+                dict(
+                    columns_value=parse_index(out_dtypes.index, store_data=True),
+                    dtypes=out_dtypes,
+                    shape=(groupby.shape[0], len(out_dtypes)),
+                )
+            )
+        else:
+            name, dtype = out_dtypes
+            kw.update(dtype=dtype, name=name, shape=(groupby.shape[0],))
+        return self.new_tileable([groupby], **kw)
+
+    @classmethod
+    def tile(cls, op):
+        in_groupby = op.inputs[0]
+        out_df = op.outputs[0]
+
+        chunks = []
+        for c in in_groupby.chunks:
+            new_op = op.copy().reset_key()
+
+            new_index = parse_index(pd.RangeIndex(-1), c.key)
+            if op.output_types[0] == OutputType.dataframe:
+                chunks.append(
+                    new_op.new_chunk(
+                        [c],
+                        index=c.index,
+                        shape=(np.nan, len(out_df.dtypes)),
+                        dtypes=out_df.dtypes,
+                        columns_value=out_df.columns_value,
+                        index_value=new_index,
+                    )
+                )
+            else:
+                chunks.append(
+                    new_op.new_chunk(
+                        [c],
+                        index=(c.index[0],),
+                        shape=(np.nan,),
+                        dtype=out_df.dtype,
+                        index_value=new_index,
+                        name=out_df.name,
+                    )
+                )
+        new_op = op.copy().reset_key()
+        kw = out_df.params.copy()
+        kw["chunks"] = chunks
+        if op.output_types[0] == OutputType.dataframe:
+            kw["nsplits"] = ((np.nan,) * len(chunks), (len(out_df.dtypes),))
+        else:
+            kw["nsplits"] = ((np.nan,) * len(chunks),)
+        return new_op.new_tileables([in_groupby], **kw)
+
+    @classmethod
+    def execute(cls, ctx, op: "GroupByFillOperand"):
+        in_data = ctx[op.inputs[0].key]
+        out_chunk = op.outputs[0]
+
+        if not in_data or in_data.empty:
+            ctx[out_chunk.key] = (
+                build_empty_df(out_chunk.dtypes)
+                if op.output_types[0] == OutputType.dataframe
+                else build_empty_series(out_chunk.dtype, name=out_chunk.name)
+            )
+            return
+
+        func_name = getattr(op, "_func_name")
+        if func_name == "fillna":
+            ctx[out_chunk.key] = in_data.fillna(
+                value=op.value,
+                method=op.method,
+                axis=op.axis,
+                limit=op.limit,
+                downcast=op.downcast,
+            )
+        else:
+            result = getattr(in_data, func_name)(limit=op.limit)
+            if result.ndim == 2:
+                ctx[out_chunk.key] = result.astype(out_chunk.dtypes, copy=False)
+            else:
+                ctx[out_chunk.key] = result.astype(out_chunk.dtype, copy=False)
+
+
+class GroupByFFill(GroupByFillOperand):
+    _op_type_ = opcodes.FILL_NA
+    _func_name = "ffill"
+
+
+class GroupByBFill(GroupByFillOperand):
+    _op_type = opcodes.FILL_NA
+    _func_name = "bfill"
+
+
+class GroupByFillNa(GroupByFillOperand):
+    _op_type = opcodes.FILL_NA
+    _func_name = "fillna"
+
+
+def ffill(groupby, limit=None):
+    """
+    Forward fill the values.
+
+    limit:  int, default None
+            Limit number of values to fill
+
+    return: Series or DataFrame
+    """
+    op = GroupByFFill(limit=limit)
+    return op(groupby)
+
+
+def bfill(groupby, limit=None):
+    """
+    Backward fill the values.
+
+    limit:  int, default None
+            Limit number of values to fill
+
+    return: Series or DataFrame
+    """
+    op = GroupByBFill(limit=limit)
+    return op(groupby)
+
+
+def fillna(groupby, value=None, method=None, axis=None, limit=None, downcast=None):
+    """
+    Fill NA/NaN values using the specified method
+
+    value:  scalar, dict, Series, or DataFrame
+            Value to use to fill holes (e.g. 0), alternately a dict/Series/DataFrame
+            of values specifying which value to use for each index (for a Series) or
+            column (for a DataFrame). Values not in the dict/Series/DataFrame
+            will not be filled. This value cannot be a list.
+    method: {'backfill','bfill','ffill',None}, default None
+    axis:   {0 or 'index', 1 or 'column'}
+    limit:  int, default None
+            If method is specified, this is the maximum number of consecutive
+            NaN values to forward/backward fill
+    downcast:   dict, default None
+                A dict of item->dtype of what to downcast if possible,
+                or the string ‘infer’ which will try to downcast to an appropriate equal type
+
+    return: DataFrame or None
+    """
+    op = GroupByFillNa(
+        value=value, method=method, axis=axis, limit=limit, downcast=downcast
+    )
+    return op(groupby)


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
@@ -363,3 +363,101 @@ def test_groupby_cum():
         assert len(r.chunks) == 4
         assert r.shape == (len(series1),)
         assert r.chunks[0].shape == (np.nan,)
+
+
+def test_groupby_fill():
+    df1 = pd.DataFrame(
+        [
+            [1, 1, 10],
+            [1, 1, np.nan],
+            [1, 1, np.nan],
+            [1, 2, np.nan],
+            [1, 2, 20],
+            [1, 2, np.nan],
+            [1, 3, np.nan],
+            [1, 3, np.nan],
+        ],
+        columns=["one", "two", "three"],
+    )
+    mdf = md.DataFrame(df1, chunk_size=3)
+
+    r = tile(getattr(mdf.groupby(["one", "two"]), "ffill")())
+    assert r.op.output_types[0] == OutputType.dataframe
+    assert r.shape == (len(df1), 1)
+    assert len(r.chunks) == 3
+    assert r.chunks[0].shape == (np.nan, 1)
+    assert r.dtypes.index.tolist() == ["three"]
+
+    r = tile(getattr(mdf.groupby(["two"]), "bfill")())
+    assert r.op.output_types[0] == OutputType.dataframe
+    assert r.shape == (len(df1), 2)
+    assert len(r.chunks) == 3
+    assert r.chunks[0].shape == (np.nan, 2)
+    assert r.dtypes.index.tolist() == ["one", "three"]
+
+    r = tile(getattr(mdf.groupby(["two"]), "backfill")())
+    assert r.op.output_types[0] == OutputType.dataframe
+    assert r.shape == (len(df1), 2)
+    assert len(r.chunks) == 3
+    assert r.chunks[0].shape == (np.nan, 2)
+    assert r.dtypes.index.tolist() == ["one", "three"]
+
+    r = tile(getattr(mdf.groupby(["one"]), "fillna")(5))
+    assert r.op.output_types[0] == OutputType.dataframe
+    assert r.shape == (len(df1), 2)
+    assert len(r.chunks) == 3
+    assert r.chunks[0].shape == (np.nan, 2)
+    assert r.dtypes.index.tolist() == ["two", "three"]
+
+    s1 = pd.Series([4, 3, 9, np.nan, np.nan, 7, 10, 8, 1, 6])
+    ms1 = md.Series(s1, chunk_size=3)
+    r = tile(getattr(ms1.groupby(lambda x: x % 2), "ffill")())
+    assert r.op.output_types[0] == OutputType.series
+    assert len(r.chunks) == 4
+    assert r.shape == (len(s1),)
+    assert r.chunks[0].shape == (np.nan,)
+
+    r = tile(getattr(ms1.groupby(lambda x: x % 2), "bfill")())
+    assert r.op.output_types[0] == OutputType.series
+    assert len(r.chunks) == 4
+    assert r.shape == (len(s1),)
+    assert r.chunks[0].shape == (np.nan,)
+
+    r = tile(getattr(ms1.groupby(lambda x: x % 2), "backfill")())
+    assert r.op.output_types[0] == OutputType.series
+    assert len(r.chunks) == 4
+    assert r.shape == (len(s1),)
+    assert r.chunks[0].shape == (np.nan,)
+
+    r = tile(getattr(ms1.groupby(lambda x: x % 2), "fillna")(5))
+    assert r.op.output_types[0] == OutputType.series
+    assert len(r.chunks) == 4
+    assert r.shape == (len(s1),)
+    assert r.chunks[0].shape == (np.nan,)
+
+    s1 = pd.Series([4, 3, 9, np.nan, np.nan, 7, 10, 8, 1, 6])
+    ms1 = md.Series(s1, chunk_size=3)
+
+    r = tile(getattr(ms1.groupby(lambda x: x % 2), "ffill")())
+    assert r.op.output_types[0] == OutputType.series
+    assert len(r.chunks) == 4
+    assert r.shape == (len(s1),)
+    assert r.chunks[0].shape == (np.nan,)
+
+    r = tile(getattr(ms1.groupby(lambda x: x % 2), "bfill")())
+    assert r.op.output_types[0] == OutputType.series
+    assert len(r.chunks) == 4
+    assert r.shape == (len(s1),)
+    assert r.chunks[0].shape == (np.nan,)
+
+    r = tile(getattr(ms1.groupby(lambda x: x % 2), "backfill")())
+    assert r.op.output_types[0] == OutputType.series
+    assert len(r.chunks) == 4
+    assert r.shape == (len(s1),)
+    assert r.chunks[0].shape == (np.nan,)
+
+    r = tile(getattr(ms1.groupby(lambda x: x % 2), "fillna")(5))
+    assert r.op.output_types[0] == OutputType.series
+    assert len(r.chunks) == 4
+    assert r.shape == (len(s1),)
+    assert r.chunks[0].shape == (np.nan,)
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
@@ -810,6 +810,67 @@ def test_groupby_cum(setup):
         )
 
 
+def test_groupby_fill(setup):
+    df1 = pd.DataFrame(
+        [
+            [1, 1, 10],
+            [1, 1, np.nan],
+            [1, 1, np.nan],
+            [1, 2, np.nan],
+            [1, 2, 20],
+            [1, 2, np.nan],
+            [1, 3, np.nan],
+            [1, 3, np.nan],
+        ],
+        columns=["one", "two", "three"],
+    )
+    mdf = md.DataFrame(df1, chunk_size=3)
+    r1 = getattr(mdf.groupby(["one", "two"]), "ffill")()
+    pd.testing.assert_frame_equal(
+        r1.execute().fetch().sort_index(),
+        getattr(df1.groupby(["one", "two"]), "ffill")().sort_index(),
+    )
+
+    r2 = getattr(mdf.groupby("two"), "bfill")()
+    pd.testing.assert_frame_equal(
+        r2.execute().fetch().sort_index(),
+        getattr(df1.groupby("two"), "bfill")().sort_index(),
+    )
+
+    r3 = getattr(mdf.groupby("one"), "fillna")(5)
+    pd.testing.assert_frame_equal(
+        r3.execute().fetch().sort_index(),
+        getattr(df1.groupby("one"), "fillna")(5).sort_index(),
+    )
+
+    r4 = getattr(mdf.groupby("two"), "backfill")()
+    pd.testing.assert_frame_equal(
+        r4.execute().fetch().sort_index(),
+        getattr(df1.groupby("two"), "backfill")().sort_index(),
+    )
+
+    s1 = pd.Series([4, 3, 9, np.nan, np.nan, 7, 10, 8, 1, 6])
+    ms1 = md.Series(s1, chunk_size=3)
+
+    r1 = getattr(ms1.groupby(lambda x: x % 2), "ffill")()
+    pd.testing.assert_series_equal(
+        r1.execute().fetch().sort_index(),
+        getattr(s1.groupby(lambda x: x % 2), "ffill")().sort_index(),
+    )
+
+    r2 = getattr(ms1.groupby(lambda x: x % 2), "bfill")()
+    pd.testing.assert_series_equal(
+        r2.execute().fetch().sort_index(),
+        getattr(s1.groupby(lambda x: x % 2), "bfill")().sort_index(),
+    )
+
+    r4 = getattr(ms1.groupby(lambda x: x % 2), "backfill")()
+    pd.testing.assert_series_equal(
+        r4.execute().fetch().sort_index(),
+        getattr(s1.groupby(lambda x: x % 2), "backfill")().sort_index(),
+    )
+
+
 def test_groupby_head(setup):
     df1 = pd.DataFrame(
         {