test: add tests for pyarrow \d replacement

Also tests for ambiguous reference \d + another literal as digit.

Author: Alvaro-Kothe

pandas/tests/strings/test_find_replace.py

@@ -9,6 +9,7 @@
 import pandas as pd
 from pandas import (
     Series,
+    StringDtype,
     _testing as tm,
 )
 from pandas.tests.strings import (
@@ -584,6 +585,10 @@ def test_replace_callable_raises(any_string_dtype, repl):
             r"\g<three> \g<two> \g<one>",
             ["Three Two One", "Baz Bar Foo"],
         ),
+        (
+            r"\3 \2 \1",
+            ["Three Two One", "Baz Bar Foo"],
+        ),
         (
             r"\g<3> \g<2> \g<1>",
             ["Three Two One", "Baz Bar Foo"],
@@ -599,6 +604,7 @@ def test_replace_callable_raises(any_string_dtype, repl):
     ],
     ids=[
         "named_groups_full_swap",
+        "numbered_groups_no_g_full_swap",
         "numbered_groups_full_swap",
         "single_group_with_literal",
         "mixed_group_reference_with_literal",
@@ -623,33 +629,79 @@ def test_replace_named_groups_regex_swap(
     [
         r"\g<20>",
         r"\20",
+        r"\40",
+        r"\4",
     ],
 )
 @pytest.mark.parametrize("use_compile", [True, False])
 def test_replace_named_groups_regex_swap_expected_fail(
-    any_string_dtype, repl, use_compile
+    any_string_dtype, repl, use_compile, request
 ):
     # GH#57636
+    if (
+        not use_compile
+        and r"\g" not in repl
+        and isinstance(any_string_dtype, StringDtype)
+        and any_string_dtype.storage == "pyarrow"
+    ):
+        # calls pyarrow method directly
+        if repl == r"\20":
+            mark = pytest.mark.xfail(reason="PyArrow interprets as group + literal")
+            request.applymarker(mark)
+
+        pa = pytest.importorskip("pyarrow")
+        error_type = pa.ArrowInvalid
+        error_msg = r"only has \d parenthesized subexpressions"
+    else:
+        error_type = re.error
+        error_msg = "invalid group reference"
+
     pattern = r"(?P<one>\w+) (?P<two>\w+) (?P<three>\w+)"
     if use_compile:
         pattern = re.compile(pattern)
     ser = Series(["One Two Three", "Foo Bar Baz"], dtype=any_string_dtype)
 
-    with pytest.raises(re.error, match="invalid group reference"):
+    with pytest.raises(error_type, match=error_msg):
         ser.str.replace(pattern, repl, regex=True)
 
 
-@pytest.mark.parametrize("use_compile", [True, False])
-def test_replace_non_named_group(any_string_dtype, use_compile):
-    ser = Series(["var.one[0]", "var.two[1]", "var.three[2]"], dtype=any_string_dtype)
-    pattern = r"\[(\d+)\]"
-    if use_compile:
-        pattern = re.compile(pattern)
-    repl = r"(\1)"
+@pytest.mark.parametrize(
+    "pattern, repl",
+    [
+        (r"(\w+) (\w+) (\w+)", r"\20"),
+        (r"(?P<one>\w+) (?P<two>\w+) (?P<three>\w+)", r"\20"),
+    ],
+)
+def test_pyarrow_ambiguous_group_references(pyarrow_string_dtype, pattern, repl):
+    # GH#62653
+    ser = Series(["One Two Three", "Foo Bar Baz"], dtype=pyarrow_string_dtype)
+
     result = ser.str.replace(pattern, repl, regex=True)
-    expected = Series(
-        ["var.one(0)", "var.two(1)", "var.three(2)"], dtype=any_string_dtype
+    expected = Series(["Two0", "Bar0"], dtype=pyarrow_string_dtype)
+    tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "pattern, repl, expected_list",
+    [
+        (
+            r"\[(?P<one>\d+)\]",
+            r"(\1)",
+            ["var.one(0)", "var.two(1)", "var.three(2)"],
+        ),
+        (
+            r"\[(\d+)\]",
+            r"(\1)",
+            ["var.one(0)", "var.two(1)", "var.three(2)"],
+        ),
+    ],
+)
+def test_pyarrow_backend_group_replacement(pattern, repl, expected_list):
+    ser = Series(["var.one[0]", "var.two[1]", "var.three[2]"]).convert_dtypes(
+        dtype_backend="pyarrow"
     )
+    result = ser.str.replace(pattern, repl, regex=True)
+    expected = Series(expected_list).convert_dtypes(dtype_backend="pyarrow")
     tm.assert_series_equal(result, expected)