test_drop.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190
  1. import numpy as np
  2. import pytest
  3. from pandas.errors import PerformanceWarning
  4. import pandas as pd
  5. from pandas import (
  6. Index,
  7. MultiIndex,
  8. )
  9. import pandas._testing as tm
  10. def test_drop(idx):
  11. dropped = idx.drop([("foo", "two"), ("qux", "one")])
  12. index = MultiIndex.from_tuples([("foo", "two"), ("qux", "one")])
  13. dropped2 = idx.drop(index)
  14. expected = idx[[0, 2, 3, 5]]
  15. tm.assert_index_equal(dropped, expected)
  16. tm.assert_index_equal(dropped2, expected)
  17. dropped = idx.drop(["bar"])
  18. expected = idx[[0, 1, 3, 4, 5]]
  19. tm.assert_index_equal(dropped, expected)
  20. dropped = idx.drop("foo")
  21. expected = idx[[2, 3, 4, 5]]
  22. tm.assert_index_equal(dropped, expected)
  23. index = MultiIndex.from_tuples([("bar", "two")])
  24. with pytest.raises(KeyError, match=r"^\('bar', 'two'\)$"):
  25. idx.drop([("bar", "two")])
  26. with pytest.raises(KeyError, match=r"^\('bar', 'two'\)$"):
  27. idx.drop(index)
  28. with pytest.raises(KeyError, match=r"^'two'$"):
  29. idx.drop(["foo", "two"])
  30. # partially correct argument
  31. mixed_index = MultiIndex.from_tuples([("qux", "one"), ("bar", "two")])
  32. with pytest.raises(KeyError, match=r"^\('bar', 'two'\)$"):
  33. idx.drop(mixed_index)
  34. # error='ignore'
  35. dropped = idx.drop(index, errors="ignore")
  36. expected = idx[[0, 1, 2, 3, 4, 5]]
  37. tm.assert_index_equal(dropped, expected)
  38. dropped = idx.drop(mixed_index, errors="ignore")
  39. expected = idx[[0, 1, 2, 3, 5]]
  40. tm.assert_index_equal(dropped, expected)
  41. dropped = idx.drop(["foo", "two"], errors="ignore")
  42. expected = idx[[2, 3, 4, 5]]
  43. tm.assert_index_equal(dropped, expected)
  44. # mixed partial / full drop
  45. dropped = idx.drop(["foo", ("qux", "one")])
  46. expected = idx[[2, 3, 5]]
  47. tm.assert_index_equal(dropped, expected)
  48. # mixed partial / full drop / error='ignore'
  49. mixed_index = ["foo", ("qux", "one"), "two"]
  50. with pytest.raises(KeyError, match=r"^'two'$"):
  51. idx.drop(mixed_index)
  52. dropped = idx.drop(mixed_index, errors="ignore")
  53. expected = idx[[2, 3, 5]]
  54. tm.assert_index_equal(dropped, expected)
  55. def test_droplevel_with_names(idx):
  56. index = idx[idx.get_loc("foo")]
  57. dropped = index.droplevel(0)
  58. assert dropped.name == "second"
  59. index = MultiIndex(
  60. levels=[Index(range(4)), Index(range(4)), Index(range(4))],
  61. codes=[
  62. np.array([0, 0, 1, 2, 2, 2, 3, 3]),
  63. np.array([0, 1, 0, 0, 0, 1, 0, 1]),
  64. np.array([1, 0, 1, 1, 0, 0, 1, 0]),
  65. ],
  66. names=["one", "two", "three"],
  67. )
  68. dropped = index.droplevel(0)
  69. assert dropped.names == ("two", "three")
  70. dropped = index.droplevel("two")
  71. expected = index.droplevel(1)
  72. assert dropped.equals(expected)
  73. def test_droplevel_list():
  74. index = MultiIndex(
  75. levels=[Index(range(4)), Index(range(4)), Index(range(4))],
  76. codes=[
  77. np.array([0, 0, 1, 2, 2, 2, 3, 3]),
  78. np.array([0, 1, 0, 0, 0, 1, 0, 1]),
  79. np.array([1, 0, 1, 1, 0, 0, 1, 0]),
  80. ],
  81. names=["one", "two", "three"],
  82. )
  83. dropped = index[:2].droplevel(["three", "one"])
  84. expected = index[:2].droplevel(2).droplevel(0)
  85. assert dropped.equals(expected)
  86. dropped = index[:2].droplevel([])
  87. expected = index[:2]
  88. assert dropped.equals(expected)
  89. msg = (
  90. "Cannot remove 3 levels from an index with 3 levels: "
  91. "at least one level must be left"
  92. )
  93. with pytest.raises(ValueError, match=msg):
  94. index[:2].droplevel(["one", "two", "three"])
  95. with pytest.raises(KeyError, match="'Level four not found'"):
  96. index[:2].droplevel(["one", "four"])
  97. def test_drop_not_lexsorted():
  98. # GH 12078
  99. # define the lexsorted version of the multi-index
  100. tuples = [("a", ""), ("b1", "c1"), ("b2", "c2")]
  101. lexsorted_mi = MultiIndex.from_tuples(tuples, names=["b", "c"])
  102. assert lexsorted_mi._is_lexsorted()
  103. # and the not-lexsorted version
  104. df = pd.DataFrame(
  105. columns=["a", "b", "c", "d"], data=[[1, "b1", "c1", 3], [1, "b2", "c2", 4]]
  106. )
  107. df = df.pivot_table(index="a", columns=["b", "c"], values="d")
  108. df = df.reset_index()
  109. not_lexsorted_mi = df.columns
  110. assert not not_lexsorted_mi._is_lexsorted()
  111. # compare the results
  112. tm.assert_index_equal(lexsorted_mi, not_lexsorted_mi)
  113. with tm.assert_produces_warning(PerformanceWarning):
  114. tm.assert_index_equal(lexsorted_mi.drop("a"), not_lexsorted_mi.drop("a"))
  115. def test_drop_with_nan_in_index(nulls_fixture):
  116. # GH#18853
  117. mi = MultiIndex.from_tuples([("blah", nulls_fixture)], names=["name", "date"])
  118. msg = r"labels \[Timestamp\('2001-01-01 00:00:00'\)\] not found in level"
  119. with pytest.raises(KeyError, match=msg):
  120. mi.drop(pd.Timestamp("2001"), level="date")
  121. @pytest.mark.filterwarnings("ignore::pandas.errors.PerformanceWarning")
  122. def test_drop_with_non_monotonic_duplicates():
  123. # GH#33494
  124. mi = MultiIndex.from_tuples([(1, 2), (2, 3), (1, 2)])
  125. result = mi.drop((1, 2))
  126. expected = MultiIndex.from_tuples([(2, 3)])
  127. tm.assert_index_equal(result, expected)
  128. def test_single_level_drop_partially_missing_elements():
  129. # GH 37820
  130. mi = MultiIndex.from_tuples([(1, 2), (2, 2), (3, 2)])
  131. msg = r"labels \[4\] not found in level"
  132. with pytest.raises(KeyError, match=msg):
  133. mi.drop(4, level=0)
  134. with pytest.raises(KeyError, match=msg):
  135. mi.drop([1, 4], level=0)
  136. msg = r"labels \[nan\] not found in level"
  137. with pytest.raises(KeyError, match=msg):
  138. mi.drop([np.nan], level=0)
  139. with pytest.raises(KeyError, match=msg):
  140. mi.drop([np.nan, 1, 2, 3], level=0)
  141. mi = MultiIndex.from_tuples([(np.nan, 1), (1, 2)])
  142. msg = r"labels \['a'\] not found in level"
  143. with pytest.raises(KeyError, match=msg):
  144. mi.drop([np.nan, 1, "a"], level=0)
  145. def test_droplevel_multiindex_one_level():
  146. # GH#37208
  147. index = MultiIndex.from_tuples([(2,)], names=("b",))
  148. result = index.droplevel([])
  149. expected = Index([2], name="b")
  150. tm.assert_index_equal(result, expected)