Skip to content

Commit

Permalink
Backport PR #54974 on branch 2.1.x (Include pyarrow_numpy string in e…
Browse files Browse the repository at this point in the history
…fficient merge implementation) (#55021)

Backport PR #54974: Include pyarrow_numpy string in efficient merge implementation

Co-authored-by: Patrick Hoefler <61934744+phofl@users.noreply.github.com>
  • Loading branch information
meeseeksmachine and phofl committed Sep 6, 2023
1 parent eab54ee commit b736912
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 5 deletions.
3 changes: 2 additions & 1 deletion pandas/core/reshape/merge.py
Original file line number Diff line number Diff line change
Expand Up @@ -2417,7 +2417,8 @@ def _factorize_keys(

elif isinstance(lk, ExtensionArray) and lk.dtype == rk.dtype:
if (isinstance(lk.dtype, ArrowDtype) and is_string_dtype(lk.dtype)) or (
isinstance(lk.dtype, StringDtype) and lk.dtype.storage == "pyarrow"
isinstance(lk.dtype, StringDtype)
and lk.dtype.storage in ["pyarrow", "pyarrow_numpy"]
):
import pyarrow as pa
import pyarrow.compute as pc
Expand Down
8 changes: 4 additions & 4 deletions pandas/tests/reshape/merge/test_merge.py
Original file line number Diff line number Diff line change
Expand Up @@ -2872,13 +2872,13 @@ def test_merge_ea_int_and_float_numpy():
tm.assert_frame_equal(result, expected.astype("float64"))


def test_merge_arrow_string_index():
def test_merge_arrow_string_index(any_string_dtype):
# GH#54894
pytest.importorskip("pyarrow")
left = DataFrame({"a": ["a", "b"]}, dtype="string[pyarrow]")
right = DataFrame({"b": 1}, index=Index(["a", "c"], dtype="string[pyarrow]"))
left = DataFrame({"a": ["a", "b"]}, dtype=any_string_dtype)
right = DataFrame({"b": 1}, index=Index(["a", "c"], dtype=any_string_dtype))
result = left.merge(right, left_on="a", right_index=True, how="left")
expected = DataFrame(
{"a": Series(["a", "b"], dtype="string[pyarrow]"), "b": [1, np.nan]}
{"a": Series(["a", "b"], dtype=any_string_dtype), "b": [1, np.nan]}
)
tm.assert_frame_equal(result, expected)

0 comments on commit b736912

Please sign in to comment.