From aff2b88efa53c02afda2a90c3594e89a48777b51 Mon Sep 17 00:00:00 2001 From: Kai Muehlbauer Date: Thu, 12 Nov 2020 08:01:46 +0100 Subject: [PATCH 1/3] Fix: make copy of dask_gufunc_kwargs before changing content (in apply_ufunc), add test --- xarray/core/computation.py | 2 ++ xarray/tests/test_computation.py | 26 ++++++++++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index 7b62c2c705f..9251edf1cb8 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -646,6 +646,8 @@ def apply_variable_ufunc( if dask_gufunc_kwargs is None: dask_gufunc_kwargs = {} + else: + dask_gufunc_kwargs = dask_gufunc_kwargs.copy() allow_rechunk = dask_gufunc_kwargs.get("allow_rechunk", None) if allow_rechunk is None: diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 63bedfaf280..1922977fdeb 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -796,6 +796,32 @@ def func(x): assert_identical(expected, actual) +@requires_dask +def test_apply_dask_new_output_sizes(): + ds = xr.Dataset({"foo": (["lon", "lat"], np.arange(10 * 10).reshape((10, 10)))}) + ds["bar"] = ds["foo"] + newdims = {"lon_new": 3, "lat_new": 6} + + def extract(obj): + def func(da): + return da[1:4, 1:7] + + return apply_ufunc( + func, + obj, + dask="parallelized", + input_core_dims=[["lon", "lat"]], + output_core_dims=[["lon_new", "lat_new"]], + dask_gufunc_kwargs=dict(output_sizes=newdims), + ) + + expected = extract(ds) + + actual = extract(ds.chunk()) + assert actual.dims == {"lon_new": 3, "lat_new": 6} + assert_identical(expected.chunk(), actual) + + def pandas_median(x): return pd.Series(x).median() From d99ce9b915048783cf55c9eeca43a6084e204359 Mon Sep 17 00:00:00 2001 From: Kai Muehlbauer Date: Thu, 12 Nov 2020 08:51:32 +0100 Subject: [PATCH 2/3] DOC: add entry to whats-new.rst --- doc/whats-new.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index dd3120f73fc..c650bf35374 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -75,6 +75,7 @@ Bug fixes ``DataArray`` objects, previously only the global attributes were retained (:issue:`4497`, :pull:`4510`). By `Mathias Hauser `_. - Improve performance where reading small slices from huge dimensions was slower than necessary (:pull:`4560`). By `Dion Häfner `_. +- Fix bug where ``dask_gufunc_kwargs`` where silently changed in :py:func:`apply_ufunc` (:pull:`4576`). By `Kai Mühlbauer `_. Documentation ~~~~~~~~~~~~~ From 3ad4665a11a8c6dbf4cd66e8a7780d6bc9a443eb Mon Sep 17 00:00:00 2001 From: Kai Muehlbauer Date: Thu, 12 Nov 2020 09:52:27 +0100 Subject: [PATCH 3/3] DOC: fix type in whats-new.rst [skip-ci] --- doc/whats-new.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index c650bf35374..9460fc08478 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -75,7 +75,7 @@ Bug fixes ``DataArray`` objects, previously only the global attributes were retained (:issue:`4497`, :pull:`4510`). By `Mathias Hauser `_. - Improve performance where reading small slices from huge dimensions was slower than necessary (:pull:`4560`). By `Dion Häfner `_. -- Fix bug where ``dask_gufunc_kwargs`` where silently changed in :py:func:`apply_ufunc` (:pull:`4576`). By `Kai Mühlbauer `_. +- Fix bug where ``dask_gufunc_kwargs`` was silently changed in :py:func:`apply_ufunc` (:pull:`4576`). By `Kai Mühlbauer `_. Documentation ~~~~~~~~~~~~~