Added testcases for StringArray addition and fixes

Author: eicchen

doc/source/whatsnew/v3.0.0.rst

@@ -200,6 +200,7 @@ Other enhancements
 - :class:`Holiday` has gained the constructor argument and field ``exclude_dates`` to exclude specific datetimes from a custom holiday calendar (:issue:`54382`)
 - :class:`Rolling` and :class:`Expanding` now support ``nunique`` (:issue:`26958`)
 - :class:`Rolling` and :class:`Expanding` now support aggregations ``first`` and ``last`` (:issue:`33155`)
+- :class:`StringDtype` now supports addition to Series/DataFrame with floats, ints, and strings (:issue:`61581`)
 - :func:`read_parquet` accepts ``to_pandas_kwargs`` which are forwarded to :meth:`pyarrow.Table.to_pandas` which enables passing additional keywords to customize the conversion to pandas, such as ``maps_as_pydicts`` to read the Parquet map data type as python dictionaries (:issue:`56842`)
 - :meth:`.DataFrameGroupBy.transform`, :meth:`.SeriesGroupBy.transform`, :meth:`.DataFrameGroupBy.agg`, :meth:`.SeriesGroupBy.agg`, :meth:`.SeriesGroupBy.apply`, :meth:`.DataFrameGroupBy.apply` now support ``kurt`` (:issue:`40139`)
 - :meth:`DataFrame.apply` supports using third-party execution engines like the Bodo.ai JIT compiler (:issue:`60668`)
@@ -227,7 +228,6 @@ Other enhancements
 - Support reading Stata 102-format (Stata 1) dta files (:issue:`58978`)
 - Support reading Stata 110-format (Stata 7) dta files (:issue:`47176`)
 - Switched wheel upload to **PyPI Trusted Publishing** (OIDC) for release-tag pushes in ``wheels.yml``. (:issue:`61718`)
--
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_300.notable_bug_fixes:

pandas/core/arrays/arrow/array.py

@@ -890,7 +890,17 @@ def _op_method_error_message(self, other, op) -> str:
     def _evaluate_op_method(self, other, op, arrow_funcs) -> Self:
         pa_type = self._pa_array.type
         other_original = other
-        other = self._box_pa(other)
+        try:
+            other = self._box_pa(other)
+        except pa.lib.ArrowTypeError:
+            # was expecting time dtype but received non-temporal dtype (time offset)
+            from pandas.core.tools.timedeltas import to_timedelta
+
+            other = self._box_pa(to_timedelta(other))
+        except ValueError as err:
+            raise TypeError(
+                "Incompatible type when converting to PyArrow dtype for operation."
+            ) from err
 
         if (
             pa.types.is_string(pa_type)
@@ -903,19 +913,31 @@ def _evaluate_op_method(self, other, op, arrow_funcs) -> Self:
                     pa.types.is_integer(other.type)
                     or pa.types.is_floating(other.type)
                     or pa.types.is_null(other.type)
+                    or pa.types.is_string(other.type)
+                    or pa.types.is_large_string(other.type)
+                    or pa.types.is_binary(other.type)
                 ):
                     other = other.cast(pa_type)
-                sep = pa.scalar("", type=pa_type)
-                try:
-                    if op is operator.add:
-                        result = pc.binary_join_element_wise(self._pa_array, other, sep)
-                    elif op is roperator.radd:
-                        result = pc.binary_join_element_wise(other, self._pa_array, sep)
-                except pa.ArrowNotImplementedError as err:
+                    sep = pa.scalar("", type=pa_type)
+                    try:
+                        if op is operator.add:
+                            result = pc.binary_join_element_wise(
+                                self._pa_array, other, sep
+                            )
+                        elif op is roperator.radd:
+                            result = pc.binary_join_element_wise(
+                                other, self._pa_array, sep
+                            )
+                    except pa.ArrowNotImplementedError as err:
+                        raise TypeError(
+                            self._op_method_error_message(other_original, op)
+                        ) from err
+                    return self._from_pyarrow_array(result)
+                else:
                     raise TypeError(
-                        self._op_method_error_message(other_original, op)
-                    ) from err
-                return self._from_pyarrow_array(result)
+                        "Can only add string arrays to dtypes "
+                        "null, int, float, str, and binary."
+                    )
             elif op in [operator.mul, roperator.rmul]:
                 binary = self._pa_array
                 integral = other

pandas/core/arrays/string_.py

@@ -1111,23 +1111,27 @@ def _cmp_method(self, other, op):
         if op.__name__ in ops.ARITHMETIC_BINOPS:
             result = np.empty_like(self._ndarray, dtype="object")
             result[mask] = self.dtype.na_value
-            try:
-                result[valid] = op(self._ndarray[valid], other)
-                if isinstance(other, Path):
-                    # GH#61940
-                    return result
-            except TypeError:
-                if is_array_like(other):
-                    if is_float_dtype(other.dtype):
-                        # Shorten whole numbers to be ints to match pyarrow behavior
+            if op.__name__ in ["add", "radd"]:
+                if isinstance(other, str) or is_string_dtype(other):
+                    pass
+                elif is_float_dtype(other) or is_integer_dtype(other):
+                    if is_float_dtype(other):
+                        # Shorten whole number floats to match pyarrow behavior
                         other = [
                             str(int(x)) if x.is_integer() else str(x) for x in other
                         ]
                     else:
                         other = other.astype(str)
-                    result[valid] = op(self._ndarray[valid], other)
                 else:
-                    raise
+                    raise TypeError(
+                        f"Only supports op({op.__name__}) between StringArray and "
+                        "dtypes int, float, and str."
+                    )
+
+            result[valid] = op(self._ndarray[valid], other)
+            if isinstance(other, Path):
+                # GH#61940
+                return result
 
             return self._from_backing_data(result)
         else:

pandas/tests/arrays/string_/test_string.py

@@ -4,6 +4,7 @@
 """
 
 import operator
+from re import escape
 
 import numpy as np
 import pytest
@@ -249,6 +250,32 @@ def test_mul(dtype):
     tm.assert_extension_array_equal(result, expected)
 
 
+def test_add_series(dtype):
+    arr = pd.array(["a", "b", "c", "d"], dtype=dtype)
+    df = pd.Series(["t", "y", "v", "w"], dtype=object)
+
+    result = arr + df
+    expected = pd.Series(["at", "by", "cv", "dw"]).astype(dtype)
+    tm.assert_series_equal(result, expected)
+
+    result = df + arr
+    expected = pd.Series(["ta", "yb", "vc", "wd"]).astype(dtype)
+    tm.assert_series_equal(result, expected)
+
+
+def test_add_series_float(dtype):
+    arr = pd.array(["a", "b", "c", "d"], dtype=dtype)
+    df = pd.Series([1, 2.0, 3.5, 4])
+
+    result = arr + df
+    expected = pd.Series(["a1", "b2", "c3.5", "d4"]).astype(dtype)
+    tm.assert_series_equal(result, expected)
+
+    result = df + arr
+    expected = pd.Series(["1a", "2b", "3.5c", "4d"]).astype(dtype)
+    tm.assert_series_equal(result, expected)
+
+
 def test_add_strings(dtype):
     arr = pd.array(["a", "b", "c", "d"], dtype=dtype)
     df = pd.DataFrame([["t", "y", "v", "w"]], dtype=object)
@@ -278,19 +305,58 @@ def test_add_frame(dtype):
     tm.assert_frame_equal(result, expected, check_dtype=False)
 
 
-def test_add_frame_mixed_type(dtype):
-    arr = pd.array(["a", "bc", 3, np.nan], dtype=dtype)
-    df = pd.DataFrame([[1, 2, 3.3, 4]])
+def test_add_frame_int(dtype):
+    arr = pd.array(["a", "b", "c", np.nan], dtype=dtype)
+    df = pd.DataFrame([[1, np.nan, 3, np.nan]])
 
     result = arr + df
-    expected = pd.DataFrame([["a1", "bc2", "33.3", np.nan]]).astype(dtype)
+    expected = pd.DataFrame([["a1", np.nan, "c3", np.nan]]).astype(dtype)
     tm.assert_frame_equal(result, expected, check_dtype=False)
 
     result = df + arr
-    expected = pd.DataFrame([["1a", "2bc", "3.33", np.nan]]).astype(dtype)
+    expected = pd.DataFrame([["1a", np.nan, "3c", np.nan]]).astype(dtype)
     tm.assert_frame_equal(result, expected, check_dtype=False)
 
 
+@pytest.mark.parametrize(
+    "invalid",
+    [
+        pd.Timedelta(hours=31),
+        pd.Timestamp("2021-01-01"),
+        np.datetime64("NaT", "ns"),
+        pd.NaT,
+        True,
+        pd.Period("2025-09"),
+        pd.Categorical(["test"]),
+        pd.offsets.Minute(3),
+        pd.Interval(1, 2, closed="right"),
+    ],
+)
+def test_add_frame_invalid(dtype, invalid):
+    arr = pd.array(["a", np.nan], dtype=dtype)
+    df = pd.DataFrame([[invalid, invalid]])
+
+    if dtype.storage == "pyarrow":
+        if invalid == pd.Categorical(["test"]):
+            msg = (
+                "Incompatible type found when converting "
+                "to PyArrow dtype for operation."
+            )
+        else:
+            msg = (
+                "Can only add string arrays to dtypes "
+                "null, int, float, str, and binary."
+            )
+        with pytest.raises(TypeError, match=msg):
+            arr + df
+    else:
+        msg = escape(
+            "Only supports op(add) between StringArray and dtypes int, float, and str."
+        )
+        with pytest.raises(TypeError, match=msg):
+            arr + df
+
+
 def test_comparison_methods_scalar(comparison_op, dtype):
     op_name = f"__{comparison_op.__name__}__"
     a = pd.array(["a", None, "c"], dtype=dtype)

pandas/tests/extension/base/ops.py

@@ -152,6 +152,9 @@ def test_arith_series_with_array(self, data, all_arithmetic_operators):
         # ndarray & other series
         op_name = all_arithmetic_operators
         ser = pd.Series(data)
+        if op_name in ["__add__", "__radd__"]:
+            pytest.mark.xfail(reason="Failed: DID NOT RAISE <class 'TypeError'>")
+
         self.check_opname(ser, op_name, pd.Series([ser.iloc[0]] * len(ser)))
 
     def test_divmod(self, data):