ENH: allow per_axis, per_level multiindex setting

TST: better error messages when levels are not sorted with core/index/get_locs

ENH: add boolean indexer support on per_axis/per_level

BUG: handle a multi-level indexed series passed like with a nested tuple of selectors
     e.g. something like: s.loc['A1':'A3',:,['C1','C3']]

Author: jreback

pandas/core/index.py

@@ -3231,34 +3231,52 @@ def partial_selection(key):
                            if key[i] != slice(None, None)]
                 return indexer, _maybe_drop_levels(indexer, ilevels,
                                                    drop_level)
-        elif isinstance(key, slice):
-            # handle a passed slice for this level
-            start = self._get_level_indexer(key.start,level=level)
-            stop  = self._get_level_indexer(key.stop,level=level)
-            step = key.step
-            indexer = slice(start.start,stop.start,step)
-            return indexer, _maybe_drop_levels(indexer, [level], drop_level)
         else:
             indexer = self._get_level_indexer(key, level=level)
-            new_index = _maybe_drop_levels(indexer, [level], drop_level)
-            return indexer, new_index
+            return indexer, _maybe_drop_levels(indexer, [level], drop_level)
 
     def _get_level_indexer(self, key, level=0):
+        # return a boolean indexer or a slice showing where the key is
+        # in the totality of values
+
         level_index = self.levels[level]
-        loc = level_index.get_loc(key)
         labels = self.labels[level]
 
-        if level > 0 or self.lexsort_depth == 0:
-            return np.array(labels == loc,dtype=bool)
+        if isinstance(key, slice):
+            # handle a slice, returnig a slice if we can
+            # otherwise a boolean indexer
+
+            start = level_index.get_loc(key.start)
+            stop  = level_index.get_loc(key.stop)
+            step = key.step
+
+            if level > 0 or self.lexsort_depth == 0:
+                # need to have like semantics here to right
+                # searching as when we are using a slice
+                # so include the stop+1 (so we include stop)
+                m = np.zeros(len(labels),dtype=bool)
+                m[np.in1d(labels,np.arange(start,stop+1,step))] = True
+                return m
+            else:
+                # sorted, so can return slice object -> view
+                i = labels.searchsorted(start, side='left')
+                j = labels.searchsorted(stop, side='right')
+                return slice(i, j, step)
+
         else:
-            # sorted, so can return slice object -> view
-            i = labels.searchsorted(loc, side='left')
-            j = labels.searchsorted(loc, side='right')
-            return slice(i, j)
+
+            loc = level_index.get_loc(key)
+            if level > 0 or self.lexsort_depth == 0:
+                return np.array(labels == loc,dtype=bool)
+            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_locs(self, tup):
         """
-        Given a tuple of slices/lists/labels to a level-wise spec
+        Given a tuple of slices/lists/labels/boolean indexer to a level-wise spec
         produce an indexer to extract those locations
 
         Parameters
@@ -3272,8 +3290,11 @@ def get_locs(self, tup):
         """
 
         # must be lexsorted to at least as many levels
-        assert self.is_lexsorted_for_tuple(tup)
-        assert self.is_unique
+        if not self.is_lexsorted_for_tuple(tup):
+            raise KeyError('MultiIndex Slicing requires the index to be fully lexsorted'
+                           ' tuple len ({0}), lexsort depth ({1})'.format(len(tup), self.lexsort_depth))
+        if not self.is_unique:
+            raise ValueError('MultiIndex Slicing requires a unique index')
 
         def _convert_indexer(r):
             if isinstance(r, slice):
@@ -3285,7 +3306,14 @@ def _convert_indexer(r):
         ranges = []
         for i,k in enumerate(tup):
 
-            if com.is_list_like(k):
+            if com._is_bool_indexer(k):
+                # a boolean indexer, must be the same length!
+                k = np.asarray(k)
+                if len(k) != len(self):
+                    raise ValueError("cannot index with a boolean indexer that is"
+                                     " not the same length as the index")
+                ranges.append(k)
+            elif com.is_list_like(k):
                 # a collection of labels to include from this level (these are or'd)
                 ranges.append(reduce(
                     np.logical_or,[ _convert_indexer(self._get_level_indexer(x, level=i)
@@ -3294,10 +3322,8 @@ def _convert_indexer(r):
                 # include all from this level
                 pass
             elif isinstance(k,slice):
-                start = self._get_level_indexer(k.start,level=i)
-                stop  = self._get_level_indexer(k.stop,level=i)
-                step = k.step
-                ranges.append(slice(start.start,stop.start,step))
+                # a slice, include BOTH of the labels
+                ranges.append(self._get_level_indexer(k,level=i))
             else:
                 # a single label
                 ranges.append(self.get_loc_level(k,level=i,drop_level=False)[0])

pandas/core/indexing.py

@@ -99,6 +99,7 @@ def _slice(self, obj, axis=0, raise_on_error=False, typ=None):
                                typ=typ)
 
     def __setitem__(self, key, value):
+
         # kludgetastic
         ax = self.obj._get_axis(0)
         if isinstance(ax, MultiIndex):
@@ -131,6 +132,11 @@ def _has_valid_tuple(self, key):
                 raise ValueError("Location based indexing can only have [%s] "
                                  "types" % self._valid_types)
 
+    def _is_nested_tuple_indexer(self, tup):
+        if any([ isinstance(ax, MultiIndex) for ax in self.obj.axes ]):
+            return any([ _is_nested_tuple(tup,ax) for ax in self.obj.axes ])
+        return False
+
     def _convert_tuple(self, key, is_setter=False):
         keyidx = []
         for i, k in enumerate(key):
@@ -716,9 +722,8 @@ def _handle_lowerdim_multi_index_axis0(self, tup):
     def _getitem_lowerdim(self, tup):
 
         # we may have a nested tuples indexer here
-        if any([ isinstance(ax, MultiIndex) for ax in self.obj.axes ]):
-            if any([ _is_nested_tuple(tup,ax) for ax in self.obj.axes ]):
-                return self._getitem_nested_tuple(tup)
+        if self._is_nested_tuple_indexer(tup):
+            return self._getitem_nested_tuple(tup)
 
         # we maybe be using a tuple to represent multiple dimensions here
         ax0 = self.obj._get_axis(0)
@@ -772,7 +777,12 @@ def _getitem_nested_tuple(self, tup):
         # multi-index dimension, try to see if we have something like
         # a tuple passed to a series with a multi-index
         if len(tup) > self.ndim:
-            return self._handle_lowerdim_multi_index_axis0(tup)
+            result = self._handle_lowerdim_multi_index_axis0(tup)
+            if result is not None:
+                return result
+
+            # this is a series with a multi-index specified a tuple of selectors
+            return self._getitem_axis(tup, axis=0, validate_iterable=True)
 
         # handle the multi-axis by taking sections and reducing
         # this is iterative
@@ -983,6 +993,8 @@ def _convert_to_indexer(self, obj, axis=0, is_setter=False):
         if isinstance(obj, slice):
             return self._convert_slice_indexer(obj, axis)
 
+        elif _is_nested_tuple(obj, labels):
+            return labels.get_locs(obj)
         elif _is_list_like(obj):
             if com._is_bool_indexer(obj):
                 obj = _check_bool_indexer(labels, obj)

pandas/tests/test_indexing.py

@@ -1072,7 +1072,13 @@ def mklbl(prefix,n):
         ix = MultiIndex.from_product([mklbl('A',5),mklbl('B',7),mklbl('C',4),mklbl('D',2)])
         df = DataFrame(np.arange(len(ix.get_values())),index=ix)
         result = df.loc[(slice('A1','A3'),slice(None), ['C1','C3']),:]
-        expected = df.loc[[ tuple([a,b,c,d]) for a,b,c,d in df.index.values if (a == 'A1' or a == 'A2') and (c == 'C1' or c == 'C3')]]
+        expected = df.loc[[ tuple([a,b,c,d]) for a,b,c,d in df.index.values if (
+            a == 'A1' or a == 'A2' or a == 'A3') and (c == 'C1' or c == 'C3')]]
+        assert_frame_equal(result, expected)
+
+        expected = df.loc[[ tuple([a,b,c,d]) for a,b,c,d in df.index.values if (
+            a == 'A1' or a == 'A2' or a == 'A3') and (c == 'C1' or c == 'C2' or c == 'C3')]]
+        result = df.loc[(slice('A1','A3'),slice(None), slice('C1','C3')),:]
         assert_frame_equal(result, expected)
 
         # test multi-index slicing with per axis and per index controls
@@ -1121,13 +1127,164 @@ def mklbl(prefix,n):
         expected = df.iloc[[0,1,3]]
         assert_frame_equal(result, expected)
 
+        # multi-level series
+        s = Series(np.arange(len(ix.get_values())),index=ix)
+        result = s.loc['A1':'A3', :, ['C1','C3']]
+        expected = s.loc[[ tuple([a,b,c,d]) for a,b,c,d in s.index.values if (
+            a == 'A1' or a == 'A2' or a == 'A3') and (c == 'C1' or c == 'C3')]]
+        assert_series_equal(result, expected)
+
+        # boolean indexers
+        result = df.loc[(slice(None),df.loc[:,('a','bar')]>5),:]
+        expected = df.iloc[[2,3]]
+        assert_frame_equal(result, expected)
+
+        def f():
+            df.loc[(slice(None),np.array([True,False])),:]
+        self.assertRaises(ValueError, f)
+
         # ambiguous cases
         # these can be multiply interpreted
         # but we can catch this in some cases
         def f():
             df.loc[(slice(None),[1])]
         self.assertRaises(KeyError, f)
 
+    def test_per_axis_per_level_getitem_doc_examples(self):
+
+        # from indexing.rst / advanced
+        def mklbl(prefix,n):
+            return ["%s%s" % (prefix,i)  for i in range(n)]
+
+        index = MultiIndex.from_product([mklbl('A',4),
+                                         mklbl('B',2),
+                                         mklbl('C',4),
+                                         mklbl('D',2)])
+        columns = MultiIndex.from_tuples([('a','foo'),('a','bar'),
+                                          ('b','foo'),('b','bah')],
+                                         names=['lvl0', 'lvl1'])
+        df = DataFrame(np.arange(len(index)*len(columns)).reshape((len(index),len(columns))),
+                       index=index,
+                       columns=columns)
+        result = df.loc[(slice('A1','A3'),slice(None), ['C1','C3']),:]
+        expected = df.loc[[ tuple([a,b,c,d]) for a,b,c,d in df.index.values if (
+            a == 'A1' or a == 'A2' or a == 'A3') and (c == 'C1' or c == 'C3')]]
+        assert_frame_equal(result, expected)
+
+        result = df.loc[(slice(None),slice(None), ['C1','C3']),:]
+        expected = df.loc[[ tuple([a,b,c,d]) for a,b,c,d in df.index.values if (
+            c == 'C1' or c == 'C3')]]
+        assert_frame_equal(result, expected)
+
+        # not sorted
+        def f():
+            df.loc['A1',(slice(None),'foo')]
+        self.assertRaises(KeyError, f)
+        df = df.sortlevel(axis=1)
+
+        df.loc['A1',(slice(None),'foo')]
+        df.loc[(slice(None),slice(None), ['C1','C3']),(slice(None),'foo')]
+
+    def test_per_axis_per_level_setitem(self):
+
+        # test multi-index slicing with per axis and per index controls
+        index = MultiIndex.from_tuples([('A',1),('A',2),('A',3),('B',1)],
+                                       names=['one','two'])
+        columns = MultiIndex.from_tuples([('a','foo'),('a','bar'),('b','foo'),('b','bah')],
+                                         names=['lvl0', 'lvl1'])
+
+        df_orig = DataFrame(np.arange(16).reshape(4, 4), index=index, columns=columns)
+        df_orig = df_orig.sortlevel(axis=0).sortlevel(axis=1)
+
+        # identity
+        df = df_orig.copy()
+        df.loc[(slice(None),slice(None)),:] = 100
+        expected = df_orig.copy()
+        expected.iloc[:,:] = 100
+        assert_frame_equal(df, expected)
+
+        df = df_orig.copy()
+        df.loc[(slice(None),slice(None)),(slice(None),slice(None))] = 100
+        expected = df_orig.copy()
+        expected.iloc[:,:] = 100
+        assert_frame_equal(df, expected)
+
+        df = df_orig.copy()
+        df.loc[:,(slice(None),slice(None))] = 100
+        expected = df_orig.copy()
+        expected.iloc[:,:] = 100
+        assert_frame_equal(df, expected)
+
+        # index
+        df = df_orig.copy()
+        df.loc[(slice(None),[1]),:] = 100
+        expected = df_orig.copy()
+        expected.iloc[[0,3]] = 100
+        assert_frame_equal(df, expected)
+
+        df = df_orig.copy()
+        df.loc[(slice(None),1),:] = 100
+        expected = df_orig.copy()
+        expected.iloc[[0,3]] = 100
+        assert_frame_equal(df, expected)
+
+        # columns
+        df = df_orig.copy()
+        df.loc[:,(slice(None),['foo'])] = 100
+        expected = df_orig.copy()
+        expected.iloc[:,[1,3]] = 100
+        assert_frame_equal(df, expected)
+
+        # both
+        df = df_orig.copy()
+        df.loc[(slice(None),1),(slice(None),['foo'])] = 100
+        expected = df_orig.copy()
+        expected.iloc[[0,3],[1,3]] = 100
+        assert_frame_equal(df, expected)
+
+        df = df_orig.copy()
+        df.loc['A','a'] = 100
+        expected = df_orig.copy()
+        expected.iloc[0:3,0:2] = 100
+        assert_frame_equal(df, expected)
+
+        # setting with a list-like
+        df = df_orig.copy()
+        df.loc[(slice(None),1),(slice(None),['foo'])] = np.array([[100, 100], [100, 100]],dtype='int64')
+        expected = df_orig.copy()
+        expected.iloc[[0,3],[1,3]] = 100
+        assert_frame_equal(df, expected)
+
+        # not enough values
+        df = df_orig.copy()
+        def f():
+            df.loc[(slice(None),1),(slice(None),['foo'])] = np.array([[100], [100, 100]],dtype='int64')
+        self.assertRaises(ValueError, f)
+        def f():
+            df.loc[(slice(None),1),(slice(None),['foo'])] = np.array([100, 100, 100, 100],dtype='int64')
+        self.assertRaises(ValueError, f)
+
+        # with an alignable rhs
+        df = df_orig.copy()
+        df.loc[(slice(None),1),(slice(None),['foo'])] = df.loc[(slice(None),1),(slice(None),['foo'])] * 5
+        expected = df_orig.copy()
+        expected.iloc[[0,3],[1,3]] = expected.iloc[[0,3],[1,3]] * 5
+        assert_frame_equal(df, expected)
+
+        df = df_orig.copy()
+        df.loc[(slice(None),1),(slice(None),['foo'])] *= df.loc[(slice(None),1),(slice(None),['foo'])]
+        expected = df_orig.copy()
+        expected.iloc[[0,3],[1,3]] *= expected.iloc[[0,3],[1,3]]
+        assert_frame_equal(df, expected)
+
+        rhs = df_orig.loc[(slice(None),1),(slice(None),['foo'])].copy()
+        rhs.loc[:,('c','bah')] = 10
+        df = df_orig.copy()
+        df.loc[(slice(None),1),(slice(None),['foo'])] *= rhs
+        expected = df_orig.copy()
+        expected.iloc[[0,3],[1,3]] *= expected.iloc[[0,3],[1,3]]
+        assert_frame_equal(df, expected)
+
     def test_getitem_multiindex(self):
 
         # GH 5725