From 6c63055f8cfaf9088a1900b34b89002d986e574b Mon Sep 17 00:00:00 2001 From: Adam Klein Date: Mon, 9 Jan 2012 14:41:44 -0500 Subject: [PATCH] ENH: enable unstack to create Series --- pandas/core/reshape.py | 5 ++++- pandas/tests/test_frame.py | 24 ++++++++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/pandas/core/reshape.py b/pandas/core/reshape.py index b9f00320d8970..6fef3a5ee3c3e 100644 --- a/pandas/core/reshape.py +++ b/pandas/core/reshape.py @@ -283,7 +283,10 @@ def _slow_pivot(index, columns, values): def unstack(obj, level): if isinstance(obj, DataFrame): - return _unstack_frame(obj, level) + if isinstance(obj.index, MultiIndex): + return _unstack_frame(obj, level) + else: + return obj.transpose().stack(dropna=False) else: unstacker = _Unstacker(obj.values, obj.index, level=level) return unstacker.get_result() diff --git a/pandas/tests/test_frame.py b/pandas/tests/test_frame.py index 6bf13e0acac11..160e213279351 100644 --- a/pandas/tests/test_frame.py +++ b/pandas/tests/test_frame.py @@ -3889,6 +3889,30 @@ def test_stack_unstack(self): assert_frame_equal(unstacked_cols.T, self.frame) assert_frame_equal(unstacked_cols_df['bar'].T, self.frame) + def test_unstack_to_series(self): + # check reversibility + data = self.frame.unstack() + self.assertTrue(isinstance(data, Series)) + undo = data.unstack().transpose() + assert_frame_equal(undo, self.frame) + + # check NA handling + data = DataFrame({'x': [1, 2, np.NaN], 'y': [3.0, 4, np.NaN]}) + data.index = Index(['a','b','c']) + result = data.unstack() + + midx = MultiIndex(levels=[['x','y'],['a','b','c']], + labels=[[0,0,0,1,1,1],[0,1,2,0,1,2]]) + expected = Series([1,2,np.NaN,3,4,np.NaN], index=midx) + + assert_series_equal(result, expected) + + # check composability of unstack + old_data = data.copy() + for _ in xrange(4): + data = data.unstack() + assert_frame_equal(old_data, data) + def test_reset_index(self): stacked = self.frame.stack()[::2] stacked = DataFrame({'foo' : stacked, 'bar' : stacked})