Zarr Python 3 tracking issue

[jhamman]

What is your issue?

Zarr-Python 3.0 is getting close to a full release. This issue tracks the integration of the 3.0 release with Xarray.

Here's a running list of issues we're solving upstream related to integration with Xarray:

• Zarr-v3 Consolidated Metadata zarr-developers/zarr-python#2113
• Make MemoryStore serializable zarr-developers/zarr-python#2204
• [v3] fix: zarr v2 compatibility fixes for Dask zarr-developers/zarr-python#2186
• ...?

Special shout out to @TomAugspurger has been front running a lot of this 🙌.

[TomAugspurger]

Testing this out

There multiple active branches right now, but you can get a usable xarray + zarr-python 3.x with these two branches:

• https://github.com/TomAugspurger/zarr-python/tree/xarray-compat
• https://github.com/TomAugspurger/xarray/tree/fix/zarr-v3

You can install these with:

pip install git+https://github.com/TomAugspurger/zarr-python@xarray-compat git+https://github.com/TomAugspurger/xarray/@fix/zarr-v3

Work Items

We require some changes on both zarr-python and xarray. I'm pushing the zarr ones to TomAugspurger/zarr-python@xarary-compat and the xarray ones to `TomAugspurger/xarray@fix/zarr-v3.

zarr-python

• Consolidated metadata: Zarr-v3 Consolidated Metadata zarr-developers/zarr-python#2113 (in progress)
• Array v2 metadata deserialization: Fixed ArrayV2Metadata parameter names zarr-developers/zarr-python#2270 (in progress)
• Support for more dtypes: [v3]: TypeError when updating Array.attrs, thanks to incorrect inferred fill_value zarr-developers/zarr-python#2153 (lots of work to do)
• Change default fill values?: Change default fill values zarr-developers/zarr-python#2265 (in discussion. related to Is _FillValue really the same as zarr's fill_value? #5475)
• base64encoded fill_value for certain data types in zarr v2: base64 encode fill value for some dtypes with zarr_format=2 zarr-developers/zarr-python#2286 (in progress)

xarray

Most of these are in my PR at #9552

• _FillValue and Zarr's fill_value causing valid values to be cast to NaN: Is _FillValue really the same as zarr's fill_value? #5475 (this could use feedback from the xarray maintainers).
• Pass through zarr_format=2/3/None in all public read / write APIs.
• Update zarr.Blosc imports to numcodecs.Blosc (can be done anytime)
• Update for filters / compressor -> codecs change: Backwards compatibility for reading Array.filters and Array.codecs zarr-developers/zarr-python#2194

Fixed issues

• basic zarr-python 2.x compatibility: Basic Zarr-python 2.x compatibility changes zarr-developers/zarr-python#2098

• Attributes.asdict: Added Attributes.asdict zarr-developers/zarr-python#2221

• Fixed codec pipeline for zarr-v2: Update V2 codec pipeline to use concrete classes zarr-developers/zarr-python#2244

• Allow mode casting for Stores zarr-developers/zarr-python#2249

• Create intermediate Groups when creating a nested node: Ensure parents are created when creating a node zarr-developers/zarr-python#2262

• ignore extra keys in V2 metadata: Ignore extra keys in v2 metadata zarr-developers/zarr-python#2297 (in progress)

Things to investigate:

• separate store / chunk_store
• writing a subset of regions

[dcherian]

@TomAugspurger are you able to open a WIP PR with in-progress work. It'd be nice to see what's needed

[TomAugspurger]

Sure, #9552 has that.

[TomAugspurger]

Question for the group: does anyone object to xarray continuing to write Zarr V2 datasets by default? I hesitate to have xarray's default be different from zarr-python's, but that would relive some pressure to address #5475 quickly, since v2 datasets should be round-tripable.

[TomAugspurger]

I think that support for reading Zarr V2 datasets with zarr-python v3 is close to being ready. I updated #9515 (comment) with some instructions on how to install two branches if anyone is able to test that out:

In [4]: xr.open_dataset("abfs://daymet-zarr/annual/hi.zarr", engine="zarr", storage_options={"account_name": "daymeteuwest"})
Out[4]:
<xarray.Dataset> Size: 137MB
Dimensions:                  (y: 584, x: 284, time: 41, nv: 2)
Coordinates:
    lat                      (y, x) float32 663kB ...
    lon                      (y, x) float32 663kB ...
  * time                     (time) datetime64[ns] 328B 1980-07-01T12:00:00 ....
  * x                        (x) float32 1kB -5.802e+06 ... -5.519e+06
  * y                        (y) float32 2kB -3.9e+04 -4e+04 ... -6.22e+05
...
    start_year:        1980

In [5]: xr.open_dataset("s3://cmip6-pds/CMIP6/ScenarioMIP/AS-RCEC/TaiESM1/ssp126/r1i1p1f1/Amon/clt/gn/v20201124", engine="zarr", storage_options={"anon": True})
Out[5]:
<xarray.Dataset> Size: 228MB
Dimensions:    (time: 1032, lat: 192, lon: 288, bnds: 2)
Coordinates:
  * lat        (lat) float64 2kB -90.0 -89.06 -88.12 -87.17 ... 88.12 89.06 90.0
...
    variant_label:             r1i1p1f1

[maxrjones]

Thank you for all your work on Zarr V3 in Xarray! Super excited about the progress here 🎉

I've been using the https://github.com/pydata/xarray/tree/zarr-v3 branch to test out icechunk, and thought I'd share some odd behavior when writing data to S3 just in case it's not a known issue. I ran this code multiple times - the first time only the attributes were written and the second the col and row variables were also written, but the S variable would never get written. All data variables were written if I used a local filesystem rather than object storage:

testv3_store = "s3://nasa-veda-scratch/test-weight-caching/scratch/test-v3.zarr"
s = np.ones(10, dtype='float64')
c = np.ones(10, dtype='int32')
r = np.ones(10, dtype='int32')
ds = xr.Dataset(
    data_vars=dict(
        S = (["n_s"], s),
        col = (["n_s"], c),
        row = (["n_s"], r)
    ),
    attrs=dict(n_in = 100, n_out = 200)
)
print(f"Original data vars: {ds.data_vars}")
ds.to_zarr(testv3_store, zarr_format=3, mode="w")
ds2 = xr.open_zarr(testv3_store, zarr_format=3).load()
print(f"Round tripped data vars: {ds2.data_vars}")

Original data vars: Data variables:
    S        (n_s) float64 80B 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0
    col      (n_s) int32 40B 1 1 1 1 1 1 1 1 1 1
    row      (n_s) int32 40B 1 1 1 1 1 1 1 1 1 1
Round tripped data vars: Data variables:
    col      (n_s) int32 40B 1 1 1 1 1 1 1 1 1 1
    row      (n_s) int32 40B 1 1 1 1 1 1 1 1 1 1

p.s., this issue didn't seem to happen when I was writing/reading from an icechunk store

[TomAugspurger]

Thanks for the testing it out and reporting that issue. I'll see if I can reproduce it.