Repository: pandas-dev/pandas
Pull request: ENH: Allow multi values for index and columns in df.pivot
- [x] closes #21425
- [x] tests added / passed
- [x] passes `black pandas`
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [x] whatsnew entry


Problem statement:
BUG: DataFrame.pivot fails on multiple columns to set as index
I am not sure anymore if `DataFrame.pivot` actually supports multiple columns to set as `index`/`columns` (it should), but in any case this error is very confusing:

```
In [1]: df = pd.DataFrame({'lev1': [1, 1, 1, 1,2, 2, 2,2], 'lev2': [1, 1, 2, 2, 1, 1, 2, 2], 'lev3': [1, 2, 1, 2, 1, 2, 1, 2], 'values': range(8)})

In [2]: df
Out[2]: 
   lev1  lev2  lev3  values
0     1     1     1       0
1     1     1     2       1
2     1     2     1       2
3     1     2     2       3
4     2     1     1       4
5     2     1     2       5
6     2     2     1       6
7     2     2     2       7

In [3]: df.pivot(index=['lev1', 'lev2'], columns='lev3', values='values')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-3-2fef29f9fd39> in <module>()
----> 1 df.pivot(index=['lev1', 'lev2'], columns='lev3', values='values')

~/scipy/pandas/pandas/core/frame.py in pivot(self, index, columns, values)
   5191         """
   5192         from pandas.core.reshape.reshape import pivot
-> 5193         return pivot(self, index=index, columns=columns, values=values)
   5194 
   5195     _shared_docs['pivot_table'] = """

~/scipy/pandas/pandas/core/reshape/reshape.py in pivot(self, index, columns, values)
    406         else:
    407             indexed = self._constructor_sliced(self[values].values,
--> 408                                                index=index)
    409     return indexed.unstack(columns)
    410 

~/scipy/pandas/pandas/core/series.py in __init__(self, data, index, dtype, name, copy, fastpath)
    260                             'Length of passed values is {val}, '
    261                             'index implies {ind}'
--> 262                             .format(val=len(data), ind=len(index)))
    263                 except TypeError:
    264                     pass

ValueError: Length of passed values is 8, index implies 2

```

Hints from issue discussion:
Shouldn't pd.pivot_table be used if multiple columns for the index is desired, or if a pivoted index-column pair has more than 1 value and thus requires aggregation? 
Why does pivot exist actually? 
Is there anything pivot can do that pivot_table cannot?
I'd like to complete this (I had this issue the other day, and just switched to `pd.pivot_table()`). Is it ok if I work on this since the other PRs are closed?

Looks like just a few changes needed in thoo's branch.
Will I be duplicating someone else's efforts if I begin work on this issue?
Since it is already some weeks that @srcole commented about this (btw, sorry for not answering to your comment back then!), I think you are welcome to do a PR.

Commit message: fix pep8

Patch to review:
diff --git a/doc/source/whatsnew/v1.1.0.rst b/doc/source/whatsnew/v1.1.0.rst
index 1cd325dad9f07..dee5b2bee2b40 100644
--- a/doc/source/whatsnew/v1.1.0.rst
+++ b/doc/source/whatsnew/v1.1.0.rst
@@ -144,6 +144,7 @@ Reshaping
 - Bug in :meth:`DataFrame.pivot_table` when only MultiIndexed columns is set (:issue:`17038`)
 - Fix incorrect error message in :meth:`DataFrame.pivot` when ``columns`` is set to ``None``. (:issue:`30924`)
 - Bug in :func:`crosstab` when inputs are two Series and have tuple names, the output will keep dummy MultiIndex as columns. (:issue:`18321`)
+- :meth:`DataFrame.pivot` can now take lists for ``index`` and ``columns`` arguments (:issue:`21425`)
 
 
 Sparse
diff --git a/pandas/core/frame.py b/pandas/core/frame.py
index 1a49388d81243..dad71b7e8ba69 100644
--- a/pandas/core/frame.py
+++ b/pandas/core/frame.py
@@ -5847,11 +5847,19 @@ def groupby(
 
         Parameters
         ----------%s
-        index : str or object, optional
+        index : str or object or a list of str, optional
             Column to use to make new frame's index. If None, uses
             existing index.
-        columns : str or object
+
+            .. versionchanged:: 1.0.0
+               Also accept list of index names.
+
+        columns : str or object or a list of str
             Column to use to make new frame's columns.
+
+            .. versionchanged:: 1.0.0
+               Also accept list of columns names.
+
         values : str, object or a list of the previous, optional
             Column(s) to use for populating new frame's values. If not
             specified, all remaining columns will be used and the result will
@@ -5918,6 +5926,38 @@ def groupby(
         one   1  2  3   x  y  z
         two   4  5  6   q  w  t
 
+        You could also assign a list of column names or a list of index names.
+
+        >>> df = pd.DataFrame({
+        ...        "lev1": [1, 1, 1, 2, 2, 2],
+        ...        "lev2": [1, 1, 2, 1, 1, 2],
+        ...        "lev3": [1, 2, 1, 2, 1, 2],
+        ...        "lev4": [1, 2, 3, 4, 5, 6],
+        ...        "values": [0, 1, 2, 3, 4, 5]})
+        >>> df
+            lev1 lev2 lev3 lev4 values
+        0   1    1    1    1    0
+        1   1    1    2    2    1
+        2   1    2    1    3    2
+        3   2    1    2    4    3
+        4   2    1    1    5    4
+        5   2    2    2    6    5
+
+        >>> df.pivot(index="lev1", columns=["lev2", "lev3"],values="values")
+        lev2    1         2
+        lev3    1    2    1    2
+        lev1
+        1     0.0  1.0  2.0  NaN
+        2     4.0  3.0  NaN  5.0
+
+        >>> df.pivot(index=["lev1", "lev2"], columns=["lev3"],values="values")
+              lev3    1    2
+        lev1  lev2
+           1     1  0.0  1.0
+                 2  2.0  NaN
+           2     1  4.0  3.0
+                 2  NaN  5.0
+
         A ValueError is raised if there are any duplicates.
 
         >>> df = pd.DataFrame({"foo": ['one', 'one', 'two', 'two'],
diff --git a/pandas/core/reshape/pivot.py b/pandas/core/reshape/pivot.py
index 930ff5f454a7b..df3b5df3c8999 100644
--- a/pandas/core/reshape/pivot.py
+++ b/pandas/core/reshape/pivot.py
@@ -431,17 +431,31 @@ def _convert_by(by):
 def pivot(data: "DataFrame", index=None, columns=None, values=None) -> "DataFrame":
     if columns is None:
         raise TypeError("pivot() missing 1 required argument: 'columns'")
+    columns = columns if is_list_like(columns) else [columns]
 
     if values is None:
-        cols = [columns] if index is None else [index, columns]
+        cols: List[str] = []
+        if index is None:
+            pass
+        elif is_list_like(index):
+            cols = list(index)
+        else:
+            cols = [index]
+        cols.extend(columns)
+
         append = index is None
         indexed = data.set_index(cols, append=append)
     else:
         if index is None:
-            index = data.index
+            index = [Series(data.index, name=data.index.name)]
+        elif is_list_like(index):
+            index = [data[idx] for idx in index]
         else:
-            index = data[index]
-        index = MultiIndex.from_arrays([index, data[columns]])
+            index = [data[index]]
+
+        data_columns = [data[col] for col in columns]
+        index.extend(data_columns)
+        index = MultiIndex.from_arrays(index)
 
         if is_list_like(values) and not isinstance(values, tuple):
             # Exclude tuple because it is seen as a single column name
diff --git a/pandas/tests/reshape/test_pivot.py b/pandas/tests/reshape/test_pivot.py
index 44073f56abfa1..0124d96230be5 100644
--- a/pandas/tests/reshape/test_pivot.py
+++ b/pandas/tests/reshape/test_pivot.py
@@ -1007,6 +1007,192 @@ def test_pivot_no_level_overlap(self):
         expected = grouped.unstack("b").unstack("c").dropna(axis=1, how="all")
         tm.assert_frame_equal(table, expected)
 
+    @pytest.mark.parametrize(
+        "input_index, input_columns, input_values, "
+        "expected_values, expected_columns, expected_index",
+        [
+            (
+                ["lev4"],
+                "lev3",
+                "values",
+                [
+                    [0.0, np.nan],
+                    [np.nan, 1.0],
+                    [2.0, np.nan],
+                    [np.nan, 3.0],
+                    [4.0, np.nan],
+                    [np.nan, 5.0],
+                    [6.0, np.nan],
+                    [np.nan, 7.0],
+                ],
+                Index([1, 2], name="lev3"),
+                Index([1, 2, 3, 4, 5, 6, 7, 8], name="lev4"),
+            ),
+            (
+                ["lev4"],
+                "lev3",
+                None,
+                [
+                    [1.0, np.nan, 1.0, np.nan, 0.0, np.nan],
+                    [np.nan, 1.0, np.nan, 1.0, np.nan, 1.0],
+                    [1.0, np.nan, 2.0, np.nan, 2.0, np.nan],
+                    [np.nan, 1.0, np.nan, 2.0, np.nan, 3.0],
+                    [2.0, np.nan, 1.0, np.nan, 4.0, np.nan],
+                    [np.nan, 2.0, np.nan, 1.0, np.nan, 5.0],
+                    [2.0, np.nan, 2.0, np.nan, 6.0, np.nan],
+                    [np.nan, 2.0, np.nan, 2.0, np.nan, 7.0],
+                ],
+                MultiIndex.from_tuples(
+                    [
+                        ("lev1", 1),
+                        ("lev1", 2),
+                        ("lev2", 1),
+                        ("lev2", 2),
+                        ("values", 1),
+                        ("values", 2),
+                    ],
+                    names=[None, "lev3"],
+                ),
+                Index([1, 2, 3, 4, 5, 6, 7, 8], name="lev4"),
+            ),
+            (
+                ["lev1", "lev2"],
+                "lev3",
+                "values",
+                [[0, 1], [2, 3], [4, 5], [6, 7]],
+                Index([1, 2], name="lev3"),
+                MultiIndex.from_tuples(
+                    [(1, 1), (1, 2), (2, 1), (2, 2)], names=["lev1", "lev2"]
+                ),
+            ),
+            (
+                ["lev1", "lev2"],
+                "lev3",
+                None,
+                [[1, 2, 0, 1], [3, 4, 2, 3], [5, 6, 4, 5], [7, 8, 6, 7]],
+                MultiIndex.from_tuples(
+                    [("lev4", 1), ("lev4", 2), ("values", 1), ("values", 2)],
+                    names=[None, "lev3"],
+                ),
+                MultiIndex.from_tuples(
+                    [(1, 1), (1, 2), (2, 1), (2, 2)], names=["lev1", "lev2"]
+                ),
+            ),
+        ],
+    )
+    def test_pivot_list_like_index(
+        self,
+        input_index,
+        input_columns,
+        input_values,
+        expected_values,
+        expected_columns,
+        expected_index,
+    ):
+        # GH 21425, test when index is given a list
+        df = pd.DataFrame(
+            {
+                "lev1": [1, 1, 1, 1, 2, 2, 2, 2],
+                "lev2": [1, 1, 2, 2, 1, 1, 2, 2],
+                "lev3": [1, 2, 1, 2, 1, 2, 1, 2],
+                "lev4": [1, 2, 3, 4, 5, 6, 7, 8],
+                "values": [0, 1, 2, 3, 4, 5, 6, 7],
+            }
+        )
+
+        result = df.pivot(index=input_index, columns=input_columns, values=input_values)
+        expected = pd.DataFrame(
+            expected_values, columns=expected_columns, index=expected_index
+        )
+        tm.assert_frame_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "input_index, input_columns, input_values, "
+        "expected_values, expected_columns, expected_index",
+        [
+            (
+                "lev4",
+                ["lev3"],
+                "values",
+                [
+                    [0.0, np.nan],
+                    [np.nan, 1.0],
+                    [2.0, np.nan],
+                    [np.nan, 3.0],
+                    [4.0, np.nan],
+                    [np.nan, 5.0],
+                    [6.0, np.nan],
+                    [np.nan, 7.0],
+                ],
+                Index([1, 2], name="lev3"),
+                Index([1, 2, 3, 4, 5, 6, 7, 8], name="lev4"),
+            ),
+            (
+                ["lev1", "lev2"],
+                ["lev3"],
+                "values",
+                [[0, 1], [2, 3], [4, 5], [6, 7]],
+                Index([1, 2], name="lev3"),
+                MultiIndex.from_tuples(
+                    [(1, 1), (1, 2), (2, 1), (2, 2)], names=["lev1", "lev2"]
+                ),
+            ),
+            (
+                ["lev1"],
+                ["lev2", "lev3"],
+                "values",
+                [[0, 1, 2, 3], [4, 5, 6, 7]],
+                MultiIndex.from_tuples(
+                    [(1, 1), (1, 2), (2, 1), (2, 2)], names=["lev2", "lev3"]
+                ),
+                Index([1, 2], name="lev1"),
+            ),
+            (
+                ["lev1", "lev2"],
+                ["lev3", "lev4"],
+                "values",
+                [
+                    [0.0, 1.0, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
+                    [np.nan, np.nan, 2.0, 3.0, np.nan, np.nan, np.nan, np.nan],
+                    [np.nan, np.nan, np.nan, np.nan, 4.0, 5.0, np.nan, np.nan],
+                    [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 6.0, 7.0],
+                ],
+                MultiIndex.from_tuples(
+                    [(1, 1), (2, 2), (1, 3), (2, 4), (1, 5), (2, 6), (1, 7), (2, 8)],
+                    names=["lev3", "lev4"],
+                ),
+                MultiIndex.from_tuples(
+                    [(1, 1), (1, 2), (2, 1), (2, 2)], names=["lev1", "lev2"]
+                ),
+            ),
+        ],
+    )
+    def test_pivot_list_like_columns(
+        self,
+        input_index,
+        input_columns,
+        input_values,
+        expected_values,
+        expected_columns,
+        expected_index,
+    ):
+        # GH 21425, test when columns is given a list
+        df = pd.DataFrame(
+            {
+                "lev1": [1, 1, 1, 1, 2, 2, 2, 2],
+                "lev2": [1, 1, 2, 2, 1, 1, 2, 2],
+                "lev3": [1, 2, 1, 2, 1, 2, 1, 2],
+                "lev4": [1, 2, 3, 4, 5, 6, 7, 8],
+                "values": [0, 1, 2, 3, 4, 5, 6, 7],
+            }
+        )
+
+        result = df.pivot(index=input_index, columns=input_columns, values=input_values)
+        expected = pd.DataFrame(
+            expected_values, columns=expected_columns, index=expected_index
+        )
+        tm.assert_frame_equal(result, expected)
+
     def test_pivot_columns_lexsorted(self):
 
         n = 10000


Code review comments:
diff_hunk: @@ -5847,11 +5847,19 @@ def groupby(
 
         Parameters
         ----------%s
-        index : str or object, optional
+        index : str or object or a list of str, optional
             Column to use to make new frame's index. If None, uses
             existing index.
-        columns : str or object
+
+            .. versionchanged:: 1.0.0
path: pandas/core/frame.py
text: @user1:
```suggestion
            .. versionchanged:: 1.1.0
```

@author:
nice catch!! my bad 😅