Ensure KeyError raised for zarr datasets missing dim names

[oliverwm1]

Add check to ensure that a KeyError is raised for zarr datasets that were not created by xarray.

• Closes ValueError raised when opening a 'plain' zarr with xarray+zarr 3 #9970
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst

[welcome]

Thank you for opening this pull request! It may take us a few days to respond here, so thank you for being patient.
If you have questions, some answers may be found in our contributing guidelines.

[oliverwm1]

Locally I am getting a failure I don't understand for default_zarr_format=2 case, related to the NCZarr code path.

/Users/oliverwm/repos/xarray/xarray/backends/zarr.py:1495: in open_zarr
    ds = open_dataset(
/Users/oliverwm/repos/xarray/xarray/backends/api.py:685: in open_dataset
    backend_ds = backend.open_dataset(
/Users/oliverwm/repos/xarray/xarray/backends/zarr.py:1585: in open_dataset
    ds = store_entrypoint.open_dataset(
/Users/oliverwm/repos/xarray/xarray/backends/store.py:44: in open_dataset
    vars, attrs = filename_or_obj.load()
/Users/oliverwm/repos/xarray/xarray/backends/common.py:312: in load
    (_decode_variable_name(k), v) for k, v in self.get_variables().items()
/Users/oliverwm/repos/xarray/xarray/backends/zarr.py:861: in get_variables
    return FrozenDict((k, self.open_store_variable(k)) for k in self.array_keys())
/Users/oliverwm/repos/xarray/xarray/core/utils.py:457: in FrozenDict
    return Frozen(dict(*args, **kwargs))
/Users/oliverwm/repos/xarray/xarray/backends/zarr.py:861: in <genexpr>
    return FrozenDict((k, self.open_store_variable(k)) for k in self.array_keys())
/Users/oliverwm/repos/xarray/xarray/backends/zarr.py:819: in open_store_variable
    dimensions, attributes = _get_zarr_dims_and_attrs(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

zarr_obj = <Array file:///private/var/folders/kn/7wksmxbj6zdgztrgmk6y12w00000gp/T/pytest-of-oliverwm/pytest-221/test_raises_key_error_on_inval0/tmp.zarr/bar shape=(3, 5) dtype=float32>
dimension_key = '_ARRAY_DIMENSIONS', try_nczarr = True

    def _get_zarr_dims_and_attrs(zarr_obj, dimension_key, try_nczarr):
        # Zarr V3 explicitly stores the dimension names in the metadata
        try:
            # if this exists, we are looking at a Zarr V3 array
            # convert None to empty tuple
            dimensions = zarr_obj.metadata.dimension_names or ()
        except AttributeError:
            # continue to old code path
            pass
        else:
            attributes = dict(zarr_obj.attrs)
            if len(zarr_obj.shape) != len(dimensions):
                raise KeyError(
                    "Zarr object is missing the `dimension_names` metadata which is "
                    "required for xarray to determine variable dimensions."
                )
            return dimensions, attributes
    
        # Zarr arrays do not have dimensions. To get around this problem, we add
        # an attribute that specifies the dimension. We have to hide this attribute
        # when we send the attributes to the user.
        # zarr_obj can be either a zarr group or zarr array
        try:
            # Xarray-Zarr
            dimensions = zarr_obj.attrs[dimension_key]
        except KeyError as e:
            if not try_nczarr:
                raise KeyError(
                    f"Zarr object is missing the attribute `{dimension_key}`, which is "
                    "required for xarray to determine variable dimensions."
                ) from e
    
            # NCZarr defines dimensions through metadata in .zarray
            zarray_path = os.path.join(zarr_obj.path, ".zarray")
>           zarray = json.loads(zarr_obj.store[zarray_path])
E           TypeError: 'LocalStore' object is not subscriptable

/Users/oliverwm/repos/xarray/xarray/backends/zarr.py:401: TypeError
==================================================================== short test summary info =====================================================================
FAILED test_backends.py::test_raises_key_error_on_invalid_zarr_store[2] - TypeError: 'LocalStore' object is not subscriptable
========================================================== 1 failed, 1 passed, 2277 deselected in 0.57s ==========================================================

@jhamman

[oliverwm1]

Test still failing for case zarr-python>=3 and default_zarr_format=2. With latest change, test fails because LocalStore.get is an async function in zarr-python 3:

/home/runner/work/xarray/xarray/xarray/tests/test_backends.py:6143: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/home/runner/work/xarray/xarray/xarray/backends/zarr.py:1495: in open_zarr
    ds = open_dataset(
/home/runner/work/xarray/xarray/xarray/backends/api.py:685: in open_dataset
    backend_ds = backend.open_dataset(
/home/runner/work/xarray/xarray/xarray/backends/zarr.py:1585: in open_dataset
    ds = store_entrypoint.open_dataset(
/home/runner/work/xarray/xarray/xarray/backends/store.py:44: in open_dataset
    vars, attrs = filename_or_obj.load()
/home/runner/work/xarray/xarray/xarray/backends/common.py:312: in load
    (_decode_variable_name(k), v) for k, v in self.get_variables().items()
/home/runner/work/xarray/xarray/xarray/backends/zarr.py:861: in get_variables
    return FrozenDict((k, self.open_store_variable(k)) for k in self.array_keys())
/home/runner/work/xarray/xarray/xarray/core/utils.py:457: in FrozenDict
    return Frozen(dict(*args, **kwargs))
/home/runner/work/xarray/xarray/xarray/backends/zarr.py:861: in <genexpr>
    return FrozenDict((k, self.open_store_variable(k)) for k in self.array_keys())
/home/runner/work/xarray/xarray/xarray/backends/zarr.py:819: in open_store_variable
    dimensions, attributes = _get_zarr_dims_and_attrs(
/home/runner/work/xarray/xarray/xarray/backends/zarr.py:401: in _get_zarr_dims_and_attrs
    zarray = json.loads(zarr_obj.store.get(zarray_path))

E               TypeError: the JSON object must be str, bytes or bytearray, not coroutine

/home/runner/micromamba/envs/xarray-tests/lib/python3.12/json/__init__.py:339: TypeError

Not sure best solution—should we check zarr version here and handle async function appropriately? Is there a way to get this info that has same API in zarr-python 2 and 3? Or can we just not test with this combination of zarr-python version and default zarr format?

@jhamman
@malmans2 tagging you since it is nczarr-related code raising this issue.

[d-v-b]

would you find it useful if zarr let you do this with a single function? because we could totally add that

[oliverwm1]

I've not needed to access the 'zarray' info before, so that wouldn't make a big difference to me. But it seems nczarr does use it and I was just trying to get this existing code to work with zarr-python 3.

[joshmoore]

If there's no check previously guaranteeing that this is zarr_format=2, an additional check for zarray_paths ending with zarr.json will likely be needed.

[oliverwm1]

I believe the first try/except block in this function is intended to ensure we have zarr_format=2 here:

https://github.com/oliverwm1/xarray/blob/c1a1c70310f256a1125b08f089e5688075f89d78/xarray/backends/zarr.py#L368-L375

[dcherian]

Thanks! I'll merge once we fix our test failures.

[welcome]

Congratulations on completing your first pull request! Welcome to Xarray! We are proud of you, and hope to see you again!