Skip to content

Commit

Permalink
multi-index column support
Browse files Browse the repository at this point in the history
  • Loading branch information
jreback committed Jul 17, 2019
1 parent d38f7d8 commit fc66d10
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 5 deletions.
8 changes: 4 additions & 4 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
import itertools
import sys
from textwrap import dedent
from typing import FrozenSet, List, Optional, Set, Type, Union
from typing import FrozenSet, List, Optional, Set, Tuple, Type, Union
import warnings

import numpy as np
Expand Down Expand Up @@ -6233,7 +6233,7 @@ def stack(self, level=-1, dropna=True):
else:
return stack(self, level, dropna=dropna)

def explode(self, column: str) -> "DataFrame":
def explode(self, column: Union[str, Tuple]) -> "DataFrame":
"""
Transform each element of a list-like to a row, replicating the
index values.
Expand All @@ -6242,7 +6242,7 @@ def explode(self, column: str) -> "DataFrame":
Parameters
----------
column : str
column : str or tuple
Returns
-------
Expand Down Expand Up @@ -6290,7 +6290,7 @@ def explode(self, column: str) -> "DataFrame":
3 4 1
"""

if not is_scalar(column):
if not (is_scalar(column) or isinstance(column, tuple)):
raise ValueError("column must be a scalar")
if not self.columns.is_unique:
raise ValueError("columns must be unique")
Expand Down
21 changes: 20 additions & 1 deletion pandas/tests/frame/test_explode.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ def test_basic():
tm.assert_frame_equal(result, expected)


def test_multi_index():
def test_multi_index_rows():
df = pd.DataFrame(
{"A": np.array([[0, 1, 2], np.nan, [], (3, 4)], dtype=object), "B": 1},
index=pd.MultiIndex.from_tuples([("a", 1), ("a", 2), ("b", 1), ("b", 2)]),
Expand Down Expand Up @@ -63,6 +63,25 @@ def test_multi_index():
tm.assert_frame_equal(result, expected)


def test_multi_index_columns():
df = pd.DataFrame(
{("A", 1): np.array([[0, 1, 2], np.nan, [], (3, 4)], dtype=object), ("A", 2): 1}
)

result = df.explode(("A", 1))
expected = pd.DataFrame(
{
("A", 1): pd.Series(
[0, 1, 2, np.nan, np.nan, 3, 4],
index=pd.Index([0, 0, 0, 1, 2, 3, 3]),
dtype=object,
),
("A", 2): 1,
}
)
tm.assert_frame_equal(result, expected)


def test_usecase():
# explode a single column
# gh-10511
Expand Down

0 comments on commit fc66d10

Please sign in to comment.