xarray: Aggregating a dimension using the Quantiles method with `skipna=True` is very slow

AFAIK, performance increasing backends such as numbagg and bottleneck are optional dependencies. So by default numpy.nanquantile will be used. In cases where users have numbagg installed but want to disable it, they could do so using xarray.set_options:

with xr.set_options(use_numbagg=False):
    # do something

Here is a test with a real-world example (5 year time series of Sentinel-1 RTC data):

This is the numbagg wrapper I’ve used (same as I’ve posted above).

Yeah, having it run in parallel is going to generally be much faster. And avoiding multiple processes / dask is going to generally save on overhead.

A welcome PR would be to add numbagg’s nanpercentile to xarray, in the same way that rolling methods run through numbagg by default if it’s installed.

(or if we had benchmarks showing another approach was more performant, happy to consider that)

Recently, I create a library for the same purpose, maybe it can help other people. It can speed up quantile computations a lot. To install, just run: pip install fastnanquantile It’s compatible with dask. More info: https://github.com/lbferreira/fastnanquantile

Example below:

import numpy as np
import xarray as xr
from fastnanquantile import xrcompat

da = xr.DataArray(
    np.random.rand(10, 1000, 1000),
    coords={"time": np.arange(10), "x": np.arange(1000), "y": np.arange(1000)},
)

# Xarray quantile (time to run: ~25s)
result_xr = da.quantile(q=0.6, dim="time")
# fastnanquantile (time to run: <1s)
result_fnq = xrcompat.xr_apply_nanquantile(da, q=0.6, dim="time")
# Check if results are equal (If results are different, an error will be raised)
np.testing.assert_almost_equal(result_fnq.values, result_xr.values, decimal=4)

Hi all, I just created a simple workaround, which might be useful for others:
https://gist.github.com/maawoo/0b34d371c3cc1960a1589ccaded868c2

It uses the _nan_quantile method of xclim and works fine for my applications. Here is a quick comparison using the same example data as in my initial post:

EDIT: I’ve updated the code to use numbagg instead of xclim.

Thanks @arongergely! I have mentioned the numpy issue in my post above (FYI, for anyone looking for it). I was really surprised to see that it’s over 2 years old and that this is now the first Xarray issue referencing it. If it’s really a “well known” issue, I think it should have been somehow mentioned in the Xarray quantiles method.

I have seen the blog post and tried to use the workaround with apply_ufunc and Dask but ran into some problems. I’ll revisit that when I have some time and will also check xclim. Seems to be very promising, thanks!

Happy holidays! 🎄