| 113 | tm.assert_index_equal(dropped.index, expected) |
| 114 | |
| 115 | def test_drop(self): |
| 116 | simple = DataFrame({"A": [1, 2, 3, 4], "B": [0, 1, 2, 3]}) |
| 117 | tm.assert_frame_equal(simple.drop("A", axis=1), simple[["B"]]) |
| 118 | tm.assert_frame_equal(simple.drop(["A", "B"], axis="columns"), simple[[]]) |
| 119 | tm.assert_frame_equal(simple.drop([0, 1, 3], axis=0), simple.loc[[2], :]) |
| 120 | tm.assert_frame_equal(simple.drop([0, 3], axis="index"), simple.loc[[1, 2], :]) |
| 121 | |
| 122 | with pytest.raises(KeyError, match=r"\[5\] not found in axis"): |
| 123 | simple.drop(5) |
| 124 | with pytest.raises(KeyError, match=r"\['C'\] not found in axis"): |
| 125 | simple.drop("C", axis=1) |
| 126 | with pytest.raises(KeyError, match=r"\[5\] not found in axis"): |
| 127 | simple.drop([1, 5]) |
| 128 | with pytest.raises(KeyError, match=r"\['C'\] not found in axis"): |
| 129 | simple.drop(["A", "C"], axis=1) |
| 130 | |
| 131 | # GH 42881 |
| 132 | with pytest.raises(KeyError, match=r"\['C', 'D', 'F'\] not found in axis"): |
| 133 | simple.drop(["C", "D", "F"], axis=1) |
| 134 | |
| 135 | # errors = 'ignore' |
| 136 | tm.assert_frame_equal(simple.drop(5, errors="ignore"), simple) |
| 137 | tm.assert_frame_equal( |
| 138 | simple.drop([0, 5], errors="ignore"), simple.loc[[1, 2, 3], :] |
| 139 | ) |
| 140 | tm.assert_frame_equal(simple.drop("C", axis=1, errors="ignore"), simple) |
| 141 | tm.assert_frame_equal( |
| 142 | simple.drop(["A", "C"], axis=1, errors="ignore"), simple[["B"]] |
| 143 | ) |
| 144 | |
| 145 | # non-unique - wheee! |
| 146 | nu_df = DataFrame( |
| 147 | list(zip(range(3), range(-3, 1), list("abc"))), columns=["a", "a", "b"] |
| 148 | ) |
| 149 | tm.assert_frame_equal(nu_df.drop("a", axis=1), nu_df[["b"]]) |
| 150 | tm.assert_frame_equal(nu_df.drop("b", axis="columns"), nu_df["a"]) |
| 151 | tm.assert_frame_equal(nu_df.drop([]), nu_df) # GH 16398 |
| 152 | |
| 153 | nu_df = nu_df.set_index(Index(["X", "Y", "X"])) |
| 154 | nu_df.columns = list("abc") |
| 155 | tm.assert_frame_equal(nu_df.drop("X", axis="rows"), nu_df.loc[["Y"], :]) |
| 156 | tm.assert_frame_equal(nu_df.drop(["X", "Y"], axis=0), nu_df.loc[[], :]) |
| 157 | |
| 158 | # inplace cache issue |
| 159 | # GH#5628 |
| 160 | df = DataFrame( |
| 161 | np.random.default_rng(2).standard_normal((10, 3)), columns=list("abc") |
| 162 | ) |
| 163 | expected = df[~(df.b > 0)] |
| 164 | return_value = df.drop(labels=df[df.b > 0].index, inplace=True) |
| 165 | assert return_value is None |
| 166 | tm.assert_frame_equal(df, expected) |
| 167 | |
| 168 | def test_drop_multiindex_not_lexsorted(self, performance_warning): |
| 169 | # GH#11640 |