BUG: Fix/test SparseSeries/SparseDataFrame stack/unstack

Author: kernc

doc/source/whatsnew/v0.21.0.txt

@@ -588,7 +588,7 @@ Sparse
 
 - Bug in ``SparseSeries`` raises ``AttributeError`` when a dictionary is passed in as data (:issue:`16905`)
 - Bug in :func:`SparseDataFrame.fillna` not filling all NaNs when frame was instantiated from SciPy sparse matrix (:issue:`16112`)
-
+- Bug in :func:`SparseSeries.unstack` and :func:`SparseDataFrame.stack` (:issue:`16614`, :issue:`15045`)
 
 Reshaping
 ^^^^^^^^^

pandas/core/reshape/reshape.py

@@ -10,7 +10,7 @@
 from pandas.core.dtypes.common import (
     _ensure_platform_int,
     is_list_like, is_bool_dtype,
-    needs_i8_conversion)
+    needs_i8_conversion, is_sparse)
 from pandas.core.dtypes.cast import maybe_promote
 from pandas.core.dtypes.missing import notna
 import pandas.core.dtypes.concat as _concat
@@ -75,10 +75,15 @@ def __init__(self, values, index, level=-1, value_columns=None,
                  fill_value=None):
 
         self.is_categorical = None
+        self.is_sparse = is_sparse(values)
         if values.ndim == 1:
             if isinstance(values, Categorical):
                 self.is_categorical = values
                 values = np.array(values)
+            elif self.is_sparse:
+                # XXX: Makes SparseArray *dense*, but it's supposedly
+                # a single column at a time, so it's "doable"
+                values = values.values
             values = values[:, np.newaxis]
         self.values = values
         self.value_columns = value_columns
@@ -177,7 +182,8 @@ def get_result(self):
                                   ordered=ordered)
                       for i in range(values.shape[-1])]
 
-        return DataFrame(values, index=index, columns=columns)
+        klass = SparseDataFrame if self.is_sparse else DataFrame
+        return klass(values, index=index, columns=columns)
 
     def get_new_values(self):
         values = self.values
@@ -472,15 +478,15 @@ def _unstack_frame(obj, level, fill_value=None):
     from pandas.core.internals import BlockManager, make_block
 
     if obj._is_mixed_type:
-        unstacker = _Unstacker(np.empty(obj.shape, dtype=bool),  # dummy
+        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 = []
+        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,
@@ -490,15 +496,25 @@ def _unstack_frame(obj, level, fill_value=None):
             new_placement = new_columns.get_indexer(new_items)
             new_values, mask = bunstacker.get_new_values()
 
-            mblk = make_block(mask.T, placement=new_placement)
-            mask_blocks.append(mblk)
+            mask_blocks[new_placement] = mask.any(0)
 
-            newb = make_block(new_values.T, placement=new_placement)
-            new_blocks.append(newb)
+            # 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
+            else:
+                new_placement = [new_placement]
+                new_values = [new_values.T]
+
+            for cols, placement in zip(new_values, new_placement):
+                newb = blk.make_block_same_class(cols, placement=placement)
+                new_blocks.append(newb)
 
-        result = DataFrame(BlockManager(new_blocks, new_axes))
-        mask_frame = DataFrame(BlockManager(mask_blocks, new_axes))
-        return result.loc[:, mask_frame.sum(0) > 0]
+        klass = type(obj)
+        assert klass in (SparseDataFrame, DataFrame), klass
+        result = klass(BlockManager(new_blocks, new_axes))
+        return result.loc[:, mask_blocks]
     else:
         unstacker = _Unstacker(obj.values, obj.index, level=level,
                                value_columns=obj.columns,
@@ -559,7 +575,9 @@ def factorize(index):
         mask = notna(new_values)
         new_values = new_values[mask]
         new_index = new_index[mask]
-    return Series(new_values, index=new_index)
+
+    klass = type(frame)._constructor_sliced
+    return klass(new_values, index=new_index)
 
 
 def stack_multiple(frame, level, dropna=True):

pandas/tests/test_multilevel.py

@@ -2388,6 +2388,29 @@ def test_iloc_mi(self):
         tm.assert_frame_equal(result, expected)
 
 
+class TestSparse(object):
+    def setup_method(self, method):
+        self.sdf = pd.SparseDataFrame({0: {0: 1}, 1: {1: 1}, 2: {2: 1}})  # eye
+        self.mi = MultiIndex.from_tuples([(0, 0), (1, 1), (2, 2)])
+
+    def test_sparse_frame_stack(self):
+        ss = self.sdf.stack()
+        expected = pd.SparseSeries(np.ones(3), index=self.mi)
+        tm.assert_sp_series_equal(ss, expected)
+
+    def test_sparse_frame_unstack(self):
+        mi = pd.MultiIndex.from_tuples([(0, 0), (1, 0), (1, 2)])
+        sdf = self.sdf
+        sdf.index = mi
+        df = pd.DataFrame(np.eye(3), index=mi).replace(0, np.nan)
+
+        tm.assert_numpy_array_equal(df.unstack().values, sdf.unstack().values)
+
+    def test_sparse_series_unstack(self):
+        frame = pd.SparseSeries(np.ones(3), index=self.mi).unstack()
+        tm.assert_sp_frame_equal(frame, self.sdf)
+
+
 class TestSorted(Base):
     """ everthing you wanted to test about sorting """