ENH: support 'out' keyword in Categorical.__array_ufunc__ (#45381)

Author: jbrockmendel

pandas/core/arrays/categorical.py

@@ -1517,6 +1517,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_numpy_ufuncs_out
+            return arraylike.dispatch_ufunc_with_out(
+                self, ufunc, method, *inputs, **kwargs
+            )
+
         if method == "reduce":
             # e.g. TestCategoricalAnalytics::test_min_max_ordered
             result = arraylike.dispatch_reduction_ufunc(

pandas/core/arrays/numpy_.py

@@ -141,6 +141,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_unary
+            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/tests/arrays/test_numpy.py

@@ -237,6 +237,11 @@ def test_ufunc_unary(ufunc):
     expected = PandasArray(ufunc(arr._ndarray))
     tm.assert_extension_array_equal(result, expected)
 
+    # same thing but with the 'out' keyword
+    out = PandasArray(np.array([-9.0, -9.0, -9.0]))
+    ufunc(arr, out=out)
+    tm.assert_extension_array_equal(out, expected)
+
 
 def test_ufunc():
     arr = PandasArray(np.array([-1.0, 0.0, 1.0]))

pandas/tests/extension/decimal/array.py

@@ -119,6 +119,18 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
         ):
             return NotImplemented
 
+        result = arraylike.maybe_dispatch_ufunc_to_dunder_op(
+            self, ufunc, method, *inputs, **kwargs
+        )
+        if result is not NotImplemented:
+            # e.g. test_array_ufunc_series_scalar_other
+            return result
+
+        if "out" in kwargs:
+            return arraylike.dispatch_ufunc_with_out(
+                self, ufunc, method, *inputs, **kwargs
+            )
+
         inputs = tuple(x._data if isinstance(x, DecimalArray) else x for x in inputs)
         result = getattr(ufunc, method)(*inputs, **kwargs)
 

pandas/tests/indexes/test_numpy_compat.py

@@ -18,6 +18,20 @@
 from pandas.core.indexes.datetimelike import DatetimeIndexOpsMixin
 
 
+def test_numpy_ufuncs_out(index):
+    result = index == index
+
+    out = np.empty(index.shape, dtype=bool)
+    np.equal(index, index, out=out)
+    tm.assert_numpy_array_equal(out, result)
+
+    if not index._is_multi:
+        # same thing on the ExtensionArray
+        out = np.empty(index.shape, dtype=bool)
+        np.equal(index.array, index.array, out=out)
+        tm.assert_numpy_array_equal(out, result)
+
+
 @pytest.mark.parametrize(
     "func",
     [
@@ -91,6 +105,10 @@ def test_numpy_ufuncs_other(index, func, request):
             # numpy 1.18 changed isinf and isnan to not raise on dt64/td64
             result = func(index)
             assert isinstance(result, np.ndarray)
+
+            out = np.empty(index.shape, dtype=bool)
+            func(index, out=out)
+            tm.assert_numpy_array_equal(out, result)
         else:
             with tm.external_error_raised(TypeError):
                 func(index)
@@ -109,7 +127,15 @@ def test_numpy_ufuncs_other(index, func, request):
             assert isinstance(result, BooleanArray)
         else:
             assert isinstance(result, np.ndarray)
-        assert not isinstance(result, Index)
+
+        out = np.empty(index.shape, dtype=bool)
+        func(index, out=out)
+
+        if not isinstance(index.dtype, np.dtype):
+            tm.assert_numpy_array_equal(out, result._data)
+        else:
+            tm.assert_numpy_array_equal(out, result)
+
     else:
         if len(index) == 0:
             pass