Skip to content

Commit

Permalink
ENH: unstack/stack multiple levels per #370, use Series index per not…
Browse files Browse the repository at this point in the history
…e in #373
  • Loading branch information
wesm committed Nov 17, 2011
1 parent b1fb3bc commit d345967
Show file tree
Hide file tree
Showing 5 changed files with 66 additions and 15 deletions.
34 changes: 24 additions & 10 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
from itertools import izip
from StringIO import StringIO
import csv
import gc
import operator
import sys

Expand Down Expand Up @@ -241,6 +240,10 @@ def _init_ndarray(self, values, index, columns, dtype=None,
if isinstance(values, Series) and values.name is not None:
if columns is None:
columns = [values.name]
if index is None:
index = values.index
else:
values = values.reindex(index)

values = _prep_ndarray(values, copy=copy)

Expand Down Expand Up @@ -1958,24 +1961,31 @@ def stack(self, level=-1, dropna=True):
Parameters
----------
level : int or string, default last level
Level to stack, can pass level name
level : int, string, or list of these, default last level
Level(s) to stack, can pass level name
Returns
-------
stacked : Series
"""
from pandas.core.reshape import stack
return stack(self, level=level, dropna=dropna)

if isinstance(level, (tuple, list)):
result = self
for lev in level:
result = stack(result, lev, dropna=dropna)
return result
else:
return stack(self, level, dropna=dropna)

def unstack(self, level=-1):
"""
"Unstack" level from MultiLevel index to produce reshaped DataFrame
Parameters
----------
level : int or string, default last level
Level to unstack, can pass level name
level : int, string, or list of these, default last level
Level(s) to unstack, can pass level name
Examples
--------
Expand All @@ -1999,10 +2009,14 @@ def unstack(self, level=-1):
-------
unstacked : DataFrame
"""
from pandas.core.reshape import _Unstacker
unstacker = _Unstacker(self.values, self.index, level=level,
value_columns=self.columns)
return unstacker.get_result()
from pandas.core.reshape import unstack
if isinstance(level, (tuple, list)):
result = self
for lev in level:
result = unstack(result, lev)
return result
else:
return unstack(self, level)

def delevel(self):
"""
Expand Down
10 changes: 10 additions & 0 deletions pandas/core/reshape.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from pandas.core.common import notnull
from pandas.core.index import MultiIndex


class ReshapeError(Exception):
pass

Expand Down Expand Up @@ -280,6 +281,15 @@ def _slow_pivot(index, columns, values):

return DataFrame(tree)

def unstack(obj, level):
if isinstance(obj, DataFrame):
columns = obj.columns
else:
columns = None
unstacker = _Unstacker(obj.values, obj.index, level=level,
value_columns=columns)
return unstacker.get_result()

def stack(frame, level=-1, dropna=True):
"""
Convert DataFrame to Series with multi-level Index. Columns become the
Expand Down
15 changes: 10 additions & 5 deletions pandas/core/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -1331,8 +1331,8 @@ def unstack(self, level=-1):
Parameters
----------
level : int, default last level
Level to unstack
level : int, string, or list of these, default last level
Level(s) to unstack, can pass level name
Examples
--------
Expand All @@ -1356,9 +1356,14 @@ def unstack(self, level=-1):
-------
unstacked : DataFrame
"""
from pandas.core.reshape import _Unstacker
unstacker = _Unstacker(self.values, self.index, level=level)
return unstacker.get_result()
from pandas.core.reshape import unstack
if isinstance(level, (tuple, list)):
result = self
for lev in level:
result = unstack(result, lev)
return result
else:
return unstack(self, level)

#----------------------------------------------------------------------
# function application
Expand Down
1 change: 1 addition & 0 deletions pandas/tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -1147,6 +1147,7 @@ def test_constructor_Series_named(self):
a = Series([1,2,3], index=['a','b','c'], name='x')
df = DataFrame(a)
self.assert_(df.columns[0] == 'x')
self.assert_(df.index.equals(a.index))

def test_astype(self):
casted = self.frame.astype(int)
Expand Down
21 changes: 21 additions & 0 deletions pandas/tests/test_multilevel.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ def setUp(self):
self.tdf = tm.makeTimeDataFrame()
self.ymd = self.tdf.groupby([lambda x: x.year, lambda x: x.month,
lambda x: x.day]).sum()
self.ymd.index.names = ['year', 'month', 'day']

def test_append(self):
a, b = self.frame[:5], self.frame[5:]
Expand Down Expand Up @@ -450,6 +451,26 @@ def test_stack_level_name(self):
expected = self.frame.stack()
assert_series_equal(result, expected)

def test_stack_unstack_multiple(self):
unstacked = self.ymd.unstack(['year', 'month'])
expected = self.ymd.unstack('year').unstack('month')
assert_frame_equal(unstacked, expected)
self.assertEquals(unstacked.columns.names,
expected.columns.names)

# series
s = self.ymd['A']
s_unstacked = s.unstack(['year', 'month'])
assert_frame_equal(s_unstacked, expected['A'])

restacked = unstacked.stack(['year', 'month'])
restacked = restacked.swaplevel(0, 1).swaplevel(1, 2)
restacked = restacked.sortlevel(0)

assert_frame_equal(restacked, self.ymd)
self.assertEquals(restacked.index.names,
self.ymd.index.names)

def test_groupby_transform(self):
s = self.frame['A']
grouper = s.index.get_level_values(0)
Expand Down

0 comments on commit d345967

Please sign in to comment.