Repository: dask/dask
Pull request #11473 (2024-10-30 12:44:30)

Implement non gil-blocking quantile method
- [ ] ref https://github.com/pydata/xarray/issues/9692
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     13 files  ±  0       13 suites  ±0   3h 44m 42s ⏱️ + 5m 33s
 13 271 tests + 29   12 209 ✅ + 29   1 062 💤 ±0  0 ❌ ±0 
138 470 runs  +319  119 447 ✅ +319  19 023 💤 ±0  0 ❌ ±0 

Results for commit d87a41ab. ± Comparison against base commit 4f830509.

[test-results]:data:application/gzip;base64,H4sIAPQ0ImcC/03Myw6DIBCF4VcxrLtgRhDoyzQ4aELqpUFYNX334g1d/t9kzpf1fugW9qygflRsST5e5VKw0c/T1kJjpnyN+x0VnP1aEtGKiNxc+PafFXlTHl+99UM2XqALYQ6HhDTt01qos8s0GCFUwWPacKxPu01vfV+meRx9zMGcVlaAbTtNTa8BpHMSCFrQjUHgJLlCSa1lvz/IDzp/GgEAAA==


Patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
index 5ffbf73f9..7a8d2e767 100644
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -124,10 +124,12 @@ try:
         nanmedian,
         nanmin,
         nanprod,
+        nanquantile,
         nanstd,
         nansum,
         nanvar,
         prod,
+        quantile,
         reduction,
         std,
         sum,
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
index 1d9faa02b..a0dfb3d4f 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -6,9 +6,10 @@ import math
 import operator
 import warnings
 from collections.abc import Iterable
-from functools import partial
+from functools import partial, reduce
 from itertools import product, repeat
 from numbers import Integral, Number
+from operator import mul
 
 import numpy as np
 from tlz import accumulate, compose, drop, get, partition_all, pluck
@@ -27,7 +28,7 @@ from dask.array.core import (
 )
 from dask.array.creation import arange, diagonal
 from dask.array.dispatch import divide_lookup, nannumel_lookup, numel_lookup
-from dask.array.numpy_compat import ComplexWarning
+from dask.array.numpy_compat import NUMPY_GE_200, ComplexWarning
 from dask.array.utils import (
     array_safe,
     asarray_safe,
@@ -1861,3 +1862,221 @@ def nanmedian(a, axis=None, keepdims=False, out=None):
 
     result = handle_out(out, result)
     return result
+
+
+@derived_from(np)
+def quantile(
+    a,
+    q,
+    axis=None,
+    out=None,
+    overwrite_input=False,
+    method="linear",
+    keepdims=False,
+    *,
+    weights=None,
+    interpolation=None,
+):
+    """
+    This works by automatically chunking the reduced axes to a single chunk if necessary
+    and then calling ``numpy.quantile`` function across the remaining dimensions
+    """
+    if axis is None:
+        if any(n_blocks > 1 for n_blocks in a.numblocks):
+            raise NotImplementedError(
+                "The da.quantile function only works along an axis.  "
+                "The full algorithm is difficult to do in parallel"
+            )
+        else:
+            axis = tuple(range(len(a.shape)))
+
+    if not isinstance(axis, Iterable):
+        axis = (axis,)
+
+    axis = [ax + a.ndim if ax < 0 else ax for ax in axis]
+
+    # rechunk if reduced axes are not contained in a single chunk
+    if builtins.any(a.numblocks[ax] > 1 for ax in axis):
+        a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})
+
+    if NUMPY_GE_200:
+        kwargs = {"weights": weights}
+    else:
+        kwargs = {}
+
+    result = a.map_blocks(
+        np.quantile,
+        q=q,
+        method=method,
+        interpolation=interpolation,
+        axis=axis,
+        keepdims=keepdims,
+        drop_axis=axis if not keepdims else None,
+        new_axis=0 if isinstance(q, Iterable) else None,
+        chunks=_get_quantile_chunks(a, q, axis, keepdims),
+        **kwargs,
+    )
+
+    result = handle_out(out, result)
+    return result
+
+
+def _span_indexers(a):
+    # We have to create an indexer so that we can index into every quantile slice
+    # The method relies on the quantile axis being the last axis, which means that
+    # we have to calculate reduce(mul, list(a.shape)[:-1]) number of quantiles
+    # We create an indexer combination for each of these quantiles
+
+    shapes = reduce(mul, list(a.shape)[1:-1])
+    original_shapes = shapes * a.shape[0]
+    indexers = [tuple(np.repeat(np.arange(a.shape[0]), shapes))]
+
+    for i in range(1, len(a.shape) - 1):
+        indexer = np.repeat(np.arange(a.shape[i]), shapes // a.shape[i])
+        indexers.append(tuple(np.tile(indexer, original_shapes // shapes)))
+        shapes //= a.shape[i]
+    return indexers
+
+
+def _custom_quantile(
+    a,
+    q,
+    axis=None,
+    method="linear",
+    interpolation=None,
+    keepdims=False,
+    **kwargs,
+):
+    if (
+        not {method, interpolation}.issubset({"linear", None})
+        or len(axis) != 1
+        or axis[0] != len(a.shape) - 1
+        or len(a.shape) == 1
+        or a.shape[-1] > 1000
+    ):
+        # bail to nanquantile. Assumptions are pretty strict for now but we
+        # do cover the xarray.quantile case.
+        return np.nanquantile(
+            a,
+            q,
+            axis=axis,
+            method=method,
+            interpolation=interpolation,
+            keepdims=keepdims,
+            **kwargs,
+        )
+    # nanquantile in NumPy is pretty slow if the quantile axis is slow because
+    # each quantile has overhead.
+    # This method works around this by calculating the quantile manually.
+    # Steps:
+    # 1. Sort the array along the quantile axis (this is the most expensive step
+    # 2. Calculate which positions are the quantile positions
+    #    (respecting NaN values, so each quantile can have a different indexer)
+    # 3. Get the neighboring values of the quantile positions
+    # 4. Perform linear interpolation between the neighboring values
+    #
+    # The main advantage is that we get rid of the overhead, removing GIL blockage
+    # and just generally making things faster.
+
+    sorted_arr = np.sort(a, axis=-1)
+    indexers = _span_indexers(a)
+    nr_quantiles = len(indexers[0])
+
+    is_scalar = False
+    if not isinstance(q, Iterable):
+        is_scalar = True
+        q = [q]
+
+    quantiles = []
+    reshape_shapes = (1,) + tuple(sorted_arr.shape[:-1]) + ((1,) if keepdims else ())
+    for single_q in list(q):
+        i = (
+            np.ones(nr_quantiles) * (a.shape[-1] - 1)
+            - np.isnan(sorted_arr).sum(axis=-1).reshape(-1)
+        ) * single_q
+        lower_value, higher_value = np.floor(i).astype(int), np.ceil(i).astype(int)
+
+        # Get neighboring values
+        lower = sorted_arr[tuple(indexers) + (tuple(lower_value),)]
+        higher = sorted_arr[tuple(indexers) + (tuple(higher_value),)]
+
+        # Perform linear interpolation
+        factor_higher = i - lower_value
+        factor_higher = np.where(factor_higher == 0.0, 1.0, factor_higher)
+        factor_lower = higher_value - i
+
+        quantiles.append(
+            (higher * factor_higher + lower * factor_lower).reshape(*reshape_shapes)
+        )
+
+    if is_scalar:
+        return quantiles[0].squeeze(axis=0)
+    else:
+        return np.concatenate(quantiles, axis=0)
+
+
+@derived_from(np)
+def nanquantile(
+    a,
+    q,
+    axis=None,
+    out=None,
+    overwrite_input=False,
+    method="linear",
+    keepdims=False,
+    *,
+    weights=None,
+    interpolation=None,
+):
+    """
+    This works by automatically chunking the reduced axes to a single chunk
+    and then calling ``numpy.nanquantile`` function across the remaining dimensions
+    """
+    if axis is None:
+        if any(n_blocks > 1 for n_blocks in a.numblocks):
+            raise NotImplementedError(
+                "The da.nanquantile function only works along an axis.  "
+                "The full algorithm is difficult to do in parallel"
+            )
+        else:
+            axis = tuple(range(len(a.shape)))
+
+    if not isinstance(axis, Iterable):
+        axis = (axis,)
+
+    axis = [ax + a.ndim if ax < 0 else ax for ax in axis]
+
+    # rechunk if reduced axes are not contained in a single chunk
+    if builtins.any(a.numblocks[ax] > 1 for ax in axis):
+        a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})
+
+    if NUMPY_GE_200:
+        kwargs = {"weights": weights}
+    else:
+        kwargs = {}
+
+    result = a.map_blocks(
+        _custom_quantile,
+        q=q,
+        method=method,
+        interpolation=interpolation,
+        axis=axis,
+        keepdims=keepdims,
+        drop_axis=axis if not keepdims else None,
+        new_axis=0 if isinstance(q, Iterable) else None,
+        chunks=_get_quantile_chunks(a, q, axis, keepdims),
+        **kwargs,
+    )
+
+    result = handle_out(out, result)
+    return result
+
+
+def _get_quantile_chunks(a, q, axis, keepdims):
+    quantile_chunk = [len(q)] if isinstance(q, Iterable) else []
+    if keepdims:
+        return quantile_chunk + [
+            1 if ax in axis else c for ax, c in enumerate(a.chunks)
+        ]
+    else:
+        return quantile_chunk + [c for ax, c in enumerate(a.chunks) if ax not in axis]
diff --git a/docs/source/array-api.rst b/docs/source/array-api.rst
index 7c6f4897e..0f5b579bb 100644
--- a/docs/source/array-api.rst
+++ b/docs/source/array-api.rst
@@ -165,6 +165,7 @@ Top level functions
    nanmedian
    nanmin
    nanprod
+   nanquantile
    nanstd
    nansum
    nanvar
@@ -185,6 +186,7 @@ Top level functions
    power
    prod
    ptp
+   quantile
    rad2deg
    radians
    ravel


Test patch:
diff --git a/dask/array/tests/test_reductions.py b/dask/array/tests/test_reductions.py
index 4bb6a033a..0c43c13fb 100644
--- a/dask/array/tests/test_reductions.py
+++ b/dask/array/tests/test_reductions.py
@@ -908,3 +908,67 @@ def test_weighted_reduction():
     # Non-broadcastable weights (too many dims)
     with pytest.raises(ValueError):
         da.reduction(dx, w_sum, np.sum, weights=[[[2]]])
+
+
+@pytest.mark.parametrize("axis", [3, 0, [1, 3]])
+@pytest.mark.parametrize("q", [0.75, [0.75], [0.75, 0.4]])
+@pytest.mark.parametrize("rechunk", [True, False])
+def test_nanquantile(rechunk, q, axis):
+    shape = 10, 15, 10, 15
+    arr = np.random.randn(*shape)
+    indexer = np.random.randint(0, 10, size=shape)
+    arr[indexer >= 8] = np.nan
+
+    darr = da.from_array(arr, chunks=(2, 3, 4, (5 if rechunk else -1)))
+    assert_eq(da.nanquantile(darr, q, axis=axis), np.nanquantile(arr, q, axis=axis))
+    assert_eq(
+        da.nanquantile(darr, q, axis=axis, keepdims=True),
+        np.nanquantile(arr, q, axis=axis, keepdims=True),
+    )
+
+
+@pytest.mark.parametrize("axis", [3, [1, 3]])
+@pytest.mark.parametrize("q", [0.75, [0.75]])
+@pytest.mark.parametrize("rechunk", [True, False])
+def test_quantile(rechunk, q, axis):
+    shape = 10, 15, 20, 15
+    arr = np.random.randn(*shape)
+    indexer = np.random.randint(0, 10, size=shape)
+    arr[indexer >= 8] = np.nan
+
+    darr = da.from_array(arr, chunks=(2, 3, 4, (5 if rechunk else -1)))
+    assert_eq(da.quantile(darr, q, axis=axis), np.quantile(arr, q, axis=axis))
+    assert_eq(
+        da.quantile(darr, q, axis=axis, keepdims=True),
+        np.quantile(arr, q, axis=axis, keepdims=True),
+    )
+
+
+def test_nanquantile_all_nan():
+    shape = 10, 15, 20, 15
+    arr = np.random.randn(*shape)
+    arr[:] = np.nan
+    darr = da.from_array(arr, chunks=(2, 3, 4, -1))
+    da.nanquantile(darr, 0.75, axis=-1).compute()
+    with pytest.raises(RuntimeWarning):
+        assert_eq(
+            da.nanquantile(darr, 0.75, axis=-1), np.nanquantile(arr, 0.75, axis=-1)
+        )
+
+
+def test_nanquantile_method():
+    shape = 10, 15, 20, 15
+    arr = np.random.randn(*shape)
+    indexer = np.random.randint(0, 10, size=shape)
+    arr[indexer >= 8] = np.nan
+    darr = da.from_array(arr, chunks=(2, 3, 4, -1))
+    assert_eq(
+        da.nanquantile(darr, 0.75, axis=-1, method="weibull"),
+        np.nanquantile(arr, 0.75, axis=-1, method="weibull"),
+    )
+
+
+def test_nanquantile_one_dim():
+    arr = np.random.randn(10)
+    darr = da.from_array(arr, chunks=(2,))
+    assert_eq(da.nanquantile(darr, 0.75, axis=-1), np.nanquantile(arr, 0.75, axis=-1))