Skip to content

Commit e40496d

Browse files
committed
BUG: Fix MultiIndex lookup with mixed datetime types GH#55969
1 parent 8be8439 commit e40496d

File tree

2 files changed

+56
-4
lines changed

2 files changed

+56
-4
lines changed

pandas/core/indexes/multi.py

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2957,8 +2957,7 @@ def sortlevel(
29572957
# error: Item "Hashable" of "Union[Hashable, Sequence[Hashable]]" has
29582958
# no attribute "__iter__" (not iterable)
29592959
level = [
2960-
self._get_level_number(lev)
2961-
for lev in level # type: ignore[union-attr]
2960+
self._get_level_number(lev) for lev in level # type: ignore[union-attr]
29622961
]
29632962
sortorder = None
29642963

@@ -3264,9 +3263,9 @@ def _get_loc_single_level_index(self, level_index: Index, key: Hashable) -> int:
32643263
else:
32653264
return level_index.get_loc(key)
32663265

3267-
def get_loc(self, key):
3266+
def get_loc(self, key, method=None):
32683267
"""
3269-
Get location for a label or a tuple of labels. The location is returned \
3268+
Get location for a label or a tuple of labels. The location is returned
32703269
as an integer/slice or boolean mask.
32713270
32723271
This method returns the integer location, slice object, or boolean mask
@@ -3310,6 +3309,25 @@ def get_loc(self, key):
33103309
>>> mi.get_loc(("b", "e"))
33113310
1
33123311
"""
3312+
# --- FIX GH#55969 START ---
3313+
# If the key contains np.datetime64 but the level is object-dtype (python objects),
3314+
# strict lookups (and binary search) can fail. Convert to python objects to match.
3315+
if isinstance(key, tuple):
3316+
new_key = list(key)
3317+
modified = False
3318+
for i, (k, level) in enumerate(zip(new_key, self.levels)):
3319+
if isinstance(k, np.datetime64) and level.dtype == object:
3320+
try:
3321+
new_key[i] = k.item()
3322+
modified = True
3323+
except (ValueError, TypeError):
3324+
pass
3325+
if modified:
3326+
key = tuple(new_key)
3327+
3328+
if method is not None:
3329+
return Index.get_loc(self, key, method=method)
3330+
33133331
self._check_indexing_error(key)
33143332

33153333
def _maybe_to_slice(loc):
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
import numpy as np
2+
import pytest
3+
4+
from pandas import (
5+
DataFrame,
6+
MultiIndex,
7+
Timestamp,
8+
)
9+
import pandas._testing as tm
10+
11+
12+
def test_mixed_datetime_types_lookup():
13+
14+
import datetime as dt
15+
16+
dates = [dt.date(2023, 11, 1), dt.date(2023, 11, 1), dt.date(2023, 11, 2)]
17+
t1 = ["A", "B", "C"]
18+
t2 = ["C", "D", "E"]
19+
vals = [10, 20, 30]
20+
21+
df = DataFrame({"dates": dates, "t1": t1, "t2": t2, "vals": vals}).set_index(
22+
["dates", "t1", "t2"]
23+
)
24+
25+
date_np = np.datetime64("2023-11-01")
26+
27+
result = df.loc[(date_np, "A")]
28+
expected_val = 10
29+
assert len(result) == 1
30+
assert result["vals"].iloc[0] == expected_val
31+
32+
msg = "'C'"
33+
with pytest.raises(KeyError, match=msg):
34+
df.loc[(date_np, "C")]

0 commit comments

Comments
 (0)