Add nanmedian for dask arrays (pydata#3604)

* Add nanmedian for dask arrays

Close pydata#2999

* Fix tests.

* fix import

* 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).

* fix tests.

* Update whats-new.rst

Author: dcherian

doc/whats-new.rst

@@ -25,6 +25,9 @@ Breaking changes
 
 New Features
 ~~~~~~~~~~~~
+- Implement :py:func:`median` and :py:func:`nanmedian` for dask arrays. This works by rechunking
+  to a single chunk along all reduction axes. (:issue:`2999`).
+  By `Deepak Cherian <https://github.com/dcherian>`_.
 - :py:func:`xarray.concat` now preserves attributes from the first Variable.
   (:issue:`2575`, :issue:`2060`, :issue:`1614`)
   By `Deepak Cherian <https://github.com/dcherian>`_.

xarray/core/dask_array_compat.py

@@ -1,8 +1,14 @@
 from distutils.version import LooseVersion
+from typing import Iterable
 
-import dask.array as da
 import numpy as np
-from dask import __version__ as dask_version
+
+try:
+    import dask.array as da
+    from dask import __version__ as dask_version
+except ImportError:
+    dask_version = "0.0.0"
+    da = None
 
 if LooseVersion(dask_version) >= LooseVersion("2.0.0"):
     meta_from_array = da.utils.meta_from_array
@@ -89,3 +95,76 @@ def meta_from_array(x, ndim=None, dtype=None):
             meta = meta.astype(dtype)
 
         return meta
+
+
+if LooseVersion(dask_version) >= LooseVersion("2.8.1"):
+    median = da.median
+else:
+    # Copied from dask v2.8.1
+    # Used under the terms of Dask's license, see licenses/DASK_LICENSE.
+    def median(a, axis=None, keepdims=False):
+        """
+        This works by automatically chunking the reduced axes to a single chunk
+        and then calling ``numpy.median`` function across the remaining dimensions
+        """
+
+        if axis is None:
+            raise NotImplementedError(
+                "The da.median function only works along an axis.  "
+                "The full algorithm is difficult to do in parallel"
+            )
+
+        if not isinstance(axis, Iterable):
+            axis = (axis,)
+
+        axis = [ax + a.ndim if ax < 0 else ax for ax in axis]
+
+        a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})
+
+        result = a.map_blocks(
+            np.median,
+            axis=axis,
+            keepdims=keepdims,
+            drop_axis=axis if not keepdims else None,
+            chunks=[1 if ax in axis else c for ax, c in enumerate(a.chunks)]
+            if keepdims
+            else None,
+        )
+
+        return result
+
+
+if LooseVersion(dask_version) > LooseVersion("2.9.0"):
+    nanmedian = da.nanmedian
+else:
+
+    def nanmedian(a, axis=None, keepdims=False):
+        """
+        This works by automatically chunking the reduced axes to a single chunk
+        and then calling ``numpy.nanmedian`` function across the remaining dimensions
+        """
+
+        if axis is None:
+            raise NotImplementedError(
+                "The da.nanmedian function only works along an axis.  "
+                "The full algorithm is difficult to do in parallel"
+            )
+
+        if not isinstance(axis, Iterable):
+            axis = (axis,)
+
+        axis = [ax + a.ndim if ax < 0 else ax for ax in axis]
+
+        a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})
+
+        result = a.map_blocks(
+            np.nanmedian,
+            axis=axis,
+            keepdims=keepdims,
+            drop_axis=axis if not keepdims else None,
+            chunks=[1 if ax in axis else c for ax, c in enumerate(a.chunks)]
+            if keepdims
+            else None,
+        )
+
+        return result

xarray/core/duck_array_ops.py

@@ -11,7 +11,7 @@
 import numpy as np
 import pandas as pd
 
-from . import dask_array_ops, dtypes, npcompat, nputils
+from . import dask_array_ops, dask_array_compat, dtypes, npcompat, nputils
 from .nputils import nanfirst, nanlast
 from .pycompat import dask_array_type
 
@@ -284,7 +284,7 @@ def _ignore_warnings_if(condition):
         yield
 
 
-def _create_nan_agg_method(name, coerce_strings=False):
+def _create_nan_agg_method(name, dask_module=dask_array, coerce_strings=False):
     from . import nanops
 
     def f(values, axis=None, skipna=None, **kwargs):
@@ -301,7 +301,7 @@ def f(values, axis=None, skipna=None, **kwargs):
             nanname = "nan" + name
             func = getattr(nanops, nanname)
         else:
-            func = _dask_or_eager_func(name)
+            func = _dask_or_eager_func(name, dask_module=dask_module)
 
         try:
             return func(values, axis=axis, **kwargs)
@@ -337,7 +337,7 @@ def f(values, axis=None, skipna=None, **kwargs):
 std.numeric_only = True
 var = _create_nan_agg_method("var")
 var.numeric_only = True
-median = _create_nan_agg_method("median")
+median = _create_nan_agg_method("median", dask_module=dask_array_compat)
 median.numeric_only = True
 prod = _create_nan_agg_method("prod")
 prod.numeric_only = True

xarray/core/nanops.py

@@ -6,8 +6,10 @@
 
 try:
     import dask.array as dask_array
+    from . import dask_array_compat
 except ImportError:
     dask_array = None
+    dask_array_compat = None  # type: ignore
 
 
 def _replace_nan(a, val):
@@ -141,7 +143,15 @@ def nanmean(a, axis=None, dtype=None, out=None):
 
 
 def nanmedian(a, axis=None, out=None):
-    return _dask_or_eager_func("nanmedian", eager_module=nputils)(a, axis=axis)
+    # 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(np.atleast_1d(axis)) == a.ndim:
+        axis = None
+    return _dask_or_eager_func(
+        "nanmedian", dask_module=dask_array_compat, eager_module=nputils
+    )(a, axis=axis)
 
 
 def _nanvar_object(value, axis=None, ddof=0, keepdims=False, **kwargs):

xarray/tests/test_dask.py

@@ -216,8 +216,10 @@ def test_reduce(self):
         self.assertLazyAndAllClose(u.argmin(dim="x"), actual)
         self.assertLazyAndAllClose((u > 1).any(), (v > 1).any())
         self.assertLazyAndAllClose((u < 1).all("x"), (v < 1).all("x"))
-        with raises_regex(NotImplementedError, "dask"):
+        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)