diff --git a/pandas/core/reshape/reshape.py b/pandas/core/reshape/reshape.py index 3d9e84954a63b..b22d143544b7b 100644 --- a/pandas/core/reshape/reshape.py +++ b/pandas/core/reshape/reshape.py @@ -395,15 +395,29 @@ def pivot(self, index=None, columns=None, values=None): See DataFrame.pivot """ if values is None: - cols = [columns] if index is None else [index, columns] + if index is None: + cols = [columns] + else: + if is_list_like(index): + cols = [column for column in index] + else: + cols = [index] + cols.append(columns) append = index is None indexed = self.set_index(cols, append=append) + else: if index is None: index = self.index + index = MultiIndex.from_arrays([index, self[columns]]) + elif is_list_like(index): + # Iterating through the list of multiple columns of an index + indexes = [self[column] for column in index] + indexes.append(self[columns]) + index = MultiIndex.from_arrays(indexes) else: index = self[index] - index = MultiIndex.from_arrays([index, self[columns]]) + index = MultiIndex.from_arrays([index, self[columns]]) 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 7e7e081408534..4474c61dddfe3 100644 --- a/pandas/tests/reshape/test_pivot.py +++ b/pandas/tests/reshape/test_pivot.py @@ -301,6 +301,34 @@ def test_pivot_multi_functions(self): expected = concat([means, stds], keys=['mean', 'std'], axis=1) tm.assert_frame_equal(result, expected) + def test_pivot_multiple_columns_as_index(self): + # adding the test case for multiple columns as index (#21425) + df = 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': [0, 1, 2, 3, 4, 5, 6, 7]}) + result = df.pivot(index=['lev1', 'lev2'], + columns='lev3', + values='values') + result_no_values = df.pivot(index=['lev1', 'lev2'], + columns='lev3') + data = [[0, 1], [2, 3], [4, 5], [6, 7]] + exp_index = pd.MultiIndex.from_product([[1, 2], [1, 2]], + names=['lev1', 'lev2']) + exp_columns_1 = Index([1, 2], name='lev3') + expected_1 = DataFrame(data=data, index=exp_index, + columns=exp_columns_1) + + exp_columns_2 = MultiIndex(levels=[['values'], [1, 2]], + labels=[[0, 0], [0, 1]], + names=[None, 'lev3']) + + expected_2 = DataFrame(data=data, index=exp_index, + columns=exp_columns_2) + + tm.assert_frame_equal(result, expected_1) + tm.assert_frame_equal(result_no_values, expected_2) + def test_pivot_index_with_nan(self): # GH 3588 nan = np.nan