Skip to content

Commit b0945c9

Browse files
phoflyehoshuadimarsky
authored andcommitted
REGR: drop raising with ea index and duplicates (pandas-dev#45983)
1 parent 064c3f7 commit b0945c9

File tree

4 files changed

+30
-2
lines changed

4 files changed

+30
-2
lines changed

doc/source/whatsnew/v1.4.2.rst

+1-1
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ including other versions of pandas.
1414

1515
Fixed regressions
1616
~~~~~~~~~~~~~~~~~
17-
-
17+
- Fixed regression in :meth:`DataFrame.drop` and :meth:`Series.drop` when :class:`Index` had extension dtype and duplicates (:issue:`45820`)
1818
-
1919

2020
.. ---------------------------------------------------------------------------

pandas/core/generic.py

+4
Original file line numberDiff line numberDiff line change
@@ -4342,6 +4342,10 @@ def _drop_axis(
43424342
if errors == "raise" and labels_missing:
43434343
raise KeyError(f"{labels} not found in axis")
43444344

4345+
if is_extension_array_dtype(mask.dtype):
4346+
# GH#45860
4347+
mask = mask.to_numpy(dtype=bool)
4348+
43454349
indexer = mask.nonzero()[0]
43464350
new_axis = axis.take(indexer)
43474351

pandas/tests/frame/methods/test_drop.py

+12
Original file line numberDiff line numberDiff line change
@@ -537,3 +537,15 @@ def test_drop_level_missing_label_multiindex(self):
537537
df = DataFrame(index=MultiIndex.from_product([range(3), range(3)]))
538538
with pytest.raises(KeyError, match="labels \\[5\\] not found in level"):
539539
df.drop(5, level=0)
540+
541+
@pytest.mark.parametrize("idx, level", [(["a", "b"], 0), (["a"], None)])
542+
def test_drop_index_ea_dtype(self, any_numeric_ea_dtype, idx, level):
543+
# GH#45860
544+
df = DataFrame(
545+
{"a": [1, 2, 2, pd.NA], "b": 100}, dtype=any_numeric_ea_dtype
546+
).set_index(idx)
547+
result = df.drop(Index([2, pd.NA]), level=level)
548+
expected = DataFrame(
549+
{"a": [1], "b": 100}, dtype=any_numeric_ea_dtype
550+
).set_index(idx)
551+
tm.assert_frame_equal(result, expected)

pandas/tests/series/methods/test_drop.py

+13-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
import pytest
22

3-
from pandas import Series
3+
from pandas import (
4+
Index,
5+
Series,
6+
)
47
import pandas._testing as tm
58

69

@@ -98,3 +101,12 @@ def test_drop_pos_args_deprecation():
98101
result = ser.drop(1, 0)
99102
expected = Series([1, 3], index=[0, 2])
100103
tm.assert_series_equal(result, expected)
104+
105+
106+
def test_drop_index_ea_dtype(any_numeric_ea_dtype):
107+
# GH#45860
108+
df = Series(100, index=Index([1, 2, 2], dtype=any_numeric_ea_dtype))
109+
idx = Index([df.index[1]])
110+
result = df.drop(idx)
111+
expected = Series(100, index=Index([1], dtype=any_numeric_ea_dtype))
112+
tm.assert_series_equal(result, expected)

0 commit comments

Comments
 (0)