ENH: Allow multi values for index and columns in df.pivot

doc/source/whatsnew/v1.1.0.rst

@@ -192,8 +192,10 @@ Reshaping
 - Bug in :meth:`DataFrame.pivot_table` when ``margin`` is ``True`` and only ``column`` is defined (:issue:`31016`)
 - 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`)
 - Bug in :func:`concat` where the resulting indices are not copied when ``copy=True`` (:issue:`29879`)
 
+
 Sparse
 ^^^^^^
 

pandas/core/frame.py

@@ -5896,11 +5896,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.1.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.1.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
@@ -5967,6 +5975,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'],

pandas/core/reshape/pivot.py

@@ -425,17 +425,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

pandas/tests/reshape/test_pivot_multilevel.py

@@ -0,0 +1,192 @@
+import numpy as np
+import pytest
+
+import pandas as pd
+from pandas import Index, MultiIndex
+import pandas._testing as tm
+
+
+@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(
+    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(
+    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)