From 9569bddafc9c977aead295f95f9c169318d6e662 Mon Sep 17 00:00:00 2001 From: Patrick Hoefler <61934744+phofl@users.noreply.github.com> Date: Tue, 5 Nov 2024 21:59:11 +0100 Subject: [PATCH 1/3] Dispatch to Dask is nanquantile is available --- xarray/core/variable.py | 3 ++- xarray/tests/test_variable.py | 14 +++++++++++--- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 16e292796a7..b2035b8cfbd 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -46,6 +46,7 @@ ) from xarray.namedarray.core import NamedArray, _raise_if_any_duplicate_dimensions from xarray.namedarray.pycompat import integer_types, is_0d_dask_array, to_duck_array +from xarray.namedarray.utils import module_available from xarray.util.deprecation_helpers import deprecate_dims NON_NUMPY_SUPPORTED_ARRAY_TYPES = ( @@ -1948,7 +1949,7 @@ def _wrapper(npa, **kwargs): output_core_dims=[["quantile"]], output_dtypes=[np.float64], dask_gufunc_kwargs=dict(output_sizes={"quantile": len(q)}), - dask="parallelized", + dask="allowed" if module_available("dask", "2024.11.0") else "parallelized", kwargs=kwargs, ) diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 1d430b6b27e..00efa507eac 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -30,6 +30,7 @@ from xarray.core.utils import NDArrayMixin from xarray.core.variable import as_compatible_data, as_variable from xarray.namedarray.pycompat import array_type +from xarray.namedarray.utils import module_available from xarray.tests import ( assert_allclose, assert_array_equal, @@ -1871,9 +1872,16 @@ def test_quantile_interpolation_deprecation(self, method) -> None: def test_quantile_chunked_dim_error(self): v = Variable(["x", "y"], self.d).chunk({"x": 2}) - # this checks for ValueError in dask.array.apply_gufunc - with pytest.raises(ValueError, match=r"consists of multiple chunks"): - v.quantile(0.5, dim="x") + if module_available("dask", "2024.10.0"): + # Dask rechunks + np.testing.assert_allclose( + v.compute().quantile(0.5, dim="x"), v.quantile(0.5, dim="x") + ) + + else: + # this checks for ValueError in dask.array.apply_gufunc + with pytest.raises(ValueError, match=r"consists of multiple chunks"): + v.quantile(0.5, dim="x") @pytest.mark.parametrize("compute_backend", ["numbagg", None], indirect=True) @pytest.mark.parametrize("q", [-0.1, 1.1, [2], [0.25, 2]]) From d4a5011749601cf2f5513c28624a37beab6e408b Mon Sep 17 00:00:00 2001 From: Patrick Hoefler <61934744+phofl@users.noreply.github.com> Date: Tue, 5 Nov 2024 22:12:30 +0100 Subject: [PATCH 2/3] Fixup --- xarray/tests/test_variable.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 00efa507eac..9946ea78d6d 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -1872,7 +1872,7 @@ def test_quantile_interpolation_deprecation(self, method) -> None: def test_quantile_chunked_dim_error(self): v = Variable(["x", "y"], self.d).chunk({"x": 2}) - if module_available("dask", "2024.10.0"): + if module_available("dask", "2024.11.0"): # Dask rechunks np.testing.assert_allclose( v.compute().quantile(0.5, dim="x"), v.quantile(0.5, dim="x") From 61b7edff63c07f498a2a00a4a8d09a963f50d9f8 Mon Sep 17 00:00:00 2001 From: Patrick Hoefler <61934744+phofl@users.noreply.github.com> Date: Tue, 5 Nov 2024 22:16:51 +0100 Subject: [PATCH 3/3] Change test --- xarray/tests/__init__.py | 1 + xarray/tests/test_variable.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index 7293a6fd931..5ed334e61dd 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -107,6 +107,7 @@ def _importorskip( has_h5netcdf, requires_h5netcdf = _importorskip("h5netcdf") has_cftime, requires_cftime = _importorskip("cftime") has_dask, requires_dask = _importorskip("dask") +has_dask_ge_2024_11_0, requires_dask_ge_2024_11_0 = _importorskip("dask", "2024.11.0") with warnings.catch_warnings(): warnings.filterwarnings( "ignore", diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 9946ea78d6d..2d5fcbe0049 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -30,13 +30,13 @@ from xarray.core.utils import NDArrayMixin from xarray.core.variable import as_compatible_data, as_variable from xarray.namedarray.pycompat import array_type -from xarray.namedarray.utils import module_available from xarray.tests import ( assert_allclose, assert_array_equal, assert_equal, assert_identical, assert_no_warnings, + has_dask_ge_2024_11_0, has_pandas_3, raise_if_dask_computes, requires_bottleneck, @@ -1872,7 +1872,7 @@ def test_quantile_interpolation_deprecation(self, method) -> None: def test_quantile_chunked_dim_error(self): v = Variable(["x", "y"], self.d).chunk({"x": 2}) - if module_available("dask", "2024.11.0"): + if has_dask_ge_2024_11_0: # Dask rechunks np.testing.assert_allclose( v.compute().quantile(0.5, dim="x"), v.quantile(0.5, dim="x")