Compatibility for zarr-python 3.x

[TomAugspurger]

This PR begins the process of adding compatibility with zarr-python 3.x. It's intended to be run against zarr-python v3 + the open PRs referenced in #9515.

All of the zarr test cases should be parameterized by zarr_format=[2, 3] with zarr-python 3.x to exercise reading and writing both formats.

This is currently passing with zarr-python==2.18.3. zarr-python 3.x has about 61 failures, all of which are related to data types that aren't yet implemented in zarr-python 3.x.

I'll also note that #5475 is going to become a larger issue once people start writing Zarr-V3 datasets.

• Closes Zarr Python 3 tracking issue #9515
• Closes Is _FillValue really the same as zarr's fill_value? #5475
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst

[TomAugspurger]

This set of changes should be backwards compatible and work with zarr-python 2.x (so reading and writing zarr v2 data).

I'll work through zarr-python 3.x now. I think we might want to parametrize most of these tests by zarr_version=[2, 3] to confirm that we can read / write zarr v2 data with zarr-python 3.x

[TomAugspurger]

Maybe this instead?

Suggested change

	x.to_dict() for x in zarr_array.metadata.codecs
	zarr_array.metadata.to_dict()["codecs"]

A bit wasteful since everything has to be serialized, but presumably zarr knows better how to serialize the codec pipeline than we do here?

[jhamman]

Great progress here @TomAugspurger. I'm impressed by how little you've changed in the backend itself and I'm noting the pain around testing (I felt some of that w/ dask as well).

[TomAugspurger]

I just pushed a commit reverting the changes to avoid values equal to the fill_value in the test cases, which aren't needed after Ryan's changes to fill value handling.

I think this is ready to go once CI finishes. I expect upstream-ci to fail on the xarray/tests/test_groupby.py::test_gappy_resample_reductions tests, but those should be unrelated.

[TomAugspurger]

There's one typing failure we might want to address:

xarray/tests/test_backends_datatree.py: note: In member "test_to_zarr_zip_store" of class "TestZarrDatatreeIO":
xarray/tests/test_backends_datatree.py:289: error: Argument 1 to "open_datatree" has incompatible type "ZipStore"; expected "str | PathLike[Any] | BufferedIOBase | AbstractDataStore"  [arg-type]

I'll do some reading about how best to handle type annotations when the proper type depends on the version of a dependency.

Edit: a complication here is that this is in open_datatree, which I think supports other backends like NetCDF too. It's not immediately clear to me whether the signature can of open_datatree needs to match for every implementation. This can maybe be addressed once we fully support datatree with zarr-python 3?

[TomNicholas]

Edit: a complication here is that this is in open_datatree, which I think supports other backends like NetCDF too. It's not immediately clear to me whether the signature can of open_datatree needs to match for every implementation. This can maybe be addressed once we fully support datatree with zarr-python 3?

I don't see why the typing of open_datatree(<zarr-store>) would need to be any different to that of open_dataset(<zarr-store>). But I think it's fine to ignore this for now as we know we need to come back to it in another PR anyway.

[TomAugspurger]

Good catch, this affects both.

I was hoping something like this would work:

from pathlib import Path

try:
    from zarr.storage import StoreLike as _StoreLike

except ImportError:
    _StoreLike = str | Path

StoreLike = type[_StoreLike]


def f(x: StoreLike) -> StoreLike:
    return x

but mypy doesn't like that

test.py:7: error: Cannot assign multiple types to name "_StoreLike" without an explicit "Type[...]" annotation  [misc]
Found 1 error in 1 file (checked 1 source file)

[jhamman]

but mypy doesn't like that

my 2 cents... we should not get hung up on this right now. a) there are plenty of other failures in the upstream-dev-mypy check unrelated to this PR and b) its probably not worth hacking something in here when there are bigger issues with the upstream zarr implementation to sort out.

[dcherian]

Thanks @TomAugspurger et al. This looks good. I have some minor comments, which I can address later today.

[dcherian]

Can this be removed at some point in the future? If so, it would be good to add a TODO

[TomAugspurger]

I'll look more closely later, but for now I think this will be required, following a deliberate change in zarr v3 consolidated metadata.

With v2 metadata, I think that consolidated happened at the store-level, and was all-or-nothing. If you have two Groups with Arrays, the consolidated metadata will be placed at the store root, and will contain everything:

# zarr v2

In [1]: import json, xarray as xr

In [2]: store = {}

In [3]: a = xr.tutorial.load_dataset("air_temperature")

In [4]: b = xr.tutorial.load_dataset("rasm")

In [5]: a.to_zarr(store=store, group="A")
/Users/tom/gh/zarr-developers/zarr-v2/.direnv/python-3.10/lib/python3.10/site-packages/xarray/core/dataset.py:2562: SerializationWarning: saving variable None with floating point data as an integer dtype without any _FillValue to use for NaNs
  return to_zarr(  # type: ignore[call-overload,misc]
Out[5]: <xarray.backends.zarr.ZarrStore at 0x11113edc0>

In [6]: b.to_zarr(store=store, group="B")
Out[6]: <xarray.backends.zarr.ZarrStore at 0x10cab2440>

In [7]: list(json.loads(store['.zmetadata'])['metadata'])
Out[7]:  # contains nodes from both A and B
['.zgroup',
 'A/.zattrs',
 'A/.zgroup',
 'A/air/.zarray',
 'A/air/.zattrs',
 'A/lat/.zarray',
 'A/lat/.zattrs',
 'A/lon/.zarray',
 'A/lon/.zattrs',
 'A/time/.zarray',
 'A/time/.zattrs',
 'B/.zattrs',
 'B/.zgroup',
 'B/Tair/.zarray',
 'B/Tair/.zattrs',
 'B/time/.zarray',
 'B/time/.zattrs',
 'B/xc/.zarray',
 'B/xc/.zattrs',
 'B/yc/.zarray',
 'B/yc/.zattrs']

With v3, consolidated metadata is scoped to a Group, so we can provide the group we want to consolidated (the zarr-python API does support "consolidate everything in the store at the root", but I don't think we want that because you'd need to open it at the root when reading, and I think it's kinda where for ds.to_zarr(group="A") to be reading / writing stuff outside of the A prefix).

[shoyer]

Potentially it would make sense to have two versions of consolidated metadata:

• Everything at a specific group/node level
• Everything in a group and all of its subgroups (i.e., for DataTree)

[TomAugspurger]

Agreed. zarr-developers/zarr-specs#309 has some discussion on adding a depth field to the spec for consolidated metadata. That's currently implicitly depth=None, which is everything below a group. depth=0 or 1 would be just the immediate children. That's not standardized or implemented anywhere yet, but the current implementation is forwards compatible and it shouldn't be a ton of effort.

[dcherian]

Let's get this in by the end of the week.

[jhamman]

👏 Thanks all! Especially @TomAugspurger for doing the lion's share of the work here.