Repository: dask/dask
Pull request: Implement {Series,DataFrame}GroupBy `fillna` methods
- [x] Closes #8708
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`

TODO before merging:
- [x] Add docstring
- [x] Open new issue for implementing value = dict/Series/DataFrame (done: #8922)

Problem statement:
Missing Built-in methods for SeriesGroupBy `ffill`
Some of the "built-in" methods for dask SeriesGroupBy objects are missing.

In pandas we could do something like...

```
import numpy as np
import pandas as pd

df = pd.DataFrame({
    'a_cat': list('aabbbaccc'),
    'b_num': [np.nan if i%3!=0 else (i+1) for i in range(9)]
})

df.groupby('a_cat')['b_num'].ffill()
```

But when I try to do it in Dask..

```
import dask.dataframe as dd

ddf = dd.from_pandas(df)
ddf.groupby('a_cat')['b_num'].ffill()
```
I receive an `AttributeError: 'SeriesGroupBy' object has no attribute 'ffill'`

---

I'm hoping this can be implemented along with the other `fillna` methods. Thanks.

Hints from issue discussion:
Thanks for opening this issue @CJC-ds if you would like to work on it please feel free to open a pull request!
I'd like to take a pass at this, I'll open a PR soon :)

Commit message: implement axis=1 using apply()

Patch to review:
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
index 780ccf36298..f8b372c8c3d 100644
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -1023,6 +1023,15 @@ def _cumcount_aggregate(a, b, fill_value=None):
     return a.add(b, fill_value=fill_value) + 1
 
 
+def _fillna_groups(groups, axis, **kwargs):
+    if axis == 0:
+        return groups.fillna(axis=axis, **kwargs)
+    else:
+        return groups.drop(columns=groups.columns[0], axis=1).fillna(
+            axis=axis, **kwargs
+        )
+
+
 class _GroupBy:
     """Superclass for DataFrameGroupBy and SeriesGroupBy
 
@@ -1984,6 +1993,45 @@ def rolling(self, window, min_periods=None, center=False, win_type=None, axis=0)
             axis=axis,
         )
 
+    def fillna(self, value=None, method=None, limit=None, axis=None):
+        if not np.isscalar(value) and value is not None:
+            raise NotImplementedError(
+                "groupby-fillna with value=dict/Series/DataFrame is currently not supported"
+            )
+        axis = DataFrame._validate_axis(axis)
+        if value is not None or axis == 0:
+            meta = self._meta_nonempty.transform(
+                _fillna_groups, value=value, method=method, limit=limit, axis=0
+            )
+            return self.transform(
+                _fillna_groups,
+                value=value,
+                method=method,
+                limit=limit,
+                axis=0,
+                meta=meta,
+            )
+        else:
+            meta = self._meta_nonempty.apply(
+                _fillna_groups, value=value, method=method, limit=limit, axis=1
+            )
+            return self.apply(
+                _fillna_groups,
+                value=value,
+                method=method,
+                limit=limit,
+                axis=1,
+                meta=meta,
+            )
+
+    @derived_from(pd.core.groupby.GroupBy)
+    def ffill(self, limit=None):
+        return self.fillna(method="ffill", limit=limit)
+
+    @derived_from(pd.core.groupby.GroupBy)
+    def bfill(self, limit=None):
+        return self.fillna(method="bfill", limit=limit)
+
 
 class DataFrameGroupBy(_GroupBy):
     _token_prefix = "dataframe-groupby-"
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
index f8b3096a55a..1d8c42d567d 100644
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -1091,6 +1091,84 @@ def test_aggregate_dask():
                 assert len(other.dask) == len(result2.dask)
 
 
+@pytest.mark.parametrize("axis", [0, 1])
+@pytest.mark.parametrize("method", ["ffill", "bfill"])
+@pytest.mark.parametrize("limit", [None, 1, 4])
+def test_fillna(axis, method, limit):
+    df = pd.DataFrame(
+        {
+            "A": [1, 1, 2, 2],
+            "B": [np.nan, 3, np.nan, 4],
+            "C": [4, np.nan, 5, np.nan],
+            "D": [6, np.nan, 7, np.nan],
+        }
+    )
+    ddf = dd.from_pandas(df, npartitions=2)
+    assert_eq(
+        df.groupby("A").fillna(0, axis=axis),
+        ddf.groupby("A").fillna(0, axis=axis),
+    )
+    assert_eq(
+        df.groupby("A").B.fillna(0),
+        ddf.groupby("A").B.fillna(0),
+    )
+    assert_eq(
+        df.groupby("A").fillna(method=method, limit=limit, axis=axis),
+        ddf.groupby("A").fillna(method=method, limit=limit, axis=axis),
+    )
+
+    with pytest.raises(NotImplementedError):
+        ddf.groupby("A").fillna({"A": 0})
+
+    with pytest.raises(NotImplementedError):
+        ddf.groupby("A").fillna(pd.Series(dtype=int))
+
+    with pytest.raises(NotImplementedError):
+        ddf.groupby("A").fillna(pd.DataFrame)
+
+
+@pytest.mark.parametrize("limit", [None, 1, 4])
+def test_ffill(limit):
+    df = pd.DataFrame(
+        {
+            "A": [1, 1, 2, 2],
+            "B": [np.nan, 3, np.nan, np.nan],
+            "C": [4, np.nan, 5, np.nan],
+            "D": [6, np.nan, 7, np.nan],
+        }
+    )
+    ddf = dd.from_pandas(df, npartitions=2)
+    assert_eq(
+        df.groupby("A").ffill(limit=limit),
+        ddf.groupby("A").ffill(limit=limit),
+    )
+    assert_eq(
+        df.groupby("A").B.ffill(limit=limit),
+        ddf.groupby("A").B.ffill(limit=limit),
+    )
+
+
+@pytest.mark.parametrize("limit", [None, 1, 4])
+def test_bfill(limit):
+    df = pd.DataFrame(
+        {
+            "A": [1, 1, 2, 2],
+            "B": [np.nan, 3, np.nan, np.nan],
+            "C": [np.nan, 4, np.nan, 5],
+            "D": [np.nan, 6, np.nan, 7],
+        }
+    )
+    ddf = dd.from_pandas(df, npartitions=2)
+    assert_eq(
+        df.groupby("A").bfill(limit=limit),
+        ddf.groupby("A").bfill(limit=limit),
+    )
+    assert_eq(
+        df.groupby("A").B.bfill(limit=limit),
+        ddf.groupby("A").B.bfill(limit=limit),
+    )
+
+
 @pytest.mark.parametrize(
     "grouper",
     [


Code review comments:
diff_hunk: @@ -1023,6 +1023,15 @@ def _cumcount_aggregate(a, b, fill_value=None):
     return a.add(b, fill_value=fill_value) + 1
 
 
+def _fillna_groups(groups, axis, **kwargs):
+    if axis == 0:
+        return groups.fillna(axis=axis, **kwargs)
+    else:
+        return groups.drop(columns=groups.columns[0], axis=1).fillna(
path: dask/dataframe/groupby.py
text: @author:
@user1-r-rose Is this way of dropping columns alright? I noticed there is a `drop_columns` dataframe-method, would it make sense to use that instead?

https://github.com/dask/dask/blob/1b1a96b4ddf31d0879e718be463a731af7f1d086/dask/dataframe/methods.py#L379-L382

@author:
And, I'm also not sure about bundling drop-column logic in the top-level `def _fillna_groups()`. is this OK, or would it be better to do this in `fillna()` (line 2014)?

@author:
Just to note: I'm yet to add doctrings and comments :)

@user2:
I think using `df.drop` directly is okay. It looks like the `drop_columns` utility is primarily used for `map_partitions`, which doesn't really apply here.

However, I think always dropping the zeroth column doesn't quite do what we want:
1. The column we want to drop might not be the first one.
2. The user might be grouping by multiple columns

What if instead we also included a `by` argument in `_fillna_group` for the column(s) by which we are grouping, the dataframe can pass in `self.by`, and then we can drop them?

@author:
> What if instead we also included a by argument in _fillna_group for the column(s) by which we are grouping, the dataframe can pass in self.by, and then we can drop them?

Yes! Thanks for this suggestion!

diff_hunk: @@ -1984,6 +1993,45 @@ def rolling(self, window, min_periods=None, center=False, win_type=None, axis=0)
             axis=axis,
         )
 
+    def fillna(self, value=None, method=None, limit=None, axis=None):
+        if not np.isscalar(value) and value is not None:
+            raise NotImplementedError(
+                "groupby-fillna with value=dict/Series/DataFrame is currently not supported"
+            )
+        axis = DataFrame._validate_axis(axis)
+        if value is not None or axis == 0:
path: dask/dataframe/groupby.py
text: @author:
If we use `value` with `fillna` the `axis` doesn't really matter because the output remains the same.

But, that said, I feel if a user specifies `axis=1` they probably assume Dask to be working row-wise internally, so maybe we should use the `apply` implementation whenever `axis=1` irrespective of `value`?

@user1:
I think your choice here is fine, I don't want to belabor the point of `apply` here, it's more of an implementation detail, and shouldn't really affect the user much.

diff_hunk: @@ -1023,6 +1023,15 @@ def _cumcount_aggregate(a, b, fill_value=None):
     return a.add(b, fill_value=fill_value) + 1
 
 
+def _fillna_groups(groups, axis, **kwargs):
path: dask/dataframe/groupby.py
text: @user1:
Nit: this is applying to a single group, so I think `group` would be a more appropriate name

@user1:
Rather than a `**kwargs`, I'd just include the rest of the args explicitly

@author:
Sounds good to me, shall do!