Skip to content

Commit 361e1af

Browse files
committed
ENH: pd.MultiIndex.get_loc(np.nan) (pandas-dev#19132)
MultiIndex.get_loc could not find nan with values including missing values as a input
1 parent d5fa16b commit 361e1af

File tree

3 files changed

+66
-5
lines changed

3 files changed

+66
-5
lines changed

doc/source/whatsnew/v1.0.0.rst

+1
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,7 @@ Indexing
229229
- Bug in reindexing a :meth:`PeriodIndex` with another type of index that contained a `Period` (:issue:`28323`) (:issue:`28337`)
230230
- Fix assignment of column via `.loc` with numpy non-ns datetime type (:issue:`27395`)
231231
- Bug in :meth:`Float64Index.astype` where ``np.inf`` was not handled properly when casting to an integer dtype (:issue:`28475`)
232+
- When index is ``MultiIndex``, Using ``.get_loc`` can't find ``nan`` with values including missing values as input (:issue:`19132`)
232233

233234
Missing
234235
^^^^^^^

pandas/core/indexes/multi.py

+38-5
Original file line numberDiff line numberDiff line change
@@ -2591,7 +2591,7 @@ def _partial_tup_index(self, tup, side="left"):
25912591
for k, (lab, lev, labs) in enumerate(zipped):
25922592
section = labs[start:end]
25932593

2594-
if lab not in lev:
2594+
if lab not in lev and not isna(lab):
25952595
if not lev.is_type_compatible(lib.infer_dtype([lab], skipna=False)):
25962596
raise TypeError("Level type mismatch: %s" % lab)
25972597

@@ -2601,13 +2601,45 @@ def _partial_tup_index(self, tup, side="left"):
26012601
loc -= 1
26022602
return start + section.searchsorted(loc, side=side)
26032603

2604-
idx = lev.get_loc(lab)
2604+
idx = self._get_loc_single_level_index_wrapper(lev, lab)
26052605
if k < n - 1:
26062606
end = start + section.searchsorted(idx, side="right")
26072607
start = start + section.searchsorted(idx, side="left")
26082608
else:
26092609
return start + section.searchsorted(idx, side=side)
26102610

2611+
def _get_loc_single_level_index_wrapper(self, level_index: Index, key) -> int:
2612+
"""
2613+
Wrapper function for MultiIndex.
2614+
If key is NA value, location of index unify as -1.
2615+
2616+
Parameters
2617+
----------
2618+
level_index: Index
2619+
Single level index in MultiIndex
2620+
key : label
2621+
2622+
Returns
2623+
-------
2624+
loc : int
2625+
If key is NA value, loc is -1
2626+
Else, location of key in index
2627+
2628+
See Also
2629+
--------
2630+
Index.get_loc : The get_loc method for (single-level) index.
2631+
2632+
Notes
2633+
-----
2634+
Depending on Index type, Handling NA value is different using get_loc.
2635+
But in MultiIndex, NA values is denoted as missing by -1.
2636+
"""
2637+
2638+
if is_scalar(key) and isna(key):
2639+
return -1
2640+
else:
2641+
return level_index.get_loc(key)
2642+
26112643
def get_loc(self, key, method=None):
26122644
"""
26132645
Get location for a label or a tuple of labels as an integer, slice or
@@ -2707,7 +2739,9 @@ def _maybe_to_slice(loc):
27072739
loc = np.arange(start, stop, dtype="int64")
27082740

27092741
for i, k in enumerate(follow_key, len(lead_key)):
2710-
mask = self.codes[i][loc] == self.levels[i].get_loc(k)
2742+
mask = self.codes[i][loc] == self._get_loc_single_level_index_wrapper(
2743+
self.levels[i], k
2744+
)
27112745
if not mask.all():
27122746
loc = loc[mask]
27132747
if not len(loc):
@@ -2934,8 +2968,7 @@ def convert_indexer(start, stop, step, indexer=indexer, codes=level_codes):
29342968
return slice(i, j, step)
29352969

29362970
else:
2937-
2938-
code = level_index.get_loc(key)
2971+
code = self._get_loc_single_level_index_wrapper(level_index, key)
29392972

29402973
if level > 0 or self.lexsort_depth == 0:
29412974
# Desired level is not sorted

pandas/tests/indexes/multi/test_indexing.py

+27
Original file line numberDiff line numberDiff line change
@@ -439,3 +439,30 @@ def test_timestamp_multiindex_indexer():
439439
)
440440
should_be = pd.Series(data=np.arange(24, len(qidx) + 24), index=qidx, name="foo")
441441
tm.assert_series_equal(result, should_be)
442+
443+
444+
def test_get_loc_with_values_including_missing_values():
445+
# issue 19132
446+
idx = MultiIndex.from_product([[np.nan, 1]] * 2)
447+
expected = slice(0, 2, None)
448+
assert idx.get_loc(np.nan) == expected
449+
450+
idx = MultiIndex.from_arrays([[np.nan, 1, 2, np.nan]])
451+
expected = np.array([True, False, False, True])
452+
tm.assert_numpy_array_equal(idx.get_loc(np.nan), expected)
453+
454+
idx = MultiIndex.from_product([[np.nan, 1]] * 3)
455+
expected = slice(2, 4, None)
456+
assert idx.get_loc((np.nan, 1)) == expected
457+
458+
459+
def test_get_indexer_with_missing_value():
460+
# issue 19132
461+
idx = MultiIndex.from_arrays([[1, np.nan, 2], [3, 4, 5]])
462+
result = idx.get_indexer([1, np.nan, 2])
463+
expected = np.array([-1, -1, -1], dtype="int64")
464+
tm.assert_numpy_array_equal(result.astype("int64"), expected)
465+
466+
result = idx.get_indexer([(np.nan, 4)])
467+
expected = np.array([1], dtype="int64")
468+
tm.assert_numpy_array_equal(result.astype("int64"), expected)

0 commit comments

Comments
 (0)