BUG: fix MultiIndex.remove_unused_levels() when index contains NaNs

closes pandas-dev#18417

Author: toobaz

doc/source/whatsnew/v0.22.0.txt

@@ -119,7 +119,7 @@ Indexing
 
 - Bug in :func:`Series.truncate` which raises ``TypeError`` with a monotonic ``PeriodIndex`` (:issue:`17717`)
 - Bug in :func:`DataFrame.groupby` where tuples were interpreted as lists of keys rather than as keys (:issue:`17979`, :issue:`18249`)
--
+- Bug in :func:`MultiIndex.remove_unused_levels`` which would fill nan values (:issue:`18417`)
 -
 
 I/O

pandas/core/indexes/multi.py

@@ -1366,6 +1366,11 @@ def remove_unused_levels(self):
 
         changed = False
         for lev, lab in zip(self.levels, self.labels):
+            null_mask = lab == -1
+            if null_mask.any():
+                lab = lab[~null_mask]
+            else:
+                null_mask = None
 
             uniques = algos.unique(lab)
 
@@ -1380,8 +1385,14 @@ def remove_unused_levels(self):
             # labels get mapped from uniques to 0:len(uniques)
             label_mapping = np.zeros(len(lev))
             label_mapping[uniques] = np.arange(len(uniques))
+
             lab = label_mapping[lab]
 
+            if null_mask is not None:
+                full_lab = - np.ones(null_mask.shape, dtype=lab.dtype)
+                full_lab[~null_mask] = lab
+                lab = full_lab
+
             # new levels are simple
             lev = lev.take(uniques)
 

pandas/tests/indexes/test_multi.py

@@ -2629,6 +2629,16 @@ def test_reconstruct_remove_unused(self):
         tm.assert_index_equal(result2, expected)
         assert result2.is_(result)
 
+    @pytest.mark.parametrize('level', [['a', 'd', 'b'],
+                                       ['a', 'd', 'b', 'unused']])
+    def test_remove_unused_nan(self, level):
+        # GH 18417
+        df = pd.MultiIndex(levels=[level, ['w', 'x', 'y', 'z']],
+                           labels=[[0, 2, -1, 1, 1], [0, 1, 2, 3, 2]])
+
+        result = df.remove_unused_levels()
+        tm.assert_index_equal(result, df)
+
     @pytest.mark.parametrize('first_type,second_type', [
         ('int64', 'int64'),
         ('datetime64[D]', 'str')])