fix: use contextual arrayspec when validating / evolving codecs (zarr-developers#3941)

* fix: use contextual arrayspec when validating / evolving codecs

* docs: changelog

* fix: gate tests depending on presence of compiled backend

Author: d-v-b

changes/3938.bugfix.md

@@ -0,0 +1,4 @@
+Fixed a `CastValue` validation bug where the "can we use an out-of-range mode" check
+inspected the source dtype instead of the target dtype. This meant arrays with a
+float source dtype and an integer target dtype incorrectly raised a `ValueError`
+when configured with a `wrap` out-of-range mode.

changes/3941.bugfix.md

@@ -0,0 +1,6 @@
+Fixed a bug where the codec pipeline evolved each codec against the original
+array spec instead of the spec produced by upstream array-to-array codecs. This
+caused failures whenever an upstream codec changed the dtype between codec
+boundaries — e.g. arrays using `CastValue` to convert a single-byte source dtype
+(`int8`) to a multi-byte target dtype (`int16`) raised a `ValueError` from
+`BytesCodec` about a missing `endian` configuration.

src/zarr/core/codec_pipeline.py

@@ -187,7 +187,18 @@ class BatchedCodecPipeline(CodecPipeline):
     batch_size: int
 
     def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
-        return type(self).from_codecs(c.evolve_from_array_spec(array_spec=array_spec) for c in self)
+        # Each codec must be evolved against the spec it will actually see
+        # at run-time, not the original array spec. Earlier array->array
+        # codecs may transform the dtype (e.g. cast_value), so the spec
+        # threaded into later codecs (the array->bytes serializer and any
+        # bytes->bytes filters) must reflect those transformations.
+        evolved: list[Codec] = []
+        spec = array_spec
+        for codec in self:
+            evolved_codec = codec.evolve_from_array_spec(array_spec=spec)
+            evolved.append(evolved_codec)
+            spec = evolved_codec.resolve_metadata(spec)
+        return type(self).from_codecs(evolved)
 
     @classmethod
     def from_codecs(cls, codecs: Iterable[Codec], *, batch_size: int | None = None) -> Self:

src/zarr/core/metadata/v3.py

@@ -503,7 +503,21 @@ def __init__(
             config=ArrayConfig.from_dict({}),  # TODO: config is not needed here.
             prototype=default_buffer_prototype(),  # TODO: prototype is not needed here.
         )
-        codecs_parsed = tuple(c.evolve_from_array_spec(array_spec) for c in codecs_parsed_partial)
+        # Thread the spec through evolution: each codec must be evolved against
+        # the spec it will actually see at run-time, not the original array spec.
+        # Earlier array->array codecs may transform the dtype (e.g. cast_value),
+        # so the spec passed to later codecs must reflect those transformations.
+        # Per-codec validate() must run before resolve_metadata(), since the
+        # latter may rely on invariants the former checks (e.g. cast_value
+        # rejects complex source dtypes that would otherwise crash _do_cast).
+        evolved: list[Codec] = []
+        spec = array_spec
+        for c in codecs_parsed_partial:
+            evolved_codec = c.evolve_from_array_spec(spec)
+            evolved_codec.validate(shape=spec.shape, dtype=spec.dtype, chunk_grid=chunk_grid_parsed)
+            evolved.append(evolved_codec)
+            spec = evolved_codec.resolve_metadata(spec)
+        codecs_parsed = tuple(evolved)
         validate_codecs(codecs_parsed_partial, data_type)
 
         object.__setattr__(self, "shape", shape_parsed)

tests/test_codec_pipeline.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import numpy as np
 import pytest
 
 import zarr
+from zarr.codecs import BytesCodec, CastValue
 from zarr.core.array import _get_chunk_spec
 from zarr.core.buffer.core import default_buffer_prototype
 from zarr.core.indexing import BasicIndexer
@@ -70,3 +72,51 @@ async def test_read_returns_get_results(
     assert len(results) == len(expected_statuses)
     for result, expected_status in zip(results, expected_statuses, strict=True):
         assert result["status"] == expected_status
+
+
+try:
+    import cast_value_rs  # noqa: F401
+
+    _HAS_CAST_VALUE_RS = True
+except ModuleNotFoundError:
+    _HAS_CAST_VALUE_RS = False
+
+requires_cast_value_rs = pytest.mark.skipif(
+    not _HAS_CAST_VALUE_RS, reason="cast-value-rs not installed"
+)
+
+
+@requires_cast_value_rs
+@pytest.mark.parametrize(
+    ("source_dtype", "target_dtype"),
+    [
+        # Source is single-byte (no endianness); target is multi-byte (has endianness).
+        # Without the fix, BytesCodec.evolve_from_array_spec sees the source dtype,
+        # strips its `endian` to None, and then chokes when the chunk_spec dtype
+        # gets transformed to the multi-byte target before bytes-decoding.
+        ("int8", "int16"),
+        ("uint8", "int32"),
+        ("int8", "float32"),
+        # Source is multi-byte; target is single-byte (the reverse direction also
+        # exercises the spec-threading logic).
+        ("int16", "int8"),
+    ],
+)
+def test_codec_pipeline_threads_dtype_through_evolve(source_dtype: str, target_dtype: str) -> None:
+    """Regression for #3937: each codec must be evolved against the spec it
+    will see at runtime, not the original array spec. cast_value transforms
+    the dtype between AA codecs and the array->bytes serializer."""
+    arr = zarr.create_array(
+        store={},
+        shape=(4,),
+        chunks=(4,),
+        dtype=source_dtype,
+        fill_value=0,
+        filters=[CastValue(data_type=target_dtype)],
+        serializer=BytesCodec(endian="little"),
+        compressors=[],
+        zarr_format=3,
+        overwrite=True,
+    )
+    arr[:] = np.asarray([0, 1, 2, 3], dtype=source_dtype)
+    np.testing.assert_array_equal(arr[:], np.asarray([0, 1, 2, 3], dtype=source_dtype))

tests/test_codecs/test_cast_value.py

@@ -165,6 +165,7 @@ def test_validation_rejects_invalid(case: ExpectErr[dict[str, Any]]) -> None:
         )
 
 
+@requires_cast_value_rs
 @pytest.mark.parametrize(
     ("source_dtype", "target_dtype"),
     [