zarr-python v3 compatibility

kerchunk/hdf.py

@@ -1,7 +1,7 @@
 import base64
 import io
 import logging
-from typing import Union, BinaryIO, Any, cast
+from typing import Union, BinaryIO
 from packaging.version import Version
 
 import fsspec.core
@@ -11,7 +11,7 @@
 import numcodecs
 
 from .codecs import FillStringsCodec
-from .utils import _encode_for_JSON
+from .utils import _encode_for_JSON, encode_fill_value
 
 try:
     import h5py
@@ -22,12 +22,6 @@
         "for more details."
     )
 
-# try:
-#     from zarr.meta import encode_fill_value
-# except ModuleNotFoundError:
-#     # https://github.com/zarr-developers/zarr-python/issues/2021
-#     from zarr.v2.meta import encode_fill_value
-
 lggr = logging.getLogger("h5-to-zarr")
 _HIDDEN_ATTRS = {  # from h5netcdf.attrs
     "REFERENCE_LIST",
@@ -504,7 +498,6 @@ def _translator(
                 lggr.debug(f"Created Zarr array: {za}")
                 self._transfer_attrs(h5obj, za)
 
-                # za.attrs["_ARRAY_DIMENSIONS"] = adims
                 lggr.debug(f"_ARRAY_DIMENSIONS = {adims}")
 
                 if "data" in kwargs:
@@ -705,43 +698,3 @@ def _is_netcdf_variable(dataset: h5py.Dataset):
 def has_visititems_links():
     return hasattr(h5py.Group, "visititems_links")
 
-
-def encode_fill_value(v: Any, dtype: np.dtype, object_codec: Any = None) -> Any:
-    # early out
-    if v is None:
-        return v
-    if dtype.kind == "V" and dtype.hasobject:
-        if object_codec is None:
-            raise ValueError("missing object_codec for object array")
-        v = object_codec.encode(v)
-        v = str(base64.standard_b64encode(v), "ascii")
-        return v
-    if dtype.kind == "f":
-        if np.isnan(v):
-            return "NaN"
-        elif np.isposinf(v):
-            return "Infinity"
-        elif np.isneginf(v):
-            return "-Infinity"
-        else:
-            return float(v)
-    elif dtype.kind in "ui":
-        return int(v)
-    elif dtype.kind == "b":
-        return bool(v)
-    elif dtype.kind in "c":
-        c = cast(np.complex128, np.dtype(complex).type())
-        v = (
-            encode_fill_value(v.real, c.real.dtype, object_codec),
-            encode_fill_value(v.imag, c.imag.dtype, object_codec),
-        )
-        return v
-    elif dtype.kind in "SV":
-        v = str(base64.standard_b64encode(v), "ascii")
-        return v
-    elif dtype.kind == "U":
-        return v
-    elif dtype.kind in "mM":
-        return int(v.view("i8"))
-    else:
-        return v

kerchunk/netCDF3.py

@@ -1,4 +1,5 @@
 from functools import reduce
+from packaging.version import Version
 from operator import mul
 
 import numpy as np
@@ -167,7 +168,13 @@ def translate(self):
         import zarr
 
         out = self.out
-        z = zarr.open(out, mode="w", zarr_format=2)
+        if Version(zarr.__version__) < Version("3.0.0.a0"):
+            store = zarr.storage.KVStore(out)
+            z = zarr.group(store=store, overwrite=True)
+        else:
+            store = zarr.storage.MemoryStore(mode="a", store_dict=out)
+            z = zarr.open(store, mode="w", zarr_format=2)
+
         for dim, var in self.variables.items():
             if dim in self.chunks:
                 shape = self.chunks[dim][-1]
@@ -197,7 +204,7 @@ def translate(self):
                     dtype=var.data.dtype,
                     fill_value=fill,
                     chunks=shape,
-                    compression=None,
+                    compressor=None,
                 )
                 part = ".".join(["0"] * len(shape)) or "0"
                 k = f"{dim}/{part}"
@@ -251,7 +258,7 @@ def translate(self):
                     dtype=base,
                     fill_value=fill,
                     chunks=(1,) + dtype.shape,
-                    compression=None,
+                    compressor=None,
                 )
                 arr.attrs.update(
                     {

kerchunk/tests/test_hdf.py

@@ -1,5 +1,6 @@
-import asyncio
+from typing import Any
 import fsspec
+import json
 import os.path as osp
 
 import fsspec.implementations
@@ -23,25 +24,29 @@ async def list_dir(store, path):
     [x async for x in store.list_dir(path)]
 
 
-def create_store(test_dict: dict):
+def create_store(test_dict: dict, remote_options: Any = None):
     if Version(zarr.__version__) < Version("3.0.0.a0"):
         return fsspec.get_mapper(
-            "reference://", fo=test_dict, remote_protocol="s3", remote_options=so
+            "reference://", fo=test_dict, remote_protocol="s3", remote_options=remote_options
         )
     else:
-        fs = fsspec.implementations.reference.ReferenceFileSystem(fo=test_dict)
+        fs = fsspec.implementations.reference.ReferenceFileSystem(fo=test_dict, remote_options=remote_options)
         return zarr.storage.RemoteStore(fs, mode="r")
 
 
 def test_single():
     """Test creating references for a single HDF file"""
-    url = "s3://noaa-nwm-retro-v2.0-pds/full_physics/2017/201704010000.CHRTOUT_DOMAIN1.comp"
+    #url = "s3://noaa-nwm-retro-v2.0-pds/full_physics/2017/201704010000.CHRTOUT_DOMAIN1.comp"
+    url = "s3://noaa-nos-ofs-pds/ngofs2/netcdf/202410/ngofs2.t03z.20241001.2ds.f020.nc"
     so = dict(anon=True, default_fill_cache=False, default_cache_type="none")
 
     with fsspec.open(url, **so) as f:
         h5chunks = SingleHdf5ToZarr(f, url, storage_options=so)
         test_dict = h5chunks.translate()
 
+    with open("test_dict.json", "w") as f:
+        json.dump(test_dict, f)
+
     store = create_store(test_dict)
 
     ds = xr.open_dataset(store, engine="zarr", zarr_format=2, backend_kwargs=dict(consolidated=False))
@@ -97,6 +102,7 @@ def test_multizarr(generate_mzz):
     """Test creating a combined reference file with MultiZarrToZarr"""
     mzz = generate_mzz
     test_dict = mzz.translate()
+
     store = create_store(test_dict)
     ds = xr.open_dataset(store, engine="zarr", zarr_format=2, backend_kwargs=dict(consolidated=False))
 

kerchunk/tests/test_netcdf.py

@@ -1,4 +1,5 @@
 import os
+from typing import Any
 
 
 import fsspec
@@ -7,6 +8,8 @@
 import pytest
 from kerchunk import netCDF3
 
+import zarr
+
 xr = pytest.importorskip("xarray")
 
 
@@ -24,16 +27,29 @@
 )
 
 
+def create_store(test_dict: dict, remote_options: Any = None):
+    if Version(zarr.__version__) < Version("3.0.0.a0"):
+        return fsspec.get_mapper(
+            "reference://", fo=test_dict, remote_protocol="s3", remote_options=remote_options
+        )
+    else:
+        fs = fsspec.implementations.reference.ReferenceFileSystem(fo=test_dict, remote_options=remote_options)
+        return zarr.storage.RemoteStore(fs, mode="r")
+
+
 def test_one(m):
     m.pipe("data.nc3", bdata)
     h = netCDF3.netcdf_recording_file("memory://data.nc3")
     out = h.translate()
+
+    store = create_store(out, remote_options={"remote_protocol": "memory"})
+
     ds = xr.open_dataset(
-        "reference://",
+        store,
         engine="zarr",
         backend_kwargs={
             "consolidated": False,
-            "storage_options": {"fo": out, "remote_protocol": "memory"},
+            "zarr_format": "2",
         },
     )
     assert (ds.data == data).all()

kerchunk/utils.py

@@ -1,11 +1,13 @@
 import base64
 import copy
 import itertools
+from typing import Any, cast
 import warnings
 
 import ujson
 
 import fsspec
+import numpy as np
 import zarr
 
 
@@ -134,6 +136,48 @@ def _encode_for_JSON(store):
     return store
 
 
+
+def encode_fill_value(v: Any, dtype: np.dtype, object_codec: Any = None) -> Any:
+    # early out
+    if v is None:
+        return v
+    if dtype.kind == "V" and dtype.hasobject:
+        if object_codec is None:
+            raise ValueError("missing object_codec for object array")
+        v = object_codec.encode(v)
+        v = str(base64.standard_b64encode(v), "ascii")
+        return v
+    if dtype.kind == "f":
+        if np.isnan(v):
+            return "NaN"
+        elif np.isposinf(v):
+            return "Infinity"
+        elif np.isneginf(v):
+            return "-Infinity"
+        else:
+            return float(v)
+    elif dtype.kind in "ui":
+        return int(v)
+    elif dtype.kind == "b":
+        return bool(v)
+    elif dtype.kind in "c":
+        c = cast(np.complex128, np.dtype(complex).type())
+        v = (
+            encode_fill_value(v.real, c.real.dtype, object_codec),
+            encode_fill_value(v.imag, c.imag.dtype, object_codec),
+        )
+        return v
+    elif dtype.kind in "SV":
+        v = str(base64.standard_b64encode(v), "ascii")
+        return v
+    elif dtype.kind == "U":
+        return v
+    elif dtype.kind in "mM":
+        return int(v.view("i8"))
+    else:
+        return v
+
+
 def do_inline(store, threshold, remote_options=None, remote_protocol=None):
     """Replace short chunks with the value of that chunk and inline metadata