ENH: indexing object with MultiIndex works if indexing the top level. address GH #120

Author: wesm

pandas/core/frame.py

@@ -793,12 +793,16 @@ def _slice(self, slobj, axis=0):
 
     def _getitem_multilevel(self, key):
         loc = self.columns.get_loc(key)
-        if isinstance(loc, slice):
+        if isinstance(loc, (slice, np.ndarray)):
             new_columns = self.columns[loc]
-            new_columns = _maybe_droplevels(new_columns, key)
-            new_values = self.values[:, loc]
-            result = DataFrame(new_values, index=self.index,
-                               columns=new_columns)
+            result_columns = _maybe_droplevels(new_columns, key)
+            if self._is_mixed_type:
+                result = self.reindex(columns=new_columns)
+                result.columns = result_columns
+            else:
+                new_values = self.values[:, loc]
+                result = DataFrame(new_values, index=self.index,
+                                   columns=result_columns)
             return result
         else:
             return self._getitem_single(key)

pandas/core/index.py

@@ -965,11 +965,13 @@ def get_loc(self, key):
             labels = self.labels[0]
             loc = level.get_loc(key)
 
-            assert(self.lexsort_depth >= 1)
-
-            i = labels.searchsorted(loc, side='left')
-            j = labels.searchsorted(loc, side='right')
-            return slice(i, j)
+            if self.lexsort_depth == 0:
+                return labels == loc
+            else:
+                # sorted, so can return slice object -> view
+                i = labels.searchsorted(loc, side='left')
+                j = labels.searchsorted(loc, side='right')
+                return slice(i, j)
 
     def _get_tuple_loc(self, tup):
         indexer = self._get_label_key(tup)

pandas/core/internals.py

@@ -438,7 +438,8 @@ def xs(self, key, axis=1, copy=True):
 
         new_axes = list(self.axes)
 
-        if isinstance(loc, slice):
+        # could be an array indexer!
+        if isinstance(loc, (slice, np.ndarray)):
             new_axes[axis] = new_axes[axis][loc]
         else:
             new_axes.pop(axis)

pandas/core/series.py

@@ -271,7 +271,7 @@ def _multilevel_index(self, key):
         values = self.values
         try:
             loc = self.index.get_loc(key)
-            if isinstance(loc, slice):
+            if isinstance(loc, (slice, np.ndarray)):
                 # TODO: what if a level contains tuples??
                 new_index = self.index[loc]
                 new_index = _maybe_droplevels(new_index, key)

pandas/tests/test_multilevel.py

@@ -276,7 +276,39 @@ def test_frame_getitem_view(self):
     def test_frame_getitem_not_sorted(self):
         df = self.frame.T
         df['foo', 'four'] = 'foo'
-        self.assertRaises(Exception, df.__getitem__, 'foo')
+
+        arrays = [np.array(x) for x in zip(*df.columns.get_tuple_index())]
+
+        result = df['foo']
+        result2 = df.ix[:, 'foo']
+        expected = df.reindex(columns=df.columns[arrays[0] == 'foo'])
+        expected.columns = expected.columns.droplevel(0)
+        assert_frame_equal(result, expected)
+        assert_frame_equal(result2, expected)
+
+        df = df.T
+        result = df.xs('foo')
+        result2 = df.ix['foo']
+        expected = df.reindex(df.index[arrays[0] == 'foo'])
+        expected.index = expected.index.droplevel(0)
+        assert_frame_equal(result, expected)
+        assert_frame_equal(result2, expected)
+
+    def test_series_getitem_not_sorted(self):
+        arrays = [['bar', 'bar', 'baz', 'baz', 'qux', 'qux', 'foo', 'foo'],
+        ['one', 'two', 'one', 'two', 'one', 'two', 'one', 'two']]
+        tuples = zip(*arrays)
+        index = MultiIndex.from_tuples(tuples)
+        s = Series(randn(8), index=index)
+
+        arrays = [np.array(x) for x in zip(*index.get_tuple_index())]
+
+        result = s['qux']
+        result2 = s.ix['qux']
+        expected = s[arrays[0] == 'qux']
+        expected.index = expected.index.droplevel(0)
+        assert_series_equal(result, expected)
+        assert_series_equal(result2, expected)
 
 if __name__ == '__main__':