BUG: IntegerArray/FloatingArray ufunc with 'out' kwd (#45122)

Author: jbrockmendel

doc/source/whatsnew/v1.4.0.rst

@@ -980,6 +980,7 @@ ExtensionArray
 - Bug in :func:`array` failing to preserve :class:`PandasArray` (:issue:`43887`)
 - NumPy ufuncs ``np.abs``, ``np.positive``, ``np.negative`` now correctly preserve dtype when called on ExtensionArrays that implement ``__abs__, __pos__, __neg__``, respectively. In particular this is fixed for :class:`TimedeltaArray` (:issue:`43899`, :issue:`23316`)
 - NumPy ufuncs ``np.minimum.reduce`` ``np.maximum.reduce``, ``np.add.reduce``, and ``np.prod.reduce`` now work correctly instead of raising ``NotImplementedError`` on :class:`Series` with ``IntegerDtype`` or ``FloatDtype`` (:issue:`43923`, :issue:`44793`)
+- NumPy ufuncs with ``out`` keyword are now supported by arrays with ``IntegerDtype`` and ``FloatingDtype`` (:issue:`??`)
 - Avoid raising ``PerformanceWarning`` about fragmented DataFrame when using many columns with an extension dtype (:issue:`44098`)
 - Bug in :class:`IntegerArray` and :class:`FloatingArray` construction incorrectly coercing mismatched NA values (e.g. ``np.timedelta64("NaT")``) to numeric NA (:issue:`44514`)
 - Bug in :meth:`BooleanArray.__eq__` and :meth:`BooleanArray.__ne__` raising ``TypeError`` on comparison with an incompatible type (like a string). This caused :meth:`DataFrame.replace` to sometimes raise a ``TypeError`` if a nullable boolean column was included (:issue:`44499`)

pandas/core/arrays/masked.py

@@ -433,6 +433,12 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
         if result is not NotImplemented:
             return result
 
+        if "out" in kwargs:
+            # e.g. test_ufunc_with_out
+            return arraylike.dispatch_ufunc_with_out(
+                self, ufunc, method, *inputs, **kwargs
+            )
+
         if method == "reduce":
             result = arraylike.dispatch_reduction_ufunc(
                 self, ufunc, method, *inputs, **kwargs

pandas/core/missing.py

@@ -92,14 +92,15 @@ def mask_missing(arr: ArrayLike, values_to_mask) -> npt.NDArray[np.bool_]:
             # GH#29553 prevent numpy deprecation warnings
             pass
         else:
-            mask |= arr == x
+            new_mask = arr == x
+            if not isinstance(new_mask, np.ndarray):
+                # usually BooleanArray
+                new_mask = new_mask.to_numpy(dtype=bool, na_value=False)
+            mask |= new_mask
 
     if na_mask.any():
         mask |= isna(arr)
 
-    if not isinstance(mask, np.ndarray):
-        # e.g. if arr is IntegerArray, then mask is BooleanArray
-        mask = mask.to_numpy(dtype=bool, na_value=False)
     return mask
 
 

pandas/tests/arrays/masked_shared.py

@@ -1,6 +1,8 @@
 """
 Tests shared by MaskedArray subclasses.
 """
+import numpy as np
+import pytest
 
 import pandas as pd
 import pandas._testing as tm
@@ -102,3 +104,35 @@ def test_compare_to_string(self, dtype):
         expected = pd.Series([False, pd.NA], dtype="boolean")
 
         self.assert_series_equal(result, expected)
+
+    def test_ufunc_with_out(self, dtype):
+        arr = pd.array([1, 2, 3], dtype=dtype)
+        arr2 = pd.array([1, 2, pd.NA], dtype=dtype)
+
+        mask = arr == arr
+        mask2 = arr2 == arr2
+
+        result = np.zeros(3, dtype=bool)
+        result |= mask
+        # If MaskedArray.__array_ufunc__ handled "out" appropriately,
+        #  `result` should still be an ndarray.
+        assert isinstance(result, np.ndarray)
+        assert result.all()
+
+        # result |= mask worked because mask could be cast lossslessly to
+        #  boolean ndarray. mask2 can't, so this raises
+        result = np.zeros(3, dtype=bool)
+        msg = "Specify an appropriate 'na_value' for this dtype"
+        with pytest.raises(ValueError, match=msg):
+            result |= mask2
+
+        # addition
+        res = np.add(arr, arr2)
+        expected = pd.array([2, 4, pd.NA], dtype=dtype)
+        tm.assert_extension_array_equal(res, expected)
+
+        # when passing out=arr, we will modify 'arr' inplace.
+        res = np.add(arr, arr2, out=arr)
+        assert res is arr
+        tm.assert_extension_array_equal(res, expected)
+        tm.assert_extension_array_equal(arr, expected)