Import dask and cubed only once

[eendebakpt]

We recently found a performance bottleneck in the xarray groupby method. Part of the bottleneck was repeated attempted imports of cubed. In this PR we refactor the code to import cubed and dask only once.

[dcherian]

Nice, thanks. Looks like there's one more mypy thing to fix. I'm happy to immediately release once that is fixed.

We recently found a performance bottleneck in the xarray groupby method.

I'm curious to learn if you've found other things. I spent quite some time optimizing this stuff a couple of years ago.

[eendebakpt]

Nice, thanks. Looks like there's one more mypy thing to fix. I'm happy to immediately release once that is fixed.

On my local system mypy is giving different output, so I am doing a bit of trial and error here.

We recently found a performance bottleneck in the xarray groupby method.

I'm curious to learn if you've found other things. I spent quite some time optimizing this stuff a couple of years ago.

Not something specific. Our use case has relatively small arrays, maybe the optimizations have been tested on larger datasets?
Here is a minimal example I extracted:

import xarray as xr
import time
import numpy as np


def method(ds: xr.Dataset):
    pauli_groups = ds.groupby("pauli")
    for _, ds in pauli_groups:
        da = ds.groupby("clifford").mean().counts


k = 12
mult = 30
index = np.arange(k*mult)
bitindex = np.arange(4)
counts = np.random.randint(0, 10, size=(4, index.size))
c = xr.DataArray(counts, coords={'bit_index': bitindex, 'index': index})
ds = xr.Dataset({'counts': c})
ds = ds.assign_coords({'clifford': xr.DataArray(
    np.tile(np.arange(k), mult), dims='index')})
p = xr.DataArray(np.tile(np.arange(4), (k//4) * mult), dims='index')
for ii in range(16):  # the value here has a large impact on the running time
    p.data[ii] = 1000 + ii
ds = ds.assign_coords({'pauli': p})
print(ds)

nn = 20
for ii in range(4):
    t0 = time.time()
    for ii in range(nn):
        method(ds)
    dt = time.time()-t0
    print(f'time {1e3*dt/nn:.2f} [ms / iteration]')

On my system a single call to method takes about 50 ms. I suspect that the equivalent calculation with just numpy arrays is much faster (not a fair comparison, xarray has to keep track of much more data!). Looking at profiling results shows that the call to mean() is actually taking most of the time.