From a4202283c1e6a187a27fbceeea935b296bbe9140 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Fri, 16 Oct 2026 08:44:08 -0300 Subject: [PATCH 1/2] Read zarr v2 NWB files with zarr-python 3 Route zarr v2 NWB files, local or streamed with stream_mode="zarr", to pynwb, where hdmf-zarr's NWBZarrV2IO keeps the backward compatibility, so the in-house reader only needs to read zarr v3. --- .../extractors/nwbextractors.py | 95 +++++++++++++++---- 1 file changed, 76 insertions(+), 19 deletions(-) diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index 0d1946d39a..d5b99334c3 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/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,25 @@ 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): + # 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 _get_zarr_format(file_path, stream_mode, storage_options) == 2: + if not HAVE_PYNWB: + raise ImportError( + "To read zarr v2 NWB files, install pynwb and hdmf-zarr: \n\n pip install pynwb hdmf-zarr\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 +1066,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 +1378,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 +1785,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( From 102d5f6087aad3bab39d03148a9425c649f847cd Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Fri, 16 Oct 2026 09:58:10 -0300 Subject: [PATCH 2/2] Route zarr v2 NWB files to pynwb only with zarr-python 3 With zarr-python 2 the in-house reader reads zarr v2 files, and the pynwb streaming path passes an FSStore that hdmf-zarr<0.14 rejects. Raise a clear error when hdmf-zarr is missing, and stop skipping the zarr streaming test with zarr-python 3. --- src/spikeinterface/extractors/nwbextractors.py | 15 ++++++++++----- .../tests/test_nwbextractors_streaming.py | 7 ------- 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index d5b99334c3..5c6f2ea6ce 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/src/spikeinterface/extractors/nwbextractors.py @@ -449,15 +449,20 @@ def _storage_backend(file_path=None, file=None, stream_mode=None): @staticmethod def _reading_method(use_pynwb, file_path=None, file=None, stream_mode=None, storage_options=None): - # 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. + # 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 _get_zarr_format(file_path, stream_mode, storage_options) == 2: - if not HAVE_PYNWB: + 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, install pynwb and hdmf-zarr: \n\n pip install pynwb hdmf-zarr\n\n" + "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 " diff --git a/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py b/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py index 9a7b2c0c67..161040cbaf 100644 --- a/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py +++ b/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/"