Read zarr v2 NWB files with zarr-python 3

src/spikeinterface/extractors/nwbextractors.py

@@ -137,6 +137,25 @@ def read_nwbfile(
     if file_path is None and file is None:
         raise ValueError("Provide either file_path or file")
 
+    if (
+        backend == "zarr"
+        and file is None
+        and _get_zarr_format(file_path, stream_mode, storage_options) == 2
+        and _has_zarr_python_3()
+    ):
+        from hdmf_zarr import NWBZarrV2IO
+
+        # hdmf-zarr>=0.14 reads zarr v2 files with a separate read-only IO. Files written by hdmf-zarr<0.14
+        # pickle their object references (e.g. the "group" column of the electrodes table), so pickle is allowed
+        io = NWBZarrV2IO(
+            path=str(file_path),
+            mode="r",
+            load_namespaces=True,
+            allow_pickle=True,
+            storage_options=storage_options,
+        )
+        return io.read()
+
     open_file = read_file_from_backend(
         file_path=file_path,
         file=file,
@@ -303,19 +322,45 @@ def _get_backend_from_local_file(file_path: str | Path) -> str:
         else:
             raise RuntimeError(f"{file_path} is not a valid HDF5 file!")
     elif file_path.is_dir():
-        try:
-            import zarr
-
-            # `mode` is keyword-only in zarr v3, and its groups are not context managers
-            zarr.open(file_path, mode="r")
-            backend = "zarr"
-        except:
+        # The root metadata file is checked instead of opening the folder, because zarr>=3 cannot open
+        # the consolidated metadata of the zarr v2 files written by hdmf-zarr<0.14
+        if _get_zarr_format(file_path) is None:
             raise RuntimeError(f"{file_path} is not a valid Zarr folder!")
+        backend = "zarr"
     else:
         raise RuntimeError(f"File {file_path} is not an existing file or folder!")
     return backend
 
 
+def _get_zarr_format(file_path: str | Path, stream_mode=None, storage_options: dict | None = None) -> int | None:
+    """
+    Return the zarr format (2 or 3) of a zarr folder, local or streamed with stream_mode="zarr", from its
+    root metadata file, or None if the folder is not a zarr group.
+    """
+    if stream_mode == "zarr":
+        import fsspec
+
+        file_system, root = fsspec.core.url_to_fs(str(file_path), **(storage_options or {}))
+        root = root.rstrip("/")
+        exists = lambda name: file_system.exists(f"{root}/{name}")
+    elif stream_mode is None:
+        folder_path = Path(file_path)
+        exists = lambda name: (folder_path / name).is_file()
+    else:
+        return None
+    if exists("zarr.json"):
+        return 3
+    if exists(".zgroup"):
+        return 2
+    return None
+
+
+def _has_zarr_python_3() -> bool:
+    import zarr
+
+    return int(zarr.__version__.split(".")[0]) >= 3
+
+
 def _find_neurodata_type_from_backend(group, path="", result=None, neurodata_type="ElectricalSeries", backend="hdf5"):
     """
     Recursively searches for groups with the specified neurodata_type hdf5 or zarr object,
@@ -402,6 +447,30 @@ def _storage_backend(file_path=None, file=None, stream_mode=None):
             return _get_backend_from_local_file(file_path)
         return "zarr" if stream_mode == "zarr" else "hdf5"
 
+    @staticmethod
+    def _reading_method(use_pynwb, file_path=None, file=None, stream_mode=None, storage_options=None):
+        # With zarr-python 3, the in-house zarr reader only reads zarr v3 files. Zarr v2 files (written by
+        # hdmf-zarr<0.14) are read with pynwb, where hdmf-zarr keeps the backward compatibility code for them.
+        if use_pynwb:
+            return "use_pynwb"
+        backend = _NWBReader._storage_backend(file_path, file, stream_mode)
+        if (
+            backend == "zarr"
+            and _has_zarr_python_3()
+            and _get_zarr_format(file_path, stream_mode, storage_options) == 2
+        ):
+            if not HAVE_PYNWB or importlib.util.find_spec("hdmf_zarr") is None:
+                raise ImportError(
+                    "To read zarr v2 NWB files with zarr-python 3, install pynwb and hdmf-zarr>=0.14, which requires "
+                    'Python 3.12 or newer: \n\n pip install pynwb "hdmf-zarr>=0.14"\n\n'
+                )
+            warnings.warn(
+                f"{file_path} is a zarr v2 NWB file, which is read with pynwb. This can be slower than "
+                "the default reader, which only reads zarr v3 files."
+            )
+            return "use_pynwb"
+        return f"use_{backend}"
+
     def _open_handle(self, *, file_path, file, stream_mode, cache, stream_cache_path, storage_options=None):
         # Open the file handle: a pynwb NWBFile, or a raw h5py.File / zarr group.
         if self.reading_method == "use_pynwb":
@@ -1002,10 +1071,7 @@ def __init__(
         if use_pynwb and not HAVE_PYNWB:
             raise ImportError(self.installation_mesg)
 
-        if use_pynwb:
-            reading_method = "use_pynwb"  # the reader detects hdf5 vs zarr itself for pynwb
-        else:
-            reading_method = f"use_{_NWBReader._storage_backend(file_path, file, self.stream_mode)}"
+        reading_method = _NWBReader._reading_method(use_pynwb, file_path, file, self.stream_mode, storage_options)
 
         self._reader = _NWBReader(
             reading_method=reading_method,
@@ -1317,10 +1383,9 @@ def __init__(
         if use_pynwb and not HAVE_PYNWB:
             raise ImportError(self.installation_mesg)
 
-        if use_pynwb:
-            reading_method = "use_pynwb"
-        else:
-            reading_method = f"use_{_NWBReader._storage_backend(file_path, stream_mode=stream_mode)}"
+        reading_method = _NWBReader._reading_method(
+            use_pynwb, file_path, stream_mode=stream_mode, storage_options=storage_options
+        )
 
         self._reader = _NWBReader(
             reading_method=reading_method,
@@ -1725,10 +1790,7 @@ def __init__(
         if use_pynwb and not HAVE_PYNWB:
             raise ImportError(self.installation_mesg)
 
-        if use_pynwb:
-            reading_method = "use_pynwb"
-        else:
-            reading_method = f"use_{_NWBReader._storage_backend(file_path, file, stream_mode)}"
+        reading_method = _NWBReader._reading_method(use_pynwb, file_path, file, stream_mode, storage_options)
 
         self._reader = _NWBReader(reading_method=reading_method, timeseries_path=timeseries_path)
         self._reader.load_timeseries(

src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py

@@ -3,18 +3,12 @@
 
 import pytest
 import numpy as np
-import zarr
 
 from spikeinterface import load
 from spikeinterface.core.testing import check_recordings_equal
 from spikeinterface.core.testing import check_recordings_equal, check_sortings_equal
 from spikeinterface.extractors.extractor_classes import NwbRecordingExtractor, NwbSortingExtractor
 
-# TODO: remove once streaming zarr is supported with zarr>=3
-skip_zarr_3 = pytest.mark.skipif(
-    int(zarr.__version__.split(".")[0]) >= 3, reason="Streaming zarr is not supported yet with zarr>=3"
-)
-
 
 @pytest.mark.streaming_extractors
 @pytest.mark.parametrize("cache", [True, False])  # Test with and without cache
@@ -191,7 +185,6 @@ def test_sorting_s3_nwb_remfile(tmp_path):
 
 
 @pytest.mark.streaming_extractors
-@skip_zarr_3
 def test_sorting_s3_nwb_zarr(tmp_path):
     file_path = (
         "s3://aind-open-data/ecephys_625749_2022-08-03_15-15-06_nwb_2023-05-16_16-34-55/"