test_indexing.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360
  1. """
  2. test_indexing tests the following Index methods:
  3. __getitem__
  4. get_loc
  5. get_value
  6. __contains__
  7. take
  8. where
  9. get_indexer
  10. get_indexer_for
  11. slice_locs
  12. asof_locs
  13. The corresponding tests.indexes.[index_type].test_indexing files
  14. contain tests for the corresponding methods specific to those Index subclasses.
  15. """
  16. import numpy as np
  17. import pytest
  18. from pandas.errors import InvalidIndexError
  19. from pandas.core.dtypes.common import (
  20. is_float_dtype,
  21. is_scalar,
  22. )
  23. from pandas import (
  24. NA,
  25. DatetimeIndex,
  26. Index,
  27. IntervalIndex,
  28. MultiIndex,
  29. NaT,
  30. PeriodIndex,
  31. TimedeltaIndex,
  32. )
  33. import pandas._testing as tm
  34. class TestTake:
  35. def test_take_invalid_kwargs(self, index):
  36. indices = [1, 2]
  37. msg = r"take\(\) got an unexpected keyword argument 'foo'"
  38. with pytest.raises(TypeError, match=msg):
  39. index.take(indices, foo=2)
  40. msg = "the 'out' parameter is not supported"
  41. with pytest.raises(ValueError, match=msg):
  42. index.take(indices, out=indices)
  43. msg = "the 'mode' parameter is not supported"
  44. with pytest.raises(ValueError, match=msg):
  45. index.take(indices, mode="clip")
  46. def test_take(self, index):
  47. indexer = [4, 3, 0, 2]
  48. if len(index) < 5:
  49. # not enough elements; ignore
  50. return
  51. result = index.take(indexer)
  52. expected = index[indexer]
  53. assert result.equals(expected)
  54. if not isinstance(index, (DatetimeIndex, PeriodIndex, TimedeltaIndex)):
  55. # GH 10791
  56. msg = r"'(.*Index)' object has no attribute 'freq'"
  57. with pytest.raises(AttributeError, match=msg):
  58. index.freq
  59. def test_take_indexer_type(self):
  60. # GH#42875
  61. integer_index = Index([0, 1, 2, 3])
  62. scalar_index = 1
  63. msg = "Expected indices to be array-like"
  64. with pytest.raises(TypeError, match=msg):
  65. integer_index.take(scalar_index)
  66. def test_take_minus1_without_fill(self, index):
  67. # -1 does not get treated as NA unless allow_fill=True is passed
  68. if len(index) == 0:
  69. # Test is not applicable
  70. return
  71. result = index.take([0, 0, -1])
  72. expected = index.take([0, 0, len(index) - 1])
  73. tm.assert_index_equal(result, expected)
  74. class TestContains:
  75. @pytest.mark.parametrize(
  76. "index,val",
  77. [
  78. (Index([0, 1, 2]), 2),
  79. (Index([0, 1, "2"]), "2"),
  80. (Index([0, 1, 2, np.inf, 4]), 4),
  81. (Index([0, 1, 2, np.nan, 4]), 4),
  82. (Index([0, 1, 2, np.inf]), np.inf),
  83. (Index([0, 1, 2, np.nan]), np.nan),
  84. ],
  85. )
  86. def test_index_contains(self, index, val):
  87. assert val in index
  88. @pytest.mark.parametrize(
  89. "index,val",
  90. [
  91. (Index([0, 1, 2]), "2"),
  92. (Index([0, 1, "2"]), 2),
  93. (Index([0, 1, 2, np.inf]), 4),
  94. (Index([0, 1, 2, np.nan]), 4),
  95. (Index([0, 1, 2, np.inf]), np.nan),
  96. (Index([0, 1, 2, np.nan]), np.inf),
  97. # Checking if np.inf in int64 Index should not cause an OverflowError
  98. # Related to GH 16957
  99. (Index([0, 1, 2], dtype=np.int64), np.inf),
  100. (Index([0, 1, 2], dtype=np.int64), np.nan),
  101. (Index([0, 1, 2], dtype=np.uint64), np.inf),
  102. (Index([0, 1, 2], dtype=np.uint64), np.nan),
  103. ],
  104. )
  105. def test_index_not_contains(self, index, val):
  106. assert val not in index
  107. @pytest.mark.parametrize(
  108. "index,val", [(Index([0, 1, "2"]), 0), (Index([0, 1, "2"]), "2")]
  109. )
  110. def test_mixed_index_contains(self, index, val):
  111. # GH#19860
  112. assert val in index
  113. @pytest.mark.parametrize(
  114. "index,val", [(Index([0, 1, "2"]), "1"), (Index([0, 1, "2"]), 2)]
  115. )
  116. def test_mixed_index_not_contains(self, index, val):
  117. # GH#19860
  118. assert val not in index
  119. def test_contains_with_float_index(self, any_real_numpy_dtype):
  120. # GH#22085
  121. dtype = any_real_numpy_dtype
  122. data = [0, 1, 2, 3] if not is_float_dtype(dtype) else [0.1, 1.1, 2.2, 3.3]
  123. index = Index(data, dtype=dtype)
  124. if not is_float_dtype(index.dtype):
  125. assert 1.1 not in index
  126. assert 1.0 in index
  127. assert 1 in index
  128. else:
  129. assert 1.1 in index
  130. assert 1.0 not in index
  131. assert 1 not in index
  132. def test_contains_requires_hashable_raises(self, index):
  133. if isinstance(index, MultiIndex):
  134. return # TODO: do we want this to raise?
  135. msg = "unhashable type: 'list'"
  136. with pytest.raises(TypeError, match=msg):
  137. [] in index
  138. msg = "|".join(
  139. [
  140. r"unhashable type: 'dict'",
  141. r"must be real number, not dict",
  142. r"an integer is required",
  143. r"\{\}",
  144. r"pandas\._libs\.interval\.IntervalTree' is not iterable",
  145. ]
  146. )
  147. with pytest.raises(TypeError, match=msg):
  148. {} in index._engine
  149. class TestGetLoc:
  150. def test_get_loc_non_hashable(self, index):
  151. # MultiIndex and Index raise TypeError, others InvalidIndexError
  152. with pytest.raises((TypeError, InvalidIndexError), match="slice"):
  153. index.get_loc(slice(0, 1))
  154. def test_get_loc_non_scalar_hashable(self, index):
  155. # GH52877
  156. from enum import Enum
  157. class E(Enum):
  158. X1 = "x1"
  159. assert not is_scalar(E.X1)
  160. exc = KeyError
  161. msg = "<E.X1: 'x1'>"
  162. if isinstance(
  163. index,
  164. (
  165. DatetimeIndex,
  166. TimedeltaIndex,
  167. PeriodIndex,
  168. IntervalIndex,
  169. ),
  170. ):
  171. # TODO: make these more consistent?
  172. exc = InvalidIndexError
  173. msg = "E.X1"
  174. with pytest.raises(exc, match=msg):
  175. index.get_loc(E.X1)
  176. def test_get_loc_generator(self, index):
  177. exc = KeyError
  178. if isinstance(
  179. index,
  180. (
  181. DatetimeIndex,
  182. TimedeltaIndex,
  183. PeriodIndex,
  184. IntervalIndex,
  185. MultiIndex,
  186. ),
  187. ):
  188. # TODO: make these more consistent?
  189. exc = InvalidIndexError
  190. with pytest.raises(exc, match="generator object"):
  191. # MultiIndex specifically checks for generator; others for scalar
  192. index.get_loc(x for x in range(5))
  193. def test_get_loc_masked_duplicated_na(self):
  194. # GH#48411
  195. idx = Index([1, 2, NA, NA], dtype="Int64")
  196. result = idx.get_loc(NA)
  197. expected = np.array([False, False, True, True])
  198. tm.assert_numpy_array_equal(result, expected)
  199. class TestGetIndexer:
  200. def test_get_indexer_base(self, index):
  201. if index._index_as_unique:
  202. expected = np.arange(index.size, dtype=np.intp)
  203. actual = index.get_indexer(index)
  204. tm.assert_numpy_array_equal(expected, actual)
  205. else:
  206. msg = "Reindexing only valid with uniquely valued Index objects"
  207. with pytest.raises(InvalidIndexError, match=msg):
  208. index.get_indexer(index)
  209. with pytest.raises(ValueError, match="Invalid fill method"):
  210. index.get_indexer(index, method="invalid")
  211. def test_get_indexer_consistency(self, index):
  212. # See GH#16819
  213. if index._index_as_unique:
  214. indexer = index.get_indexer(index[0:2])
  215. assert isinstance(indexer, np.ndarray)
  216. assert indexer.dtype == np.intp
  217. else:
  218. msg = "Reindexing only valid with uniquely valued Index objects"
  219. with pytest.raises(InvalidIndexError, match=msg):
  220. index.get_indexer(index[0:2])
  221. indexer, _ = index.get_indexer_non_unique(index[0:2])
  222. assert isinstance(indexer, np.ndarray)
  223. assert indexer.dtype == np.intp
  224. def test_get_indexer_masked_duplicated_na(self):
  225. # GH#48411
  226. idx = Index([1, 2, NA, NA], dtype="Int64")
  227. result = idx.get_indexer_for(Index([1, NA], dtype="Int64"))
  228. expected = np.array([0, 2, 3], dtype=result.dtype)
  229. tm.assert_numpy_array_equal(result, expected)
  230. class TestConvertSliceIndexer:
  231. def test_convert_almost_null_slice(self, index):
  232. # slice with None at both ends, but not step
  233. key = slice(None, None, "foo")
  234. if isinstance(index, IntervalIndex):
  235. msg = "label-based slicing with step!=1 is not supported for IntervalIndex"
  236. with pytest.raises(ValueError, match=msg):
  237. index._convert_slice_indexer(key, "loc")
  238. else:
  239. msg = "'>=' not supported between instances of 'str' and 'int'"
  240. with pytest.raises(TypeError, match=msg):
  241. index._convert_slice_indexer(key, "loc")
  242. class TestPutmask:
  243. def test_putmask_with_wrong_mask(self, index):
  244. # GH#18368
  245. if not len(index):
  246. return
  247. fill = index[0]
  248. msg = "putmask: mask and data must be the same size"
  249. with pytest.raises(ValueError, match=msg):
  250. index.putmask(np.ones(len(index) + 1, np.bool_), fill)
  251. with pytest.raises(ValueError, match=msg):
  252. index.putmask(np.ones(len(index) - 1, np.bool_), fill)
  253. with pytest.raises(ValueError, match=msg):
  254. index.putmask("foo", fill)
  255. @pytest.mark.parametrize(
  256. "idx", [Index([1, 2, 3]), Index([0.1, 0.2, 0.3]), Index(["a", "b", "c"])]
  257. )
  258. def test_getitem_deprecated_float(idx):
  259. # https://github.com/pandas-dev/pandas/issues/34191
  260. msg = "Indexing with a float is no longer supported"
  261. with pytest.raises(IndexError, match=msg):
  262. idx[1.0]
  263. @pytest.mark.parametrize(
  264. "idx,target,expected",
  265. [
  266. ([np.nan, "var1", np.nan], [np.nan], np.array([0, 2], dtype=np.intp)),
  267. (
  268. [np.nan, "var1", np.nan],
  269. [np.nan, "var1"],
  270. np.array([0, 2, 1], dtype=np.intp),
  271. ),
  272. (
  273. np.array([np.nan, "var1", np.nan], dtype=object),
  274. [np.nan],
  275. np.array([0, 2], dtype=np.intp),
  276. ),
  277. (
  278. DatetimeIndex(["2020-08-05", NaT, NaT]),
  279. [NaT],
  280. np.array([1, 2], dtype=np.intp),
  281. ),
  282. (["a", "b", "a", np.nan], [np.nan], np.array([3], dtype=np.intp)),
  283. (
  284. np.array(["b", np.nan, float("NaN"), "b"], dtype=object),
  285. Index([np.nan], dtype=object),
  286. np.array([1, 2], dtype=np.intp),
  287. ),
  288. ],
  289. )
  290. def test_get_indexer_non_unique_multiple_nans(idx, target, expected):
  291. # GH 35392
  292. axis = Index(idx)
  293. actual = axis.get_indexer_for(target)
  294. tm.assert_numpy_array_equal(actual, expected)
  295. def test_get_indexer_non_unique_nans_in_object_dtype_target(nulls_fixture):
  296. idx = Index([1.0, 2.0])
  297. target = Index([1, nulls_fixture], dtype="object")
  298. result_idx, result_missing = idx.get_indexer_non_unique(target)
  299. tm.assert_numpy_array_equal(result_idx, np.array([0, -1], dtype=np.intp))
  300. tm.assert_numpy_array_equal(result_missing, np.array([1], dtype=np.intp))