Move .unstack() logic onto BlockManager and Block

Author: kernc

pandas/core/internals.py

@@ -1312,6 +1312,10 @@ def equals(self, other):
             return False
         return array_equivalent(self.values, other.values)
 
+    def _unstack(self, new_values, new_placement):
+        """Return a list of unstacked blocks of self"""
+        return [make_block(new_values, placement=new_placement)]
+
     def quantile(self, qs, interpolation='linear', axis=0, mgr=None):
         """
         compute the quantiles of the
@@ -1554,6 +1558,12 @@ def _slice(self, slicer):
     def _try_cast_result(self, result, dtype=None):
         return result
 
+    def _unstack(self, new_values, new_placement):
+        # NonConsolidatable blocks can have a single item only, so we return
+        # one block per item
+        return [self.make_block_same_class(vals, [place])
+                for vals, place in zip(new_values.T, new_placement)]
+
 
 class NumericBlock(Block):
     __slots__ = ()
@@ -4066,6 +4076,28 @@ def canonicalize(block):
         return all(block.equals(oblock)
                    for block, oblock in zip(self_blocks, other_blocks))
 
+    def unstack(self, unstacker):
+        """Return blockmanager with all blocks unstacked"""
+        dummy = unstacker(np.empty((0, 0)), value_columns=self.items)
+        new_columns = dummy.get_new_columns()
+        new_index = dummy.get_new_index()
+        new_blocks = []
+        mask_blocks = np.zeros_like(new_columns, dtype=bool)
+
+        for blk in self.blocks:
+            bunstacker = unstacker(
+                blk.values.T, value_columns=self.items[blk.mgr_locs.indexer])
+            new_items = bunstacker.get_new_columns()
+            new_values, mask = bunstacker.get_new_values()
+            new_placement = new_columns.get_indexer(new_items)
+            mask_blocks[new_placement] = mask.any(0)
+            new_blocks.extend(blk._unstack(new_values, new_placement))
+
+        new_blocks = [b for b, keep in zip(new_blocks, mask_blocks) if keep]
+        new_axes = [new_columns[mask_blocks], new_index]
+        bm = BlockManager(new_blocks, new_axes)
+        return bm
+
 
 class SingleBlockManager(BlockManager):
     """ manage a single block with """

pandas/core/reshape/reshape.py

@@ -2,6 +2,7 @@
 # pylint: disable=W0703,W0622,W0613,W0201
 from pandas.compat import range, zip
 from pandas import compat
+from functools import partial
 import itertools
 import re
 
@@ -466,48 +467,12 @@ def unstack(obj, level, fill_value=None):
 
 
 def _unstack_frame(obj, level, fill_value=None):
-    from pandas.core.internals import BlockManager, make_block as _make_block
-
     if obj._is_mixed_type:
-        unstacker = _Unstacker(np.empty((0, 0)),  # dummy
-                               obj.index, level=level,
-                               value_columns=obj.columns)
-        new_columns = unstacker.get_new_columns()
-        new_index = unstacker.get_new_index()
-        new_axes = [new_columns, new_index]
-
-        new_blocks = []
-        mask_blocks = np.zeros_like(new_columns, dtype=bool)
-        for blk in obj._data.blocks:
-            blk_items = obj._data.items[blk.mgr_locs.indexer]
-            bunstacker = _Unstacker(blk.values.T, obj.index, level=level,
-                                    value_columns=blk_items,
-                                    fill_value=fill_value)
-            new_items = bunstacker.get_new_columns()
-            new_placement = new_columns.get_indexer(new_items)
-            new_values, mask = bunstacker.get_new_values()
-
-            mask_blocks[new_placement] = mask.any(0)
-
-            # BlockManager can't handle SparseBlocks with multiple items,
-            # so lets make one block for each item
-            if is_sparse(blk.values):
-                new_placement = [[i] for i in new_placement]
-                new_values = new_values.T
-                make_block = blk.make_block_same_class
-            else:
-                new_placement = [new_placement]
-                new_values = [new_values.T]
-                make_block = _make_block
-
-            for cols, placement in zip(new_values, new_placement):
-                newb = make_block(cols, placement=placement)
-                new_blocks.append(newb)
-
+        unstacker = partial(_Unstacker, index=obj.index,
+                            level=level, fill_value=fill_value)
+        blocks = obj._data.unstack(unstacker)
         klass = type(obj)
-        assert klass in (SparseDataFrame, DataFrame), klass
-        result = klass(BlockManager(new_blocks, new_axes))
-        return result.loc[:, mask_blocks]
+        return klass(blocks)
     else:
         unstacker = _Unstacker(obj.values, obj.index, level=level,
                                value_columns=obj.columns,