Icechunk store

[tomwhite]

Fixes #628

This is a proof of concept/draft for a store_icechunk function that uses Cubed callbacks to do the Icechunk store merging. The tests pass locally, but not in CI yet. I'm not sure where this code should live.

[dcherian]

could definitely address this later but regions is quite an important kwarg.

[dcherian]

@paraseba shall we commit to icechunk.distributed.merge_stores as public API?

https://github.com/earth-mover/icechunk/blob/a8c33b81d329c32a8e34e7151a1a71620967067a/icechunk-python/python/icechunk/distributed.py#L12-L16

[dcherian]

@tomwhite It looks like callbacks are "accumulated" on every worker. Is that right?

[paraseba]

@dcherian My only concern with it is the implementation of merge_stores is probably very slow, using no parallelism, but I don't see any issues with making it public

[dcherian]

Right we can fix that, but note that dask and presumably cubed are accumulating on remote workers already, so there is already some parallelism in how it is used.

Moreover it'd be nice not to have store.change_set_bytes be the public API :)

[tomwhite]

It looks like callbacks are "accumulated" on every worker. Is that right?

No, they are being accumulated on the client so there is no parallelism. I don't know what Icechunk is doing here, but would it be possible to merge a batch in one go rather than one at a time? Could that be more efficient?

[dcherian]

but would it be possible to merge a batch in one go rather than one at a time

we'll have to build some rust API, we'll get to this eventually.

No, they are being accumulated on the client so there is no parallelism.

ah ok. In that case, how about calling reduction on each array that was written? That way you parallelize the merge across blocks for each array, and then the only serial bit is the merging across arrays, which will be a lot smaller. I considered this approach for dask, but then just wrote out a tree reduction across all chunks.

EDIT: or is the reduction approach not viable because you need to serialize to Zarr at some point?

[tomwhite]

or is the reduction approach not viable because you need to serialize to Zarr at some point?

Yes, that's basically it - Cubed also separates the data paths for array manipulations (contents of the blocks) from the metadata operations (block IDs and - for Icechunk - the changesets). So I think merging in batches would be more feasible.

we'll have to build some rust API, we'll get to this eventually.

+1

[dcherian]

Yes, that's basically it

:( I was afraid so.

[tomwhite]

I'm not sure where this code should live.

After chatting to @rabernat, @jhamman, @TomNicholas we think Icechunk integration can live in this repo - as an optional dependency. There may be further changes once Xarray supports pluggable writers, but this PR can go in once tests are passing.

[rabernat]

This is amazing!

Note that the Python API is getting a major refactor in earth-mover/icechunk#470 which will almost certainly break this.

[tomwhite]

Note that the Python API is getting a major refactor in earth-mover/icechunk#470 which will almost certainly break this.

@rabernat thanks for the heads up. I'll keep an eye on that PR.

I think the best place to run the Icechunk tests here is in the Zarr v3 tests workflow in CI.

[tomwhite]

This is working well now with the latest zarr and icechunk releases. I plan to merge it soon.

[TomNicholas]

Awesome @tomwhite !