MAINT: implement median

dask/array/__init__.py

@@ -211,7 +211,7 @@
         nancumsum,
         reduction,
     )
-    from .percentile import percentile
+    from .percentile import percentile, median
     from . import ma
     from . import random, linalg, overlap, fft, backends
     from .overlap import map_overlap

dask/array/core.py

@@ -1997,6 +1997,14 @@ def mean(self, axis=None, dtype=None, keepdims=False, split_every=None, out=None
             out=out,
         )
 
+    def median(self):
+        """
+        Implements an approximate version of the median function.
+        See :func:`percentile` for more detail.
+        """
+        from .percentile import median
+        return median(self)
+
     @derived_from(np.ndarray)
     def std(
         self, axis=None, dtype=None, keepdims=False, ddof=0, split_every=None, out=None

dask/array/percentile.py

@@ -149,6 +149,14 @@ def percentile(a, q, interpolation="linear", method="default"):
     return Array(graph, name2, chunks=((len(q),),), dtype=dtype)
 
 
+def median(x, interpolation="linear"):
+    """
+    Implements an approximate version of the median function.
+    See :func:`percentile` for more detail.
+    """
+    return percentile(x, 50, interpolation=interpolation, method="tdigest")
+
+
 def merge_percentiles(finalq, qs, vals, interpolation="lower", Ns=None):
     """ Combine several percentile calculations of different data.
 

dask/array/tests/test_percentiles.py

@@ -52,6 +52,21 @@ def test_percentile(method):
         )
 
 
+@pytest.mark.skipif(not crick, reason="Requires crick")
+def test_median():
+    rng = np.random.RandomState(42)
+    n = int(10e3)
+    x = rng.rand(n)
+    d = da.from_array(x, chunks=n // 10)
+
+    assert_eq(d.median(), da.median(d))
+    assert_eq(da.percentile(d, 50, method="tdigest"), da.median(d))
+    diff = np.abs(d.median().compute() - np.median(x))
+    rel_error = diff / np.median(x)
+    assert rel_error < 2e-3
+    assert diff < 1e-3
+
+
 @pytest.mark.skip
 def test_percentile_with_categoricals():
     try:

dask/dataframe/core.py

@@ -2755,6 +2755,9 @@ def nlargest(self, n=5, split_every=None):
             n=n,
         )
 
+    def median(self):
+        return da.median(self)
+
     @derived_from(pd.Series)
     def nsmallest(self, n=5, split_every=None):
         return aca(

dask/dataframe/tests/test_dataframe.py

@@ -4111,3 +4111,11 @@ def test_pop():
     assert s.name == "y"
     assert ddf.columns == ["x"]
     assert_eq(ddf, df[["x"]])
+
+
+def test_median():
+    N = int(10e3)
+    rng = np.random.RandomState(42)
+    s = pd.Series(rng.rand(N))
+    ds = dd.from_pandas(s, npartitions=N // 10)
+    assert ds.median().compute() == pytest.approx(s.median(), rel=0.002)