Patch ops

Author: jbrockmendel

pandas/core/arrays/arrow/array.py

@@ -50,10 +50,16 @@
     pandas_dtype,
 )
 from pandas.core.dtypes.dtypes import DatetimeTZDtype
+from pandas.core.dtypes.generic import (
+    ABCDataFrame,
+    ABCIndex,
+    ABCSeries,
+)
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import (
     algorithms as algos,
+    arraylike,
     missing,
     ops,
     roperator,
@@ -752,6 +758,39 @@ def __array__(
 
         return self.to_numpy(dtype=dtype, copy=copy)
 
+    def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
+        if any(
+            isinstance(other, (ABCSeries, ABCIndex, ABCDataFrame)) for other in inputs
+        ):
+            return NotImplemented
+
+        result = arraylike.maybe_dispatch_ufunc_to_dunder_op(
+            self, ufunc, method, *inputs, **kwargs
+        )
+        if result is not NotImplemented:
+            return result
+
+        if "out" in kwargs:
+            return arraylike.dispatch_ufunc_with_out(
+                self, ufunc, method, *inputs, **kwargs
+            )
+
+        if method == "reduce":
+            result = arraylike.dispatch_reduction_ufunc(
+                self, ufunc, method, *inputs, **kwargs
+            )
+            if result is not NotImplemented:
+                return result
+
+        if self.dtype.kind == "f":
+            # e.g. test_log_arrow_backed_missing_value
+            new_inputs = [
+                x if x is not self else x.to_numpy(na_value=np.nan) for x in inputs
+            ]
+            return getattr(ufunc, method)(*new_inputs, **kwargs)
+
+        return arraylike.default_array_ufunc(self, ufunc, method, *inputs, **kwargs)
+
     def __invert__(self) -> Self:
         # This is a bit wise op for integer types
         if pa.types.is_integer(self._pa_array.type):
@@ -923,7 +962,13 @@ def _logical_method(self, other, op) -> Self:
             return self._evaluate_op_method(other, op, ARROW_LOGICAL_FUNCS)
 
     def _arith_method(self, other, op) -> Self:
-        return self._evaluate_op_method(other, op, ARROW_ARITHMETIC_FUNCS)
+        result = self._evaluate_op_method(other, op, ARROW_ARITHMETIC_FUNCS)
+        if is_nan_na() and result.dtype.kind == "f":
+            parr = result._pa_array
+            mask = pc.is_nan(parr).to_numpy()
+            arr = pc.replace_with_mask(parr, mask, pa.scalar(None, type=parr.type))
+            result = type(self)(arr)
+        return result
 
     def equals(self, other) -> bool:
         if not isinstance(other, ArrowExtensionArray):

pandas/core/arrays/base.py

@@ -2539,14 +2539,6 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
             if result is not NotImplemented:
                 return result
 
-        # TODO: putting this here is hacky as heck
-        if self.dtype == "float64[pyarrow]":
-            # e.g. test_log_arrow_backed_missing_value
-            new_inputs = [
-                x if x is not self else x.to_numpy(na_value=np.nan) for x in inputs
-            ]
-            return getattr(ufunc, method)(*new_inputs, **kwargs)
-
         return arraylike.default_array_ufunc(self, ufunc, method, *inputs, **kwargs)
 
     def map(self, mapper, na_action: Literal["ignore"] | None = None):

pandas/tests/extension/test_arrow.py

@@ -3589,3 +3589,26 @@ def test_timestamp_dtype_matches_to_datetime():
     expected = pd.Series([ts], dtype=dtype1).convert_dtypes(dtype_backend="pyarrow")
 
     tm.assert_series_equal(result, expected)
+
+
+def test_ops_with_nan_is_na(using_nan_is_na):
+    # GH#61732
+    ser = pd.Series([-1, 0, 1], dtype="int64[pyarrow]")
+
+    result = ser - np.nan
+    if using_nan_is_na:
+        assert result.isna().all()
+    else:
+        assert not result.isna().any()
+
+    result = ser * np.nan
+    if using_nan_is_na:
+        assert result.isna().all()
+    else:
+        assert not result.isna().any()
+
+    result = ser / 0
+    if using_nan_is_na:
+        assert result.isna()[1]
+    else:
+        assert not result.isna()[1]

pandas/tests/series/test_npfuncs.py

@@ -38,7 +38,7 @@ def test_numpy_argwhere(index):
 
 
 @td.skip_if_no("pyarrow")
-def test_log_arrow_backed_missing_value():
+def test_log_arrow_backed_missing_value(using_nan_is_na):
     # GH#56285
     ser = Series([1, 2, None], dtype="float64[pyarrow]")
     result = np.log(ser)