Icechunk store - #633
Conversation
| *, | ||
| sources: Union["Array", Sequence["Array"]], | ||
| targets: List[zarr.Array], | ||
| executor=None, |
There was a problem hiding this comment.
could definitely address this later but regions is quite an important kwarg.
| if self.store is None: | ||
| self.store = store | ||
| else: | ||
| self.store.merge(store.change_set_bytes()) |
There was a problem hiding this comment.
@paraseba shall we commit to icechunk.distributed.merge_stores as public API?
There was a problem hiding this comment.
@tomwhite It looks like callbacks are "accumulated" on every worker. Is that right?
There was a problem hiding this comment.
@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
There was a problem hiding this comment.
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 :)
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
or is the
reductionapproach 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
There was a problem hiding this comment.
Yes, that's basically it
:( I was afraid so.
a4a84c7 to
471fc9e
Compare
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. |
|
This is amazing! 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. |
cb7a1c1 to
c21a446
Compare
719f83c to
7cc3b8e
Compare
Don't include icechunk in coverage
5fc8970 to
d5a6351
Compare
|
This is working well now with the latest zarr and icechunk releases. I plan to merge it soon. |
|
Awesome @tomwhite ! |
5 participants
Fixes #628
This is a proof of concept/draft for a
store_icechunkfunction 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.