Skip to content

MultiIndex Support for DataFrame.pivot #25330

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 21 commits into from
Closed
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion doc/source/whatsnew/v0.25.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ Other Enhancements
^^^^^^^^^^^^^^^^^^

- :meth:`Timestamp.replace` now supports the ``fold`` argument to disambiguate DST transition times (:issue:`25017`)
-
- :meth:`DataFrame.pivot` now supports multiple indexs.
-

.. _whatsnew_0250.api_breaking:
Expand Down
18 changes: 16 additions & 2 deletions pandas/core/reshape/pivot.py
Original file line number Diff line number Diff line change
Expand Up @@ -368,15 +368,29 @@ def _convert_by(by):
@Appender(_shared_docs['pivot'], indents=1)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also think you will need to update the docstring

def pivot(data, index=None, columns=None, values=None):
if values is None:
cols = [columns] if index is None else [index, columns]
if index is None:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you add comments here delineating the cases

cols = [columns]
else:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you make this an if/elsif/else

if is_list_like(index):
cols = [column for column in index]
else:
cols = [index]
cols.append(columns)
append = index is None
indexed = data.set_index(cols, append=append)

else:
if index is None:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same here

index = data.index
index = MultiIndex.from_arrays([index, data[columns]])
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do you need to pass names?

elif is_list_like(index):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would help readability if the conditions here were refactored to be in the same sequence as the conditions in the branch above, as IIUC we essentially evaluate the same conditions in both branches

# Iterating through the list of multiple columns of an index
indexes = [data[column] for column in index]
indexes.append(data[columns])
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What does this do?

index = MultiIndex.from_arrays(indexes)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

names?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think so but let me know if I should.

else:
index = data[index]
index = MultiIndex.from_arrays([index, data[columns]])
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can't you construct "the first draft" of index and then call index = index = MultiIndex.from_arrays([index, data[columns]]) outside of the if-else block?

index = MultiIndex.from_arrays([index, data[columns]])

if is_list_like(values) and not isinstance(values, tuple):
# Exclude tuple because it is seen as a single column name
Expand Down
28 changes: 28 additions & 0 deletions pandas/tests/reshape/test_pivot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you do the test caess 1 after the other and use the standard

result=
expected=
assert_frame_equal(result, expected)

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]],
codes=[[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)

@pytest.mark.parametrize('method', [True, False])
def test_pivot_index_with_nan(self, method):
# GH 3588
Expand Down