BUG: reimplement MultiIndex.remove_unused_levels

* Add a large random test case for remove_unused_levels that failed the
previous implementation

* Fix #16556, a performance issue with the previous implementation

* Add inplace functionality

* Always return (if not inplace) at least a view instead of the original
index

Author: rhendric

pandas/core/indexes/multi.py

@@ -1252,7 +1252,7 @@ def _sort_levels_monotonic(self):
                           names=self.names, sortorder=self.sortorder,
                           verify_integrity=False)
 
-    def remove_unused_levels(self):
+    def remove_unused_levels(self, inplace=False):
         """
         create a new MultiIndex from the current that removing
         unused levels, meaning that they are not expressed in the labels
@@ -1263,6 +1263,11 @@ def remove_unused_levels(self):
 
         .. versionadded:: 0.20.0
 
+        Parameters
+        ----------
+        inplace : bool
+            if True, mutates in place
+
         Returns
         -------
         MultiIndex
@@ -1290,42 +1295,42 @@ def remove_unused_levels(self):
         new_levels = []
         new_labels = []
 
-        changed = np.ones(self.nlevels, dtype=bool)
-        for i, (lev, lab) in enumerate(zip(self.levels, self.labels)):
+        changed = False
+        for lev, lab in zip(self.levels, self.labels):
 
             uniques = algos.unique(lab)
 
             # nothing unused
             if len(uniques) == len(lev):
                 new_levels.append(lev)
                 new_labels.append(lab)
-                changed[i] = False
                 continue
 
-            # set difference, then reverse sort
-            diff = Index(np.arange(len(lev))).difference(uniques)
-            unused = diff.sort_values(ascending=False)
+            changed = True
+
+            # 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 labels, we remove the unsued
-            # by decrementing the labels for that value
-            # prob a better way
-            for u in unused:
-
-                lab = np.where(lab > u, lab - 1, lab)
-
             new_levels.append(lev)
             new_labels.append(lab)
+        
+        if inplace:
+            idx = self
+        else:
+            idx = self._shallow_copy()
 
-        # nothing changed
-        if not changed.any():
-            return self
+        if changed:
+            idx._reset_identity()
+            idx._set_levels(new_levels, validate=False)
+            idx._set_labels(new_labels, validate=False)
 
-        return MultiIndex(new_levels, new_labels,
-                          names=self.names, sortorder=self.sortorder,
-                          verify_integrity=False)
+        if not inplace:
+            return idx
 
     @property
     def nlevels(self):

pandas/tests/indexes/test_multi.py

@@ -2489,7 +2489,27 @@ def test_reconstruct_remove_unused(self):
         # idempotent
         result2 = result.remove_unused_levels()
         tm.assert_index_equal(result2, expected)
-        assert result2 is result
+        assert result2.is_(result)
+
+    def test_remove_unused_levels_large(self):
+        # because tests should be deterministic:
+        rng = np.random.RandomState(4) # chosen by fair dice roll. guaranteed to be random.
+
+        size = 1<<16
+        df = DataFrame(dict(first=rng.randint(0, 1<<13, size),
+                            second=rng.randint(0, 1<<10, size),
+                            third=rng.rand(size)))
+        df = df.groupby(['first', 'second']).sum()
+        df = df[df.third < 0.1]
+
+        result = df.index.remove_unused_levels()
+        assert len(result.levels[0]) < len(df.index.levels[0])
+        assert len(result.levels[1]) < len(df.index.levels[1])
+        assert result.equals(df.index)
+
+        # in place
+        df.index.remove_unused_levels(inplace=True)
+        tm.assert_index_equal(df.index, result)
 
     def test_isin(self):
         values = [('foo', 2), ('bar', 3), ('quux', 4)]