|
6 | 6 | from pandas.errors import InvalidIndexError
|
7 | 7 |
|
8 | 8 | from pandas import (
|
| 9 | + NA, |
9 | 10 | CategoricalIndex,
|
10 | 11 | Interval,
|
11 | 12 | IntervalIndex,
|
| 13 | + NaT, |
12 | 14 | Timedelta,
|
13 | 15 | date_range,
|
14 | 16 | timedelta_range,
|
@@ -168,6 +170,20 @@ def test_get_loc_non_scalar_errors(self, key):
|
168 | 170 | with pytest.raises(InvalidIndexError, match=msg):
|
169 | 171 | idx.get_loc(key)
|
170 | 172 |
|
| 173 | + def test_get_indexer_with_nans(self): |
| 174 | + # GH#41831 |
| 175 | + index = IntervalIndex([np.nan, Interval(1, 2), np.nan]) |
| 176 | + |
| 177 | + expected = np.array([True, False, True]) |
| 178 | + for key in [None, np.nan, NA]: |
| 179 | + assert key in index |
| 180 | + result = index.get_loc(key) |
| 181 | + tm.assert_numpy_array_equal(result, expected) |
| 182 | + |
| 183 | + for key in [NaT, np.timedelta64("NaT", "ns"), np.datetime64("NaT", "ns")]: |
| 184 | + with pytest.raises(KeyError, match=str(key)): |
| 185 | + index.get_loc(key) |
| 186 | + |
171 | 187 |
|
172 | 188 | class TestGetIndexer:
|
173 | 189 | @pytest.mark.parametrize(
|
@@ -326,6 +342,17 @@ def test_get_indexer_non_monotonic(self):
|
326 | 342 | expected = np.array([1, 2], dtype=np.intp)
|
327 | 343 | tm.assert_numpy_array_equal(result, expected)
|
328 | 344 |
|
| 345 | + def test_get_indexer_with_nans(self): |
| 346 | + # GH#41831 |
| 347 | + index = IntervalIndex([np.nan, np.nan]) |
| 348 | + other = IntervalIndex([np.nan]) |
| 349 | + |
| 350 | + assert not index._index_as_unique |
| 351 | + |
| 352 | + result = index.get_indexer_for(other) |
| 353 | + expected = np.array([0, 1], dtype=np.intp) |
| 354 | + tm.assert_numpy_array_equal(result, expected) |
| 355 | + |
329 | 356 |
|
330 | 357 | class TestSliceLocs:
|
331 | 358 | def test_slice_locs_with_interval(self):
|
|
0 commit comments