Apply function to groupby result using dask

[iuryt]

I could manage to make xarray use dask for groupby by chunking data first, but I couldn't find a way to group in a lazy mode.
That means that I couldn't find a way to parallelize the function I apply for the groups.

import xarray as xr
import numpy as np

n = int(1e7)
ds = xr.Dataset(
    dict(
        x = ("obs",np.random.randn(n)),
        y = ("obs",np.random.randint(1,20+1,n))
    )
)

ds.chunk(obs=1000).groupby(y).min()

What happens is that when I do ds.chunk(obs=1000).groupby(y), it computes the groups, which means it won't apply the min function to the groups in parallel.

I thought about defining a dask.delayed function, but it takes some time to run a loop through the groups.

Any idea?

[dcherian]

use flox for now. I'm slowing working to fix this.

[iuryt]

Thanks @dcherian!

I think this is going to work for one of my examples.
If I understood correctly, flox have predefined functions to apply to the groups.
Do you know how we could apply any arbitrary function.
In my case, I wanted to group the data by year, do some reshaping and saving for each year.

• "sum"
• "nansum"
• "prod"
• "nanprod"
• "count"
• "mean"
• "nanmean"
• "var"
• "nanvar"
• "std"
• "nanstd"
• "argmin"
• "argmax"
• "first"
• "last"

[dcherian]

I wanted to group the data by year, do some reshaping and saving for each year.

An example would be useful, but grouping by year should be doable with just normal xarray no? The time vector is usually small

[iuryt]

It is not as easy to reproduce because I have a data that has only one dimension obs and I can have many multiple observations in a day and consequently, in a year. But I think flox works fine for the example above, so I think we can mark this as an answer.

import xarray as xr
import numpy as np
from flox.xarray import xarray_reduce

n = int(1e7)
dsi = xr.Dataset(
    dict(
        x = ("obs",np.random.randn(n)),
        y = ("obs",np.random.randint(1,20+1,n))
    )
).chunk(obs=1000)

expected = np.unique(dsi.y)
group_min = xarray_reduce(dsi.x,dsi.y,expected_groups=expected,func="min").load()