Specify type for pa.chunked_array

Author: jbrockmendel

pandas/core/arrays/_arrow_string_mixins.py

@@ -11,6 +11,7 @@
 
 from pandas.compat import (
     pa_version_under10p1,
+    pa_version_under11p0,
     pa_version_under13p0,
     pa_version_under17p0,
 )
@@ -103,6 +104,25 @@ def _str_get(self, i: int) -> Self:
         result = pc.if_else(not_out_of_bounds, selected, null_value)
         return type(self)(result)
 
+    def _str_slice(
+        self, start: int | None = None, stop: int | None = None, step: int | None = None
+    ) -> Self:
+        if pa_version_under11p0:
+            # GH#59724
+            result = self._apply_elementwise(lambda val: val[start:stop:step])
+            return type(self)(pa.chunked_array(result, type=self._pa_array.type))
+        if start is None:
+            if step is not None and step < 0:
+                # GH#59710
+                start = -1
+            else:
+                start = 0
+        if step is None:
+            step = 1
+        return type(self)(
+            pc.utf8_slice_codeunits(self._pa_array, start=start, stop=stop, step=step)
+        )
+
     def _str_slice_replace(
         self, start: int | None = None, stop: int | None = None, repl: str | None = None
     ) -> Self:

pandas/core/arrays/arrow/array.py

@@ -2393,25 +2393,6 @@ def _str_rpartition(self, sep: str, expand: bool) -> Self:
         result = self._apply_elementwise(predicate)
         return type(self)(pa.chunked_array(result))
 
-    def _str_slice(
-        self, start: int | None = None, stop: int | None = None, step: int | None = None
-    ) -> Self:
-        if pa_version_under11p0:
-            # GH#59724
-            result = self._apply_elementwise(lambda val: val[start:stop:step])
-            return type(self)(pa.chunked_array(result))
-        if start is None:
-            if step is not None and step < 0:
-                # GH#59710
-                start = -1
-            else:
-                start = 0
-        if step is None:
-            step = 1
-        return type(self)(
-            pc.utf8_slice_codeunits(self._pa_array, start=start, stop=stop, step=step)
-        )
-
     def _str_len(self) -> Self:
         return type(self)(pc.utf8_length(self._pa_array))
 

pandas/core/arrays/string_arrow.py

@@ -293,7 +293,7 @@ def astype(self, dtype, copy: bool = True):
     _str_startswith = ArrowStringArrayMixin._str_startswith
     _str_endswith = ArrowStringArrayMixin._str_endswith
     _str_pad = ArrowStringArrayMixin._str_pad
-    _str_slice = ArrowExtensionArray._str_slice
+    _str_slice = ArrowStringArrayMixin._str_slice
 
     def _str_contains(
         self, pat, case: bool = True, flags: int = 0, na=np.nan, regex: bool = True

pandas/tests/extension/test_arrow.py

@@ -2036,12 +2036,7 @@ def test_str_join_string_type():
         [None, 2, None, ["ab", None]],
         [None, 2, 1, ["ab", None]],
         [1, 3, 1, ["bc", None]],
-        (
-            None,
-            None,
-            -1,
-            ["dcba", None],
-        ),
+        (None, None, -1, ["dcba", None]),
     ],
 )
 def test_str_slice(start, stop, step, exp):

pandas/tests/strings/test_string_array.py

@@ -13,7 +13,7 @@
 
 
 @pytest.mark.filterwarnings("ignore:Falling back")
-def test_string_array(nullable_string_dtype, any_string_method, request):
+def test_string_array(nullable_string_dtype, any_string_method):
     method_name, args, kwargs = any_string_method
 
     data = ["a", "bb", np.nan, "ccc"]

pandas/tests/strings/test_strings.py

@@ -403,7 +403,7 @@ def test_pipe_failures(any_string_dtype):
         (3, 0, -1, ["ofa", "aba", np.nan, "aba"]),
     ],
 )
-def test_slice(start, stop, step, expected, any_string_dtype, request):
+def test_slice(start, stop, step, expected, any_string_dtype):
     ser = Series(["aafootwo", "aabartwo", np.nan, "aabazqux"], dtype=any_string_dtype)
     result = ser.str.slice(start, stop, step)
     expected = Series(expected, dtype=any_string_dtype)