Skip to content

Commit fbbd654

Browse files
Backport PR #45983: REGR: drop raising with ea index and duplicates (#46020)
Co-authored-by: Patrick Hoefler <[email protected]>
1 parent 14fa026 commit fbbd654

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
@@ -4351,6 +4351,10 @@ def _drop_axis(
43514351
if errors == "raise" and labels_missing:
43524352
raise KeyError(f"{labels} not found in axis")
43534353

4354+
if is_extension_array_dtype(mask.dtype):
4355+
# GH#45860
4356+
mask = mask.to_numpy(dtype=bool)
4357+
43544358
indexer = mask.nonzero()[0]
43554359
new_axis = axis.take(indexer)
43564360

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

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

0 commit comments

Comments
 (0)