reindex multi-index at level with reordered labels

Author: behzadnouri

doc/source/whatsnew/v0.15.2.txt

@@ -119,6 +119,7 @@ Bug Fixes
 - Bug in ``BlockManager`` where setting values with different type would break block integrity (:issue:`8850`)
 - Bug in ``DatetimeIndex`` when using ``time`` object as key (:issue:`8667`)
 - Bug in ``merge`` where ``how='left'`` and ``sort=False`` would not preserve left frame order (:issue:`7331`)
+- Bug in ``MultiIndex.reindex`` where reindexing at level would not reorder labels (:issue:`4088`)
 
 - Fix negative step support for label-based slices (:issue:`8753`)
 

pandas/core/index.py

@@ -1828,13 +1828,41 @@ def _join_non_unique(self, other, how='left', return_indexers=False):
         else:
             return join_index
 
-    def _join_level(self, other, level, how='left', return_indexers=False):
+    def _join_level(self, other, level, how='left',
+                    return_indexers=False,
+                    keep_order=True):
         """
         The join method *only* affects the level of the resulting
         MultiIndex. Otherwise it just exactly aligns the Index data to the
-        labels of the level in the MultiIndex. The order of the data indexed by
-        the MultiIndex will not be changed (currently)
-        """
+        labels of the level in the MultiIndex. If `keep_order` == True, the
+        order of the data indexed by the MultiIndex will not be changed;
+        otherwise, it will tie out with `other`.
+        """
+        from pandas.algos import groupsort_indexer
+
+        def _get_leaf_sorter(labels):
+            '''
+            returns sorter for the inner most level while preserving the
+            order of higher levels
+            '''
+            if labels[0].size == 0:
+                return np.empty(0, dtype='int64')
+
+            if len(labels) == 1:
+                lab = com._ensure_int64(labels[0])
+                sorter, _ = groupsort_indexer(lab, 1 + lab.max())
+                return sorter
+
+            # find indexers of begining of each set of
+            # same-key labels w.r.t all but last level
+            tic = labels[0][:-1] != labels[0][1:]
+            for lab in labels[1:-1]:
+                tic |= lab[:-1] != lab[1:]
+
+            starts = np.hstack(([True], tic, [True])).nonzero()[0]
+            lab = com._ensure_int64(labels[-1])
+            return lib.get_level_sorter(lab, starts)
+
         if isinstance(self, MultiIndex) and isinstance(other, MultiIndex):
             raise TypeError('Join on level between two MultiIndex objects '
                             'is ambiguous')
@@ -1849,33 +1877,69 @@ def _join_level(self, other, level, how='left', return_indexers=False):
         level = left._get_level_number(level)
         old_level = left.levels[level]
 
+        if not right.is_unique:
+            raise NotImplementedError('Index._join_level on non-unique index '
+                                      'is not implemented')
+
         new_level, left_lev_indexer, right_lev_indexer = \
             old_level.join(right, how=how, return_indexers=True)
 
-        if left_lev_indexer is not None:
+        if left_lev_indexer is None:
+            if keep_order or len(left) == 0:
+                left_indexer = None
+                join_index = left
+            else:  # sort the leaves
+                left_indexer = _get_leaf_sorter(left.labels[:level + 1])
+                join_index = left[left_indexer]
+
+        else:
             left_lev_indexer = com._ensure_int64(left_lev_indexer)
             rev_indexer = lib.get_reverse_indexer(left_lev_indexer,
                                                   len(old_level))
 
             new_lev_labels = com.take_nd(rev_indexer, left.labels[level],
                                          allow_fill=False)
-            omit_mask = new_lev_labels != -1
 
             new_labels = list(left.labels)
             new_labels[level] = new_lev_labels
 
-            if not omit_mask.all():
-                new_labels = [lab[omit_mask] for lab in new_labels]
-
             new_levels = list(left.levels)
             new_levels[level] = new_level
 
-            join_index = MultiIndex(levels=new_levels, labels=new_labels,
-                                    names=left.names, verify_integrity=False)
-            left_indexer = np.arange(len(left))[new_lev_labels != -1]
-        else:
-            join_index = left
-            left_indexer = None
+            if keep_order:  # just drop missing values. o.w. keep order
+                left_indexer = np.arange(len(left))
+                mask = new_lev_labels != -1
+                if not mask.all():
+                    new_labels = [lab[mask] for lab in new_labels]
+                    left_indexer = left_indexer[mask]
+
+            else:  # tie out the order with other
+                if level == 0:  # outer most level, take the fast route
+                    ngroups = 1 + new_lev_labels.max()
+                    left_indexer, counts = groupsort_indexer(new_lev_labels,
+                                                             ngroups)
+                    # missing values are placed first; drop them!
+                    left_indexer = left_indexer[counts[0]:]
+                    new_labels = [lab[left_indexer] for lab in new_labels]
+
+                else:  # sort the leaves
+                    mask = new_lev_labels != -1
+                    mask_all = mask.all()
+                    if not mask_all:
+                        new_labels = [lab[mask] for lab in new_labels]
+
+                    left_indexer = _get_leaf_sorter(new_labels[:level + 1])
+                    new_labels = [lab[left_indexer] for lab in new_labels]
+
+                    # left_indexers are w.r.t masked frame.
+                    # reverse to original frame!
+                    if not mask_all:
+                        left_indexer = mask.nonzero()[0][left_indexer]
+
+            join_index = MultiIndex(levels=new_levels,
+                                    labels=new_labels,
+                                    names=left.names,
+                                    verify_integrity=False)
 
         if right_lev_indexer is not None:
             right_indexer = com.take_nd(right_lev_indexer,
@@ -3925,7 +3989,8 @@ def reindex(self, target, method=None, level=None, limit=None):
             else:
                 target = _ensure_index(target)
             target, indexer, _ = self._join_level(target, level, how='right',
-                                                  return_indexers=True)
+                                                  return_indexers=True,
+                                                  keep_order=False)
         else:
             if self.equals(target):
                 indexer = None

pandas/lib.pyx

@@ -1138,6 +1138,27 @@ def row_bool_subset_object(ndarray[object, ndim=2] values,
 
     return out
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def get_level_sorter(ndarray[int64_t, ndim=1] label,
+                     ndarray[int64_t, ndim=1] starts):
+    """
+    argsort for a single level of a multi-index, keeping the order of higher
+    levels unchanged. `starts` points to starts of same-key indices w.r.t
+    to leading levels; equivalent to:
+        np.hstack([label[starts[i]:starts[i+1]].argsort(kind='mergesort')
+            + starts[i] for i in range(len(starts) - 1)])
+    """
+    cdef:
+        int64_t l, r
+        Py_ssize_t i
+        ndarray[int64_t, ndim=1] out = np.empty(len(label), dtype=np.int64)
+
+    for i in range(len(starts) - 1):
+        l, r = starts[i], starts[i + 1]
+        out[l:r] = l + label[l:r].argsort(kind='mergesort')
+
+    return out
 
 def group_count(ndarray[int64_t] values, Py_ssize_t size):
     cdef:

pandas/tests/test_frame.py

@@ -1897,6 +1897,66 @@ def test_reversed_reindex_ffill_raises(self):
         self.assertRaises(ValueError, df.reindex, dr[::-1], method='ffill')
         self.assertRaises(ValueError, df.reindex, dr[::-1], method='bfill')
 
+    def test_reindex_level(self):
+        from itertools import permutations
+        icol = ['jim', 'joe', 'jolie']
+
+        def verify_first_level(df, level, idx):
+            f = lambda val: np.nonzero(df[level] == val)[0]
+            i = np.concatenate(list(map(f, idx)))
+            left = df.set_index(icol).reindex(idx, level=level)
+            right = df.iloc[i].set_index(icol)
+            assert_frame_equal(left, right)
+
+        def verify(df, level, idx, indexer):
+            left = df.set_index(icol).reindex(idx, level=level)
+            right = df.iloc[indexer].set_index(icol)
+            assert_frame_equal(left, right)
+
+        df = pd.DataFrame({'jim':list('B' * 4 + 'A' * 2 + 'C' * 3),
+                           'joe':list('abcdeabcd')[::-1],
+                           'jolie':[10, 20, 30] * 3,
+                           'joline': np.random.randint(0, 1000, 9)})
+
+        target = [['C', 'B', 'A'], ['F', 'C', 'A', 'D'], ['A'], ['D', 'F'],
+                  ['A', 'B', 'C'], ['C', 'A', 'B'], ['C', 'B'], ['C', 'A'],
+                  ['A', 'B'], ['B', 'A', 'C'], ['A', 'C', 'B']]
+
+        for idx in target:
+            verify_first_level(df, 'jim', idx)
+
+        verify(df, 'joe', list('abcde'), [3, 2, 1, 0, 5, 4, 8, 7, 6])
+        verify(df, 'joe', list('abcd'),  [3, 2, 1, 0, 5, 8, 7, 6])
+        verify(df, 'joe', list('abc'),   [3, 2, 1, 8, 7, 6])
+        verify(df, 'joe', list('eca'),   [1, 3, 4, 6, 8])
+        verify(df, 'joe', list('edc'),   [0, 1, 4, 5, 6])
+        verify(df, 'joe', list('eadbc'), [3, 0, 2, 1, 4, 5, 8, 7, 6])
+        verify(df, 'joe', list('edwq'),  [0, 4, 5])
+        verify(df, 'joe', list('wq'),    [])
+
+        df = DataFrame({'jim':['mid'] * 5 + ['btm'] * 8 + ['top'] * 7,
+                        'joe':['3rd'] * 2 + ['1st'] * 3 + ['2nd'] * 3 +
+                              ['1st'] * 2 + ['3rd'] * 3 + ['1st'] * 2 +
+                              ['3rd'] * 3 + ['2nd'] * 2,
+                        'jolie':np.random.randint(0, 1000, 20),
+                        'joline': np.random.randn(20).round(3) * 10})
+
+        for idx in permutations(df['jim'].unique()):
+            for i in range(3):
+                verify_first_level(df, 'jim', idx[:i+1])
+
+        i = [2,3,4,0,1,8,9,5,6,7,10,11,12,13,14,18,19,15,16,17]
+        verify(df, 'joe', ['1st', '2nd', '3rd'], i)
+
+        i = [0,1,2,3,4,10,11,12,5,6,7,8,9,15,16,17,18,19,13,14]
+        verify(df, 'joe', ['3rd', '2nd', '1st'], i)
+
+        i = [0,1,5,6,7,10,11,12,18,19,15,16,17]
+        verify(df, 'joe', ['2nd', '3rd'], i)
+
+        i = [0,1,2,3,4,10,11,12,8,9,15,16,17,13,14]
+        verify(df, 'joe', ['3rd', '1st'], i)
+
     def test_getitem_ix_float_duplicates(self):
         df = pd.DataFrame(np.random.randn(3, 3),
                           index=[0.1, 0.2, 0.2], columns=list('abc'))