ENH: Index.reindex refactoring and getting drop to work with MultiIndex. address GH #101

Author: wesm

pandas/core/generic.py

@@ -157,17 +157,9 @@ def drop(self, labels, axis=0):
         """
         axis_name = self._get_axis_name(axis)
         axis = self._get_axis(axis)
-
-        labels = np.asarray(list(labels), dtype=object)
-
-        indexer, mask = axis.get_indexer(labels)
-        if not mask.all():
-            raise ValueError('labels %s not contained in axis' % labels[-mask])
-
-        new_axis = np.delete(np.asarray(axis), indexer)
+        new_axis = axis.drop(labels)
         return self.reindex(**{axis_name : new_axis})
 
-
 class NDFrame(PandasObject):
     """
 

pandas/core/index.py

@@ -268,6 +268,10 @@ def get_indexer(self, target, method=None):
                                             target.indexMap, method)
         return indexer, mask
 
+    def reindex(self, target, method=None):
+        indexer, mask = self.get_indexer(target, method=method)
+        return target, indexer, mask
+
     def slice_locs(self, start=None, end=None):
         """
 
@@ -300,6 +304,13 @@ def delete(self, loc):
         arr = np.delete(np.asarray(self), loc)
         return Index(arr)
 
+    def drop(self, labels):
+        labels = np.asarray(list(labels), dtype=object)
+        indexer, mask = self.get_indexer(labels)
+        if not mask.all():
+            raise ValueError('labels %s not contained in axis' % labels[-mask])
+        return self.delete(indexer)
+
 class DateIndex(Index):
     pass
 
@@ -486,6 +497,25 @@ def take(self, *args, **kwargs):
         new_labels = [lab.take(*args, **kwargs) for lab in self.labels]
         return MultiIndex(levels=self.levels, labels=new_labels)
 
+    def drop(self, labels):
+        try:
+            arr = np.asarray(list(labels), dtype=object)
+            indexer, mask = self.get_indexer(arr)
+            if not mask.all():
+                raise ValueError('labels %s not contained in axis' % arr[-mask])
+        except Exception:
+            pass
+
+        inds = []
+        for label in labels:
+            loc = self.get_loc(label)
+            if isinstance(loc, int):
+                inds.append(loc)
+            else:
+                inds.extend(range(loc.start, loc.stop))
+
+        return self.delete(inds)
+
     def droplevel(self, level=0):
         """
         Return Index with requested level removed. If MultiIndex has only 2
@@ -556,17 +586,32 @@ def get_indexer(self, target, method=None):
         }
         method = aliases.get(method, method)
 
-        if not isinstance(target, MultiIndex):
-            raise TypeError('Can only align with other MultiIndex objects')
+        if isinstance(target, MultiIndex):
+            target_index = target.get_tuple_index()
+        else:
+            if len(target) > 0:
+                val = target[0]
+                if not isinstance(val, tuple) or len(val) != self.nlevels:
+                    raise ValueError('can only pass MultiIndex or '
+                                     'array of tuples')
 
-        self_index = self.get_tuple_index()
-        target_index = target.get_tuple_index()
+            target_index = target
 
+        self_index = self.get_tuple_index()
         indexer, mask = _tseries.getFillVec(self_index, target_index,
                                             self_index.indexMap,
-                                            target_index.indexMap, method)
+                                            target.indexMap, method)
         return indexer, mask
 
+    def reindex(self, target, method=None):
+        indexer, mask = self.get_indexer(target, method=method)
+
+        # hopefully?
+        if not isinstance(target, MultiIndex):
+            target = MultiIndex.from_tuples(target)
+
+        return target, indexer, mask
+
     def get_tuple_index(self):
         return Index(list(self))
 
@@ -699,14 +744,18 @@ def equals(self, other):
         if len(self) != len(other):
             return False
 
-        # if not self.equal_levels(other):
-        #     return False
-
         for i in xrange(self.nlevels):
-            if not self.levels[i].equals(other.levels[i]):
-                return False
-            if not np.array_equal(self.labels[i], other.labels[i]):
+            svalues = np.asarray(self.levels[i]).take(self.labels[i])
+            ovalues = np.asarray(other.levels[i]).take(other.labels[i])
+
+            if not np.array_equal(svalues, ovalues):
                 return False
+
+            # if not self.levels[i].equals(other.levels[i]):
+            #     return False
+            # if not np.array_equal(self.labels[i], other.labels[i]):
+            #     return False
+
         return True
 
     def equal_levels(self, other):

pandas/core/internals.py

@@ -124,7 +124,7 @@ def reindex_items_from(self, new_ref_items):
         -------
         reindexed : Block
         """
-        indexer, mask = self.items.get_indexer(new_ref_items)
+        new_ref_items, indexer, mask = self.items.reindex(new_ref_items)
         masked_idx = indexer[mask]
         new_values = self.values.take(masked_idx, axis=0)
         new_items = self.items.take(masked_idx)
@@ -525,7 +525,10 @@ def _delete_from_block(self, i, item):
 
     def _add_new_block(self, item, value):
         # Do we care about dtype at the moment?
-        new_block = make_block(value, [item], self.items)
+
+        # hm, elaborate hack?
+        loc = self.items.get_loc(item)
+        new_block = make_block(value, self.items[loc:loc+1], self.items)
         self.blocks.append(new_block)
 
     def _find_block(self, item):
@@ -546,7 +549,7 @@ def reindex_axis(self, new_axis, method=None, axis=0):
         new_axis = _ensure_index(new_axis)
         cur_axis = self.axes[axis]
 
-        indexer, mask = cur_axis.get_indexer(new_axis, method)
+        new_axis, indexer, mask = cur_axis.reindex(new_axis, method)
 
         # TODO: deal with length-0 case? or does it fall out?
         notmask = -mask
@@ -572,16 +575,16 @@ def reindex_items(self, new_items):
             data = data.consolidate()
             return data.reindex_items(new_items)
 
+        # TODO: this part could be faster (!)
+        new_items, _, mask = self.items.reindex(new_items)
+        notmask = -mask
+
         new_blocks = []
         for block in self.blocks:
             newb = block.reindex_items_from(new_items)
             if len(newb.items) > 0:
                 new_blocks.append(newb)
 
-        # TODO: this part could be faster (!)
-        _, mask = self.items.get_indexer(new_items)
-        notmask = -mask
-
         if notmask.any():
             extra_items = new_items[notmask]
 

pandas/core/series.py

@@ -1061,7 +1061,7 @@ def reindex(self, index=None, method=None):
         if len(self.index) == 0:
             return Series(nan, index=index)
 
-        fill_vec, mask = self.index.get_indexer(index, method=method)
+        new_index, fill_vec, mask = self.index.reindex(index, method=method)
         new_values = self.values.take(fill_vec)
 
         notmask = -mask
@@ -1073,7 +1073,7 @@ def reindex(self, index=None, method=None):
 
             np.putmask(new_values, notmask, nan)
 
-        return Series(new_values, index=index)
+        return Series(new_values, index=new_index)
 
     def reindex_like(self, other, method=None):
         """

pandas/tests/test_index.py

@@ -496,8 +496,13 @@ def test_get_indexer(self):
         assert_almost_equal(r2, rbfill2)
 
         # pass non-MultiIndex
+        r1, r2 = idx1.get_indexer(idx2.get_tuple_index())
+        rexp1, rexp2 = idx1.get_indexer(idx2)
+        assert_almost_equal(r1, rexp1)
+        assert_almost_equal(r2, rexp2)
+
         self.assertRaises(Exception, idx1.get_indexer,
-                          idx2.get_tuple_index())
+                          list(zip(*idx2.get_tuple_index())[0]))
 
     def test_format(self):
         self.index.format()

pandas/tests/test_multilevel.py

@@ -1,10 +1,6 @@
 # pylint: disable-msg=W0612,E1101
-from copy import deepcopy
-from datetime import datetime, timedelta
 from cStringIO import StringIO
-import cPickle as pickle
 import operator
-import os
 import unittest
 
 from numpy import random, nan
@@ -52,6 +48,11 @@ def _test_roundtrip(frame):
         _test_roundtrip(self.ymd)
         _test_roundtrip(self.ymd.T)
 
+    def test_reindex(self):
+        reindexed = self.frame.ix[[('foo', 'one'), ('bar', 'one')]]
+        expected = self.frame.ix[[0, 3]]
+        assert_frame_equal(reindexed, expected)
+
     def test_repr_to_string(self):
         repr(self.frame)
         repr(self.ymd)
@@ -160,7 +161,8 @@ def test_sortlevel_mixed(self):
         dft['foo', 'three'] = 'bar'
 
         sorted_after = dft.sortlevel(1, axis=1)
-        assert_frame_equal(sorted_before, sorted_after.drop(['foo'], axis=1))
+        assert_frame_equal(sorted_before.drop([('foo', 'three')], axis=1),
+                           sorted_after.drop([('foo', 'three')], axis=1))
 
     def test_alignment(self):
         pass