test_isin.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. import numpy as np
  2. import pytest
  3. from pandas import MultiIndex
  4. import pandas._testing as tm
  5. def test_isin_nan():
  6. idx = MultiIndex.from_arrays([["foo", "bar"], [1.0, np.nan]])
  7. tm.assert_numpy_array_equal(idx.isin([("bar", np.nan)]), np.array([False, True]))
  8. tm.assert_numpy_array_equal(
  9. idx.isin([("bar", float("nan"))]), np.array([False, True])
  10. )
  11. def test_isin_missing(nulls_fixture):
  12. # GH48905
  13. mi1 = MultiIndex.from_tuples([(1, nulls_fixture)])
  14. mi2 = MultiIndex.from_tuples([(1, 1), (1, 2)])
  15. result = mi2.isin(mi1)
  16. expected = np.array([False, False])
  17. tm.assert_numpy_array_equal(result, expected)
  18. def test_isin():
  19. values = [("foo", 2), ("bar", 3), ("quux", 4)]
  20. idx = MultiIndex.from_arrays([["qux", "baz", "foo", "bar"], np.arange(4)])
  21. result = idx.isin(values)
  22. expected = np.array([False, False, True, True])
  23. tm.assert_numpy_array_equal(result, expected)
  24. # empty, return dtype bool
  25. idx = MultiIndex.from_arrays([[], []])
  26. result = idx.isin(values)
  27. assert len(result) == 0
  28. assert result.dtype == np.bool_
  29. def test_isin_level_kwarg():
  30. idx = MultiIndex.from_arrays([["qux", "baz", "foo", "bar"], np.arange(4)])
  31. vals_0 = ["foo", "bar", "quux"]
  32. vals_1 = [2, 3, 10]
  33. expected = np.array([False, False, True, True])
  34. tm.assert_numpy_array_equal(expected, idx.isin(vals_0, level=0))
  35. tm.assert_numpy_array_equal(expected, idx.isin(vals_0, level=-2))
  36. tm.assert_numpy_array_equal(expected, idx.isin(vals_1, level=1))
  37. tm.assert_numpy_array_equal(expected, idx.isin(vals_1, level=-1))
  38. msg = "Too many levels: Index has only 2 levels, not 6"
  39. with pytest.raises(IndexError, match=msg):
  40. idx.isin(vals_0, level=5)
  41. msg = "Too many levels: Index has only 2 levels, -5 is not a valid level number"
  42. with pytest.raises(IndexError, match=msg):
  43. idx.isin(vals_0, level=-5)
  44. with pytest.raises(KeyError, match=r"'Level 1\.0 not found'"):
  45. idx.isin(vals_0, level=1.0)
  46. with pytest.raises(KeyError, match=r"'Level -1\.0 not found'"):
  47. idx.isin(vals_1, level=-1.0)
  48. with pytest.raises(KeyError, match="'Level A not found'"):
  49. idx.isin(vals_1, level="A")
  50. idx.names = ["A", "B"]
  51. tm.assert_numpy_array_equal(expected, idx.isin(vals_0, level="A"))
  52. tm.assert_numpy_array_equal(expected, idx.isin(vals_1, level="B"))
  53. with pytest.raises(KeyError, match="'Level C not found'"):
  54. idx.isin(vals_1, level="C")
  55. @pytest.mark.parametrize(
  56. "labels,expected,level",
  57. [
  58. ([("b", np.nan)], np.array([False, False, True]), None),
  59. ([np.nan, "a"], np.array([True, True, False]), 0),
  60. (["d", np.nan], np.array([False, True, True]), 1),
  61. ],
  62. )
  63. def test_isin_multi_index_with_missing_value(labels, expected, level):
  64. # GH 19132
  65. midx = MultiIndex.from_arrays([[np.nan, "a", "b"], ["c", "d", np.nan]])
  66. result = midx.isin(labels, level=level)
  67. tm.assert_numpy_array_equal(result, expected)
  68. def test_isin_empty():
  69. # GH#51599
  70. midx = MultiIndex.from_arrays([[1, 2], [3, 4]])
  71. result = midx.isin([])
  72. expected = np.array([False, False])
  73. tm.assert_numpy_array_equal(result, expected)
  74. def test_isin_generator():
  75. # GH#52568
  76. midx = MultiIndex.from_tuples([(1, 2)])
  77. result = midx.isin(x for x in [(1, 2)])
  78. expected = np.array([True])
  79. tm.assert_numpy_array_equal(result, expected)