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_variable.py

@@ -1604,6 +1604,12 @@ def test_reduce_funcs(self):
 
         assert_identical(v.max(), Variable([], pd.Timestamp("2000-01-03")))
 
+    @requires_dask
+    def test_median(self):
+        v = Variable(["x", "y"], self.d).chunk(2)
+        with raises_regex(NotImplementedError, "difficult to do in parallel"):
+            v.median(["x", "y"])
+
     def test_reduce_keepdims(self):
         v = Variable(["x", "y"], self.d)