diff --git a/pandas/core/indexes/multi.py b/pandas/core/indexes/multi.py index 90f710b0de4de..4bc418a9a33d6 100644 --- a/pandas/core/indexes/multi.py +++ b/pandas/core/indexes/multi.py @@ -2957,8 +2957,7 @@ def sortlevel( # error: Item "Hashable" of "Union[Hashable, Sequence[Hashable]]" has # no attribute "__iter__" (not iterable) level = [ - self._get_level_number(lev) - for lev in level # type: ignore[union-attr] + self._get_level_number(lev) for lev in level # type: ignore[union-attr] ] sortorder = None @@ -3264,52 +3263,42 @@ def _get_loc_single_level_index(self, level_index: Index, key: Hashable) -> int: else: return level_index.get_loc(key) - def get_loc(self, key): + def get_loc(self, key, method=None): """ - Get location for a label or a tuple of labels. The location is returned \ - as an integer/slice or boolean mask. - - This method returns the integer location, slice object, or boolean mask - corresponding to the specified key, which can be a single label or a tuple - of labels. The key represents a position in the MultiIndex, and the location - indicates where the key is found within the index. + Get location for a label or a tuple of labels. Parameters ---------- key : label or tuple of labels (one for each level) - A label or tuple of labels that correspond to the levels of the MultiIndex. - The key must match the structure of the MultiIndex. + The key to locate. + method : str or None, optional + Method for getting the location (see Index.get_loc). Returns ------- - int, slice object or boolean mask - If the key is past the lexsort depth, the return may be a - boolean mask array, otherwise it is always a slice or int. - - See Also - -------- - Index.get_loc : The get_loc method for (single-level) index. - MultiIndex.slice_locs : Get slice location given start label(s) and - end label(s). - MultiIndex.get_locs : Get location for a label/slice/list/mask or a - sequence of such. - - Notes - ----- - The key cannot be a slice, list of same-level labels, a boolean mask, - or a sequence of such. If you want to use those, use - :meth:`MultiIndex.get_locs` instead. - - Examples - -------- - >>> mi = pd.MultiIndex.from_arrays([list("abb"), list("def")]) + int, slice, or boolean mask + Location(s) of the key. + """ + # GH#55969: If key has np.datetime64 but level is object-dtype + # (python objects), strict lookups/binary search can fail. + # Convert to python objects to match. + if isinstance(key, tuple): + new_key = list(key) + modified = False + # Use strict=False as key len might be < levels len + for i, (k, level) in enumerate(zip(new_key, self.levels, strict=False)): + if isinstance(k, np.datetime64) and level.dtype == object: + try: + new_key[i] = k.item() + modified = True + except (ValueError, TypeError): + pass + if modified: + key = tuple(new_key) - >>> mi.get_loc("b") - slice(1, 3, None) + if method is not None: + return Index.get_loc(self, key, method=method) - >>> mi.get_loc(("b", "e")) - 1 - """ self._check_indexing_error(key) def _maybe_to_slice(loc): diff --git a/pandas/tests/indexes/multi/test_gh55969.py b/pandas/tests/indexes/multi/test_gh55969.py new file mode 100644 index 0000000000000..a792689308806 --- /dev/null +++ b/pandas/tests/indexes/multi/test_gh55969.py @@ -0,0 +1,34 @@ +import numpy as np +import pytest + +from pandas import ( + DataFrame, + MultiIndex, + Timestamp, +) +import pandas._testing as tm + + +def test_mixed_datetime_types_lookup(): + + import datetime as dt + + dates = [dt.date(2023, 11, 1), dt.date(2023, 11, 1), dt.date(2023, 11, 2)] + t1 = ["A", "B", "C"] + t2 = ["C", "D", "E"] + vals = [10, 20, 30] + + df = DataFrame({"dates": dates, "t1": t1, "t2": t2, "vals": vals}).set_index( + ["dates", "t1", "t2"] + ) + + date_np = np.datetime64("2023-11-01") + + result = df.loc[(date_np, "A")] + expected_val = 10 + assert len(result) == 1 + assert result["vals"].iloc[0] == expected_val + + msg = "'C'" + with pytest.raises(KeyError, match=msg): + df.loc[(date_np, "C")]