Skip to content

Commit 2d8b55e

Browse files
committed
Fix Index.equal
1 parent 25e6462 commit 2d8b55e

File tree

3 files changed

+58
-14
lines changed

3 files changed

+58
-14
lines changed

pandas/core/indexes/base.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5481,11 +5481,7 @@ def equals(self, other: Any) -> bool:
54815481
# quickly return if the lengths are different
54825482
return False
54835483

5484-
if (
5485-
isinstance(self.dtype, StringDtype)
5486-
and self.dtype.na_value is np.nan
5487-
and other.dtype != self.dtype
5488-
):
5484+
if isinstance(self.dtype, StringDtype) and other.dtype != self.dtype:
54895485
# TODO(infer_string) can we avoid this special case?
54905486
# special case for object behavior
54915487
return other.equals(self.astype(object))

pandas/tests/frame/test_arithmetic.py

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2183,19 +2183,18 @@ def test_enum_column_equality():
21832183
tm.assert_series_equal(result, expected)
21842184

21852185

2186-
def test_mixed_col_index_dtype(using_infer_string):
2186+
def test_mixed_col_index_dtype():
21872187
# GH 47382
21882188
df1 = DataFrame(columns=list("abc"), data=1.0, index=[0])
21892189
df2 = DataFrame(columns=list("abc"), data=0.0, index=[0])
21902190
df1.columns = df2.columns.astype("string")
21912191
result = df1 + df2
21922192
expected = DataFrame(columns=list("abc"), data=1.0, index=[0])
2193-
if using_infer_string:
2194-
# df2.columns.dtype will be "str" instead of object,
2195-
# so the aligned result will be "string", not object
2196-
if HAS_PYARROW:
2197-
dtype = "string[pyarrow]"
2198-
else:
2199-
dtype = "string"
2200-
expected.columns = expected.columns.astype(dtype)
2193+
2194+
if HAS_PYARROW:
2195+
dtype = "string[pyarrow]"
2196+
else:
2197+
dtype = "string"
2198+
expected.columns = expected.columns.astype(dtype)
2199+
22012200
tm.assert_frame_equal(result, expected)

pandas/tests/indexes/test_base.py

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
ensure_index,
4141
ensure_index_from_sequences,
4242
)
43+
from pandas.testing import assert_series_equal
4344

4445

4546
class TestIndex:
@@ -1717,3 +1718,51 @@ def test_is_monotonic_pyarrow_list_type():
17171718
idx = Index([[1], [2, 3]], dtype=pd.ArrowDtype(pa.list_(pa.int64())))
17181719
assert not idx.is_monotonic_increasing
17191720
assert not idx.is_monotonic_decreasing
1721+
1722+
1723+
@pytest.mark.parametrize(
1724+
"dtype",
1725+
[
1726+
"string[python]",
1727+
pytest.param(
1728+
pd.StringDtype(storage="pyarrow", na_value=pd.NA),
1729+
marks=td.skip_if_no("pyarrow"),
1730+
),
1731+
pytest.param(
1732+
pd.StringDtype(storage="pyarrow", na_value=np.nan),
1733+
marks=td.skip_if_no("pyarrow"),
1734+
),
1735+
],
1736+
)
1737+
def test_index_equals_different_string_dtype(dtype):
1738+
# GH 61099
1739+
idx_obj = Index(["a", "b", "c"])
1740+
idx_str = Index(["a", "b", "c"], dtype=dtype)
1741+
1742+
assert idx_obj.equals(idx_str)
1743+
assert idx_str.equals(idx_obj)
1744+
1745+
1746+
@pytest.mark.parametrize(
1747+
"dtype",
1748+
[
1749+
"string[python]",
1750+
pytest.param(
1751+
pd.StringDtype(storage="pyarrow", na_value=pd.NA),
1752+
marks=td.skip_if_no("pyarrow"),
1753+
),
1754+
pytest.param(
1755+
pd.StringDtype(storage="pyarrow", na_value=np.nan),
1756+
marks=td.skip_if_no("pyarrow"),
1757+
),
1758+
],
1759+
)
1760+
def test_index_comparison_different_string_dtype(dtype):
1761+
# GH 61099
1762+
idx = Index(["a", "b", "c"])
1763+
s_obj = Series([1, 2, 3], index=idx)
1764+
s_str = Series([4, 5, 6], index=idx.astype(dtype))
1765+
1766+
expected = Series([True, True, True], index=["a", "b", "c"])
1767+
result = s_obj < s_str
1768+
assert_series_equal(result, expected)

0 commit comments

Comments
 (0)