How do I read an LZ4 compressed zarr array written from CPU to GPU?

[ivirshup]

I'm figuring out some direct GPU loading stuff with zarr and am running into some confusing bits. If I write out some arrays on a CPU, then want to read them in, what's the API for doing this?

kvikio has a notebook that demonstrates this by modifying the .zarray json inplace:

# Define a simple, dictionary-based store. In real scenarios this can be a filesystem or some other persistent store.
store = {}
zarr.save_array(store, x, compressor=cpu_codec)

# Check that the data was written by the expected codec.
meta = json.loads(store[".zarray"])
assert meta["compressor"]["id"] == "lz4"

# Change codec to GPU/nvCOMP-based.
meta["compressor"] = {"id": NVCOMP_CODEC_ID, "algorithm": LZ4_ALGO}
store[".zarray"] = json.dumps(meta).encode()

Is there a good way to do this that doesn't (a) modify the .zarray json in-place and (b) uses a higher level API?

[jakirkham]

Maybe it is worth reading issue ( zarr-developers/numcodecs#316 )

Also would check out issue ( rapidsai/kvikio#66 )

[ivirshup]

Very useful, thanks!

[weiji14]

So kvikIO=23.10.00 released yesterday has some support for LZ4 compression/decompression via nvCOMP. See the example at https://docs.rapids.ai/api/kvikio/23.10/zarr. Specifically, you should be able to do something like:

z = kvikio.zarr.open_cupy_array(store="example.zarr", mode="r", compressor=kvikio.zarr.CompatCompressor.lz4(), path="/")

There might still be some issues with it (see rapidsai/kvikio#297), but it should be easier than having to modify the .zarray file 🤞