Woooo thanks @dcherian !

@Illviljan thanks for all your comments!

Would you (or @keewis?) be willing to approve this PR now? I would really like to merge this so that I can release a version of xarray that I can use as a dependency for cubed-xarray.

(Okay now the failures are from https://github.com/pydata/xarray/pull/7815 which I've separated out, and from https://github.com/pydata/xarray/pull/7561 being recently merged into main which is definitely not my fault :sweat_smile: https://github.com/pydata/xarray/pull/7019/commits/316c63d55f4e2c317b028842f752a40596f16c6d shows that this PR passes the tests by itself.)

I would like to merge this now please! It works, it passes the tests, including mypy.

The main feature not in this PR is using parallel=True with open_mfdataset, which is still coupled to dask.delayed - I made #7811 to track that so I could get this PR merged.

If we merge this I can start properly testing cubed with xarray (in cubed-xarray).

@shoyer @dcherian if one of you could merge this or otherwise tell me anything else you think is still required!

I'm having problems with ensuring the behaviour of the chunks='auto' option is consistent between .chunk and open_dataset

Update on this rabbit hole: This commit to dask changed the behaviour of dask's auto-chunking logic, such that if I run my little test script test_old_get_chunk.py on dask releases before and after that commit I get different chunking patterns:

python (cubed) tom@tom-XPS-9315:~/Documents/Work/Code/dask$ git checkout 2022.9.2 Previous HEAD position was 7fe622b44 Add docs on running Dask in a standalone Python script (#9513) HEAD is now at 3ef47422b bump version to 2022.9.2 (cubed) tom@tom-XPS-9315:~/Documents/Work/Code/dask$ python ../experimentation/bugs/auto_chunking/test_old_get_chunk.py {'x': (362, 138), 'y': (362, 138)} (cubed) tom@tom-XPS-9315:~/Documents/Work/Code/dask$ git checkout 2022.9.1 Previous HEAD position was 3ef47422b bump version to 2022.9.2 HEAD is now at b944abf68 bump version to 2022.9.1 (cubed) tom@tom-XPS-9315:~/Documents/Work/Code/dask$ python ../experimentation/bugs/auto_chunking/test_old_get_chunk.py {'x': (250, 250), 'y': (250, 250)} (I was absolutely tearing my hair out trying to find this bug, because after the change normalize_chunks became a pure function, but before the change it actually wasn't, so I was trying calling normalize_chunks with the exact same set of input arguments and was still not able to reproduce the bug :angry: )

Anyway what this means is as this PR vendors dask.array.core.normalize_chunks, but the behaviour of dask.array.core.normalize_chunks changed between the version in CI job min-all-deps and the other CI jobs, the single vendored function cannot possibly match both behaviours.

I think one simple way to fix this failure without should be to upgrade the minimum version of dask to >=2022.9.2 (from 2022.1.1 where it currently is).

EDIT: I tried changing the minimum version of dask-core in min-all-deps.yml but the conda solve failed. But also would updating to 2022.9.2 now violate xarray's minimum dependency versions policy?

EDIT2: Another way to fix this should be to un-vendor dask.array.core.normalize_chunks within xarray. We could still achieve the goal of running cubed without dask by making normalize_chunks the responsibility of the chunkmanager instead, as cubed's vendored version of that function is not subject to xarray's minimum dependencies requirement.

I'm having problems with ensuring the behaviour of the chunks='auto' option is consistent between .chunk and open_dataset. These problems appeared since vendoring dask.array.core.normalize_chunks. Right now the only failing tests use chunks='auto' (e.g xarray/tests/test_backends.py::test_chunking_consintency[auto] - yes there's a typo in that test's name), and they fail because xarray decides on different sizes for the automatically-chosen chunks.

What's weird is that all tests pass for me locally but these failures occur on just some of the CI jobs (and which CI jobs is not even consistent apparently???). I have no idea why this would behave differently on only some of the CI jobs, especially after double-checking that array-chunk-size is being correctly determined from the dask config variable within normalize_chunks.

I've made a bare-bones cubed-xarray package to store the CubedManager class, as well as any integration tests (yet to be written). @tomwhite you should have an invitation to be an owner of that repo. It uses the entrypoint exposed in this PR to hook in, and seems to work for me locally :grin:

Thanks for the review @dcherian! I agree with basically everything you wrote.

The main difficulty I have at this point is non-reproducible failures as described here

I would like to get to the point where you can use xarray with a chunked array without ever importing dask. I think this PR gets very close, but that would be tricky to test because cubed depends on dask (so I can't just run the test suite without dask in the environment

I just released Cubed 0.6.0 which doesn't have a dependency on Dask, so this should be possible now.

Actually testing cubed with xarray in an environment without dask is currently blocked by rechunker's explicitly dependency on dask, see https://github.com/pangeo-data/rechunker/issues/139

EDIT: We can hack around this by pip installing cubed, then pip uninstalling dask as mentioned here

Does this mean my comment https://github.com/pydata/xarray/pull/7019#discussion_r970713341 is valid again?

Yes I think it does @headtr1ck - thanks for the reminder about that.

I now want to finish this PR by exposing the "chunk manager" interface as a new entrypoint, copying the pattern used for xarray's backends. That would allow me to move the cubed-specific CubedManager code into a separate repository, have the choice of chunkmanager default to whatever is installed, but ask explicitly what to do if multiple chunkmanagers are installed. That should address your comment @headtr1ck.

Thanks @dcherian ! Once I copied that explicit indexer business I was able to get serialization to and from zarr working with cubed!

```python In [1]: import xarray as xr

In [2]: from cubed import Spec

In [3]: ds = xr.open_dataset( ...: 'airtemps.zarr', ...: chunks={}, ...: from_array_kwargs={ ...: 'manager': 'cubed', ...: 'spec': Spec(work_dir="tmp", max_mem=20e6), ...: } ...: ) /home/tom/Documents/Work/Code/xarray/xarray/backends/plugins.py:139: RuntimeWarning: 'netcdf4' fails while guessing warnings.warn(f"{engine!r} fails while guessing", RuntimeWarning) /home/tom/Documents/Work/Code/xarray/xarray/backends/plugins.py:139: RuntimeWarning: 'scipy' fails while guessing warnings.warn(f"{engine!r} fails while guessing", RuntimeWarning)

In [4]: ds['air'] Out[4]: <xarray.DataArray 'air' (time: 2920, lat: 25, lon: 53)> cubed.Array<array-004, shape=(2920, 25, 53), dtype=float32, chunks=((730, 730, 730, 730), (13, 12), (27, 26))> Coordinates: * lat (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0 * lon (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0 * time (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00 Attributes: GRIB_id: 11 ...

In [5]: ds.isel(time=slice(100, 300)).to_zarr("cubed_subset.zarr") /home/tom/Documents/Work/Code/xarray/xarray/core/dataset.py:2118: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs return to_zarr( # type: ignore Out[5]: <xarray.backends.zarr.ZarrStore at 0x7f34953033c0> ```

I tried opening a zarr store into xarray with chunking via cubed, but I got an error inside the indexing adapter classes. Somehow the type is completely wrong - would be good to type hint this part of the code, because this happens despite mypy passing now.

```python

create example zarr store

orig = xr.tutorial.open_dataset("air_temperature") orig.to_zarr('air2.zarr')

open it as a cubed array

ds = xr.open_dataset('air2.zarr', engine='zarr', chunks={}, from_array_kwargs={'manager': 'cubed'})

fails at this point

ds.load() ```

This still works fine for dask.

Thanks @tomwhite - I think it might make sense for me to remove the CubedManager class from this PR and instead put that & cubed+xarray tests into another repo. That keeps xarray's changes minimal, doesn't require putting cubed in any xarray CI envs, and hopefully allows us to merge the ChunkManager changes here earlier.

Places dask is still explicitly imported in xarray

There are a few remaining places where I haven't generalised to remove specific import dask calls either because it won't be imported at runtime unless you ask for it, cubed doesn't implement the equivalent function, that function isn't in the array API standard, or because I'm not sure if the dask concept used generalises to other parallel frameworks.

• [ ] open_mfdataset(..., parallel=True) - there is no cubed.delayed to wrap the open_dataset calls in,
• [ ] Dataset.__dask_graph__ and all the other similar dask magic methods
• [ ] dask_array_ops.rolling - uses functions from dask.array.overlap,
• [ ] dask_array_ops.least_squares - uses dask.array.apply_along_axis and dask.array.linalg.lstsq,
• [ ] dask_array_ops.push - uses dask.array.reductions.cumreduction

I would like to get to the point where you can use xarray with a chunked array without ever importing dask. I think this PR gets very close, but that would be tricky to test because cubed depends on dask (so I can't just run the test suite without dask in the environment), and there are not yet any other parallel chunk-aware frameworks I know of (ramba and arkouda don't have a chunks attribute so wouldn't require this PR).

I'm making progress with this PR, and now that @tomwhite implemented cubed.apply_gufunc I've re-routed xarray.apply_ufunc to use whatever version of apply_gufunc is defined by the chosen ChunkManager. This means many basic operations should now just work:

```python In [1]: import xarray as xr

In [2]: da = xr.DataArray([1, 2, 3], dims='x')

In [3]: da_chunked = da.chunk(from_array_kwargs={'manager': 'cubed'})

In [4]: da_chunked Out[4]: <xarray.DataArray (x: 3)> cubed.Array<array-003, shape=(3,), dtype=int64, chunks=((3,),)> Dimensions without coordinates: x

In [5]: da_chunked.mean() Out[5]: <xarray.DataArray ()> cubed.Array<array-006, shape=(), dtype=int64, chunks=()>

In [6]: da_chunked.mean().compute() [cubed.Array<array-009, shape=(), dtype=int64, chunks=()>] Out[6]: <xarray.DataArray ()> array(2) ```

(You need to install both cubed>0.5.0 and the main branch of rechunker for this to work.)

I still have a fair bit more to do on this PR (see checklist at top), but for testing should I:

1. Start making a test_cubed.py file in xarray as part of this PR with bespoke tests,
2. Put bespoke tests for xarray wrapping cubed somewhere else (e.g. the cubed repo or a new cubed-xarray repo),
3. Merge this PR without cubed-specific tests and concentrate on finishing the general duck-array testing framework in #6908 so we can implement (b) in the way we actually eventually want things to work for 3rd-party duck array libraries?

I would prefer not to have this PR grow to be thousands of lines by including tests in it, but also waiting for #6908 might take a while because that's also a fairly ambitious PR.

The fact that the tests are currently green for this PR (ignoring some mypy stuff) is evidence that the decoupling of dask from xarray is working so far.

(I have already added some tests for the ability to register custom ChunkManagers though.)

@drtodd13 mentioned today that ramba doesn't actually require explicit chunks to work, which I hadn't realised. So forcing wrapped libraries to implement an explicit chunks method might be too restrictive. Ramba could possibly work entirely through the numpy array API standard.

@drtodd13 tagging you here and linking my notes from today's distributed arrays working group meeting for the links and references to this PR.

This was more abstract than expected

Yeah I was kind of asking whether this was unnecessarily abstract, and if there was a simpler design that achieved the same flexibility.