test_indexing.py 3.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. import numpy as np
  2. import pytest
  3. from pandas import (
  4. Index,
  5. RangeIndex,
  6. )
  7. import pandas._testing as tm
  8. class TestGetIndexer:
  9. def test_get_indexer(self):
  10. index = RangeIndex(start=0, stop=20, step=2)
  11. target = RangeIndex(10)
  12. indexer = index.get_indexer(target)
  13. expected = np.array([0, -1, 1, -1, 2, -1, 3, -1, 4, -1], dtype=np.intp)
  14. tm.assert_numpy_array_equal(indexer, expected)
  15. def test_get_indexer_pad(self):
  16. index = RangeIndex(start=0, stop=20, step=2)
  17. target = RangeIndex(10)
  18. indexer = index.get_indexer(target, method="pad")
  19. expected = np.array([0, 0, 1, 1, 2, 2, 3, 3, 4, 4], dtype=np.intp)
  20. tm.assert_numpy_array_equal(indexer, expected)
  21. def test_get_indexer_backfill(self):
  22. index = RangeIndex(start=0, stop=20, step=2)
  23. target = RangeIndex(10)
  24. indexer = index.get_indexer(target, method="backfill")
  25. expected = np.array([0, 1, 1, 2, 2, 3, 3, 4, 4, 5], dtype=np.intp)
  26. tm.assert_numpy_array_equal(indexer, expected)
  27. def test_get_indexer_limit(self):
  28. # GH#28631
  29. idx = RangeIndex(4)
  30. target = RangeIndex(6)
  31. result = idx.get_indexer(target, method="pad", limit=1)
  32. expected = np.array([0, 1, 2, 3, 3, -1], dtype=np.intp)
  33. tm.assert_numpy_array_equal(result, expected)
  34. @pytest.mark.parametrize("stop", [0, -1, -2])
  35. def test_get_indexer_decreasing(self, stop):
  36. # GH#28678
  37. index = RangeIndex(7, stop, -3)
  38. result = index.get_indexer(range(9))
  39. expected = np.array([-1, 2, -1, -1, 1, -1, -1, 0, -1], dtype=np.intp)
  40. tm.assert_numpy_array_equal(result, expected)
  41. class TestTake:
  42. def test_take_preserve_name(self):
  43. index = RangeIndex(1, 5, name="foo")
  44. taken = index.take([3, 0, 1])
  45. assert index.name == taken.name
  46. def test_take_fill_value(self):
  47. # GH#12631
  48. idx = RangeIndex(1, 4, name="xxx")
  49. result = idx.take(np.array([1, 0, -1]))
  50. expected = Index([2, 1, 3], dtype=np.int64, name="xxx")
  51. tm.assert_index_equal(result, expected)
  52. # fill_value
  53. msg = "Unable to fill values because RangeIndex cannot contain NA"
  54. with pytest.raises(ValueError, match=msg):
  55. idx.take(np.array([1, 0, -1]), fill_value=True)
  56. # allow_fill=False
  57. result = idx.take(np.array([1, 0, -1]), allow_fill=False, fill_value=True)
  58. expected = Index([2, 1, 3], dtype=np.int64, name="xxx")
  59. tm.assert_index_equal(result, expected)
  60. msg = "Unable to fill values because RangeIndex cannot contain NA"
  61. with pytest.raises(ValueError, match=msg):
  62. idx.take(np.array([1, 0, -2]), fill_value=True)
  63. with pytest.raises(ValueError, match=msg):
  64. idx.take(np.array([1, 0, -5]), fill_value=True)
  65. msg = "index -5 is out of bounds for (axis 0 with )?size 3"
  66. with pytest.raises(IndexError, match=msg):
  67. idx.take(np.array([1, -5]))
  68. class TestWhere:
  69. def test_where_putmask_range_cast(self):
  70. # GH#43240
  71. idx = RangeIndex(0, 5, name="test")
  72. mask = np.array([True, True, False, False, False])
  73. result = idx.putmask(mask, 10)
  74. expected = Index([10, 10, 2, 3, 4], dtype=np.int64, name="test")
  75. tm.assert_index_equal(result, expected)
  76. result = idx.where(~mask, 10)
  77. tm.assert_index_equal(result, expected)