Zarr v3 backend configuration

pyproject.toml

@@ -40,7 +40,7 @@ dependencies = [
     "PyYAML>=5.4",
     "h5py>=3.9.0",
     "hdmf>=4.1.1",
-    "hdmf_zarr>=0.11",
+    "hdmf_zarr>=0.12.1.dev40",
     "pynwb>=3.0.0",
     "pydantic>=2.0",
     "typing_extensions>=4.1.0",
@@ -53,7 +53,7 @@ dependencies = [
     "docstring-parser",  # For building json schema from method signatures
     "packaging",  # Issue 903
     "referencing", # for the json schema references
-    "numcodecs<0.16.0", # 0.16.0 is incompatible with zarr < 3: https://github.com/zarr-developers/numcodecs/issues/721
+    "zarr>=3.0",
 ]
 
 

src/neuroconv/tools/nwb_helpers/__init__.py

@@ -20,6 +20,7 @@
 from ._configuration_models._zarr_backend import ZarrBackendConfiguration
 from ._configuration_models._zarr_dataset_io import (
     AVAILABLE_ZARR_COMPRESSION_METHODS,
+    AVAILABLE_ZARR_FILTER_METHODS,
     ZarrDatasetIOConfiguration,
 )
 from ._configure_backend import configure_backend
@@ -38,6 +39,7 @@
 __all__ = [
     "AVAILABLE_HDF5_COMPRESSION_METHODS",
     "AVAILABLE_ZARR_COMPRESSION_METHODS",
+    "AVAILABLE_ZARR_FILTER_METHODS",
     "BACKEND_CONFIGURATIONS",
     "DATASET_IO_CONFIGURATIONS",
     "BACKEND_NWB_IO",

src/neuroconv/tools/nwb_helpers/_configuration_models/_base_dataset_io.py

@@ -5,7 +5,6 @@
 from typing import Any, Literal
 
 import h5py
-import numcodecs
 import numpy as np
 import zarr
 from hdmf import Container
@@ -25,6 +24,7 @@
 from pynwb.ecephys import ElectricalSeries
 from pynwb.image import ImageSeries
 from typing_extensions import Self
+from zarr.abc.codec import BytesBytesCodec
 
 from neuroconv.tools.hdmf import get_full_data_shape
 from neuroconv.tools.iterative_write import get_electrical_series_chunk_shape
@@ -127,10 +127,8 @@ class DatasetIOConfiguration(BaseModel, ABC):
             "For optimized writing speeds and minimal RAM usage, a total size of around 1 GB is recommended."
         ),
     )
-    compression_method: str | InstanceOf[h5py._hl.filters.FilterRefBase] | InstanceOf[numcodecs.abc.Codec] | None = (
-        Field(
-            description="The specified compression method to apply to this dataset. Set to `None` to disable compression.",
-        )
+    compression_method: str | InstanceOf[h5py._hl.filters.FilterRefBase] | InstanceOf[BytesBytesCodec] | None = Field(
+        description="The specified compression method to apply to this dataset. Set to `None` to disable compression.",
     )
     compression_options: dict[str, Any] | None = Field(
         default=None, description="The optional parameters to use for the specified compression method."

src/neuroconv/tools/nwb_helpers/_configuration_models/_zarr_dataset_io.py

@@ -2,78 +2,74 @@
 
 from typing import Any, Literal
 
-import numcodecs
-import zarr
 from hdmf import Container
-from pydantic import Field, InstanceOf, model_validator
+from pydantic import Field, InstanceOf, PositiveInt, model_validator
 from typing_extensions import Self
+from zarr.abc.codec import ArrayArrayCodec, BytesBytesCodec
+from zarr.codecs import BloscCodec, GzipCodec, ZstdCodec
+from zarr.codecs.numcodecs import BZ2, LZ4, LZMA, Shuffle, Zlib
 
 from ._base_dataset_io import DatasetIOConfiguration
 
-_base_zarr_codecs = set(zarr.codec_registry.keys())
-_lossy_zarr_codecs = set(("astype", "bitround", "quantize"))
-
-# These filters do nothing for us, or are things that ought to be implemented at lower HDMF levels
-# or indirectly using HDMF data structures
-_excluded_zarr_codecs = set(
-    (
-        "json2",  # no data savings
-        "pickle",  # no data savings
-        "vlen-utf8",  # enforced by HDMF
-        "vlen-array",  # enforced by HDMF
-        "vlen-bytes",  # enforced by HDMF
-        "msgpack2",  # think more on if we want to include this for variable length string datasets
-        "adler32",  # checksum
-        "crc32",  # checksum
-        "fixedscaleoffset",  # enforced indirectly by HDMF/PyNWB data types
-        "base64",  # unsure what this would ever be used for
-        "n5_wrapper",  # different data format
-        "pcodec",  # is erroneously imported before numcodecs 0.15, see https://numcodecs.readthedocs.io/en/stable/release.html?utm_source=chatgpt.com#id9
-    )
-)
-
-# Forbidding lossy codecs for now, but they could be allowed in the future with warnings?
-# (Users can always initialize and pass explicitly via code)
-_available_zarr_codecs = set(_base_zarr_codecs - _lossy_zarr_codecs - _excluded_zarr_codecs)
+# Curated mapping of string names to zarr v3 BytesBytesCodec classes.
+# Prefer native zarr codecs where available; fall back to numcodecs wrappers otherwise.
+AVAILABLE_ZARR_COMPRESSION_METHODS: dict[str, type[BytesBytesCodec]] = {
+    "gzip": GzipCodec,
+    "blosc": BloscCodec,
+    "zstd": ZstdCodec,
+    "bz2": BZ2,
+    "lzma": LZMA,
+    "zlib": Zlib,
+    "lz4": LZ4,
+    "shuffle": Shuffle,
+}
 
-AVAILABLE_ZARR_COMPRESSION_METHODS = {
-    codec_name: zarr.codec_registry[codec_name] for codec_name in _available_zarr_codecs
+# Curated mapping of string names to zarr v3 ArrayArrayCodec classes for filters.
+AVAILABLE_ZARR_FILTER_METHODS: dict[str, type[ArrayArrayCodec]] = {
+    "delta": __import__("zarr.codecs.numcodecs", fromlist=["Delta"]).Delta,
 }
 
 
 class ZarrDatasetIOConfiguration(DatasetIOConfiguration):
     """A data model for configuring options about an object that will become a Zarr Dataset in the file."""
 
     compression_method: (
-        Literal[tuple(AVAILABLE_ZARR_COMPRESSION_METHODS.keys())] | InstanceOf[numcodecs.abc.Codec] | None
+        Literal[tuple(AVAILABLE_ZARR_COMPRESSION_METHODS.keys())] | InstanceOf[BytesBytesCodec] | None
     ) = Field(
-        default="gzip",  # TODO: would like this to be 'auto'
+        default="gzip",
         description=(
             "The specified compression method to apply to this dataset. "
             "Can be either a string that matches an available method on your system, "
-            "or an instantiated numcodec.Codec object."
+            "or an instantiated zarr BytesBytesCodec object (e.g. zarr.codecs.GzipCodec(level=5)). "
             "Set to `None` to disable compression."
         ),
     )
-    # TODO: actually provide better schematic rendering of options. Only support defaults in GUIDE for now.
-    # Looks like they'll have to be hand-typed however... Can try parsing the numpy docstrings - no annotation typing.
     compression_options: dict[str, Any] | None = Field(
         default=None, description="The optional parameters to use for the specified compression method."
     )
-    filter_methods: (
-        list[Literal[tuple(AVAILABLE_ZARR_COMPRESSION_METHODS.keys())] | InstanceOf[numcodecs.abc.Codec]] | None
-    ) = Field(
-        default=None,
-        description=(
-            "The ordered collection of filtering methods to apply to this dataset prior to compression. "
-            "Each element can be either a string that matches an available method on your system, "
-            "or an instantiated numcodec.Codec object."
-            "Set to `None` to disable filtering."
-        ),
+    filter_methods: list[Literal[tuple(AVAILABLE_ZARR_FILTER_METHODS.keys())] | InstanceOf[ArrayArrayCodec]] | None = (
+        Field(
+            default=None,
+            description=(
+                "The ordered collection of filtering methods to apply to this dataset prior to compression. "
+                "Each element can be either a string that matches an available method on your system, "
+                "or an instantiated zarr ArrayArrayCodec object (e.g. zarr.codecs.numcodecs.Delta()). "
+                "Set to `None` to disable filtering."
+            ),
+        )
     )
     filter_options: list[dict[str, Any]] | None = Field(
         default=None, description="The optional parameters to use for each specified filter method."
     )
+    shard_shape: tuple[PositiveInt, ...] | None = Field(
+        default=None,
+        description=(
+            "The specified shape to use for sharding the dataset. "
+            "Each shard contains one or more chunks. When set, each axis must be >= the corresponding "
+            "chunk_shape axis, and chunk axes must evenly divide shard axes. "
+            "Set to `None` to disable sharding (default)."
+        ),
+    )
 
     def __str__(self) -> str:  # Inherited docstring from parent. noqa: D105
         string = super().__str__()
@@ -83,6 +79,8 @@ def __str__(self) -> str:  # Inherited docstring from parent. noqa: D105
             string += f"\n  filter options : {self.filter_options}"
         if self.filter_methods is not None or self.filter_options is not None:
             string += "\n"
+        if self.shard_shape is not None:
+            string += f"\n  shard shape : {self.shard_shape}\n"
 
         return string
 
@@ -109,21 +107,46 @@ def validate_filter_methods_and_options_length_match(cls, values: dict[str, Any]
 
         return values
 
+    @model_validator(mode="after")
+    def validate_shard_shape(self) -> Self:
+        if self.shard_shape is None or self.chunk_shape is None:
+            return self
+
+        if len(self.shard_shape) != len(self.chunk_shape):
+            raise ValueError(
+                f"Length of shard_shape ({len(self.shard_shape)}) does not match "
+                f"chunk_shape ({len(self.chunk_shape)}) for dataset at location '{self.location_in_file}'!"
+            )
+
+        if any(shard_axis < chunk_axis for shard_axis, chunk_axis in zip(self.shard_shape, self.chunk_shape)):
+            raise ValueError(
+                f"Some dimensions of the shard_shape {self.shard_shape} are smaller than the "
+                f"chunk_shape {self.chunk_shape} for dataset at location '{self.location_in_file}'!"
+            )
+
+        if any(shard_axis % chunk_axis != 0 for shard_axis, chunk_axis in zip(self.shard_shape, self.chunk_shape)):
+            raise ValueError(
+                f"Some dimensions of the chunk_shape {self.chunk_shape} do not evenly divide the "
+                f"shard_shape {self.shard_shape} for dataset at location '{self.location_in_file}'!"
+            )
+
+        return self
+
     def get_data_io_kwargs(self) -> dict[str, Any]:
         filters = None
         if self.filter_methods:
             filters = list()
             all_filter_options = self.filter_options or [dict() for _ in self.filter_methods]
             for filter_method, filter_options in zip(self.filter_methods, all_filter_options):
                 if isinstance(filter_method, str):
-                    filters.append(zarr.codec_registry[filter_method](**filter_options))
-                elif isinstance(filter_method, numcodecs.abc.Codec):
+                    filters.append(AVAILABLE_ZARR_FILTER_METHODS[filter_method](**filter_options))
+                elif isinstance(filter_method, ArrayArrayCodec):
                     filters.append(filter_method)
 
         if isinstance(self.compression_method, str):
             compression_options = self.compression_options or dict()
-            compressor = zarr.codec_registry[self.compression_method](**compression_options)
-        if isinstance(self.compression_method, numcodecs.abc.Codec):
+            compressor = AVAILABLE_ZARR_COMPRESSION_METHODS[self.compression_method](**compression_options)
+        elif isinstance(self.compression_method, BytesBytesCodec):
             compressor = self.compression_method
         elif self.compression_method is None:
             compressor = False
@@ -155,8 +178,13 @@ def from_neurodata_object_with_existing(
             neurodata_object=neurodata_object,
             dataset_name=dataset_name,
         )
-        compression_method = getattr(neurodata_object, dataset_name).compressor
-        filter_methods = getattr(neurodata_object, dataset_name).filters
+        dataset = getattr(neurodata_object, dataset_name)
+        # zarr v3: .compressors is a tuple of BytesBytesCodec; take first or None
+        compressors = dataset.compressors
+        compression_method = compressors[0] if compressors else None
+        # zarr v3: .filters is a tuple of ArrayArrayCodec
+        filters = dataset.filters
+        filter_methods = list(filters) if filters else None
         return cls(
             **kwargs,
             compression_method=compression_method,

src/neuroconv/tools/nwb_helpers/_configure_backend.py

@@ -65,7 +65,10 @@ def configure_backend(
             # we wrap each neurodata_object in a DataChunkIterator in order to support changes to the I/O settings.
             # For more detail, see https://github.com/hdmf-dev/hdmf/issues/1170.
             data_chunk_iterator_class = DataChunkIterator
-            data_chunk_iterator_kwargs = dict(buffer_size=math.prod(dataset_configuration.buffer_shape))
+            data_chunk_iterator_kwargs = dict(
+                buffer_size=math.prod(dataset_configuration.buffer_shape),
+                dtype=dataset_configuration.dtype,
+            )
 
         # Table columns
         if isinstance(neurodata_object, Data):

src/neuroconv/tools/nwb_helpers/_dataset_configuration.py

@@ -49,7 +49,7 @@ def _is_dataset_written_to_file(
     ) or (
         isinstance(candidate_dataset, zarr.Array)  # If the source data is a Zarr Array
         and backend == "zarr"
-        and Path(candidate_dataset.store.path).resolve()
+        and Path(candidate_dataset.store.root).resolve()
         == normalized_existing  # If the source Zarr 'file' is the appending NWBFile
     )
 
@@ -140,6 +140,16 @@ def get_default_dataset_io_configurations(
                 # Skip datasets with any zero-length axes
                 dataset_name = "data"
                 candidate_dataset = getattr(column, dataset_name)
+
+                # Skip datasets that are plain string lists (e.g., string columns in zarr v3)
+                if (
+                    backend == "zarr"
+                    and isinstance(candidate_dataset, list)
+                    and len(candidate_dataset) > 0
+                    and isinstance(candidate_dataset[0], str)
+                ):
+                    continue
+
                 full_shape = get_data_shape(data=candidate_dataset)
                 if any(axis_length == 0 for axis_length in full_shape):
                     continue
@@ -223,13 +233,22 @@ def get_existing_dataset_io_configurations(nwbfile: NWBFile) -> Generator[Datase
                 if any(isinstance(value, Container) for value in candidate_dataset):
                     continue  # Skip
 
-                # Skip when columns whose values are a reference type
+                # Skip over columns whose values are a reference type
                 if isinstance(column, TimeSeriesReferenceVectorData):
                     continue
 
                 # Skip datasets with any zero-length axes
                 dataset_name = "data"
                 candidate_dataset = getattr(column, dataset_name)
+
+                # Skip datasets that are plain string lists (e.g., string columns in zarr v3)
+                if (
+                    isinstance(candidate_dataset, list)
+                    and len(candidate_dataset) > 0
+                    and isinstance(candidate_dataset[0], str)
+                ):
+                    continue
+
                 full_shape = get_data_shape(data=candidate_dataset)
                 if any(axis_length == 0 for axis_length in full_shape):
                     continue

src/neuroconv/tools/testing/_mock/_mock_dataset_models.py

@@ -1,12 +1,13 @@
 from typing import Any, Iterable, Literal
 
 import h5py
-import numcodecs
 import numpy as np
+from zarr.abc.codec import ArrayArrayCodec, BytesBytesCodec
 
 from ...nwb_helpers import (
     AVAILABLE_HDF5_COMPRESSION_METHODS,
     AVAILABLE_ZARR_COMPRESSION_METHODS,
+    AVAILABLE_ZARR_FILTER_METHODS,
     HDF5BackendConfiguration,
     HDF5DatasetIOConfiguration,
     ZarrBackendConfiguration,
@@ -49,12 +50,13 @@ def mock_ZarrDatasetIOConfiguration(
     dtype: np.dtype = np.dtype("int16"),
     chunk_shape: tuple[int, ...] = (78_125, 64),  # ~10 MB
     buffer_shape: tuple[int, ...] = (1_250_000, 384),  # ~1 GB
-    compression_method: Literal[tuple(AVAILABLE_ZARR_COMPRESSION_METHODS.keys())] | numcodecs.abc.Codec | None = "gzip",
+    compression_method: Literal[tuple(AVAILABLE_ZARR_COMPRESSION_METHODS.keys())] | BytesBytesCodec | None = "gzip",
     compression_options: dict[str, Any] | None = None,
     filter_methods: (
-        Iterable[Literal[tuple(AVAILABLE_ZARR_COMPRESSION_METHODS.keys())] | numcodecs.abc.Codec | None] | None
+        Iterable[Literal[tuple(AVAILABLE_ZARR_FILTER_METHODS.keys())] | ArrayArrayCodec | None] | None
     ) = None,
     filter_options: Iterable[dict[str, Any]] | None = None,
+    shard_shape: tuple[int, ...] | None = None,
 ) -> ZarrDatasetIOConfiguration:
     """Mock object of a ZarrDatasetIOConfiguration with NeuroPixel-like values to show chunk/buffer recommendations."""
     return ZarrDatasetIOConfiguration(
@@ -69,6 +71,7 @@ def mock_ZarrDatasetIOConfiguration(
         compression_options=compression_options,
         filter_methods=filter_methods,
         filter_options=filter_options,
+        shard_shape=shard_shape,
     )