ENH: unstack multiple levels in one shot per #1181

Author: wesm

RELEASE.rst

@@ -70,6 +70,8 @@ pandas 0.8.0
   - Fix segfault caused by empty groups passed to groupby (#1048)
   - Fix occasionally misbehaved reindexing in the presence of NaN labels (#522)
   - Fix imprecise logic causing weird Series results from .apply (#1183)
+  - Unstack multiple levels in one shot, avoiding empty columns in some
+    cases. Fix pivot table bug (#1181)
 
 pandas 0.7.3
 ============

pandas/core/frame.py

@@ -2992,17 +2992,7 @@ def unstack(self, level=-1):
         unstacked : DataFrame or Series
         """
         from pandas.core.reshape import unstack
-        if isinstance(level, (tuple, list)):
-            result = self
-            to_unstack = level
-            while to_unstack:
-                lev = to_unstack[0]
-                result = unstack(result, lev)
-                to_unstack = [other - 1 if other > lev else other
-                              for other in to_unstack[1:]]
-            return result
-        else:
-            return unstack(self, level)
+        return unstack(self, level)
 
     #----------------------------------------------------------------------
     # Time series-related

pandas/core/reshape.py

@@ -209,37 +209,52 @@ def _unstack_multiple(data, clocs):
 
     index = data.index
 
-    clevels, rlevels = _partition(index.levels, clocs)
-    clabels, rlabels = _partition(index.labels, clocs)
-    cnames, rnames = _partition(index.names, clocs)
+    clocs = [index._get_level_number(i) for i in clocs]
+
+    rlocs = [i for i in range(index.nlevels) if i not in clocs]
+
+    clevels = [index.levels[i] for i in clocs]
+    clabels = [index.labels[i] for i in clocs]
+    cnames = [index.names[i] for i in clocs]
+    rlevels = [index.levels[i] for i in rlocs]
+    rlabels = [index.labels[i] for i in rlocs]
+    rnames = [index.names[i] for i in rlocs]
 
     shape = [len(x) for x in clevels]
     group_index = get_group_index(clabels, shape)
 
     comp_ids, obs_ids = _compress_group_index(group_index, sort=False)
+    recons_labels = decons_group_index(obs_ids, shape)
 
     dummy_index = MultiIndex(levels=rlevels + [obs_ids],
                              labels=rlabels + [comp_ids],
                              names=rnames + ['__placeholder__'])
 
-    dummy = DataFrame(data.values, index=dummy_index,
-                      columns=data.columns)
-
-    unstacked = dummy.unstack('__placeholder__')
-
-    if isinstance(unstacked, Series):
-        unstcols = unstacked.index
+    if isinstance(data, Series):
+        dummy = Series(data.values, index=dummy_index)
+        unstacked = dummy.unstack('__placeholder__')
+        new_levels = clevels
+        new_names = cnames
+        new_labels = recons_labels
     else:
-        unstcols = unstacked.columns
+        if isinstance(data.columns, MultiIndex):
+            raise NotImplementedError('Unstacking multiple levels with '
+                                      'hierarchical columns not yet supported')
 
-    new_levels = [unstcols.levels[0]] + clevels
-    new_names = [data.columns.name] + cnames
+        dummy = DataFrame(data.values, index=dummy_index,
+                          columns=data.columns)
 
-    recons_labels = decons_group_index(obs_ids, shape)
+        unstacked = dummy.unstack('__placeholder__')
+        if isinstance(unstacked, Series):
+            unstcols = unstacked.index
+        else:
+            unstcols = unstacked.columns
+        new_levels = [unstcols.levels[0]] + clevels
+        new_names = [data.columns.name] + cnames
 
-    new_labels = [unstcols.labels[0]]
-    for rec in recons_labels:
-        new_labels.append(rec.take(unstcols.labels[-1]))
+        new_labels = [unstcols.labels[0]]
+        for rec in recons_labels:
+            new_labels.append(rec.take(unstcols.labels[-1]))
 
     new_columns = MultiIndex(levels=new_levels, labels=new_labels,
                              names=new_names)
@@ -251,22 +266,6 @@ def _unstack_multiple(data, clocs):
 
     return unstacked
 
-
-def _partition(values, inds):
-    left = []
-    right = []
-
-    set_inds = set(inds)
-
-    for i, val in enumerate(values):
-        if i in set_inds:
-            left.append(val)
-        else:
-            right.append(val)
-
-    return left, right
-
-
 def pivot(self, index=None, columns=None, values=None):
     """
     See DataFrame.pivot
@@ -351,6 +350,9 @@ def _slow_pivot(index, columns, values):
     return DataFrame(tree)
 
 def unstack(obj, level):
+    if isinstance(level, (tuple, list)):
+        return _unstack_multiple(obj, level)
+
     if isinstance(obj, DataFrame):
         if isinstance(obj.index, MultiIndex):
             return _unstack_frame(obj, level)

pandas/core/series.py

@@ -1814,13 +1814,7 @@ def unstack(self, level=-1):
         unstacked : DataFrame
         """
         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)
+        return unstack(self, level)
 
     #----------------------------------------------------------------------
     # function application

pandas/tests/test_multilevel.py

@@ -607,15 +607,15 @@ def test_unstack(self):
         # test that ints work
         unstacked = self.ymd.astype(int).unstack()
 
-    # def test_unstack_multiple_no_empty_columns(self):
-    #     index = MultiIndex.from_tuples([(0, 'foo', 0), (0, 'bar', 0),
-    #                                     (1, 'baz', 1), (1, 'qux', 1)])
+    def test_unstack_multiple_no_empty_columns(self):
+        index = MultiIndex.from_tuples([(0, 'foo', 0), (0, 'bar', 0),
+                                        (1, 'baz', 1), (1, 'qux', 1)])
 
-    #     s = Series(np.random.randn(4), index=index)
+        s = Series(np.random.randn(4), index=index)
 
-    #     unstacked = s.unstack([1, 2])
-    #     expected = unstacked.dropna(axis=1, how='all')
-    #     assert_frame_equal(unstacked, expected)
+        unstacked = s.unstack([1, 2])
+        expected = unstacked.dropna(axis=1, how='all')
+        assert_frame_equal(unstacked, expected)
 
     def test_stack(self):
         # regular roundtrip
@@ -738,12 +738,12 @@ def test_stack_unstack_multiple(self):
 
         # GH #451
         unstacked = self.ymd.unstack([1, 2])
-        expected = self.ymd.unstack(1).unstack(1)
+        expected = self.ymd.unstack(1).unstack(1).dropna(axis=1, how='all')
         assert_frame_equal(unstacked, expected)
 
         unstacked = self.ymd.unstack([2, 1])
-        expected = self.ymd.unstack(2).unstack(1)
-        assert_frame_equal(unstacked, expected)
+        expected = self.ymd.unstack(2).unstack(1).dropna(axis=1, how='all')
+        assert_frame_equal(unstacked, expected.ix[:, unstacked.columns])
 
     def test_groupby_transform(self):
         s = self.frame['A']

pandas/tools/pivot.py

@@ -98,12 +98,10 @@ def pivot_table(data, values=None, rows=None, cols=None, aggfunc='mean',
     grouped = data.groupby(keys)
     agged = grouped.agg(aggfunc)
 
-    table = _unstack_multiple(agged, range(len(rows), len(keys)))
+    to_unstack = [agged.index.names[i]
+                  for i in range(len(rows), len(keys))]
 
-    # table = agged
-    # for i in range(len(cols)):
-    #     name = table.index.names[len(rows)]
-    #     table = table.unstack(name)
+    table = agged.unstack(to_unstack)
 
     if fill_value is not None:
         table = table.fillna(value=fill_value)