Backport PR pandas-dev#45451: Add back __array_wrap__ for dask compatibility (pandas-dev#45491)

Co-authored-by: Joris Van den Bossche <[email protected]>

Author: meeseeksmachine

pandas/core/generic.py

@@ -2071,6 +2071,41 @@ def empty(self) -> bool_t:
     def __array__(self, dtype: npt.DTypeLike | None = None) -> np.ndarray:
         return np.asarray(self._values, dtype=dtype)
 
+    def __array_wrap__(
+        self,
+        result: np.ndarray,
+        context: tuple[Callable, tuple[Any, ...], int] | None = None,
+    ):
+        """
+        Gets called after a ufunc and other functions.
+
+        Parameters
+        ----------
+        result: np.ndarray
+            The result of the ufunc or other function called on the NumPy array
+            returned by __array__
+        context: tuple of (func, tuple, int)
+            This parameter is returned by ufuncs as a 3-element tuple: (name of the
+            ufunc, arguments of the ufunc, domain of the ufunc), but is not set by
+            other numpy functions.q
+
+        Notes
+        -----
+        Series implements __array_ufunc_ so this not called for ufunc on Series.
+        """
+        # Note: at time of dask 2022.01.0, this is still used by dask
+        res = lib.item_from_zerodim(result)
+        if is_scalar(res):
+            # e.g. we get here with np.ptp(series)
+            # ptp also requires the item_from_zerodim
+            return res
+        d = self._construct_axes_dict(self._AXIS_ORDERS, copy=False)
+        # error: Argument 1 to "NDFrame" has incompatible type "ndarray";
+        # expected "BlockManager"
+        return self._constructor(res, **d).__finalize__(  # type: ignore[arg-type]
+            self, method="__array_wrap__"
+        )
+
     @final
     def __array_ufunc__(
         self, ufunc: np.ufunc, method: str, *inputs: Any, **kwargs: Any

pandas/tests/base/test_misc.py

@@ -84,6 +84,16 @@ def test_ndarray_compat_properties(index_or_series_obj):
     assert Series([1]).item() == 1
 
 
+def test_array_wrap_compat():
+    # Note: at time of dask 2022.01.0, this is still used by eg dask
+    # (https://github.com/dask/dask/issues/8580).
+    # This test is a small dummy ensuring coverage
+    orig = Series([1, 2, 3], dtype="int64", index=["a", "b", "c"])
+    result = orig.__array_wrap__(np.array([2, 4, 6], dtype="int64"))
+    expected = orig * 2
+    tm.assert_series_equal(result, expected)
+
+
 @pytest.mark.skipif(PYPY, reason="not relevant for PyPy")
 def test_memory_usage(index_or_series_obj):
     obj = index_or_series_obj

pandas/tests/test_downstream.py

@@ -50,6 +50,30 @@ def test_dask(df):
         pd.set_option("compute.use_numexpr", olduse)
 
 
+@pytest.mark.filterwarnings("ignore:.*64Index is deprecated:FutureWarning")
+def test_dask_ufunc():
+    # At the time of dask 2022.01.0, dask is still directly using __array_wrap__
+    # for some ufuncs (https://github.com/dask/dask/issues/8580).
+
+    # dask sets "compute.use_numexpr" to False, so catch the current value
+    # and ensure to reset it afterwards to avoid impacting other tests
+    olduse = pd.get_option("compute.use_numexpr")
+
+    try:
+        dask = import_module("dask")  # noqa:F841
+        import dask.array as da
+        import dask.dataframe as dd
+
+        s = pd.Series([1.5, 2.3, 3.7, 4.0])
+        ds = dd.from_pandas(s, npartitions=2)
+
+        result = da.fix(ds).compute()
+        expected = np.fix(s)
+        tm.assert_series_equal(result, expected)
+    finally:
+        pd.set_option("compute.use_numexpr", olduse)
+
+
 def test_xarray(df):
 
     xarray = import_module("xarray")  # noqa:F841