Check zarr format before opening.

PiperOrigin-RevId: 791314170

Author: sylmak

xarray_tensorstore.py

@@ -21,9 +21,11 @@
 from typing import Optional, TypeVar
 
 import numpy as np
+import packaging
 import tensorstore
 import xarray
 from xarray.core import indexing
+import zarr
 
 
 __version__ = '0.1.5'  # keep in sync with setup.py
@@ -176,12 +178,12 @@ def read(xarraydata: XarrayData, /) -> XarrayData:
 _DEFAULT_STORAGE_DRIVER = 'file'
 
 
-def _zarr_spec_from_path(path: str) -> ...:
+def _zarr_spec_from_path(path: str, zarr_format: int) -> ...:
   if re.match(r'\w+\://', path):  # path is a URI
     kv_store = path
   else:
     kv_store = {'driver': _DEFAULT_STORAGE_DRIVER, 'path': path}
-  return {'driver': 'zarr', 'kvstore': kv_store}
+  return {'driver': f'zarr{zarr_format}', 'kvstore': kv_store}
 
 
 def _raise_if_mask_and_scale_used_for_data_vars(ds: xarray.Dataset):
@@ -207,6 +209,14 @@ def _raise_if_mask_and_scale_used_for_data_vars(ds: xarray.Dataset):
         )
 
 
+def _get_zarr_format(path: str) -> int:
+  """Returns the Zarr format of the given path."""
+  if packaging.version.parse(zarr.__version__).major >= 3:
+    return zarr.open_group(path, mode='r').metadata.zarr_format
+  else:
+    return 2
+
+
 def open_zarr(
     path: str,
     *,
@@ -271,7 +281,10 @@ def open_zarr(
     # incorrect data values.
     _raise_if_mask_and_scale_used_for_data_vars(ds)
 
-  specs = {k: _zarr_spec_from_path(os.path.join(path, k)) for k in ds}
+  zarr_format = _get_zarr_format(path)
+  specs = {
+      k: _zarr_spec_from_path(os.path.join(path, k), zarr_format) for k in ds
+  }
   array_futures = {
       k: tensorstore.open(spec, read=True, write=write, context=context)
       for k, spec in specs.items()

xarray_tensorstore_test.py

@@ -14,12 +14,13 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
+import packaging
 import pandas as pd
 import pytest
 import tensorstore
 import xarray
-from xarray.core import indexing
 import xarray_tensorstore
+import zarr
 
 
 class XarrayTensorstoreTest(parameterized.TestCase):
@@ -145,13 +146,19 @@ def test_open_zarr_from_uri(self):
     opened = xarray_tensorstore.open_zarr('file://' + path)
     xarray.testing.assert_identical(source, opened)
 
-  def test_read_dataset(self):
+  @parameterized.parameters(
+      {'zarr_format': 2},
+      {'zarr_format': 3},
+  )
+  def test_read_dataset(self, zarr_format):
+    if packaging.version.parse(zarr.__version__).major < 3 and zarr_format == 3:
+      self.skipTest('zarr format 3 is not supported in zarr < 3.0.0')
     source = xarray.Dataset(
         {'baz': (('x', 'y', 'z'), np.arange(24).reshape(2, 3, 4))},
         coords={'x': np.arange(2)},
     )
     path = self.create_tempdir().full_path
-    source.chunk().to_zarr(path)
+    source.chunk().to_zarr(path, zarr_format=zarr_format)
 
     opened = xarray_tensorstore.open_zarr(path)
     read = xarray_tensorstore.read(opened)
@@ -160,15 +167,21 @@ def test_read_dataset(self):
     self.assertIsNotNone(read.variables['baz']._data.future)
     xarray.testing.assert_identical(read, source)
 
-  def test_read_dataarray(self):
+  @parameterized.parameters(
+      {'zarr_format': 2},
+      {'zarr_format': 3},
+  )
+  def test_read_dataarray(self, zarr_format):
+    if packaging.version.parse(zarr.__version__).major < 3  and zarr_format == 3:
+      self.skipTest('zarr format 3 is not supported in zarr < 3.0.0')
     source = xarray.DataArray(
         np.arange(24).reshape(2, 3, 4),
         dims=('x', 'y', 'z'),
         name='baz',
         coords={'x': np.arange(2)},
     )
     path = self.create_tempdir().full_path
-    source.to_dataset().chunk().to_zarr(path)
+    source.to_dataset().chunk().to_zarr(path, zarr_format=zarr_format)
 
     opened = xarray_tensorstore.open_zarr(path)['baz']
     read = xarray_tensorstore.read(opened)