Add to_icechunk for xarray

[dcherian]

Adds icechunk.xarray.to_icechunk.

TODO:

• Add load_stored kwarg to dask.array.store. dask/dask#11465
• more docs?
• update examples/dask_write.py
• release xarray

I chose to not update dask_write.py because it does show a different pattern that uses dask as an orchestrator and not dask.array.

[dcherian]

This adds two APIs. One is "easy" and does everything in one step, the other allows extra control over initializing the store, and subsequent writes.

1. icechunk.xarray.to_icechunk(dataset, store=store, ...) which is like ds.to_zarr without the compute kwarg.
2. To "initialize" the repo, and execute a distributed writer in a separate step, the API is:

from icechunk.xarray import XarrayDatasetWriter

writer = XarrayDatasetWriter(ds, store=store)
writer.write_metadata(group="new2", mode="w") # write metadata
writer.write_eager() # write in-memory arrays
writer.write_lazy() # eagerly write dask arrays

The write_eager method could execute an async write of all the eager in-memory coordinate arrays, but does not do so right now.

[dcherian]

Think about appends & region writes here.

[TomNicholas]

Naive question: if icechunk.xarray.to_icechunk is functionally the same as ds.to_zarr(icechunkstore) then why does it need to exist?

Also should we call icechunk.xarray.to_icechunk from inside virtualizarr's .virtualize.to_icechunk accessor method? Since that can also write non-virtual variables.

cc @mpiannucci

[dcherian]

is functionally the same as

It is not. to_zarr will not work for distributed writes. It would write the chunks, but the commit wouldn't be meaningful.

[TomNicholas]

.to_zarr will not work for distributed writes. It would write the chunks, but the commit wouldn't be meaningful.

Can you expand on what this means in practice? If I have a big distributed dask array, I call to .to_zarr, then commit, what exactly would go wrong?

[dcherian]

Can you expand on what this means in practice?

The store object you are holding has no record of the distributed writes because those changes were never communicated back. So you'll receive all fill values when trying to read: #383

[mpiannucci]

🥵

[rabernat]

This is a masterpiece Deepak!

Reviewed the code and tests and didn't spot any red flags. I think I get the gist of how it works, but I'm not confident I'd be able to identify any bugs there anyway. The tests inspire great confidence.

I left a bunch of suggestions on the docs.

[rabernat]

We've done so much work in Xarray around safe chunks, region writes, etc. It's a pity that we can't reuse any of that here.