BUG: avoid unnecessary casting when unstacking index with unused levels

closes #17845

Author: toobaz

pandas/core/reshape/reshape.py

@@ -89,7 +89,7 @@ def __init__(self, values, index, level=-1, value_columns=None,
         if value_columns is None and values.shape[1] != 1:  # pragma: no cover
             raise ValueError('must pass column labels for multi-column data')
 
-        self.index = index
+        self.index = index.remove_unused_levels()
 
         if isinstance(self.index, MultiIndex):
             if index._reference_duplicate_name(level):
@@ -102,11 +102,12 @@ def __init__(self, values, index, level=-1, value_columns=None,
         # when index includes `nan`, need to lift levels/strides by 1
         self.lift = 1 if -1 in self.index.labels[self.level] else 0
 
-        self.new_index_levels = list(index.levels)
-        self.new_index_names = list(index.names)
+        self.new_index_levels = list(self.index.levels)
+        self.new_index_names = list(self.index.names)
 
         self.removed_name = self.new_index_names.pop(self.level)
         self.removed_level = self.new_index_levels.pop(self.level)
+        self.removed_level_full = index.levels[self.level]
 
         self._make_sorted_values_labels()
         self._make_selectors()
@@ -156,21 +157,10 @@ def _make_selectors(self):
         self.compressor = comp_index.searchsorted(np.arange(ngroups))
 
     def get_result(self):
-        # TODO: find a better way than this masking business
-
-        values, value_mask = self.get_new_values()
+        values, _ = self.get_new_values()
         columns = self.get_new_columns()
         index = self.get_new_index()
 
-        # filter out missing levels
-        if values.shape[1] > 0:
-            col_inds, obs_ids = compress_group_index(self.sorted_labels[-1])
-            # rare case, level values not observed
-            if len(obs_ids) < self.full_shape[1]:
-                inds = (value_mask.sum(0) > 0).nonzero()[0]
-                values = algos.take_nd(values, inds, axis=1)
-                columns = columns[inds]
-
         # may need to coerce categoricals here
         if self.is_categorical is not None:
             categories = self.is_categorical.categories
@@ -259,17 +249,23 @@ def get_new_columns(self):
         width = len(self.value_columns)
         propagator = np.repeat(np.arange(width), stride)
         if isinstance(self.value_columns, MultiIndex):
-            new_levels = self.value_columns.levels + (self.removed_level,)
+            new_levels = self.value_columns.levels + (self.removed_level_full,)
             new_names = self.value_columns.names + (self.removed_name,)
 
             new_labels = [lab.take(propagator)
                           for lab in self.value_columns.labels]
         else:
-            new_levels = [self.value_columns, self.removed_level]
+            new_levels = [self.value_columns, self.removed_level_full]
             new_names = [self.value_columns.name, self.removed_name]
             new_labels = [propagator]
 
-        new_labels.append(np.tile(np.arange(stride) - self.lift, width))
+        if len(self.removed_level_full) != len(self.removed_level):
+            repeater = self.removed_level_full.get_indexer(self.removed_level)
+            if self.lift:
+                repeater = np.insert(repeater, 0, -1)
+        else:
+            repeater = np.arange(stride) - self.lift
+        new_labels.append(np.tile(repeater, width))
         return MultiIndex(levels=new_levels, labels=new_labels,
                           names=new_names, verify_integrity=False)
 

pandas/tests/frame/test_reshape.py

@@ -536,6 +536,56 @@ def test_unstack_dtypes(self):
             assert left.shape == (3, 2)
             tm.assert_frame_equal(left, right)
 
+    def test_unstack_unused_levels(self):
+        # GH 17845: sliced columns of int DataFrame
+        idx = pd.MultiIndex.from_product([['a'], ['A', 'B', 'C', 'D']])[:-1]
+        df = pd.DataFrame([[1, 0]] * 3, index=idx)
+
+        result = df.unstack()
+        exp_col = pd.MultiIndex.from_product([[0, 1], ['A', 'B', 'C']])
+        expected = pd.DataFrame([[1, 1, 1, 0, 0, 0]], index=['a'],
+                                columns=exp_col)
+        tm.assert_frame_equal(result, expected)
+        assert((result.columns.levels[1] == idx.levels[1]).all())
+
+        # Unused items on both levels
+        levels = [[0, 1, 7], [0, 1, 2, 3]]
+        labels = [[0, 0, 1, 1], [0, 2, 0, 2]]
+        idx = pd.MultiIndex(levels, labels)
+        block = np.arange(4).reshape(2, 2)
+        df = pd.DataFrame(np.concatenate([block, block + 4]), index=idx)
+        result = df.unstack()
+        expected = pd.DataFrame(np.concatenate([block * 2, block * 2 - 1],
+                                               axis=1),
+                                columns=idx)
+        assert((result.columns.levels[1] == idx.levels[1]).all())
+
+        # With mixed dtype and NaN
+        levels = [['a', 2, 'c'], [1, 3, 5, 7]]
+        labels = [[0, -1, 1, 1], [0, 2, -1, 2]]
+        idx = pd.MultiIndex(levels, labels)
+        data = np.arange(8)
+        df = pd.DataFrame(data.reshape(4, 2), index=idx)
+
+        cases = ((0, [13, 16, 6, 9, 2, 5, 8, 11],
+                  [np.nan, 'a', 2], [np.nan, 5, 1]),
+                 (1, [8, 11, 1, 4, 12, 15, 13, 16],
+                  [np.nan, 5, 1], [np.nan, 'a', 2]))
+        for level, idces, col_level, idx_level in cases:
+            result = df.unstack(level=level)
+            exp_data = np.zeros(18) * np.nan
+            exp_data[idces] = data
+            cols = pd.MultiIndex.from_product([[0, 1], col_level])
+            expected = pd.DataFrame(exp_data.reshape(3, 6),
+                                    index=idx_level, columns=cols)
+            # Broken (GH 18455):
+            # tm.assert_frame_equal(result, expected)
+            diff = result - expected
+            assert(diff.sum().sum() == 0)
+            assert((diff + 1).sum().sum() == 8)
+
+            assert((result.columns.levels[1] == idx.levels[level]).all())
+
     def test_unstack_non_unique_index_names(self):
         idx = MultiIndex.from_tuples([('a', 'b'), ('c', 'd')],
                                      names=['c1', 'c1'])