feat(dtype): read nested data types in a Context, with their JSON location

changes/4382.bugfix.md

@@ -0,0 +1,17 @@
+Made the `fill_value` check on the write path substantially cheaper.
+`NDBuffer.all_equal` runs on every chunk write when `write_empty_chunks` is
+false, which is the default, and it dominated the cost of writing chunks.
+
+Two changes. The bit-pattern comparison used for a zero `fill_value`, which
+keeps `-0.0` distinct from `0.0`, now views the buffer as an unsigned integer
+of the same width instead of as a void dtype; NumPy has vectorised integer
+comparison loops but no void one, so the void form fell back to a generic
+elementwise path. Separately, the scan now walks slices of the leading axis,
+doubling in length, and returns on the first slice that differs, so a buffer
+that is not uniformly `fill_value` no longer has to be read in full. The
+leading axis is used rather than a flattened view because a chunk is normally
+a strided view into a larger array and cannot be flattened without a copy. A
+buffer that really is uniform is still covered in the same number of element
+comparisons.
+
+Results are unchanged, including for `-0.0`, NaN, and non-contiguous buffers.

changes/4385.bugfix.md

@@ -0,0 +1 @@
+Zarr format 2 arrays whose data type is a single-byte boolean or integer, or raw bytes, written with a `<` or `>` byte order (e.g. `<i1`, `>b1`, `<V4`, as written by netCDF-C, Zarr.jl and jzarr) can be read again, including as fields of structured data types. The byte order of these data types is not relevant, and metadata is written back with the canonical `|` spelling. Also, `datetime64` and `timedelta64` data types no longer accept an object codec in Zarr format 2, and a structured data type with an unrecognized field data type is rejected with `DataTypeValidationError`.

changes/4387.feature.md

@@ -0,0 +1 @@
+Added `zarr.registry.Context`, the extensions that Zarr metadata is read with; for now it holds the data type registry. A structured data type now resolves its field data types from the context it is read in, instead of always from the default registry. `get_data_type_from_json`, `get_data_type_from_native_dtype`, `parse_dtype`, `parse_data_type`, `Structured.from_json` and `Struct.from_json` take an optional `context` (default: `Context.default()`). A structured data type with a field whose data type matches no data type now raises `zarr.errors.NestedDataTypeValidationError` with the JSON Pointer of that field (e.g. `at /configuration/fields/1/data_type`), instead of reporting that the whole structured data type matches nothing. Data types registered through the `zarr.data_type` entry point group are now loaded; previously they were collected but never registered.

examples/custom_dtype/custom_dtype.py

@@ -22,7 +22,7 @@
 import pytest
 
 import zarr
-from zarr.dtype import ZDType, check_dtype_spec_v2, data_type_registry
+from zarr.dtype import ZDType, check_dtype_spec_no_object_codec_v2, data_type_registry
 from zarr.errors import DataTypeValidationError
 from zarr.types import JSON, DTypeConfig_V2, DTypeJSON, ZarrFormat
 
@@ -79,9 +79,7 @@ def _check_json_v2(cls, data: DTypeJSON) -> TypeGuard[DTypeConfig_V2[Literal["|b
 
         See the Zarr docs for more information about the JSON encoding for data types.
         """
-        return (
-            check_dtype_spec_v2(data) and data["name"] == "int2" and data["object_codec_id"] is None
-        )
+        return check_dtype_spec_no_object_codec_v2(data) and data["name"] == "int2"
 
     @classmethod
     def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[Literal["int2"]]:

src/zarr/core/buffer/core.py

@@ -28,6 +28,15 @@
 __all__: list[str] = []
 
 
+# Integer dtypes used to compare raw bytes when `fill_value` is zero, keyed by
+# itemsize. A view to a dtype of the same width is always permitted by numpy.
+_BITWISE_DTYPE = {1: "u1", 2: "u2", 4: "u4", 8: "u8"}
+
+# Elements in the first slice examined by `NDBuffer.all_equal`. Large enough to
+# amortise one numpy call, small enough that a mismatch in it is nearly free.
+_ALL_EQUAL_BLOCK = 1 << 14
+
+
 @runtime_checkable
 class ArrayLike(Protocol):
     """Protocol for the array-like type that underlie Buffer"""
@@ -531,32 +540,66 @@ def __len__(self) -> int:
     def __repr__(self) -> str:
         return f"<NDBuffer shape={self.shape} dtype={self.dtype} {self._data!r}>"
 
-    def all_equal(self, other: Any, equal_nan: bool = True) -> bool:
-        """Compare to `other` using np.array_equal."""
-        if other is None:
-            # Handle None fill_value for Zarr V2
-            return False
+    @staticmethod
+    def _compare_all(data: NDArrayLike, other: Any, equal_nan: bool) -> bool:
+        """Whether every element of `data` equals `other`, inspecting all of it."""
         # Handle positive and negative zero by comparing bit patterns:
         if (
             np.asarray(other).dtype.kind == "f"
             and other == 0.0
-            and self._data.dtype.kind not in ("U", "S", "T", "O", "V")
+            and data.dtype.kind not in ("U", "S", "T", "O", "V")
         ):
-            _data, other = np.broadcast_arrays(self._data, np.asarray(other, self._data.dtype))
-            void_dtype = f"V{_data.dtype.itemsize}"
-            return np.array_equal(_data.view(void_dtype), other.view(void_dtype))
+            _data, other = np.broadcast_arrays(data, np.asarray(other, data.dtype))
+            # Read the bytes as unsigned integers rather than as a void dtype.
+            # Both compare bit patterns, so both keep -0.0 distinct from 0.0,
+            # but numpy has vectorised integer comparison loops and no void
+            # one: the void form falls back to a generic elementwise path that
+            # is more than an order of magnitude slower. Widths with no integer
+            # of the same size (longdouble, complex128) keep the void form.
+            bitwise_dtype = _BITWISE_DTYPE.get(_data.dtype.itemsize, f"V{_data.dtype.itemsize}")
+            return bool(np.array_equal(_data.view(bitwise_dtype), other.view(bitwise_dtype)))
         # use array_equal to obtain equal_nan=True functionality
-        # Since fill-value is a scalar, isn't there a faster path than allocating a new array for fill value
-        # every single time we have to write data?
-        _data, other = np.broadcast_arrays(self._data, other)
-        return np.array_equal(
-            _data,
-            other,
-            equal_nan=equal_nan
-            if self._data.dtype.kind not in ("U", "S", "T", "O", "V")
-            else False,
+        _data, other = np.broadcast_arrays(data, other)
+        return bool(
+            np.array_equal(
+                _data,
+                other,
+                equal_nan=equal_nan if data.dtype.kind not in ("U", "S", "T", "O", "V") else False,
+            )
         )
 
+    def all_equal(self, other: Any, equal_nan: bool = True) -> bool:
+        """Compare to `other` using np.array_equal."""
+        if other is None:
+            # Handle None fill_value for Zarr V2
+            return False
+        data = self._data
+        # A buffer that is not uniformly `other` still has to be read in full
+        # before a whole-buffer comparison can report the first mismatch, and
+        # on the write path that is the common case. Walking it in slices of
+        # the leading axis returns as soon as one slice differs. Slices double
+        # in length, so a buffer that really is uniform is covered in O(log n)
+        # comparisons and costs the same as the single scan this replaces.
+        #
+        # The leading axis is used rather than a flattened view because a chunk
+        # is usually a strided view into a larger array, which cannot be
+        # flattened without copying it. Slicing axis 0 is a view whatever the
+        # layout.
+        if data.ndim > 0 and np.ndim(other) == 0 and data.size > _ALL_EQUAL_BLOCK:
+            leading = data.shape[0]
+            per_index = data.size // leading if leading else 0
+            if leading > 1 and per_index:
+                step = max(1, _ALL_EQUAL_BLOCK // per_index)
+                start = 0
+                while start < leading:
+                    stop = min(start + step, leading)
+                    if not self._compare_all(data[start:stop], other, equal_nan):
+                        return False
+                    start = stop
+                    step *= 2
+                return True
+        return self._compare_all(data, other, equal_nan)
+
     def fill(self, value: Any) -> None:
         self._data.fill(value)
 

src/zarr/core/context.py

@@ -0,0 +1,99 @@
+"""
+The extensions that Zarr metadata is read with, and the state of a read in progress.
+
+A `Context` is the set of extensions a read may resolve, by kind. It is a value: the same for a
+whole read, and passed in by the caller as `context=`. A `Reading` is where a read has got to: the
+context, the Zarr format of the document, and the location in the document. It changes at every
+nested step, and is threaded through the read by zarr.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from zarr.core.common import ZarrFormat
+    from zarr.core.dtype.common import DTypeJSON
+    from zarr.core.dtype.registry import DataTypeRegistry
+    from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType
+
+
+@dataclass(frozen=True, kw_only=True)
+class Context:
+    """
+    The extensions that Zarr metadata is read with, by kind.
+
+    Attributes
+    ----------
+    data_types : DataTypeRegistry
+        The data types a read may resolve.
+    """
+
+    data_types: DataTypeRegistry
+
+    @classmethod
+    def default(cls) -> Context:
+        """
+        The context of the default registries, such as `zarr.dtype.data_type_registry`.
+
+        It refers to the registries themselves, so it includes extensions registered after it is
+        created.
+        """
+        # avoid circular import
+        from zarr.core.dtype import data_type_registry
+
+        return cls(data_types=data_type_registry)
+
+
+type JSONLocation = tuple[str | int, ...]
+
+
+def json_pointer(loc: JSONLocation) -> str:
+    """
+    The JSON Pointer (RFC 6901) of a location in a JSON document, e.g. `/codecs/0/configuration`.
+    """
+    return "".join("/" + str(key).replace("~", "~0").replace("/", "~1") for key in loc)
+
+
+@dataclass(frozen=True, kw_only=True)
+class Reading:
+    """
+    The state of a read of Zarr metadata in progress.
+
+    An extension that contains other extensions, such as a structured data type, reads each of them
+    with the reading returned by `at`, so that they come from the same context, use the same Zarr
+    format, and report errors with their location.
+
+    Attributes
+    ----------
+    context : Context
+        The extensions the read may resolve.
+    zarr_format : ZarrFormat
+        The Zarr format of the document.
+    loc : tuple[str | int, ...]
+        The location in the document of the JSON being read, as keys and indices. A Zarr V2 data
+        type is located by its name as written in the document (the value of `"dtype"`), not by
+        the `{"name": ..., "object_codec_id": ...}` form zarr reads it in.
+    """
+
+    context: Context
+    zarr_format: ZarrFormat
+    loc: JSONLocation = ()
+
+    def at(self, *keys: str | int) -> Reading:
+        """
+        The reading of the JSON at `keys` within the JSON of this reading.
+        """
+        return replace(self, loc=(*self.loc, *keys))
+
+    def resolve_data_type(self, data: DTypeJSON) -> ZDType[TBaseDType, TBaseScalar]:
+        """
+        Resolve the JSON representation of a data type at the location of this reading.
+
+        Raises
+        ------
+        ValueError
+            If no data type in the context matches `data`.
+        """
+        return self.context.data_types._match_json(data, reading=self)

src/zarr/core/dtype/__init__.py

@@ -43,6 +43,7 @@
 import numpy.typing as npt
 
 from zarr.core.common import JSON
+from zarr.core.context import Context, Reading
 from zarr.core.dtype.npy.string import (
     FixedLengthUTF32,
     FixedLengthUTF32JSON_V2,
@@ -161,9 +162,12 @@
 
 
 # TODO: find a better name for this function
-def get_data_type_from_native_dtype(dtype: npt.DTypeLike) -> ZDType[TBaseDType, TBaseScalar]:
+def get_data_type_from_native_dtype(
+    dtype: npt.DTypeLike, *, context: Context | None = None
+) -> ZDType[TBaseDType, TBaseScalar]:
     """
-    Get a data type wrapper (an instance of ``ZDType``) from a native data type, e.g. a numpy dtype.
+    Get a data type wrapper (an instance of ``ZDType``) from a native data type, e.g. a numpy dtype,
+    from the data types in `context`, or `Context.default()` if it is None.
     """
     if not isinstance(dtype, np.dtype):
         na_dtype: np.dtype[np.generic]
@@ -174,23 +178,31 @@ def get_data_type_from_native_dtype(dtype: npt.DTypeLike) -> ZDType[TBaseDType,
             na_dtype = np.dtype(dtype)
     else:
         na_dtype = dtype
-    return data_type_registry.match_dtype(dtype=na_dtype)
+    return (Context.default() if context is None else context).data_types.match_dtype(
+        dtype=na_dtype
+    )
 
 
 def get_data_type_from_json(
-    dtype_spec: DTypeJSON, *, zarr_format: ZarrFormat
+    dtype_spec: DTypeJSON, *, zarr_format: ZarrFormat, context: Context | None = None
 ) -> ZDType[TBaseDType, TBaseScalar]:
     """
     Given a JSON representation of a data type and a Zarr format version,
-    attempt to create a ZDType instance from the registered ZDType classes.
+    attempt to create a ZDType instance from the data types in `context`, or `Context.default()` if
+    it is None. The data types a data type contains, such as the fields of a structured data type,
+    are resolved from the same context.
     """
-    return data_type_registry.match_json(dtype_spec, zarr_format=zarr_format)
+    reading = Reading(
+        context=Context.default() if context is None else context, zarr_format=zarr_format
+    )
+    return reading.resolve_data_type(dtype_spec)
 
 
 def parse_data_type(
     dtype_spec: ZDTypeLike,
     *,
     zarr_format: ZarrFormat,
+    context: Context | None = None,
 ) -> ZDType[TBaseDType, TBaseScalar]:
     """
     Interpret the input as a ZDType.
@@ -208,6 +220,8 @@ def parse_data_type(
         The Zarr format version. This parameter is required because this function will attempt to
         parse the JSON representation of a data type, and the JSON representation of data types
         varies between Zarr 2 and Zarr 3.
+    context : Context | None
+        The extensions to choose data types from. The default is `Context.default()`.
 
     Returns
     -------
@@ -227,13 +241,14 @@ def parse_data_type(
     # DateTime64(endianness='little', scale_factor=10, unit='s')
     ```
     """
-    return parse_dtype(dtype_spec, zarr_format=zarr_format)
+    return parse_dtype(dtype_spec, zarr_format=zarr_format, context=context)
 
 
 def parse_dtype(
     dtype_spec: ZDTypeLike,
     *,
     zarr_format: ZarrFormat,
+    context: Context | None = None,
 ) -> ZDType[TBaseDType, TBaseScalar]:
     """
     Convert the input as a ZDType.
@@ -248,6 +263,8 @@ def parse_dtype(
         The Zarr format version. This parameter is required because this function will attempt to
         parse the JSON representation of a data type, and the JSON representation of data types
         varies between Zarr 2 and Zarr 3.
+    context : Context | None
+        The extensions to choose data types from. The default is `Context.default()`.
 
     Returns
     -------
@@ -272,7 +289,7 @@ def parse_dtype(
     # First attempt to interpret the input as JSON
     if isinstance(dtype_spec, Mapping | str | Sequence):
         try:
-            return get_data_type_from_json(dtype_spec, zarr_format=zarr_format)
+            return get_data_type_from_json(dtype_spec, zarr_format=zarr_format, context=context)
         except ValueError:
             # no data type matched this JSON-like input
             pass
@@ -282,7 +299,7 @@ def parse_dtype(
         return VariableLengthUTF8()  # type: ignore[return-value]
     # otherwise, we have either a numpy dtype string, or a zarr v3 dtype string, and in either case
     # we can create a native dtype from it, and do the dtype inference from that
-    return get_data_type_from_native_dtype(dtype_spec)  # type: ignore[arg-type]
+    return get_data_type_from_native_dtype(dtype_spec, context=context)  # type: ignore[arg-type]
 
 
 def __getattr__(name: str) -> object: