How to tell if flox is actually being used?

[claytharrison]

I have a process that is using .groupby() on a very large, dask-chunked dataset, and so it seems like a good candidate for using flox to improve things. However, I don't see any difference in speed or memory use when flox is installed compared to when it's not. Now I can't tell if my process just isn't one that can be improved by flox or if xarray simply isn't detecting flox for some reason. Is there any way to tell if xarray can actually "see" flox?

Alternatively, maybe somebody can tell me if my use-case makes sense in the first place. It goes something like this:

from pathlib import Path
import xarray as xr

def ds_preprocess(dataset):
    # do stuff to dataset (transform vars, add some new vars, including "cell")
    return dataset

def cell_ds_process(dataset):
    # do stuff to dataset (transform vars, add some new vars)
    return dataset

filenames = list(Path("/path/to/files/").glob("*.nc"))
with xr.open_mfdataset(filenames, concat_dim="obs", combine="nested") as fp:
    data = ds_preprocess(fp)
    for cell, cell_ds in data.groupby("cell"):
        cell_ds = cell_ds_process(cell_ds)
        fname = get_fname(cell_ds)
        cell_ds.to_netcdf(fname)

[dcherian]

for cell, cell_ds in data.groupby("cell")

It is not being used here. flox only accelerates reductions, so things like data.groupby("cell").mean(). Depending on whats in cell_ds_process you might be able to migrate the heavy lifting to flox.