Repository: dask/dask
Pull request: Multi value pivot table
- [x] Tests added / passed
- [x] Passes `black dask` / `flake8 dask`

Closes #6008 

Problem statement:
dask.dataframe.pivot_table does not work with multiple values 
In pivot_table when I use multiple columns for values `values=['D','A']` I have the following error in dask.
`ValueError: 'values' must be the name of an existing column`
With the same scripts but with just one value for `values='A'` there is no error. 
Same scripts with `values=['D','A']` work well in pandas. 

```python
import pandas as pd
import dask.dataframe as ds
df = pd.DataFrame({'A':[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20],'B':['1th','2th','3th','4th','5th','1th','1th','3th','3th','3th','1th','2th','3th','4th','5th','1th','1th','3th','3th','3th'],
'C':['aa','bb','cc','aa','cc','aa','bb','cc','aa','cc','aa','bb','cc','aa','cc','aa','bb','cc','aa','cc'],
'D':[25,50,45,14,12,89,7,11,32,10,25,50,45,14,12,89,7,11,32,10]})
df
```
```python-traceback

	A	B	C	D
0	1	1th	aa	25
1	2	2th	bb	50
2	3	3th	cc	45
3	4	4th	aa	14
4	5	5th	cc	12
5	6	1th	aa	89
6	7	1th	bb	7
7	8	3th	cc	11
8	9	3th	aa	32
9	10	3th	cc	10
10	11	1th	aa	25
11	12	2th	bb	50
12	13	3th	cc	45
13	14	4th	aa	14
14	15	5th	cc	12
15	16	1th	aa	89
16	17	1th	bb	7
17	18	3th	cc	11
18	19	3th	aa	32
19	20	3th	cc	10
```
Pivot_table in pandas work:
```python
df.pivot_table(index='B',columns='C',values=['D','A'])
```
```python-traceback
	A	D
C	aa	bb	cc	aa	bb	cc
B						
1th	8.5	12.0	NaN	57.0	7.0	NaN
2th	NaN	7.0	NaN	NaN	50.0	NaN
3th	14.0	NaN	12.0	32.0	NaN	22.0
4th	9.0	NaN	NaN	14.0	NaN	NaN
5th	NaN	NaN	10.0	NaN	NaN	12.0

```
Pivot_table in Dask:
```python
dsdf = ds.from_pandas(df,npartitions=2)
dsdf['C'] = dsdf['C'].astype('category').cat.as_known()
print(dsdf.columns)
dsdf.pivot_table(index='B',columns='C',values=['D','A'])
```
```python-traceback

Index(['A', 'B', 'C', 'D'], dtype='object')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-217-28d2e140399f> in <module>
      1 dsdf['C'] = dsdf['C'].astype('category').cat.as_known()
      2 print(dsdf.columns)
----> 3 dsdf.pivot_table(index='B',columns='C',values=['D','A'])

~\Anaconda3\lib\site-packages\dask\dataframe\core.py in pivot_table(self, index, columns, values, aggfunc)
   4258 
   4259         return pivot_table(
-> 4260             self, index=index, columns=columns, values=values, aggfunc=aggfunc
   4261         )
   4262 

~\Anaconda3\lib\site-packages\dask\dataframe\reshape.py in pivot_table(df, index, columns, values, aggfunc)
    222         )
    223     if not is_scalar(values) or values is None:
--> 224         raise ValueError("'values' must be the name of an existing column")
    225     if not is_scalar(aggfunc) or aggfunc not in ("mean", "sum", "count"):
    226         raise ValueError("aggfunc must be either 'mean', 'sum' or 'count'")

ValueError: 'values' must be the name of an existing column
```

Hints from issue discussion:
Thanks for raising this issue! MultiIndexes are not supported in dask, so I don't think you can create a `pivot_table` passing a list of columns. That said, it would be more friendly to include a specific error for this case since it is supported in pandas. Would you be willing to open a PR to catch this case and return a more precise error?
Just to clarify, I think we do support MultiIndexes as long as there's a single partition. That's way `.groupby(['A', 'B']).count()` works.

I'm not familiar with the `pivot_table` implementation, but I suspect it results in a single-partition DataFrame so this may be doable. 
Oh! I guess I read https://github.com/dask/dask/issues/1493 too quickly. Sorry for the misinformation!
So to summarize @MSKazemi , pivot tables with multi-column values for `values` are not currently supported, but is in scope.  If you or anyone else who happens upon this issue wants to work on this that would be welcome.

Commit message: remove comment

Patch to review:
diff --git a/dask/dataframe/reshape.py b/dask/dataframe/reshape.py
index e775f06e46d..c1eb32618aa 100644
--- a/dask/dataframe/reshape.py
+++ b/dask/dataframe/reshape.py
@@ -6,6 +6,7 @@
 from .utils import is_categorical_dtype, is_scalar, has_known_categories
 from ..utils import M
 import sys
+from pandas.api.types import is_list_like
 
 ###############################################################
 # Dummies
@@ -186,7 +187,8 @@ def pivot_table(df, index=None, columns=None, values=None, aggfunc="mean"):
     """
     Create a spreadsheet-style pivot table as a DataFrame. Target ``columns``
     must have category dtype to infer result's ``columns``.
-    ``index``, ``columns``, ``values`` and ``aggfunc`` must be all scalar.
+    ``index``, ``columns``, and ``aggfunc`` must be all scalar.
+    ``values`` can be scalar or list-like.
 
     Parameters
     ----------
@@ -195,8 +197,8 @@ def pivot_table(df, index=None, columns=None, values=None, aggfunc="mean"):
         column to be index
     columns : scalar
         column to be columns
-    values : scalar
-        column to aggregate
+    values : scalar or list(scalar)
+        column(s) to aggregate
     aggfunc : {'mean', 'sum', 'count'}, default 'mean'
 
     Returns
@@ -220,14 +222,23 @@ def pivot_table(df, index=None, columns=None, values=None, aggfunc="mean"):
             "`df[columns].cat.as_known()` beforehand to ensure "
             "known categories"
         )
-    if not is_scalar(values) or values is None:
-        raise ValueError("'values' must be the name of an existing column")
+    if not (
+        is_list_like(values)
+        and all([is_scalar(v) for v in values])
+        or is_scalar(values)
+    ):
+        raise ValueError("'values' must be refer to existing an column or columns")
     if not is_scalar(aggfunc) or aggfunc not in ("mean", "sum", "count"):
         raise ValueError("aggfunc must be either 'mean', 'sum' or 'count'")
 
-    # _emulate can't work for empty data
-    # the result must have CategoricalIndex columns
-    new_columns = pd.CategoricalIndex(df[columns].cat.categories, name=columns)
+    columns_contents = pd.CategoricalIndex(df[columns].cat.categories, name=columns)
+    if is_scalar(values):
+        new_columns = columns_contents
+    else:
+        new_columns = pd.MultiIndex.from_product(
+            (values, columns_contents), names=[None, columns]
+        )
+
     meta = pd.DataFrame(
         columns=new_columns, dtype=np.float64, index=pd.Index(df._meta[index])
     )
diff --git a/dask/dataframe/tests/test_reshape.py b/dask/dataframe/tests/test_reshape.py
index 5238bcf04c5..bc4481578e1 100644
--- a/dask/dataframe/tests/test_reshape.py
+++ b/dask/dataframe/tests/test_reshape.py
@@ -170,19 +170,21 @@ def test_get_dummies_errors():
         dd.get_dummies(ddf.x)
 
 
-@pytest.mark.parametrize("aggfunc", ["mean", "sum", "count"])
-def test_pivot_table(aggfunc):
+@pytest.mark.parametrize("values", ["B", ["B"], ["B", "D"]])
+@pytest.mark.parametrize("aggfunc", ["mean"])
+def test_pivot_table(values, aggfunc):
     df = pd.DataFrame(
         {
             "A": np.random.choice(list("XYZ"), size=100),
             "B": np.random.randn(100),
             "C": pd.Categorical(np.random.choice(list("abc"), size=100)),
+            "D": np.random.randn(100),
         }
     )
     ddf = dd.from_pandas(df, 5)
 
-    res = dd.pivot_table(ddf, index="A", columns="C", values="B", aggfunc=aggfunc)
-    exp = pd.pivot_table(df, index="A", columns="C", values="B", aggfunc=aggfunc)
+    res = dd.pivot_table(ddf, index="A", columns="C", values=values, aggfunc=aggfunc)
+    exp = pd.pivot_table(df, index="A", columns="C", values=values, aggfunc=aggfunc)
     if aggfunc == "count":
         # dask result cannot be int64 dtype depending on divisions because of NaN
         exp = exp.astype(np.float64)
@@ -190,8 +192,8 @@ def test_pivot_table(aggfunc):
     assert_eq(res, exp)
 
     # method
-    res = ddf.pivot_table(index="A", columns="C", values="B", aggfunc=aggfunc)
-    exp = df.pivot_table(index="A", columns="C", values="B", aggfunc=aggfunc)
+    res = ddf.pivot_table(index="A", columns="C", values=values, aggfunc=aggfunc)
+    exp = df.pivot_table(index="A", columns="C", values=values, aggfunc=aggfunc)
     if aggfunc == "count":
         # dask result cannot be int64 dtype depending on divisions because of NaN
         exp = exp.astype(np.float64)
@@ -249,10 +251,6 @@ def test_pivot_table_errors():
     with pytest.raises(ValueError) as err:
         dd.pivot_table(ddf, index="A", columns=["C"], values="B")
     assert msg in str(err.value)
-    msg = "'values' must be the name of an existing column"
-    with pytest.raises(ValueError) as err:
-        dd.pivot_table(ddf, index="A", columns="C", values=["B"])
-    assert msg in str(err.value)
 
     msg = "aggfunc must be either 'mean', 'sum' or 'count'"
     with pytest.raises(ValueError) as err:
diff --git a/dask/dataframe/utils.py b/dask/dataframe/utils.py
index 47e92a77d6b..84c32adb38f 100644
--- a/dask/dataframe/utils.py
+++ b/dask/dataframe/utils.py
@@ -741,7 +741,7 @@ def _check_dask(dsk, check_names=True, check_dtypes=True, result=None):
             )
         elif isinstance(dsk, dd.DataFrame):
             assert "DataFrame" in type(result).__name__, type(result)
-            assert isinstance(dsk.columns, pd.Index), type(dsk.columns)
+            assert isinstance(dsk.columns, (pd.Index, pd.MultiIndex)), type(dsk.columns)
             assert type(dsk._meta) == type(result), type(dsk._meta)
             if check_names:
                 tm.assert_index_equal(dsk.columns, result.columns)


Code review comments:
diff_hunk: @@ -220,14 +222,23 @@ def pivot_table(df, index=None, columns=None, values=None, aggfunc="mean"):
             "`df[columns].cat.as_known()` beforehand to ensure "
             "known categories"
         )
-    if not is_scalar(values) or values is None:
-        raise ValueError("'values' must be the name of an existing column")
+    if not (
+        is_list_like(values)
+        and all([is_scalar(v) for v in values])
+        or is_scalar(values)
+    ):
+        raise ValueError("'values' must be refer to existing an column or columns")
path: dask/dataframe/reshape.py
text: @user1:
```suggestion
        raise ValueError("'values' must refer to existing column or columns")
```

diff_hunk: @@ -220,14 +222,23 @@ def pivot_table(df, index=None, columns=None, values=None, aggfunc="mean"):
             "`df[columns].cat.as_known()` beforehand to ensure "
             "known categories"
         )
-    if not is_scalar(values) or values is None:
-        raise ValueError("'values' must be the name of an existing column")
+    if not (
+        is_list_like(values)
+        and all([is_scalar(v) for v in values])
+        or is_scalar(values)
+    ):
+        raise ValueError("'values' must be refer to existing an column or columns")
     if not is_scalar(aggfunc) or aggfunc not in ("mean", "sum", "count"):
         raise ValueError("aggfunc must be either 'mean', 'sum' or 'count'")
 
-    # _emulate can't work for empty data
-    # the result must have CategoricalIndex columns
-    new_columns = pd.CategoricalIndex(df[columns].cat.categories, name=columns)
+    columns_contents = pd.CategoricalIndex(df[columns].cat.categories, name=columns)
+    if is_scalar(values):
+        new_columns = columns_contents
+    else:
+        new_columns = pd.MultiIndex.from_product(
+            (values, columns_contents), names=[None, columns]
path: dask/dataframe/reshape.py
text: @user1:
I think the values should be sorted before passing off to `pd.MultiIndex` otherwise the order of the `._meta` won't match the order of the output. 

Compare:
```python
import pandas as pd
import dask.dataframe as ds
df = pd.DataFrame({
    'A':[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20],
    'B':['1th','2th','3th','4th','5th','1th','1th','3th','3th','3th','1th','2th','3th','4th','5th','1th','1th','3th','3th','3th'],
    'C':['aa','bb','cc','aa','cc','aa','bb','cc','aa','cc','aa','bb','cc','aa','cc','aa','bb','cc','aa','cc'],
    'D':[25,50,45,14,12,89,7,11,32,10,25,50,45,14,12,89,7,11,32,10]
})
df.pivot_table(index='B',columns='C',values=['D','A'])
```
to 
```python
dsdf = ds.from_pandas(df,npartitions=2)
dsdf['C'] = dsdf['C'].astype('category').cat.as_known()
dsdf.pivot_table(index='B',columns='C',values=['D','A'])
```



@author:
Not sure I follow. What difference should i be looking for in those two snippets?

@user1:
Sorry shoul have included output :) If you execute them you'll see that in the pandas version the columns are sorted so the output is `A`, `D`:
```
        A                 D            
C      aa    bb    cc    aa    bb    cc
B                                      
1th   8.5  12.0   NaN  57.0   7.0   NaN
2th   NaN   7.0   NaN   NaN  50.0   NaN
3th  14.0   NaN  12.0  32.0   NaN  22.0
4th   9.0   NaN   NaN  14.0   NaN   NaN
5th   NaN   NaN  10.0   NaN   NaN  12.0
```
whereas in the dask one the columns are not sorted:

```
Dask DataFrame Structure:
                     D                          A                  
                    aa       bb       cc       aa       bb       cc
npartitions=1                                                      
               float64  float64  float64  float64  float64  float64
                   ...      ...      ...      ...      ...      ...
Dask Name: truediv, 17 tasks
```

So I am suggesting:
```suggestion
            (sorted(values), columns_contents), names=[None, columns]
```

diff_hunk: @@ -170,28 +170,30 @@ def test_get_dummies_errors():
         dd.get_dummies(ddf.x)
 
 
-@pytest.mark.parametrize("aggfunc", ["mean", "sum", "count"])
-def test_pivot_table(aggfunc):
+@pytest.mark.parametrize("values", ["B", ["B"], ["B", "D"]])
+@pytest.mark.parametrize("aggfunc", ["mean"])
path: dask/dataframe/tests/test_reshape.py
text: @user2:
I don't think we shouldn't be getting rid of tests: 

```suggestion
@user1.mark.parametrize("aggfunc", ["mean", "sum", "count"])
```

@author:
you're right! That was a mistake on my part.

diff_hunk: @@ -741,7 +741,7 @@ def _check_dask(dsk, check_names=True, check_dtypes=True, result=None):
             )
         elif isinstance(dsk, dd.DataFrame):
             assert "DataFrame" in type(result).__name__, type(result)
-            assert isinstance(dsk.columns, pd.Index), type(dsk.columns)
+            assert isinstance(dsk.columns, (pd.Index, pd.MultiIndex)), type(dsk.columns)
path: dask/dataframe/utils.py
text: @user1:
`pd.MultiIndex` inherits from `pd.Index` so this change isn't needed.

@author:
Yeah you are right. I'll revert this.

diff_hunk: @@ -170,28 +170,30 @@ def test_get_dummies_errors():
         dd.get_dummies(ddf.x)
 
 
-@pytest.mark.parametrize("aggfunc", ["mean", "sum", "count"])
-def test_pivot_table(aggfunc):
+@pytest.mark.parametrize("values", ["B", ["B"], ["B", "D"]])
+@pytest.mark.parametrize("aggfunc", ["mean"])
path: dask/dataframe/tests/test_reshape.py
text: @user1:
Any reason why sum/count were dropped ?

@author:
careless mistake, will put those back!