diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 8a08ee68fa577..cca8ed9789518 100644 --- a/pandas/core/generic.py +++ b/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 diff --git a/pandas/tests/base/test_misc.py b/pandas/tests/base/test_misc.py index f3be4749fb3aa..bcab890b6047e 100644 --- a/pandas/tests/base/test_misc.py +++ b/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 diff --git a/pandas/tests/test_downstream.py b/pandas/tests/test_downstream.py index f1db4a2fc22cb..5ba5b6b1116dc 100644 --- a/pandas/tests/test_downstream.py +++ b/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