Repository: dask/dask

Issue:
[BUG] Mean groupby aggregation fails with Pandas>25.0
As discussed in [cudf#3486](https://github.com/rapidsai/cudf/pull/3486#issuecomment-570708074), Dask's `mean` groupby aggregation is failing in certain cases for Pandas>25.0.  The problem does not show up for other aggregations, because it seems to arise during the [element-wise truediv operation](https://github.com/dask/dask/blob/8530bdd0f54d7576f05c9162bb3a835864edf016/dask/dataframe/groupby.py#L1322).  

The problem shows up for (at least) the following specific case:
- Pandas>0.25
- The groupby operation is performed on multiple columns, with one column being `pd.Categorical`
- The aggregation uses `split_out>1`

Here is the reproducer:

```python
import pandas as pd
import dask.dataframe as dd

column = ["b", "d"]
df = pd.DataFrame(
    {
        "a": [0, 1, 2, 3, 4, 5, 6, 7],
        "b": [1, 0, 0, 2, 1, 1, 2, 0],
        "c": ["dog", "cat", "cat", "dog", "dog", "dog", "cat", "bird"],
    }
).fillna(0)
df["d"] = df["c"].astype("category")
ddf = dd.from_pandas(df, npartitions=3)

df1 = ddf.groupby(column).a.sum(split_out=2)
df2 = ddf.groupby(column).a.count(split_out=2)
dfm = df1 / df2

dfm.compute().sort_values().dropna()
```

For Pandas 0.24.2, I get the following output:
```
b  d   
0  cat     1.5
1  dog     3.0
2  dog     3.0
   cat     6.0
0  bird    7.0
Name: a, dtype: float64
```
However, using Pandas 0.25.3 leads to an error...

<details>
<summary> Click here for Error/Result </summary>
<br>

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-1-79144d6d3e36> in <module>
     18 dfm = df1 / df2
     19 
---> 20 dfm.compute().sort_values().dropna()

~/workspace/ucx_cudf++_12/dask/dask/base.py in compute(self, **kwargs)
    163         dask.base.compute
    164         """
--> 165         (result,) = compute(self, traverse=False, **kwargs)
    166         return result
    167 

~/workspace/ucx_cudf++_12/dask/dask/base.py in compute(*args, **kwargs)
    435     postcomputes = [x.__dask_postcompute__() for x in collections]
    436     results = schedule(dsk, keys, **kwargs)
--> 437     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    438 
    439 

~/workspace/ucx_cudf++_12/dask/dask/base.py in <listcomp>(.0)
    435     postcomputes = [x.__dask_postcompute__() for x in collections]
    436     results = schedule(dsk, keys, **kwargs)
--> 437     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    438 
    439 

~/workspace/ucx_cudf++_12/dask/dask/dataframe/core.py in finalize(results)
     97 
     98 def finalize(results):
---> 99     return _concat(results)
    100 
    101 

~/workspace/ucx_cudf++_12/dask/dask/dataframe/core.py in _concat(args)
     93     # this seems easier. TODO: don't do this.
     94     args2 = [i for i in args if len(i)]
---> 95     return args[0] if not args2 else methods.concat(args2, uniform=True)
     96 
     97 

~/workspace/ucx_cudf++_12/dask/dask/dataframe/methods.py in concat(dfs, axis, join, uniform, filter_warning)
    371         func = concat_dispatch.dispatch(type(dfs[0]))
    372         return func(
--> 373             dfs, axis=axis, join=join, uniform=uniform, filter_warning=filter_warning
    374         )
    375 

~/workspace/ucx_cudf++_12/dask/dask/dataframe/methods.py in concat_pandas(dfs, axis, join, uniform, filter_warning)
    417     if has_categoricalindex:
    418         dfs2 = [df.reset_index(drop=True) for df in dfs]
--> 419         ind = concat([df.index for df in dfs])
    420     else:
    421         dfs2 = dfs

~/workspace/ucx_cudf++_12/dask/dask/dataframe/methods.py in concat(dfs, axis, join, uniform, filter_warning)
    371         func = concat_dispatch.dispatch(type(dfs[0]))
    372         return func(
--> 373             dfs, axis=axis, join=join, uniform=uniform, filter_warning=filter_warning
    374         )
    375 

~/workspace/ucx_cudf++_12/dask/dask/dataframe/methods.py in concat_pandas(dfs, axis, join, uniform, filter_warning)
    395                 arrays = [
    396                     concat([i._get_level_values(n) for i in dfs])
--> 397                     for n in range(first.nlevels)
    398                 ]
    399                 return pd.MultiIndex.from_arrays(arrays, names=first.names)

~/workspace/ucx_cudf++_12/dask/dask/dataframe/methods.py in <listcomp>(.0)
    395                 arrays = [
    396                     concat([i._get_level_values(n) for i in dfs])
--> 397                     for n in range(first.nlevels)
    398                 ]
    399                 return pd.MultiIndex.from_arrays(arrays, names=first.names)

~/workspace/ucx_cudf++_12/dask/dask/dataframe/methods.py in concat(dfs, axis, join, uniform, filter_warning)
    371         func = concat_dispatch.dispatch(type(dfs[0]))
    372         return func(
--> 373             dfs, axis=axis, join=join, uniform=uniform, filter_warning=filter_warning
    374         )
    375 

~/workspace/ucx_cudf++_12/dask/dask/dataframe/methods.py in concat_pandas(dfs, axis, join, uniform, filter_warning)
    386             #     if not isinstance(dfs[i], pd.CategoricalIndex):
    387             #         dfs[i] = dfs[i].astype("category")
--> 388             return pd.CategoricalIndex(union_categoricals(dfs), name=dfs[0].name)
    389         elif isinstance(dfs[0], pd.MultiIndex):
    390             first, rest = dfs[0], dfs[1:]

~/miniconda3/envs/cudf++12/lib/python3.7/site-packages/pandas/core/dtypes/concat.py in union_categoricals(to_union, sort_categories, ignore_order)
    337             raise TypeError("all components to combine must be Categorical")
    338 
--> 339     to_union = [_maybe_unwrap(x) for x in to_union]
    340     first = to_union[0]
    341 

~/miniconda3/envs/cudf++12/lib/python3.7/site-packages/pandas/core/dtypes/concat.py in <listcomp>(.0)
    337             raise TypeError("all components to combine must be Categorical")
    338 
--> 339     to_union = [_maybe_unwrap(x) for x in to_union]
    340     first = to_union[0]
    341 

~/miniconda3/envs/cudf++12/lib/python3.7/site-packages/pandas/core/dtypes/concat.py in _maybe_unwrap(x)
    335             return x
    336         else:
--> 337             raise TypeError("all components to combine must be Categorical")
    338 
    339     to_union = [_maybe_unwrap(x) for x in to_union]

TypeError: all components to combine must be Categorical
```

</details>

The truediv operation ultimately results in the concatination of a `pd.CategoricalIndex`.  To do this, dask [uses pandas' `union_categoricals`](https://github.com/dask/dask/blob/8530bdd0f54d7576f05c9162bb3a835864edf016/dask/dataframe/methods.py#L385).  However, in this specific case, the the second component of the argument is an `"object"`-based index, rather than a `pd.CategoricalIndex`.


Discussion:

Resolving patch:
diff --git a/dask/dataframe/methods.py b/dask/dataframe/methods.py
--- a/dask/dataframe/methods.py
+++ b/dask/dataframe/methods.py
@@ -382,6 +382,9 @@ def concat_pandas(dfs, axis=0, join="outer", uniform=False, filter_warning=True)
     # Support concatenating indices along axis 0
     if isinstance(dfs[0], pd.Index):
         if isinstance(dfs[0], pd.CategoricalIndex):
+            for i in range(1, len(dfs)):
+                if not isinstance(dfs[i], pd.CategoricalIndex):
+                    dfs[i] = dfs[i].astype("category")
             return pd.CategoricalIndex(union_categoricals(dfs), name=dfs[0].name)
         elif isinstance(dfs[0], pd.MultiIndex):
             first, rest = dfs[0], dfs[1:]


Test patch:
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -2281,7 +2281,7 @@ def test_groupby_dropna_cudf(dropna, by):
         dask_result = ddf.groupby(by, dropna=dropna).e.sum()
         cudf_result = df.groupby(by, dropna=dropna).e.sum()
     if by in ["c", "d"]:
-        # Loose string/category index name in cudf...
+        # Lose string/category index name in cudf...
         dask_result = dask_result.compute()
         dask_result.index.name = cudf_result.index.name
 
@@ -2296,3 +2296,28 @@ def test_rounding_negative_var():
 
     ddf = dd.from_pandas(df, npartitions=2)
     assert_eq(ddf.groupby("ids").x.std(), df.groupby("ids").x.std())
+
+
+@pytest.mark.parametrize("split_out", [2, 3])
+@pytest.mark.parametrize("column", [["b", "c"], ["b", "d"], ["b", "e"]])
+def test_groupby_split_out_multiindex(split_out, column):
+    df = pd.DataFrame(
+        {
+            "a": np.arange(8),
+            "b": [1, 0, 0, 2, 1, 1, 2, 0],
+            "c": [0, 1] * 4,
+            "d": ["dog", "cat", "cat", "dog", "dog", "dog", "cat", "bird"],
+        }
+    ).fillna(0)
+    df["e"] = df["d"].astype("category")
+    ddf = dd.from_pandas(df, npartitions=3)
+
+    ddf_result_so1 = (
+        ddf.groupby(column).a.mean(split_out=1).compute().sort_values().dropna()
+    )
+
+    ddf_result = (
+        ddf.groupby(column).a.mean(split_out=split_out).compute().sort_values().dropna()
+    )
+
+    assert_eq(ddf_result, ddf_result_so1, check_index=False)