Make sure that we don't rechunk the entire variable to one chunk

by reducing over all dimensions. Dask raises an error when axis=None
but not when axis=range(a.ndim).

Author: dcherian

xarray/core/nanops.py

@@ -143,6 +143,12 @@ def nanmean(a, axis=None, dtype=None, out=None):
 
 
 def nanmedian(a, axis=None, out=None):
+    # The dask algorithm works by rechunking to one chunk along axis
+    # Make sure we trigger the dask error when passing all dimensions
+    # so that we don't rechunk the entire array to one chunk and
+    # possibly blow memory
+    if axis is not None and len(axis) == a.ndim:
+        axis = None
     return _dask_or_eager_func(
         "nanmedian", dask_module=dask_array_compat, eager_module=nputils
     )(a, axis=axis)

xarray/tests/test_dask.py

@@ -218,6 +218,8 @@ def test_reduce(self):
         self.assertLazyAndAllClose((u < 1).all("x"), (v < 1).all("x"))
         with raises_regex(NotImplementedError, "only works along an axis"):
             v.median()
+        with raises_regex(NotImplementedError, "only works along an axis"):
+            v.median(v.dims)
         with raise_if_dask_computes():
             v.reduce(duck_array_ops.mean)