ENH/TST: Add TestBaseArithmeticOps tests for ArrowExtensionArray #47601

pandas/core/arrays/arrow/array.py

@@ -57,6 +57,34 @@
         "ge": pc.greater_equal,
     }
 
+    ARROW_LOGICAL_FUNCS = {
+        "and": pc.and_kleene,
+        "rand": lambda x, y: pc.and_kleene(y, x),
+        "or": pc.or_kleene,
+        "ror": lambda x, y: pc.or_kleene(y, x),
+        "xor": pc.xor,
+        "rxor": lambda x, y: pc.xor(y, x),
+    }
+
+    ARROW_ARITHMETIC_FUNCS = {
+        "add": pc.add_checked,
+        "radd": lambda x, y: pc.add(y, x),
+        "sub": pc.subtract_checked,
+        "rsub": lambda x, y: pc.subtract_checked(y, x),
+        "mul": pc.multiply_checked,
+        "rmul": lambda x, y: pc.multiply_checked(y, x),
+        "truediv": NotImplemented,  # pc.divide_checked,
+        "rtruediv": NotImplemented,  # lambda x, y: pc.divide_checked(y, x),
+        "floordiv": NotImplemented,
+        "rfloordiv": NotImplemented,
+        "mod": NotImplemented,
+        "rmod": NotImplemented,
+        "divmod": NotImplemented,
+        "rdivmod": NotImplemented,
+        "pow": pc.power_checked,
+        "rpow": lambda x, y: pc.power_checked(y, x),
+    }
+
 if TYPE_CHECKING:
     from pandas import Series
 
@@ -74,6 +102,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 +292,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, other)
+        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/tests/extension/test_arrow.py

@@ -179,6 +179,16 @@ def data_missing_for_sorting(data_for_grouping):
     )
 
 
+@pytest.fixture
+def data_for_twos(data):
+    """Length-100 array in which all the elements are two."""
+    pa_dtype = data.dtype.pyarrow_dtype
+    if pa.types.is_integer(pa_dtype) or pa.types.is_floating(pa_dtype):
+        return pd.array([2] * 100, dtype=data.dtype)
+    # tests will be xfailed where 2 is not a valid scalar for pa_dtype
+    return data
+
+
 @pytest.fixture
 def na_value():
     """The scalar missing value for this type. Default 'None'"""
@@ -1491,6 +1501,141 @@ def test_where_series(self, data, na_value, as_frame, request, using_array_manag
         super().test_where_series(data, na_value, as_frame)
 
 
+class TestBaseArithmeticOps(base.BaseArithmeticOpsTests):
+
+    divmod_exc = NotImplementedError
+
+    def test_arith_series_with_scalar(self, data, all_arithmetic_operators, request):
+        pa_dtype = data.dtype.pyarrow_dtype
+        if all_arithmetic_operators in {
+            "__truediv__",
+            "__rtruediv__",
+            "__floordiv__",
+            "__rfloordiv__",
+            "__mod__",
+            "__rmod__",
+        }:
+            self.series_scalar_exc = NotImplementedError
+        elif not (pa.types.is_floating(pa_dtype) or pa.types.is_integer(pa_dtype)):
+            self.series_scalar_exc = pa.ArrowNotImplementedError
+        else:
+            self.series_scalar_exc = None
+        if all_arithmetic_operators == "__rpow__" and (
+            pa.types.is_floating(pa_dtype) or pa.types.is_integer(pa_dtype)
+        ):
+            request.node.add_marker(
+                pytest.mark.xfail(
+                    reason=(
+                        f"GH 29997: 1**pandas.NA == 1 while 1**pyarrow.NA == NULL "
+                        f"for {pa_dtype}"
+                    )
+                )
+            )
+        super().test_arith_series_with_scalar(data, all_arithmetic_operators)
+
+    def test_arith_frame_with_scalar(self, data, all_arithmetic_operators, request):
+        pa_dtype = data.dtype.pyarrow_dtype
+        if all_arithmetic_operators in {
+            "__truediv__",
+            "__rtruediv__",
+            "__floordiv__",
+            "__rfloordiv__",
+            "__mod__",
+            "__rmod__",
+        }:
+            self.frame_scalar_exc = NotImplementedError
+        elif not (pa.types.is_floating(pa_dtype) or pa.types.is_integer(pa_dtype)):
+            self.frame_scalar_exc = pa.ArrowNotImplementedError
+        else:
+            self.frame_scalar_exc = None
+        if all_arithmetic_operators == "__rpow__" and (
+            pa.types.is_floating(pa_dtype) or pa.types.is_integer(pa_dtype)
+        ):
+            request.node.add_marker(
+                pytest.mark.xfail(
+                    reason=(
+                        f"GH 29997: 1**pandas.NA == 1 while 1**pyarrow.NA == NULL "
+                        f"for {pa_dtype}"
+                    )
+                )
+            )
+        super().test_arith_frame_with_scalar(data, all_arithmetic_operators)
+
+    def test_arith_series_with_array(
+        self, data, all_arithmetic_operators, request, monkeypatch
+    ):
+        pa_dtype = data.dtype.pyarrow_dtype
+        if all_arithmetic_operators in {
+            "__truediv__",
+            "__rtruediv__",
+            "__floordiv__",
+            "__rfloordiv__",
+            "__mod__",
+            "__rmod__",
+        }:
+            self.series_array_exc = NotImplementedError
+        elif not (pa.types.is_floating(pa_dtype) or pa.types.is_integer(pa_dtype)):
+            self.series_array_exc = pa.ArrowNotImplementedError
+        else:
+            self.series_array_exc = None
+        if all_arithmetic_operators == "__rpow__" and (
+            pa.types.is_floating(pa_dtype) or pa.types.is_integer(pa_dtype)
+        ):
+            request.node.add_marker(
+                pytest.mark.xfail(
+                    reason=(
+                        f"GH 29997: 1**pandas.NA == 1 while 1**pyarrow.NA == NULL "
+                        f"for {pa_dtype}"
+                    )
+                )
+            )
+        elif all_arithmetic_operators in (
+            "__sub__",
+            "__rsub__",
+        ) and pa.types.is_unsigned_integer(pa_dtype):
+            request.node.add_marker(
+                pytest.mark.xfail(
+                    raises=pa.ArrowInvalid,
+                    reason=(
+                        f"Implemented pyarrow.compute.subtract_checked "
+                        f"which raises on overflow for {pa_dtype}"
+                    ),
+                )
+            )
+        op_name = all_arithmetic_operators
+        ser = pd.Series(data)
+        # pd.Series([ser.iloc[0]] * len(ser)) may not return ArrowExtensionArray
+        # since ser.iloc[0] is a python scalar
+        other = pd.Series(pd.array([ser.iloc[0]] * len(ser), dtype=data.dtype))
+        if pa.types.is_floating(pa_dtype) or pa.types.is_integer(pa_dtype):
+            # BaseOpsUtil._combine can upcast expected dtype
+            # (because it generates expected on python scalars)
+            # while ArrowExtensionArray maintains original type
+            super_combine = TestBaseArithmeticOps._combine
+
+            def _patch_combine(self, obj, other, op):
+                expected = super_combine(self, obj, other, op)
+                if isinstance(expected, pd.Series):
+                    pa_array = pa.array(expected._values).cast(obj.dtype.pyarrow_dtype)
+                    pd_array = type(expected._values)(pa_array)
+                    expected = pd.Series(pd_array)
+                return expected
+
+            monkeypatch.setattr(TestBaseArithmeticOps, "_combine", _patch_combine)
+        self.check_opname(ser, op_name, other, exc=self.series_array_exc)
+
+    def test_add_series_with_extension_array(self, data, request):
+        pa_dtype = data.dtype.pyarrow_dtype
+        if not (pa.types.is_integer(pa_dtype) or pa.types.is_floating(pa_dtype)):
+            request.node.add_marker(
+                pytest.mark.xfail(
+                    raises=pa.ArrowNotImplementedError,
+                    reason=f"add_checked not implemented for {pa_dtype}",
+                )
+            )
+        super().test_add_series_with_extension_array(data)
+
+
 def test_arrowdtype_construct_from_string_type_with_unsupported_parameters():
     with pytest.raises(NotImplementedError, match="Passing pyarrow type"):
         ArrowDtype.construct_from_string("timestamp[s, tz=UTC][pyarrow]")