ENH/TST: Add TestBaseArithmeticOps tests for ArrowExtensionArray #47601

pandas/compat/__init__.py

@@ -27,6 +27,7 @@
     pa_version_under5p0,
     pa_version_under6p0,
     pa_version_under7p0,
+    pa_version_under8p0,
 )
 
 if TYPE_CHECKING:
@@ -158,4 +159,5 @@ def get_lzma_file() -> type[lzma.LZMAFile]:
     "pa_version_under5p0",
     "pa_version_under6p0",
     "pa_version_under7p0",
+    "pa_version_under8p0",
 ]

pandas/core/arrays/arrow/array.py

@@ -57,6 +57,76 @@
         "ge": pc.greater_equal,
     }
 
+    ARROW_LOGICAL_FUNCS = {
+        "and": NotImplemented if pa_version_under2p0 else pc.and_kleene,
+        "rand": NotImplemented
+        if pa_version_under2p0
+        else lambda x, y: pc.and_kleene(y, x),
+        "or": NotImplemented if pa_version_under2p0 else pc.or_kleene,
+        "ror": NotImplemented
+        if pa_version_under2p0
+        else lambda x, y: pc.or_kleene(y, x),
+        "xor": NotImplemented if pa_version_under2p0 else pc.xor,
+        "rxor": NotImplemented if pa_version_under2p0 else lambda x, y: pc.xor(y, x),
+    }
+
+    def cast_for_truediv(
+        arrow_array: pa.ChunkedArray, pa_object: pa.Array | pa.Scalar
+    ) -> pa.ChunkedArray:
+        # Ensure int / int -> float mirroring Python/Numpy behavior
+        # as pc.divide_checked(int, int) -> int
+        if pa.types.is_integer(arrow_array.type) and pa.types.is_integer(
+            pa_object.type
+        ):
+            return arrow_array.cast(pa.float64())
+        return arrow_array
+
+    def floordiv_compat(
+        left: pa.ChunkedArray | pa.Array | pa.Scalar,
+        right: pa.ChunkedArray | pa.Array | pa.Scalar,
+    ) -> pa.ChunkedArray:
+        # Ensure int // int -> int mirroring Python/Numpy behavior
+        # as pc.floor(pc.divide_checked(int, int)) -> float
+        result = pc.floor(pc.divide_checked(left, right))
+        if pa.types.is_integer(left.type) and pa.types.is_integer(right.type):
+            result = result.cast(left.type)
+        return result
+
+    ARROW_ARITHMETIC_FUNCS = {
+        "add": NotImplemented if pa_version_under2p0 else pc.add_checked,
+        "radd": NotImplemented
+        if pa_version_under2p0
+        else lambda x, y: pc.add_checked(y, x),
+        "sub": NotImplemented if pa_version_under2p0 else pc.subtract_checked,
+        "rsub": NotImplemented
+        if pa_version_under2p0
+        else lambda x, y: pc.subtract_checked(y, x),
+        "mul": NotImplemented if pa_version_under2p0 else pc.multiply_checked,
+        "rmul": NotImplemented
+        if pa_version_under2p0
+        else lambda x, y: pc.multiply_checked(y, x),
+        "truediv": NotImplemented
+        if pa_version_under2p0
+        else lambda x, y: pc.divide_checked(cast_for_truediv(x, y), y),
+        "rtruediv": NotImplemented
+        if pa_version_under2p0
+        else lambda x, y: pc.divide_checked(y, cast_for_truediv(x, y)),
+        "floordiv": NotImplemented
+        if pa_version_under2p0
+        else lambda x, y: floordiv_compat(x, y),
+        "rfloordiv": NotImplemented
+        if pa_version_under2p0
+        else lambda x, y: floordiv_compat(y, x),
+        "mod": NotImplemented,
+        "rmod": NotImplemented,
+        "divmod": NotImplemented,
+        "rdivmod": NotImplemented,
+        "pow": NotImplemented if pa_version_under2p0 else pc.power_checked,
+        "rpow": NotImplemented
+        if pa_version_under2p0
+        else lambda x, y: pc.power_checked(y, x),
+    }
+
 if TYPE_CHECKING:
     from pandas import Series
 
@@ -74,6 +144,7 @@ def to_pyarrow_type(
     elif isinstance(dtype, pa.DataType):
         pa_dtype = dtype
     elif dtype:
+        # Accepts python types too
         pa_dtype = pa.from_numpy_dtype(dtype)
     else:
         pa_dtype = None
@@ -263,6 +334,28 @@ def _cmp_method(self, other, op):
             result = result.to_numpy()
         return BooleanArray._from_sequence(result)
 
+    def _evaluate_op_method(self, other, op, arrow_funcs):
+        pc_func = arrow_funcs[op.__name__]
+        if pc_func is NotImplemented:
+            raise NotImplementedError(f"{op.__name__} not implemented.")
+        if isinstance(other, ArrowExtensionArray):
+            result = pc_func(self._data, other._data)
+        elif isinstance(other, (np.ndarray, list)):
+            result = pc_func(self._data, pa.array(other, from_pandas=True))
+        elif is_scalar(other):
+            result = pc_func(self._data, pa.scalar(other))
+        else:
+            raise NotImplementedError(
+                f"{op.__name__} not implemented for {type(other)}"
+            )
+        return type(self)(result)
+
+    def _logical_method(self, other, op):
+        return self._evaluate_op_method(other, op, ARROW_LOGICAL_FUNCS)
+
+    def _arith_method(self, other, op):
+        return self._evaluate_op_method(other, op, ARROW_ARITHMETIC_FUNCS)
+
     def equals(self, other) -> bool:
         if not isinstance(other, ArrowExtensionArray):
             return False

pandas/core/strings/object_array.py

@@ -360,7 +360,7 @@ def _str_get_dummies(self, sep="|"):
         arr = Series(self).fillna("")
         try:
             arr = sep + arr + sep
-        except TypeError:
+        except (TypeError, NotImplementedError):
             arr = sep + arr.astype(str) + sep
 
         tags: set[str] = set()

pandas/tests/arrays/string_/test_string.py

@@ -101,7 +101,7 @@ def test_add(dtype, request):
             "unsupported operand type(s) for +: 'ArrowStringArray' and "
             "'ArrowStringArray'"
         )
-        mark = pytest.mark.xfail(raises=TypeError, reason=reason)
+        mark = pytest.mark.xfail(raises=NotImplementedError, reason=reason)
         request.node.add_marker(mark)
 
     a = pd.Series(["a", "b", "c", None, None], dtype=dtype)
@@ -142,7 +142,7 @@ def test_add_2d(dtype, request):
 def test_add_sequence(dtype, request):
     if dtype.storage == "pyarrow":
         reason = "unsupported operand type(s) for +: 'ArrowStringArray' and 'list'"
-        mark = pytest.mark.xfail(raises=TypeError, reason=reason)
+        mark = pytest.mark.xfail(raises=NotImplementedError, reason=reason)
         request.node.add_marker(mark)
 
     a = pd.array(["a", "b", None, None], dtype=dtype)
@@ -160,7 +160,7 @@ def test_add_sequence(dtype, request):
 def test_mul(dtype, request):
     if dtype.storage == "pyarrow":
         reason = "unsupported operand type(s) for *: 'ArrowStringArray' and 'int'"
-        mark = pytest.mark.xfail(raises=TypeError, reason=reason)
+        mark = pytest.mark.xfail(raises=NotImplementedError, reason=reason)
         request.node.add_marker(mark)
 
     a = pd.array(["a", "b", None], dtype=dtype)

pandas/tests/extension/base/ops.py

@@ -67,10 +67,10 @@ class BaseArithmeticOpsTests(BaseOpsUtil):
     * divmod_exc = TypeError
     """
 
-    series_scalar_exc: type[TypeError] | None = TypeError
-    frame_scalar_exc: type[TypeError] | None = TypeError
-    series_array_exc: type[TypeError] | None = TypeError
-    divmod_exc: type[TypeError] | None = TypeError
+    series_scalar_exc: type[Exception] | None = TypeError
+    frame_scalar_exc: type[Exception] | None = TypeError
+    series_array_exc: type[Exception] | None = TypeError
+    divmod_exc: type[Exception] | None = TypeError
 
     def test_arith_series_with_scalar(self, data, all_arithmetic_operators):
         # series & scalar