From 35b338e9ca877224cb3138a9d5cea1f538353a26 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Thu, 8 Jul 2021 05:58:05 -0700 Subject: [PATCH] BUG: .loc failing to drop first level (#42435) --- pandas/core/indexes/multi.py | 10 +++++++--- pandas/core/indexing.py | 17 ++++++++++++----- pandas/tests/frame/indexing/test_xs.py | 13 +++++++------ pandas/tests/indexing/test_loc.py | 12 ++++++++++++ pandas/tests/series/test_arithmetic.py | 2 +- 5 files changed, 39 insertions(+), 15 deletions(-) diff --git a/pandas/core/indexes/multi.py b/pandas/core/indexes/multi.py index 05580b03feba1..8903d29782610 100644 --- a/pandas/core/indexes/multi.py +++ b/pandas/core/indexes/multi.py @@ -3030,16 +3030,20 @@ def maybe_mi_droplevels(indexer, levels): # everything continue else: - raise TypeError( - f"Expected label or tuple of labels, got {key}" - ) + # e.g. test_xs_IndexSlice_argument_not_implemented + k_index = np.zeros(len(self), dtype=bool) + k_index[loc_level] = True + else: k_index = loc_level elif com.is_null_slice(k): # taking everything, does not affect `indexer` below continue + else: + # FIXME: this message can be inaccurate, e.g. + # test_series_varied_multiindex_alignment raise TypeError(f"Expected label or tuple of labels, got {key}") if indexer is None: diff --git a/pandas/core/indexing.py b/pandas/core/indexing.py index e362213c52bd5..387dcca6897b7 100644 --- a/pandas/core/indexing.py +++ b/pandas/core/indexing.py @@ -882,17 +882,21 @@ def _getitem_nested_tuple(self, tup: tuple): if self.name != "loc": # This should never be reached, but lets be explicit about it raise ValueError("Too many indices") - if isinstance(self.obj, ABCSeries) and any( - isinstance(k, tuple) for k in tup - ): - # GH#35349 Raise if tuple in tuple for series - raise ValueError("Too many indices") if all(is_hashable(x) or com.is_null_slice(x) for x in tup): # GH#10521 Series should reduce MultiIndex dimensions instead of # DataFrame, IndexingError is not raised when slice(None,None,None) # with one row. with suppress(IndexingError): return self._handle_lowerdim_multi_index_axis0(tup) + elif isinstance(self.obj, ABCSeries) and any( + isinstance(k, tuple) for k in tup + ): + # GH#35349 Raise if tuple in tuple for series + # Do this after the all-hashable-or-null-slice check so that + # we are only getting non-hashable tuples, in particular ones + # that themselves contain a slice entry + # See test_loc_series_getitem_too_many_dimensions + raise ValueError("Too many indices") # this is a series with a multi-index specified a tuple of # selectors @@ -1127,6 +1131,9 @@ def _handle_lowerdim_multi_index_axis0(self, tup: tuple): return self._get_label(tup, axis=axis) except TypeError as err: # slices are unhashable + # FIXME: this raises when we have a DatetimeIndex first level and a + # string for the first tup entry + # see test_partial_slicing_with_multiindex raise IndexingError("No label returned") from err except KeyError as ek: diff --git a/pandas/tests/frame/indexing/test_xs.py b/pandas/tests/frame/indexing/test_xs.py index a76aec9ebda44..d2704876c31c5 100644 --- a/pandas/tests/frame/indexing/test_xs.py +++ b/pandas/tests/frame/indexing/test_xs.py @@ -318,12 +318,13 @@ def test_xs_IndexSlice_argument_not_implemented(self, klass): if klass is Series: obj = obj[0] - msg = ( - "Expected label or tuple of labels, got " - r"\(\('foo', 'qux', 0\), slice\(None, None, None\)\)" - ) - with pytest.raises(TypeError, match=msg): - obj.xs(IndexSlice[("foo", "qux", 0), :]) + expected = obj.iloc[-2:].droplevel(0) + + result = obj.xs(IndexSlice[("foo", "qux", 0), :]) + tm.assert_equal(result, expected) + + result = obj.loc[IndexSlice[("foo", "qux", 0), :]] + tm.assert_equal(result, expected) @pytest.mark.parametrize("klass", [DataFrame, Series]) def test_xs_levels_raises(self, klass): diff --git a/pandas/tests/indexing/test_loc.py b/pandas/tests/indexing/test_loc.py index e96b25418d408..2e0f48849530a 100644 --- a/pandas/tests/indexing/test_loc.py +++ b/pandas/tests/indexing/test_loc.py @@ -1663,6 +1663,18 @@ def test_loc_multiindex_levels_contain_values_not_in_index_anymore(self, lt_valu with pytest.raises(KeyError, match=r"\['b'\] not in index"): df.loc[df["a"] < lt_value, :].loc[["b"], :] + def test_loc_drops_level(self): + # Based on test_series_varied_multiindex_alignment, where + # this used to fail to drop the first level + mi = MultiIndex.from_product( + [list("ab"), list("xy"), [1, 2]], names=["ab", "xy", "num"] + ) + ser = Series(range(8), index=mi) + + loc_result = ser.loc["a", :, :] + expected = ser.index.droplevel(0)[:4] + tm.assert_index_equal(loc_result.index, expected) + class TestLocSetitemWithExpansion: @pytest.mark.slow diff --git a/pandas/tests/series/test_arithmetic.py b/pandas/tests/series/test_arithmetic.py index 6b7c2c698c211..4d1c75da72399 100644 --- a/pandas/tests/series/test_arithmetic.py +++ b/pandas/tests/series/test_arithmetic.py @@ -924,7 +924,7 @@ def test_series_varied_multiindex_alignment(): [1000 * i for i in range(1, 5)], index=pd.MultiIndex.from_product([list("xy"), [1, 2]], names=["xy", "num"]), ) - result = s1.loc[pd.IndexSlice["a", :, :]] + s2 + result = s1.loc[pd.IndexSlice[["a"], :, :]] + s2 expected = Series( [1000, 2001, 3002, 4003], index=pd.MultiIndex.from_tuples(