TST: tests for per_axis_per_level_getitem

ENH: add core/indexing.py/_getitem_nested_tuple to handle the nested_tuple cases for partial multi-indexing

Author: jreback

pandas/core/index.py

@@ -3249,7 +3249,7 @@ def _get_level_indexer(self, key, level=0):
             j = labels.searchsorted(loc, side='right')
             return slice(i, j)
 
-    def get_locs(self, tup):
+    def get_specs(self, tup):
         """Convert a tuple of slices/label lists/labels to a level-wise spec
 
         Parameters
@@ -3322,8 +3322,9 @@ def get_locs(self, tup):
                     stop = level.get_loc(k.stop)
             else:
                 # a single label
-                start = level.get_loc(k)
-                stop = start
+                # make this into a list of a tuple
+                ranges.append([level.get_loc(k)])
+                continue
 
             ranges.append((start,stop))
 
@@ -3337,14 +3338,14 @@ def get_locs(self, tup):
 
         return ranges
 
-    def locs_to_indexer(self, specs):
+    def specs_to_indexer(self, specs):
         """ Take a location specification to an indexer
 
         Parameters
         ----------
         self: a sufficiently lexsorted, unique/non-dupe MultIindex.
         specs: a list of 2-tuples/list of label positions. Specifically, The
-        output of _tuple_to_mi_locs.
+        output of get_specs
         len(specs) must matc ix.nlevels.
 
         Returns
@@ -3362,8 +3363,8 @@ def locs_to_indexer(self, specs):
         ('A2', 'B0')
         ('A2', 'B1')
 
-        >>> locs = mi.get_locs((slice('A0','A2'),['B0', 'B1']))
-        >>> list(mi.locs_to_indexer(locs))
+        >>> locs = mi.get_specs((slice('A0','A2'),['B0', 'B1']))
+        >>> list(mi.specs_to_indexer(locs))
         [0, 1, 2, 3]
 
         Which are all the labels having 'A0' to 'A2' (non-inclusive) at level=0
@@ -3393,33 +3394,11 @@ def _iter_vectorize(specs, i=0):
             if len(specs)-1 == i:
                 return np.array(valrange)
             else:
-                tmpl = np.array([v for v in _iter_vectorize(specs,i+1)])
+                tmpl=np.array([v for v in _iter_vectorize(specs,i+1)])
                 res=np.tile(tmpl,(len(valrange),1))
                 steps=(np.array(valrange)*step_size).reshape((len(valrange),1))
                 return (res+steps).flatten()
 
-
-        def _iter_generator(specs, i=0):
-            step_size = giant_steps[i]
-            spec=specs[i]
-            if isinstance(spec,tuple):
-                # tuples are 2-tuples of (start,stop) label indices to include
-                valrange = compat.range(*spec)
-            elif isinstance(spec,list):
-                # lists are discrete label indicies to include
-                valrange = spec
-
-            if len(specs)-1 == i:
-                # base case
-                for v in valrange:
-                    yield v
-                else:
-                    for base in valrange:
-                        base *= step_size
-                        for v in _iter_generator(specs,i+1):
-                            yield base + v
-            # validate
-
         return _iter_vectorize(specs)
 
     def truncate(self, before=None, after=None):

pandas/core/indexing.py

@@ -693,35 +693,39 @@ def _convert_for_reindex(self, key, axis=0):
 
             return keyarr
 
-    def _getitem_lowerdim(self, tup):
+    def _handle_lowerdim_multi_index_axis0(self, tup):
+        # we have an axis0 multi-index, handle or raise
 
-        ax0 = self.obj._get_axis(0)
-        # a bit kludgy
-        if isinstance(ax0, MultiIndex):
-            try:
-                # fast path for series or for tup devoid of slices
-                return self._get_label(tup, axis=0)
-            except TypeError:
-                # slices are unhashable
-                pass
-            except Exception as e1:
-                if isinstance(tup[0], (slice, Index)):
-                    raise IndexingError("Handle elsewhere")
+        try:
+            # fast path for series or for tup devoid of slices
+            return self._get_label(tup, axis=0)
+        except TypeError:
+            # slices are unhashable
+            pass
+        except Exception as e1:
+            if isinstance(tup[0], (slice, Index)):
+                raise IndexingError("Handle elsewhere")
 
-                # raise the error if we are not sorted
-                if not ax0.is_lexsorted_for_tuple(tup):
-                    raise e1
+            # raise the error if we are not sorted
+            ax0 = self.obj._get_axis(0)
+            if not ax0.is_lexsorted_for_tuple(tup):
+                raise e1
 
-                # GH911 introduced this clause, but the regression test
-                # added for it now passes even without it. Let's rock the boat.
-                # 2014/01/27
+        return None
 
-                # # should we abort, or keep going?
-                # try:
-                #     loc = ax0.get_loc(tup[0])
-                # except KeyError:
-                #     raise e1
+    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)
+
+        # we maybe be using a tuple to represent multiple dimensions here
+        ax0 = self.obj._get_axis(0)
+        if isinstance(ax0, MultiIndex):
+            result = self._handle_lowerdim_multi_index_axis0(tup)
+            if result is not None:
+                return result
 
         if len(tup) > self.obj.ndim:
             raise IndexingError("Too many indexers. handle elsewhere")
@@ -760,7 +764,31 @@ def _getitem_lowerdim(self, tup):
 
         raise IndexingError('not applicable')
 
-    def _getitem_axis(self, key, axis=0):
+    def _getitem_nested_tuple(self, tup):
+        # we have a nested tuple so have at least 1 multi-index level
+        # we should be able to match up the dimensionaility here
+
+        # we have too many indexers for our dim, but have at least 1
+        # 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)
+
+        # handle the multi-axis by taking sections and reducing
+        # this is iterative
+        obj = self.obj
+        axis = 0
+        for key in tup:
+
+            obj = getattr(obj, self.name)._getitem_axis(key, axis=axis, validate_iterable=True)
+            axis += 1
+
+            if obj.ndim < self.ndim:
+                axis -= 1
+
+        return obj
+
+    def _getitem_axis(self, key, axis=0, validate_iterable=False):
 
         self._has_valid_type(key, axis)
         labels = self.obj._get_axis(axis)
@@ -1058,7 +1086,7 @@ def __getitem__(self, key):
         else:
             return self._getitem_axis(key, axis=0)
 
-    def _getitem_axis(self, key, axis=0):
+    def _getitem_axis(self, key, axis=0, validate_iterable=False):
         raise NotImplementedError()
 
     def _getbool_axis(self, key, axis=0):
@@ -1135,6 +1163,7 @@ def _has_valid_type(self, key, axis):
             # require all elements in the index
             idx = _ensure_index(key)
             if not idx.isin(ax).all():
+
                 raise KeyError("[%s] are not in ALL in the [%s]" %
                                (key, self.obj._get_axis_name(axis)))
 
@@ -1164,7 +1193,7 @@ def error():
 
         return True
 
-    def _getitem_axis(self, key, axis=0):
+    def _getitem_axis(self, key, axis=0, validate_iterable=False):
         labels = self.obj._get_axis(axis)
 
         if isinstance(key, slice):
@@ -1178,12 +1207,15 @@ def _getitem_axis(self, key, axis=0):
             if hasattr(key, 'ndim') and key.ndim > 1:
                 raise ValueError('Cannot index with multidimensional key')
 
+            if validate_iterable:
+                self._has_valid_type(key, axis)
             return self._getitem_iterable(key, axis=axis)
-        elif isinstance(key, tuple) and isinstance(labels, MultiIndex) and \
-            any([isinstance(x,slice) for x in key]):
-            locs = labels.get_locs(key)
-            g = labels.locs_to_indexer(locs)
-            return self.obj.iloc[g]
+        elif _is_nested_tuple(key, labels):
+            specs = labels.get_specs(key)
+            g = labels.specs_to_indexer(specs)
+            indexer = [ slice(None) ] * self.ndim
+            indexer[axis] = g
+            return self.obj.iloc[tuple(indexer)]
         else:
             self._has_valid_type(key, axis)
             return self._get_label(key, axis=axis)
@@ -1256,7 +1288,7 @@ def _get_slice_axis(self, slice_obj, axis=0):
         else:
             return self.obj.take(slice_obj, axis=axis, convert=False)
 
-    def _getitem_axis(self, key, axis=0):
+    def _getitem_axis(self, key, axis=0, validate_iterable=False):
 
         if isinstance(key, slice):
             self._has_valid_type(key, axis)
@@ -1515,6 +1547,24 @@ def _maybe_convert_ix(*args):
         return args
 
 
+def _is_nested_tuple(tup, labels):
+    # check for a compatiable nested tuple and multiindexes among the axes
+
+    if not isinstance(tup, tuple):
+        return False
+
+    # are we nested tuple of: tuple,list,slice
+    for i, k in enumerate(tup):
+
+        #if i > len(axes):
+        #    raise IndexingError("invalid indxing tuple passed, has too many indexers for this object")
+        #ax = axes[i]
+        if isinstance(k, (tuple, list, slice)):
+            return isinstance(labels, MultiIndex)
+
+    return False
+
+
 def _is_null_slice(obj):
     return (isinstance(obj, slice) and obj.start is None and
             obj.stop is None and obj.step is None)

pandas/tests/test_indexing.py

@@ -1062,6 +1062,62 @@ def test_xs_multiindex(self):
         expected.columns = expected.columns.droplevel('lvl1')
         assert_frame_equal(result, expected)
 
+    def test_per_axis_per_level_getitem(self):
+
+        # GH6134
+        # example test case
+        def mklbl(prefix,n):
+            return ["%s%s" % (prefix,i)  for i in range(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')]]
+        assert_frame_equal(result, expected)
+
+        # 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','hello'),('b','world')],
+                                         names=['lvl0', 'lvl1'])
+
+        df = DataFrame(np.arange(16).reshape(4, 4), index=index, columns=columns)
+        df = df.sortlevel(axis=0).sortlevel(axis=1)
+
+        # identity
+        result = df.loc[(slice(None),slice(None)),:]
+        assert_frame_equal(result, df)
+        result = df.loc[(slice(None),slice(None)),(slice(None),slice(None))]
+        assert_frame_equal(result, df)
+        result = df.loc[:,(slice(None),slice(None))]
+        assert_frame_equal(result, df)
+
+        # index
+        result = df.loc[(slice(None),[1]),:]
+        expected = df.iloc[[0,3]]
+        assert_frame_equal(result, expected)
+
+        result = df.loc[(slice(None),1),:]
+        expected = df.iloc[[0,3]]
+        assert_frame_equal(result, expected)
+
+        # columns
+        result = df.loc[:,(slice(None),['world'])]
+        expected = df.iloc[:,[3]]
+        assert_frame_equal(result, expected)
+
+        # both
+        result = df.loc[(slice(None),1),(slice(None),['world'])]
+        expected = df.iloc[[0,3],[3]]
+        assert_frame_equal(result, expected)
+
+        # 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_getitem_multiindex(self):
 
         # GH 5725