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

closes #18417

Author: toobaz

doc/source/whatsnew/v0.22.0.txt

@@ -121,7 +121,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

@@ -1365,25 +1365,31 @@ def remove_unused_levels(self):
         new_labels = []
 
         changed = False
-        for lev, lab in zip(self.levels, self.labels):
+        for idx, (lev, lab) in enumerate(zip(self.levels, self.labels)):
+            na_idxs = np.where(lab == -1)[0]
+
+            if len(na_idxs):
+                lab = np.delete(lab, na_idxs)
 
             uniques = algos.unique(lab)
 
             # nothing unused
-            if len(uniques) == len(lev):
-                new_levels.append(lev)
-                new_labels.append(lab)
-                continue
+            if len(uniques) != len(lev):
+                changed = True
+
+                # labels get mapped from uniques to 0:len(uniques)
+                label_mapping = np.zeros(len(lev))
+                label_mapping[uniques] = np.arange(len(uniques))
 
-            changed = True
+                lab = label_mapping[lab]
 
-            # 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]
+                # new levels are simple
+                lev = lev.take(uniques)
 
-            # new levels are simple
-            lev = lev.take(uniques)
+                if len(na_idxs):
+                    lab = np.insert(lab, na_idxs - np.arange(len(na_idxs)), -1)
+            else:
+                lab = self.labels[idx]
 
             new_levels.append(lev)
             new_labels.append(lab)

pandas/tests/indexes/test_multi.py

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