support for removing unused levels (internally)

xref #2770

Author: jreback

pandas/indexes/multi.py

@@ -1177,7 +1177,7 @@ def from_product(cls, iterables, sortorder=None, names=None):
  labels = cartesian_product(labels)
  return MultiIndex(levels, labels, sortorder=sortorder, names=names)
 
- def _reconstruct(self, sort=False):
+ def _reconstruct(self, sort=False, remove_unused=False):
  """
  reconstruct the MultiIndex
 
@@ -1188,20 +1188,32 @@ def _reconstruct(self, sort=False):
  ----------
  sort: boolean, default False
  monotonically sort the levels
+ remove_unused: boolean, default False
+ remove unsued levels
 
  Returns
  -------
  MultiIndex
 
  """
+
+ if sort and remove_unused:
+ raise ValueError("only support one of sort / remove_unused")
+
+ if not (sort or remove_unused):
+ raise ValueError("must supply one of sort / remove_unsued")
+
+ levels = self.levels
+ labels = self.labels
+ new_levels = []
+ new_labels = []
+
  if sort:
 
  if self.is_monotonic:
  return self
 
- new_levels = []
- new_labels = []
- for lev, lab in zip(self.levels, self.labels):
+ for lev, lab in zip(levels, labels):
 
  if lev.is_monotonic:
  new_levels.append(lev)
@@ -1219,6 +1231,40 @@ def _reconstruct(self, sort=False):
  new_levels.append(lev)
  new_labels.append(lab)
 
+ elif remove_unused:
+
+ changed = np.zeros(self.nlevels, dtype=bool)
+ for i, (lev, lab) in enumerate(zip(levels, labels)):
+
+ uniques = np.sort(algos.unique(lab))
+
+ # nothing unused
+ if len(uniques) == len(lev):
+ new_levels.append(lev)
+ new_labels.append(lab)
+ changed[i] = True
+ continue
+
+ unused = list(reversed(sorted(set(
+ np.arange(len(lev))) - set(uniques))))
+
+ # 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)
+
+ # nothing changed
+ if not changed.any():
+ return self
+
  return MultiIndex(new_levels, new_labels,
  names=self.names, sortorder=self.sortorder,
  verify_integrity=False)

pandas/tests/indexes/test_multi.py

@@ -2411,7 +2411,19 @@ def test_is_monotonic(self):
 
  self.assertFalse(i.is_monotonic)
 
- def test_reconstruct(self):
+ def test_reconstruct_api(self):
+
+ mi = MultiIndex.from_arrays([
+ ['A', 'A', 'B', 'B', 'B'], [1, 2, 1, 2, 3]
+ ])
+
+ with pytest.raises(ValueError):
+ mi._reconstruct()
+
+ with pytest.raises(ValueError):
+ mi._reconstruct(sort=True, remove_unused=True)
+
+ def test_reconstruct_sort(self):
 
  # starts off lexsorted & monotonic
  mi = MultiIndex.from_arrays([
@@ -2456,6 +2468,35 @@ def test_reconstruct(self):
  assert mi.equals(recons)
  assert Index(mi.values).equals(Index(recons.values))
 
+ def test_reconstruct_remove_unused(self):
+ # xref to GH 2770
+ df = DataFrame([['deleteMe', 1, 9],
+ ['keepMe', 2, 9],
+ ['keepMeToo', 3, 9]],
+ columns=['first', 'second', 'third'])
+ df2 = df.set_index(['first', 'second'], drop=False)
+ df2 = df2[df2['first'] != 'deleteMe']
+
+ # removed levels are there
+ expected = MultiIndex(levels=[['deleteMe', 'keepMe', 'keepMeToo'],
+ [1, 2, 3]],
+ labels=[[1, 2], [1, 2]],
+ names=['first', 'second'])
+ result = df2.index
+ tm.assert_index_equal(result, expected)
+
+ expected = MultiIndex(levels=[['keepMe', 'keepMeToo'],
+ [2, 3]],
+ labels=[[0, 1], [0, 1]],
+ names=['first', 'second'])
+ result = df2.index._reconstruct(remove_unused=True)
+ tm.assert_index_equal(result, expected)
+
+ # idempotent
+ result2 = result._reconstruct(remove_unused=True)
+ tm.assert_index_equal(result2, expected)
+ assert result2 is result
+
  def test_isin(self):
  values = [('foo', 2), ('bar', 3), ('quux', 4)]