Zarr reader

[norlandrhagen]

WIP PR to add a Zarr reader. Thanks to @TomNicholas for the how to write a reader guide.

• Closes Add Zarr Reader(s) #262
• Tests added
• Tests passing
• Full type hint coverage
• Changes are documented in docs/releases.rst
• optimizations (e.g. using async interface to list lengths of chunks for each variable concurrently)
• New functionality has documentation
• Read v3 Zarr

Future PR(s):

• Read v2 Zarr
• sharded v3 data

To Do:

• Open the store using zarr-python v3 (behind a protected import). This should handle both v2 and v3 stores for us.
• Use zarr-python to list the variables in the store, and check that all loadable_variables are present
  For each virtual variable:
• Use zarr-python to get the attributes and the dimension names, and coordinate names (which come from the .zmetadata or zarr.json)
• Use zarr-python to also get the dtype and chunk grid info + everything else needed to create the virtualizarr.zarr.ZArray object (eventually we can skip this step and use a zarr-python array metadata class directly instead of virtualizarr.zarr.ZArray, see
  Replace VirtualiZarr.ZArray with zarr ArrayMetadata #175)
• Use the knowledge of the store location, variable name, and the zarr format to deduce which directory / S3 prefix the chunks must live in.
• List all the chunks in that directory using fsspec.ls(detail=True), as that should also return the nbytes of each chunk. Remember that chunks are allowed to be missing.
• The offset of each chunk is just 0 (ignoring sharding for now), and the length is the file size fsspec returned. The paths are just all the paths fsspec listed.
• Parse the path and length information returned by fsspec into the structure that we can pass to ChunkManifest.init
• Create a ManifestArray from our ChunkManifest and ZArray
• Wrap that ManifestArray in an xarray.Variable, using the dims and attrs we read before
  Get the loadable_variables by just using xr.open_zarr on the same store (should use drop_variables to avoid handling the virtual variables that we already have).
• Use separate_coords to set the correct variables as coordinate variables (and avoid building indexes whilst doing it)
• Merge all the variables into one xr.Dataset and return it.
• All the above should be wrapped in a virtualizarr.readers.zarr.open_virtual_dataset function, which then should be called as a method from a ZarrVirtualBackend(VirtualBackend) subclass.
• Finally add that ZarrVirtualBackend to the list of readers in virtualizarr.backend.py

[norlandrhagen]

#273

[norlandrhagen]

Bit of an update. With the help from @sharkinsspatial @abarciauskas-bgse and @maxrjones I got a Zarr loaded as a virtual dataset.

<xarray.Dataset> Size: 3kB
Dimensions:  (time: 10, lat: 9, lon: 18)
Coordinates:
    lat      (lat) float32 36B ManifestArray<shape=(9,), dtype=float32, chunk...
    lon      (lon) float32 72B ManifestArray<shape=(18,), dtype=float32, chun...
  * time     (time) datetime64[ns] 80B 2013-01-01 ... 2013-01-03T06:00:00
Data variables:
    air      (time, lat, lon) int16 3kB ManifestArray<shape=(10, 9, 18), dtyp...
Attributes:
    Conventions:  COARDS
    title:        4x daily NMC reanalysis (1948)
    description:  Data is from NMC initialized reanalysis\n(4x/day).  These a...
    platform:     Model
    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...

Next up is how to deal with fill_values.

When I try to write it to Kerchunk JSON, I’m running into some fill_value dtype issues in the Zarray.

ZArray(shape=(10,), chunks=(10,), dtype='<f4', fill_value=np.float32(nan), order='C', compressor=None, filters=None, zarr_format=2)

Where fill_value=np.float32(nan) . When I try to write these to JSON via ds.virtualize.to_kerchunk(format="dict"), I get TypeError: np.float32(nan) is not JSON serializable.

Wondering how fill_values like np.float32(nan) should be handled.

There seems to be some conversion logic in @sharkinsspatial's HDF reader for converting fill_values . It also looks like there is some fill_value handling in zarr.py.

[TomNicholas]

I got a Zarr loaded as a virtual dataset.

Amazing!

When I try to write it to Kerchunk JSON, I’m running into some fill_value dtype issues in the Zarray.

ZArray(shape=(10,), chunks=(10,), dtype='<f4', fill_value=np.float32(nan), order='C', compressor=None, filters=None, zarr_format=2)

Where fill_value=np.float32(nan) . When I try to write these to JSON via ds.virtualize.to_kerchunk(format="dict"), I get TypeError: np.float32(nan) is not JSON serializable.

Wondering how fill_values like np.float32(nan) should be handled.

This seems like an issue that should actually be orthogonal to this PR (if it weren't for the ever-present difficulty of testing). Either the problem is in the ZArray class and what types it allows, or it's in the Kerchunk writer not knowing how to serialize a valid ZArray. Either way if np.float32(nan) is a valid fill_value for a zarr array then it's not the fault of the new zarr reader.

[TomNicholas]

This is a a great start! I think the main thing here is that we don't actually need kerchunk in order to test this reader.

[TomNicholas]

get chunk size with zarr-python (zarr-developers/zarr-python#2426) instead of fsspec

I think we should just do this in this PR. We can point to Tom's PR for now in the CI, but I expect that will get merged before this does anyway. If you look at Tom's implementation it's basically what we're doing here.

[TomNicholas]

Store.getsize was just merged so we can possibly just use the same upstream env we are already using zarr-python zarr-developers/zarr-python#2426

[TomNicholas]

Thanks for your patience @norlandrhagen , and for your help @maxrjones .

I believe I spotted one small bug, but otherwise this looks great.

The implementation is very neat now!

[TomNicholas]

Do we have an issue to track learning to read zarr v2?

[norlandrhagen]

#565