import numpy as np import pytest from pandas.errors import ( IndexingError, PerformanceWarning, ) import pandas as pd from pandas import ( DataFrame, Index, MultiIndex, Series, ) import pandas._testing as tm @pytest.fixture def single_level_multiindex(): """single level MultiIndex""" return MultiIndex( levels=[["foo", "bar", "baz", "qux"]], codes=[[0, 1, 2, 3]], names=["first"] ) @pytest.fixture def frame_random_data_integer_multi_index(): levels = [[0, 1], [0, 1, 2]] codes = [[0, 0, 0, 1, 1, 1], [0, 1, 2, 0, 1, 2]] index = MultiIndex(levels=levels, codes=codes) return DataFrame(np.random.default_rng(2).standard_normal((6, 2)), index=index) class TestMultiIndexLoc: def test_loc_setitem_frame_with_multiindex(self, multiindex_dataframe_random_data): frame = multiindex_dataframe_random_data frame.loc[("bar", "two"), "B"] = 5 assert frame.loc[("bar", "two"), "B"] == 5 # with integer labels df = frame.copy() df.columns = list(range(3)) df.loc[("bar", "two"), 1] = 7 assert df.loc[("bar", "two"), 1] == 7 def test_loc_getitem_general(self, any_real_numpy_dtype): # GH#2817 dtype = any_real_numpy_dtype data = { "amount": {0: 700, 1: 600, 2: 222, 3: 333, 4: 444}, "col": {0: 3.5, 1: 3.5, 2: 4.0, 3: 4.0, 4: 4.0}, "num": {0: 12, 1: 11, 2: 12, 3: 12, 4: 12}, } df = DataFrame(data) df = df.astype({"col": dtype, "num": dtype}) df = df.set_index(keys=["col", "num"]) key = 4.0, 12 # emits a PerformanceWarning, ok with tm.assert_produces_warning(PerformanceWarning): tm.assert_frame_equal(df.loc[key], df.iloc[2:]) # this is ok return_value = df.sort_index(inplace=True) assert return_value is None res = df.loc[key] # col has float dtype, result should be float64 Index col_arr = np.array([4.0] * 3, dtype=dtype) year_arr = np.array([12] * 3, dtype=dtype) index = MultiIndex.from_arrays([col_arr, year_arr], names=["col", "num"]) expected = DataFrame({"amount": [222, 333, 444]}, index=index) tm.assert_frame_equal(res, expected) def test_loc_getitem_multiindex_missing_label_raises(self): # GH#21593 df = DataFrame( np.random.default_rng(2).standard_normal((3, 3)), columns=[[2, 2, 4], [6, 8, 10]], index=[[4, 4, 8], [8, 10, 12]], ) with pytest.raises(KeyError, match=r"^2$"): df.loc[2] def test_loc_getitem_list_of_tuples_with_multiindex( self, multiindex_year_month_day_dataframe_random_data ): ser = multiindex_year_month_day_dataframe_random_data["A"] expected = ser.reindex(ser.index[49:51]) result = ser.loc[[(2000, 3, 10), (2000, 3, 13)]] tm.assert_series_equal(result, expected) def test_loc_getitem_series(self): # GH14730 # passing a series as a key with a MultiIndex index = MultiIndex.from_product([[1, 2, 3], ["A", "B", "C"]]) x = Series(index=index, data=range(9), dtype=np.float64) y = Series([1, 3]) expected = Series( data=[0, 1, 2, 6, 7, 8], index=MultiIndex.from_product([[1, 3], ["A", "B", "C"]]), dtype=np.float64, ) result = x.loc[y] tm.assert_series_equal(result, expected) result = x.loc[[1, 3]] tm.assert_series_equal(result, expected) # GH15424 y1 = Series([1, 3], index=[1, 2]) result = x.loc[y1] tm.assert_series_equal(result, expected) empty = Series(data=[], dtype=np.float64) expected = Series( [], index=MultiIndex(levels=index.levels, codes=[[], []], dtype=np.float64), dtype=np.float64, ) result = x.loc[empty] tm.assert_series_equal(result, expected) def test_loc_getitem_array(self): # GH15434 # passing an array as a key with a MultiIndex index = MultiIndex.from_product([[1, 2, 3], ["A", "B", "C"]]) x = Series(index=index, data=range(9), dtype=np.float64) y = np.array([1, 3]) expected = Series( data=[0, 1, 2, 6, 7, 8], index=MultiIndex.from_product([[1, 3], ["A", "B", "C"]]), dtype=np.float64, ) result = x.loc[y] tm.assert_series_equal(result, expected) # empty array: empty = np.array([]) expected = Series( [], index=MultiIndex(levels=index.levels, codes=[[], []], dtype=np.float64), dtype="float64", ) result = x.loc[empty] tm.assert_series_equal(result, expected) # 0-dim array (scalar): scalar = np.int64(1) expected = Series(data=[0, 1, 2], index=["A", "B", "C"], dtype=np.float64) result = x.loc[scalar] tm.assert_series_equal(result, expected) def test_loc_multiindex_labels(self): df = DataFrame( np.random.default_rng(2).standard_normal((3, 3)), columns=[["i", "i", "j"], ["A", "A", "B"]], index=[["i", "i", "j"], ["X", "X", "Y"]], ) # the first 2 rows expected = df.iloc[[0, 1]].droplevel(0) result = df.loc["i"] tm.assert_frame_equal(result, expected) # 2nd (last) column expected = df.iloc[:, [2]].droplevel(0, axis=1) result = df.loc[:, "j"] tm.assert_frame_equal(result, expected) # bottom right corner expected = df.iloc[[2], [2]].droplevel(0).droplevel(0, axis=1) result = df.loc["j"].loc[:, "j"] tm.assert_frame_equal(result, expected) # with a tuple expected = df.iloc[[0, 1]] result = df.loc[("i", "X")] tm.assert_frame_equal(result, expected) def test_loc_multiindex_ints(self): df = DataFrame( np.random.default_rng(2).standard_normal((3, 3)), columns=[[2, 2, 4], [6, 8, 10]], index=[[4, 4, 8], [8, 10, 12]], ) expected = df.iloc[[0, 1]].droplevel(0) result = df.loc[4] tm.assert_frame_equal(result, expected) def test_loc_multiindex_missing_label_raises(self): df = DataFrame( np.random.default_rng(2).standard_normal((3, 3)), columns=[[2, 2, 4], [6, 8, 10]], index=[[4, 4, 8], [8, 10, 12]], ) with pytest.raises(KeyError, match=r"^2$"): df.loc[2] @pytest.mark.parametrize("key, pos", [([2, 4], [0, 1]), ([2], []), ([2, 3], [])]) def test_loc_multiindex_list_missing_label(self, key, pos): # GH 27148 - lists with missing labels _do_ raise df = DataFrame( np.random.default_rng(2).standard_normal((3, 3)), columns=[[2, 2, 4], [6, 8, 10]], index=[[4, 4, 8], [8, 10, 12]], ) with pytest.raises(KeyError, match="not in index"): df.loc[key] def test_loc_multiindex_too_many_dims_raises(self): # GH 14885 s = Series( range(8), index=MultiIndex.from_product([["a", "b"], ["c", "d"], ["e", "f"]]), ) with pytest.raises(KeyError, match=r"^\('a', 'b'\)$"): s.loc["a", "b"] with pytest.raises(KeyError, match=r"^\('a', 'd', 'g'\)$"): s.loc["a", "d", "g"] with pytest.raises(IndexingError, match="Too many indexers"): s.loc["a", "d", "g", "j"] def test_loc_multiindex_indexer_none(self): # GH6788 # multi-index indexer is None (meaning take all) attributes = ["Attribute" + str(i) for i in range(1)] attribute_values = ["Value" + str(i) for i in range(5)] index = MultiIndex.from_product([attributes, attribute_values]) df = 0.1 * np.random.default_rng(2).standard_normal((10, 1 * 5)) + 0.5 df = DataFrame(df, columns=index) result = df[attributes] tm.assert_frame_equal(result, df) # GH 7349 # loc with a multi-index seems to be doing fallback df = DataFrame( np.arange(12).reshape(-1, 1), index=MultiIndex.from_product([[1, 2, 3, 4], [1, 2, 3]]), ) expected = df.loc[([1, 2],), :] result = df.loc[[1, 2]] tm.assert_frame_equal(result, expected) def test_loc_multiindex_incomplete(self): # GH 7399 # incomplete indexers s = Series( np.arange(15, dtype="int64"), MultiIndex.from_product([range(5), ["a", "b", "c"]]), ) expected = s.loc[:, "a":"c"] result = s.loc[0:4, "a":"c"] tm.assert_series_equal(result, expected) result = s.loc[:4, "a":"c"] tm.assert_series_equal(result, expected) result = s.loc[0:, "a":"c"] tm.assert_series_equal(result, expected) # GH 7400 # multiindexer getitem with list of indexers skips wrong element s = Series( np.arange(15, dtype="int64"), MultiIndex.from_product([range(5), ["a", "b", "c"]]), ) expected = s.iloc[[6, 7, 8, 12, 13, 14]] result = s.loc[2:4:2, "a":"c"] tm.assert_series_equal(result, expected) def test_get_loc_single_level(self, single_level_multiindex): single_level = single_level_multiindex s = Series( np.random.default_rng(2).standard_normal(len(single_level)), index=single_level, ) for k in single_level.values: s[k] def test_loc_getitem_int_slice(self): # GH 3053 # loc should treat integer slices like label slices index = MultiIndex.from_product([[6, 7, 8], ["a", "b"]]) df = DataFrame(np.random.default_rng(2).standard_normal((6, 6)), index, index) result = df.loc[6:8, :] expected = df tm.assert_frame_equal(result, expected) index = MultiIndex.from_product([[10, 20, 30], ["a", "b"]]) df = DataFrame(np.random.default_rng(2).standard_normal((6, 6)), index, index) result = df.loc[20:30, :] expected = df.iloc[2:] tm.assert_frame_equal(result, expected) # doc examples result = df.loc[10, :] expected = df.iloc[0:2] expected.index = ["a", "b"] tm.assert_frame_equal(result, expected) result = df.loc[:, 10] expected = df[10] tm.assert_frame_equal(result, expected) @pytest.mark.parametrize( "indexer_type_1", (list, tuple, set, slice, np.ndarray, Series, Index) ) @pytest.mark.parametrize( "indexer_type_2", (list, tuple, set, slice, np.ndarray, Series, Index) ) def test_loc_getitem_nested_indexer(self, indexer_type_1, indexer_type_2): # GH #19686 # .loc should work with nested indexers which can be # any list-like objects (see `is_list_like` (`pandas.api.types`)) or slices def convert_nested_indexer(indexer_type, keys): if indexer_type == np.ndarray: return np.array(keys) if indexer_type == slice: return slice(*keys) return indexer_type(keys) a = [10, 20, 30] b = [1, 2, 3] index = MultiIndex.from_product([a, b]) df = DataFrame( np.arange(len(index), dtype="int64"), index=index, columns=["Data"] ) keys = ([10, 20], [2, 3]) types = (indexer_type_1, indexer_type_2) # check indexers with all the combinations of nested objects # of all the valid types indexer = tuple( convert_nested_indexer(indexer_type, k) for indexer_type, k in zip(types, keys) ) if indexer_type_1 is set or indexer_type_2 is set: with pytest.raises(TypeError, match="as an indexer is not supported"): df.loc[indexer, "Data"] return else: result = df.loc[indexer, "Data"] expected = Series( [1, 2, 4, 5], name="Data", index=MultiIndex.from_product(keys) ) tm.assert_series_equal(result, expected) def test_multiindex_loc_one_dimensional_tuple(self, frame_or_series): # GH#37711 mi = MultiIndex.from_tuples([("a", "A"), ("b", "A")]) obj = frame_or_series([1, 2], index=mi) obj.loc[("a",)] = 0 expected = frame_or_series([0, 2], index=mi) tm.assert_equal(obj, expected) @pytest.mark.parametrize("indexer", [("a",), ("a")]) def test_multiindex_one_dimensional_tuple_columns(self, indexer): # GH#37711 mi = MultiIndex.from_tuples([("a", "A"), ("b", "A")]) obj = DataFrame([1, 2], index=mi) obj.loc[indexer, :] = 0 expected = DataFrame([0, 2], index=mi) tm.assert_frame_equal(obj, expected) @pytest.mark.parametrize( "indexer, exp_value", [(slice(None), 1.0), ((1, 2), np.nan)] ) def test_multiindex_setitem_columns_enlarging(self, indexer, exp_value): # GH#39147 mi = MultiIndex.from_tuples([(1, 2), (3, 4)]) df = DataFrame([[1, 2], [3, 4]], index=mi, columns=["a", "b"]) df.loc[indexer, ["c", "d"]] = 1.0 expected = DataFrame( [[1, 2, 1.0, 1.0], [3, 4, exp_value, exp_value]], index=mi, columns=["a", "b", "c", "d"], ) tm.assert_frame_equal(df, expected) def test_sorted_multiindex_after_union(self): # GH#44752 midx = MultiIndex.from_product( [pd.date_range("20110101", periods=2), Index(["a", "b"])] ) ser1 = Series(1, index=midx) ser2 = Series(1, index=midx[:2]) df = pd.concat([ser1, ser2], axis=1) expected = df.copy() result = df.loc["2011-01-01":"2011-01-02"] tm.assert_frame_equal(result, expected) df = DataFrame({0: ser1, 1: ser2}) result = df.loc["2011-01-01":"2011-01-02"] tm.assert_frame_equal(result, expected) df = pd.concat([ser1, ser2.reindex(ser1.index)], axis=1) result = df.loc["2011-01-01":"2011-01-02"] tm.assert_frame_equal(result, expected) def test_loc_no_second_level_index(self): # GH#43599 df = DataFrame( index=MultiIndex.from_product([list("ab"), list("cd"), list("e")]), columns=["Val"], ) res = df.loc[np.s_[:, "c", :]] expected = DataFrame( index=MultiIndex.from_product([list("ab"), list("e")]), columns=["Val"] ) tm.assert_frame_equal(res, expected) def test_loc_multi_index_key_error(self): # GH 51892 df = DataFrame( { (1, 2): ["a", "b", "c"], (1, 3): ["d", "e", "f"], (2, 2): ["g", "h", "i"], (2, 4): ["j", "k", "l"], } ) with pytest.raises(KeyError, match=r"(1, 4)"): df.loc[0, (1, 4)] @pytest.mark.parametrize( "indexer, pos", [ ([], []), # empty ok (["A"], slice(3)), (["A", "D"], []), # "D" isn't present -> raise (["D", "E"], []), # no values found -> raise (["D"], []), # same, with single item list: GH 27148 (pd.IndexSlice[:, ["foo"]], slice(2, None, 3)), (pd.IndexSlice[:, ["foo", "bah"]], slice(2, None, 3)), ], ) def test_loc_getitem_duplicates_multiindex_missing_indexers(indexer, pos): # GH 7866 # multi-index slicing with missing indexers idx = MultiIndex.from_product( [["A", "B", "C"], ["foo", "bar", "baz"]], names=["one", "two"] ) ser = Series(np.arange(9, dtype="int64"), index=idx).sort_index() expected = ser.iloc[pos] if expected.size == 0 and indexer != []: with pytest.raises(KeyError, match=str(indexer)): ser.loc[indexer] elif indexer == (slice(None), ["foo", "bah"]): # "bah" is not in idx.levels[1], raising KeyError enforced in 2.0 with pytest.raises(KeyError, match="'bah'"): ser.loc[indexer] else: result = ser.loc[indexer] tm.assert_series_equal(result, expected) @pytest.mark.parametrize("columns_indexer", [([], slice(None)), (["foo"], [])]) def test_loc_getitem_duplicates_multiindex_empty_indexer(columns_indexer): # GH 8737 # empty indexer multi_index = MultiIndex.from_product((["foo", "bar", "baz"], ["alpha", "beta"])) df = DataFrame( np.random.default_rng(2).standard_normal((5, 6)), index=range(5), columns=multi_index, ) df = df.sort_index(level=0, axis=1) expected = DataFrame(index=range(5), columns=multi_index.reindex([])[0]) result = df.loc[:, columns_indexer] tm.assert_frame_equal(result, expected) def test_loc_getitem_duplicates_multiindex_non_scalar_type_object(): # regression from < 0.14.0 # GH 7914 df = DataFrame( [[np.mean, np.median], ["mean", "median"]], columns=MultiIndex.from_tuples([("functs", "mean"), ("functs", "median")]), index=["function", "name"], ) result = df.loc["function", ("functs", "mean")] expected = np.mean assert result == expected def test_loc_getitem_tuple_plus_slice(): # GH 671 df = DataFrame( { "a": np.arange(10), "b": np.arange(10), "c": np.random.default_rng(2).standard_normal(10), "d": np.random.default_rng(2).standard_normal(10), } ).set_index(["a", "b"]) expected = df.loc[0, 0] result = df.loc[(0, 0), :] tm.assert_series_equal(result, expected) def test_loc_getitem_int(frame_random_data_integer_multi_index): df = frame_random_data_integer_multi_index result = df.loc[1] expected = df[-3:] expected.index = expected.index.droplevel(0) tm.assert_frame_equal(result, expected) def test_loc_getitem_int_raises_exception(frame_random_data_integer_multi_index): df = frame_random_data_integer_multi_index with pytest.raises(KeyError, match=r"^3$"): df.loc[3] def test_loc_getitem_lowerdim_corner(multiindex_dataframe_random_data): df = multiindex_dataframe_random_data # test setup - check key not in dataframe with pytest.raises(KeyError, match=r"^\('bar', 'three'\)$"): df.loc[("bar", "three"), "B"] # in theory should be inserting in a sorted space???? df.loc[("bar", "three"), "B"] = 0 expected = 0 result = df.sort_index().loc[("bar", "three"), "B"] assert result == expected def test_loc_setitem_single_column_slice(): # case from https://github.com/pandas-dev/pandas/issues/27841 df = DataFrame( "string", index=list("abcd"), columns=MultiIndex.from_product([["Main"], ("another", "one")]), ) df["labels"] = "a" df.loc[:, "labels"] = df.index tm.assert_numpy_array_equal(np.asarray(df["labels"]), np.asarray(df.index)) # test with non-object block df = DataFrame( np.nan, index=range(4), columns=MultiIndex.from_tuples([("A", "1"), ("A", "2"), ("B", "1")]), ) expected = df.copy() df.loc[:, "B"] = np.arange(4) expected.iloc[:, 2] = np.arange(4) tm.assert_frame_equal(df, expected) def test_loc_nan_multiindex(using_infer_string): # GH 5286 tups = [ ("Good Things", "C", np.nan), ("Good Things", "R", np.nan), ("Bad Things", "C", np.nan), ("Bad Things", "T", np.nan), ("Okay Things", "N", "B"), ("Okay Things", "N", "D"), ("Okay Things", "B", np.nan), ("Okay Things", "D", np.nan), ] df = DataFrame( np.ones((8, 4)), columns=Index(["d1", "d2", "d3", "d4"]), index=MultiIndex.from_tuples(tups, names=["u1", "u2", "u3"]), ) result = df.loc["Good Things"].loc["C"] expected = DataFrame( np.ones((1, 4)), index=Index( [np.nan], dtype="object" if not using_infer_string else "string[pyarrow_numpy]", name="u3", ), columns=Index(["d1", "d2", "d3", "d4"]), ) tm.assert_frame_equal(result, expected) def test_loc_period_string_indexing(): # GH 9892 a = pd.period_range("2013Q1", "2013Q4", freq="Q") i = (1111, 2222, 3333) idx = MultiIndex.from_product((a, i), names=("Period", "CVR")) df = DataFrame( index=idx, columns=( "OMS", "OMK", "RES", "DRIFT_IND", "OEVRIG_IND", "FIN_IND", "VARE_UD", "LOEN_UD", "FIN_UD", ), ) result = df.loc[("2013Q1", 1111), "OMS"] alt = df.loc[(a[0], 1111), "OMS"] assert np.isnan(alt) # Because the resolution of the string matches, it is an exact lookup, # not a slice assert np.isnan(result) alt = df.loc[("2013Q1", 1111), "OMS"] assert np.isnan(alt) def test_loc_datetime_mask_slicing(): # GH 16699 dt_idx = pd.to_datetime(["2017-05-04", "2017-05-05"]) m_idx = MultiIndex.from_product([dt_idx, dt_idx], names=["Idx1", "Idx2"]) df = DataFrame( data=[[1, 2], [3, 4], [5, 6], [7, 6]], index=m_idx, columns=["C1", "C2"] ) result = df.loc[(dt_idx[0], (df.index.get_level_values(1) > "2017-05-04")), "C1"] expected = Series( [3], name="C1", index=MultiIndex.from_tuples( [(pd.Timestamp("2017-05-04"), pd.Timestamp("2017-05-05"))], names=["Idx1", "Idx2"], ), ) tm.assert_series_equal(result, expected) def test_loc_datetime_series_tuple_slicing(): # https://github.com/pandas-dev/pandas/issues/35858 date = pd.Timestamp("2000") ser = Series( 1, index=MultiIndex.from_tuples([("a", date)], names=["a", "b"]), name="c", ) result = ser.loc[:, [date]] tm.assert_series_equal(result, ser) def test_loc_with_mi_indexer(): # https://github.com/pandas-dev/pandas/issues/35351 df = DataFrame( data=[["a", 1], ["a", 0], ["b", 1], ["c", 2]], index=MultiIndex.from_tuples( [(0, 1), (1, 0), (1, 1), (1, 1)], names=["index", "date"] ), columns=["author", "price"], ) idx = MultiIndex.from_tuples([(0, 1), (1, 1)], names=["index", "date"]) result = df.loc[idx, :] expected = DataFrame( [["a", 1], ["b", 1], ["c", 2]], index=MultiIndex.from_tuples([(0, 1), (1, 1), (1, 1)], names=["index", "date"]), columns=["author", "price"], ) tm.assert_frame_equal(result, expected) def test_loc_mi_with_level1_named_0(): # GH#37194 dti = pd.date_range("2016-01-01", periods=3, tz="US/Pacific") ser = Series(range(3), index=dti) df = ser.to_frame() df[1] = dti df2 = df.set_index(0, append=True) assert df2.index.names == (None, 0) df2.index.get_loc(dti[0]) # smoke test result = df2.loc[dti[0]] expected = df2.iloc[[0]].droplevel(None) tm.assert_frame_equal(result, expected) ser2 = df2[1] assert ser2.index.names == (None, 0) result = ser2.loc[dti[0]] expected = ser2.iloc[[0]].droplevel(None) tm.assert_series_equal(result, expected) def test_getitem_str_slice(): # GH#15928 df = DataFrame( [ ["20160525 13:30:00.023", "MSFT", "51.95", "51.95"], ["20160525 13:30:00.048", "GOOG", "720.50", "720.93"], ["20160525 13:30:00.076", "AAPL", "98.55", "98.56"], ["20160525 13:30:00.131", "AAPL", "98.61", "98.62"], ["20160525 13:30:00.135", "MSFT", "51.92", "51.95"], ["20160525 13:30:00.135", "AAPL", "98.61", "98.62"], ], columns="time,ticker,bid,ask".split(","), ) df2 = df.set_index(["ticker", "time"]).sort_index() res = df2.loc[("AAPL", slice("2016-05-25 13:30:00")), :].droplevel(0) expected = df2.loc["AAPL"].loc[slice("2016-05-25 13:30:00"), :] tm.assert_frame_equal(res, expected) def test_3levels_leading_period_index(): # GH#24091 pi = pd.PeriodIndex( ["20181101 1100", "20181101 1200", "20181102 1300", "20181102 1400"], name="datetime", freq="D", ) lev2 = ["A", "A", "Z", "W"] lev3 = ["B", "C", "Q", "F"] mi = MultiIndex.from_arrays([pi, lev2, lev3]) ser = Series(range(4), index=mi, dtype=np.float64) result = ser.loc[(pi[0], "A", "B")] assert result == 0.0 class TestKeyErrorsWithMultiIndex: def test_missing_keys_raises_keyerror(self): # GH#27420 KeyError, not TypeError df = DataFrame(np.arange(12).reshape(4, 3), columns=["A", "B", "C"]) df2 = df.set_index(["A", "B"]) with pytest.raises(KeyError, match="1"): df2.loc[(1, 6)] def test_missing_key_raises_keyerror2(self): # GH#21168 KeyError, not "IndexingError: Too many indexers" ser = Series(-1, index=MultiIndex.from_product([[0, 1]] * 2)) with pytest.raises(KeyError, match=r"\(0, 3\)"): ser.loc[0, 3] def test_missing_key_combination(self): # GH: 19556 mi = MultiIndex.from_arrays( [ np.array(["a", "a", "b", "b"]), np.array(["1", "2", "2", "3"]), np.array(["c", "d", "c", "d"]), ], names=["one", "two", "three"], ) df = DataFrame(np.random.default_rng(2).random((4, 3)), index=mi) msg = r"\('b', '1', slice\(None, None, None\)\)" with pytest.raises(KeyError, match=msg): df.loc[("b", "1", slice(None)), :] with pytest.raises(KeyError, match=msg): df.index.get_locs(("b", "1", slice(None))) with pytest.raises(KeyError, match=r"\('b', '1'\)"): df.loc[("b", "1"), :] def test_getitem_loc_commutability(multiindex_year_month_day_dataframe_random_data): df = multiindex_year_month_day_dataframe_random_data ser = df["A"] result = ser[2000, 5] expected = df.loc[2000, 5]["A"] tm.assert_series_equal(result, expected) def test_loc_with_nan(): # GH: 27104 df = DataFrame( {"col": [1, 2, 5], "ind1": ["a", "d", np.nan], "ind2": [1, 4, 5]} ).set_index(["ind1", "ind2"]) result = df.loc[["a"]] expected = DataFrame( {"col": [1]}, index=MultiIndex.from_tuples([("a", 1)], names=["ind1", "ind2"]) ) tm.assert_frame_equal(result, expected) result = df.loc["a"] expected = DataFrame({"col": [1]}, index=Index([1], name="ind2")) tm.assert_frame_equal(result, expected) def test_getitem_non_found_tuple(): # GH: 25236 df = DataFrame([[1, 2, 3, 4]], columns=["a", "b", "c", "d"]).set_index( ["a", "b", "c"] ) with pytest.raises(KeyError, match=r"\(2\.0, 2\.0, 3\.0\)"): df.loc[(2.0, 2.0, 3.0)] def test_get_loc_datetime_index(): # GH#24263 index = pd.date_range("2001-01-01", periods=100) mi = MultiIndex.from_arrays([index]) # Check if get_loc matches for Index and MultiIndex assert mi.get_loc("2001-01") == slice(0, 31, None) assert index.get_loc("2001-01") == slice(0, 31, None) loc = mi[::2].get_loc("2001-01") expected = index[::2].get_loc("2001-01") assert loc == expected loc = mi.repeat(2).get_loc("2001-01") expected = index.repeat(2).get_loc("2001-01") assert loc == expected loc = mi.append(mi).get_loc("2001-01") expected = index.append(index).get_loc("2001-01") # TODO: standardize return type for MultiIndex.get_loc tm.assert_numpy_array_equal(loc.nonzero()[0], expected) def test_loc_setitem_indexer_differently_ordered(): # GH#34603 mi = MultiIndex.from_product([["a", "b"], [0, 1]]) df = DataFrame([[1, 2], [3, 4], [5, 6], [7, 8]], index=mi) indexer = ("a", [1, 0]) df.loc[indexer, :] = np.array([[9, 10], [11, 12]]) expected = DataFrame([[11, 12], [9, 10], [5, 6], [7, 8]], index=mi) tm.assert_frame_equal(df, expected) def test_loc_getitem_index_differently_ordered_slice_none(): # GH#31330 df = DataFrame( [[1, 2], [3, 4], [5, 6], [7, 8]], index=[["a", "a", "b", "b"], [1, 2, 1, 2]], columns=["a", "b"], ) result = df.loc[(slice(None), [2, 1]), :] expected = DataFrame( [[3, 4], [7, 8], [1, 2], [5, 6]], index=[["a", "b", "a", "b"], [2, 2, 1, 1]], columns=["a", "b"], ) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("indexer", [[1, 2, 7, 6, 2, 3, 8, 7], [1, 2, 7, 6, 3, 8]]) def test_loc_getitem_index_differently_ordered_slice_none_duplicates(indexer): # GH#40978 df = DataFrame( [1] * 8, index=MultiIndex.from_tuples( [(1, 1), (1, 2), (1, 7), (1, 6), (2, 2), (2, 3), (2, 8), (2, 7)] ), columns=["a"], ) result = df.loc[(slice(None), indexer), :] expected = DataFrame( [1] * 8, index=[[1, 1, 2, 1, 2, 1, 2, 2], [1, 2, 2, 7, 7, 6, 3, 8]], columns=["a"], ) tm.assert_frame_equal(result, expected) result = df.loc[df.index.isin(indexer, level=1), :] tm.assert_frame_equal(result, df) def test_loc_getitem_drops_levels_for_one_row_dataframe(): # GH#10521 "x" and "z" are both scalar indexing, so those levels are dropped mi = MultiIndex.from_arrays([["x"], ["y"], ["z"]], names=["a", "b", "c"]) df = DataFrame({"d": [0]}, index=mi) expected = df.droplevel([0, 2]) result = df.loc["x", :, "z"] tm.assert_frame_equal(result, expected) ser = Series([0], index=mi) result = ser.loc["x", :, "z"] expected = Series([0], index=Index(["y"], name="b")) tm.assert_series_equal(result, expected) def test_mi_columns_loc_list_label_order(): # GH 10710 cols = MultiIndex.from_product([["A", "B", "C"], [1, 2]]) df = DataFrame(np.zeros((5, 6)), columns=cols) result = df.loc[:, ["B", "A"]] expected = DataFrame( np.zeros((5, 4)), columns=MultiIndex.from_tuples([("B", 1), ("B", 2), ("A", 1), ("A", 2)]), ) tm.assert_frame_equal(result, expected) def test_mi_partial_indexing_list_raises(): # GH 13501 frame = DataFrame( np.arange(12).reshape((4, 3)), index=[["a", "a", "b", "b"], [1, 2, 1, 2]], columns=[["Ohio", "Ohio", "Colorado"], ["Green", "Red", "Green"]], ) frame.index.names = ["key1", "key2"] frame.columns.names = ["state", "color"] with pytest.raises(KeyError, match="\\[2\\] not in index"): frame.loc[["b", 2], "Colorado"] def test_mi_indexing_list_nonexistent_raises(): # GH 15452 s = Series(range(4), index=MultiIndex.from_product([[1, 2], ["a", "b"]])) with pytest.raises(KeyError, match="\\['not' 'found'\\] not in index"): s.loc[["not", "found"]] def test_mi_add_cell_missing_row_non_unique(): # GH 16018 result = DataFrame( [[1, 2, 5, 6], [3, 4, 7, 8]], index=["a", "a"], columns=MultiIndex.from_product([[1, 2], ["A", "B"]]), ) result.loc["c"] = -1 result.loc["c", (1, "A")] = 3 result.loc["d", (1, "A")] = 3 expected = DataFrame( [ [1.0, 2.0, 5.0, 6.0], [3.0, 4.0, 7.0, 8.0], [3.0, -1.0, -1, -1], [3.0, np.nan, np.nan, np.nan], ], index=["a", "a", "c", "d"], columns=MultiIndex.from_product([[1, 2], ["A", "B"]]), ) tm.assert_frame_equal(result, expected) def test_loc_get_scalar_casting_to_float(): # GH#41369 df = DataFrame( {"a": 1.0, "b": 2}, index=MultiIndex.from_arrays([[3], [4]], names=["c", "d"]) ) result = df.loc[(3, 4), "b"] assert result == 2 assert isinstance(result, np.int64) result = df.loc[[(3, 4)], "b"].iloc[0] assert result == 2 assert isinstance(result, np.int64) def test_loc_empty_single_selector_with_names(): # GH 19517 idx = MultiIndex.from_product([["a", "b"], ["A", "B"]], names=[1, 0]) s2 = Series(index=idx, dtype=np.float64) result = s2.loc["a"] expected = Series([np.nan, np.nan], index=Index(["A", "B"], name=0)) tm.assert_series_equal(result, expected) def test_loc_keyerror_rightmost_key_missing(): # GH 20951 df = DataFrame( { "A": [100, 100, 200, 200, 300, 300], "B": [10, 10, 20, 21, 31, 33], "C": range(6), } ) df = df.set_index(["A", "B"]) with pytest.raises(KeyError, match="^1$"): df.loc[(100, 1)] def test_multindex_series_loc_with_tuple_label(): # GH#43908 mi = MultiIndex.from_tuples([(1, 2), (3, (4, 5))]) ser = Series([1, 2], index=mi) result = ser.loc[(3, (4, 5))] assert result == 2