Skip to content

Commit 1b3ebe4

Browse files
authored
Include pyarrow_numpy string in efficient merge implementation (#54974)
Inlude pyarrow_numpy string in efficient merge implementation
1 parent e681413 commit 1b3ebe4

File tree

2 files changed

+6
-5
lines changed

2 files changed

+6
-5
lines changed

pandas/core/reshape/merge.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -2421,7 +2421,8 @@ def _factorize_keys(
24212421

24222422
elif isinstance(lk, ExtensionArray) and lk.dtype == rk.dtype:
24232423
if (isinstance(lk.dtype, ArrowDtype) and is_string_dtype(lk.dtype)) or (
2424-
isinstance(lk.dtype, StringDtype) and lk.dtype.storage == "pyarrow"
2424+
isinstance(lk.dtype, StringDtype)
2425+
and lk.dtype.storage in ["pyarrow", "pyarrow_numpy"]
24252426
):
24262427
import pyarrow as pa
24272428
import pyarrow.compute as pc

pandas/tests/reshape/merge/test_merge.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -2949,13 +2949,13 @@ def test_merge_ea_int_and_float_numpy():
29492949
tm.assert_frame_equal(result, expected.astype("float64"))
29502950

29512951

2952-
def test_merge_arrow_string_index():
2952+
def test_merge_arrow_string_index(any_string_dtype):
29532953
# GH#54894
29542954
pytest.importorskip("pyarrow")
2955-
left = DataFrame({"a": ["a", "b"]}, dtype="string[pyarrow]")
2956-
right = DataFrame({"b": 1}, index=Index(["a", "c"], dtype="string[pyarrow]"))
2955+
left = DataFrame({"a": ["a", "b"]}, dtype=any_string_dtype)
2956+
right = DataFrame({"b": 1}, index=Index(["a", "c"], dtype=any_string_dtype))
29572957
result = left.merge(right, left_on="a", right_index=True, how="left")
29582958
expected = DataFrame(
2959-
{"a": Series(["a", "b"], dtype="string[pyarrow]"), "b": [1, np.nan]}
2959+
{"a": Series(["a", "b"], dtype=any_string_dtype), "b": [1, np.nan]}
29602960
)
29612961
tm.assert_frame_equal(result, expected)

0 commit comments

Comments
 (0)