diff --git a/doc/source/whatsnew/v0.21.0.txt b/doc/source/whatsnew/v0.21.0.txt index f760d0b6359a2..243cae1ac79e5 100644 --- a/doc/source/whatsnew/v0.21.0.txt +++ b/doc/source/whatsnew/v0.21.0.txt @@ -129,7 +129,7 @@ Other Enhancements - `read_*` methods can now infer compression from non-string paths, such as ``pathlib.Path`` objects (:issue:`17206`). - :func:`pd.read_sas()` now recognizes much more of the most frequently used date (datetime) formats in SAS7BDAT files (:issue:`15871`). - :func:`DataFrame.items` and :func:`Series.items` is now present in both Python 2 and 3 and is lazy in all cases (:issue:`13918`, :issue:`17213`) - +- :func:`DataFrame.pivot` now accepts a list of values (:issue:`17160`). .. _whatsnew_0210.api_breaking: diff --git a/pandas/core/frame.py b/pandas/core/frame.py index b5b3df64d24c0..857ac43586432 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -4102,7 +4102,7 @@ def pivot(self, index=None, columns=None, values=None): existing index. columns : string or object Column name to use to make new frame's columns - values : string or object, optional + values : string, object or a list of the previous, optional Column name to use for populating new frame's values. If not specified, all remaining columns will be used and the result will have hierarchically indexed columns diff --git a/pandas/core/reshape/reshape.py b/pandas/core/reshape/reshape.py index 455da9246783c..42f5cd2ae3c22 100644 --- a/pandas/core/reshape/reshape.py +++ b/pandas/core/reshape/reshape.py @@ -367,15 +367,17 @@ def pivot(self, index=None, columns=None, values=None): cols = [columns] if index is None else [index, columns] append = index is None indexed = self.set_index(cols, append=append) - return indexed.unstack(columns) else: - if index is None: - index = self.index + index = self.index if index is None else self[index] + index = MultiIndex.from_arrays([index, self[columns]]) + if isinstance(values, list): + indexed = DataFrame(self[values].values, + index=index, + columns=values) else: - index = self[index] - indexed = Series(self[values].values, - index=MultiIndex.from_arrays([index, self[columns]])) - return indexed.unstack(columns) + indexed = Series(self[values].values, + index=index) + return indexed.unstack(columns) def pivot_simple(index, columns, values): diff --git a/pandas/tests/reshape/test_pivot.py b/pandas/tests/reshape/test_pivot.py index 879ac96680fbb..10519a0db96c2 100644 --- a/pandas/tests/reshape/test_pivot.py +++ b/pandas/tests/reshape/test_pivot.py @@ -351,6 +351,29 @@ def test_pivot_periods(self): pv = df.pivot(index='p1', columns='p2', values='data1') tm.assert_frame_equal(pv, expected) + def test_pivot_with_multi_values(self): + df = pd.DataFrame({'foo': ['one', 'one', 'one', 'two', 'two', 'two'], + 'bar': ['A', 'B', 'C', 'A', 'B', 'C'], + 'baz': [1, 2, 3, 4, 5, 6], + 'zoo': ['x', 'y', 'z', 'q', 'w', 't']}) + + results = df.pivot(index='zoo', columns='foo', values=['bar', 'baz']) + + data = [[None, 'A', None, 4], + [None, 'C', None, 6], + [None, 'B', None, 5], + ['A', None, 1, None], + ['B', None, 2, None], + ['C', None, 3, None]] + index = Index(data=['q', 't', 'w', 'x', 'y', 'z'], name='zoo') + columns = MultiIndex(levels=[['bar', 'baz'], ['one', 'two']], + labels=[[0, 0, 1, 1], [0, 1, 0, 1]], + names=[None, 'foo']) + expected = DataFrame(data=data, index=index, + columns=columns, dtype='object') + + tm.assert_frame_equal(results, expected) + def test_margins(self): def _check_output(result, values_col, index=['A', 'B'], columns=['C'],