diff --git a/changes/4510.bugfix.md b/changes/4510.bugfix.md new file mode 100644 index 0000000000..f74b190f34 --- /dev/null +++ b/changes/4510.bugfix.md @@ -0,0 +1 @@ +The `cast_value` codec now checks that the array's fill value survives a round trip through the codec, as the [`cast_value` spec](https://github.com/zarr-developers/zarr-extensions/tree/main/codecs/cast_value#data_type) requires. Previously `out_of_range="clamp"` or `"wrap"`, or rounding, could encode the fill value to something that decodes to a different value, so the unwritten part of a partly written chunk read back as that value instead of the fill value. For example, a `float32` array with fill value `1000.0` and an `int8` target with clamping read `127.0` there. Creating such an array now raises a `ValueError` that names the fill value, both data types, and the value it decodes to. Opening an array already stored with such a fill value still works: it keeps its fill value, and a warning says what its partly written chunks hold and how to fix it. The check runs once when metadata is created or read, not for every chunk, and the codecs inside a shard are checked with the array's fill value rather than the data type's default scalar, so a shard's codec chain is no longer refused for a fill value it would never be asked to encode. diff --git a/src/zarr/codecs/cast_value.py b/src/zarr/codecs/cast_value.py index 629b4e7a27..985c37389a 100644 --- a/src/zarr/codecs/cast_value.py +++ b/src/zarr/codecs/cast_value.py @@ -10,7 +10,7 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from dataclasses import dataclass, replace from importlib.metadata import PackageNotFoundError, version from typing import TYPE_CHECKING, Final, Literal, TypedDict, cast @@ -25,10 +25,11 @@ if TYPE_CHECKING: from typing import NotRequired, Self + from zarr.abc.codec import Codec from zarr.core.array_spec import ArraySpec from zarr.core.buffer import NDBuffer from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType - from zarr.core.metadata.v3 import ChunkGridMetadata + from zarr.core.metadata.v3 import ArrayV3Metadata, ChunkGridMetadata class ScalarMapJSON(TypedDict): encode: NotRequired[list[tuple[object, object]]] @@ -175,6 +176,21 @@ def _check_representable( ) from e +def _same_scalar(a: np.ndarray, b: np.ndarray) -> bool: + """Return whether two one-element arrays of the same dtype hold the same value. + + NaN equals NaN, whatever its payload, and the sign of zero must match, as + the `cast_value` spec requires signed zero to be preserved. Comparing bytes + covers both, and works for every permitted dtype, including the + `ml_dtypes` floats. + """ + a_is_nan = bool(a[0] != a[0]) + b_is_nan = bool(b[0] != b[0]) + if a_is_nan or b_is_nan: + return a_is_nan and b_is_nan + return a.tobytes() == b.tobytes() + + # --------------------------------------------------------------------------- # Codec # --------------------------------------------------------------------------- @@ -361,23 +377,85 @@ def _get_scalar_map( return None return self.scalar_map.get(direction) # type: ignore[return-value] + def _cast_fill( + self, + fill: np.ndarray, + *, + source: ZDType[TBaseDType, TBaseScalar], + direction: Literal["encode", "decode"], + ) -> np.ndarray: + """Cast a one-element array holding a fill value, as the codec casts elements. + + `encode` casts `fill`, a scalar of the `source` data type, to the codec's + data type; `decode` casts `fill`, a scalar of the codec's data type, back + to `source`. Each uses the scalar map of its direction. + + Raises + ------ + ValueError + If the cast fails, naming the fill value and both data types. + """ + from_dtype, to_dtype = ( + (source, self.dtype) if direction == "encode" else (self.dtype, source) + ) + try: + return self._do_cast( + fill, + target_dtype=to_dtype.to_native_dtype(), + scalar_map=self._get_scalar_map(direction), + ) + except (ValueError, OverflowError) as e: + what = "fill value" if direction == "encode" else "encoded fill value" + raise ValueError( + f"The {what} {fill[0]} cannot be cast from {from_dtype.to_json(zarr_format=3)} " + f"to {to_dtype.to_json(zarr_format=3)} by the cast_value codec: {e}" + ) from e + def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: """ Update the fill value of the output spec by applying casting procedure. - """ - target_zdtype = self.dtype - target_native = target_zdtype.to_native_dtype() - source_native = chunk_spec.dtype.to_native_dtype() - fill = chunk_spec.fill_value - fill_arr = np.array([fill], dtype=source_native) - - new_fill_arr = self._do_cast( - fill_arr, target_dtype=target_native, scalar_map=self._get_scalar_map("encode") + This runs for every chunk that is read or written, so it only casts the + fill value forward. Whether the fill value survives a round trip through + the codec is checked once, when array metadata is created or read, by + `_check_fill_value_round_trip`. + """ + fill = np.array([chunk_spec.fill_value], dtype=chunk_spec.dtype.to_native_dtype()) + encoded = self._cast_fill(fill, source=chunk_spec.dtype, direction="encode") + return replace( + chunk_spec, dtype=self.dtype, fill_value=self.dtype.to_native_dtype().type(encoded[0]) ) - new_fill = target_native.type(new_fill_arr[0]) - return replace(chunk_spec, dtype=target_zdtype, fill_value=new_fill) + def _check_fill_value(self, chunk_spec: ArraySpec) -> None: + """ + Check that the fill value of `chunk_spec` survives a round trip through the codec. + + Encoding the fill value and decoding the result must give back the + original fill value. Otherwise a chunk that is only partly written would + read back a different value wherever it holds the fill value, since what + is stored there is the encoded fill value. See the ``data_type`` section + of the + [`cast_value` spec](https://github.com/zarr-developers/zarr-extensions/tree/main/codecs/cast_value#data_type). + + Raises + ------ + ValueError + If the fill value cannot be cast to the target data type, cannot + be cast back, or does not decode to the original fill value. + """ + source = chunk_spec.dtype + fill = np.array([chunk_spec.fill_value], dtype=source.to_native_dtype()) + encoded = self._cast_fill(fill, source=source, direction="encode") + decoded = self._cast_fill(encoded, source=source, direction="decode") + if not _same_scalar(fill, decoded): + source_name = source.to_json(zarr_format=3) + raise ValueError( + f"The fill value {fill[0]} does not survive a round trip through the cast_value " + f"codec: it encodes from {source_name} to {encoded[0]} in " + f"{self.dtype.to_json(zarr_format=3)}, which decodes back to {decoded[0]} in " + f"{source_name}. Choose a fill value that the codec's data_type, rounding, " + "out_of_range, and scalar_map preserve." + ) def _encode_sync( self, @@ -435,3 +513,51 @@ def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> ) num_elements = input_byte_length // source_itemsize return num_elements * target_itemsize + + +def _check_fill_value_round_trip(metadata: ArrayV3Metadata) -> None: + """ + Check that the fill value of `metadata` survives each `cast_value` codec in its codecs. + + Each codec is checked with the fill value it sees when a chunk is encoded: + the array's fill value as transformed by the array-to-array codecs before + it. The codecs inside a sharding codec are checked the same way. + + Parameters + ---------- + metadata : ArrayV3Metadata + The array metadata whose fill value and codecs to check. + + Raises + ------ + ValueError + If a `cast_value` codec cannot encode its fill value, cannot decode the + encoded fill value, or decodes it to a different value. See + `CastValue._check_fill_value`. + """ + # Imported here: `zarr.core.metadata.v3` reads metadata with this check, and + # importing it at module level would import this module in a cycle. + from zarr.core.array_spec import ArrayConfig, ArraySpec + from zarr.core.buffer.core import default_buffer_prototype + + spec = ArraySpec( + shape=metadata.shape, + dtype=metadata.data_type, + fill_value=metadata.fill_value, + config=ArrayConfig.from_dict({}), + prototype=default_buffer_prototype(), + ) + _check_codecs(metadata.codecs, spec) + + +def _check_codecs(codecs: Iterable[Codec], spec: ArraySpec) -> None: + """Check the `cast_value` codecs among `codecs` against the fill value threaded + through them from `spec`, descending into the codecs of a sharding codec.""" + from zarr.codecs.sharding import ShardingCodec + + for codec in codecs: + if isinstance(codec, CastValue): + codec._check_fill_value(spec) + elif isinstance(codec, ShardingCodec): + _check_codecs(codec.codecs, spec) + spec = codec.resolve_metadata(spec) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 06f7338755..48c46174bc 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -27,6 +27,7 @@ from zarr.abc.numcodec import Numcodec, _is_numcodec from zarr.codecs._v2 import V2Codec from zarr.codecs.bytes import BytesCodec +from zarr.codecs.cast_value import _check_fill_value_round_trip from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec from zarr.codecs.zstd import ZstdCodec from zarr.core._info import ArrayInfo @@ -612,7 +613,7 @@ def _create_metadata_v3( else: fill_value_parsed = fill_value - return ArrayV3Metadata( + metadata = ArrayV3Metadata( shape=shape, data_type=dtype, chunk_grid=chunk_grid, @@ -622,6 +623,8 @@ def _create_metadata_v3( dimension_names=tuple(dimension_names) if dimension_names else None, attributes=attributes or {}, ) + _check_fill_value_round_trip(metadata) + return metadata @classmethod async def _create_v3( diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 2d5d169e6b..8164012dfd 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -41,6 +41,7 @@ from zarr.core.json_parse import parse_field from zarr.core.metadata.common import parse_attributes from zarr.core.metadata.repair import ( + Reading, mark_repaired, repair_array_document, ) @@ -723,6 +724,24 @@ def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self: extra_fields=allowed_extra_fields, storage_transformers=_data_typed.get("storage_transformers", ()), # type: ignore[arg-type] ) + # Imported here: `zarr.codecs` imports the chunk pipeline, which imports this module. + from zarr.codecs.cast_value import _check_fill_value_round_trip + + try: + _check_fill_value_round_trip(metadata) + except ValueError as e: + readings.append( + Reading( + moves_chunks=False, + warning=( + f"{e} Partly written chunks of this array store the encoded fill " + "value, so their unwritten elements read back as the decoded value, " + "not the fill value. Other implementations may refuse to open the " + "array. To fix it, copy the data to a new array with a fill value " + "the codec preserves." + ), + ) + ) return mark_repaired(metadata, data, readings, path) def to_dict(self) -> dict[str, JSON]: diff --git a/tests/test_codecs/test_cast_value.py b/tests/test_codecs/test_cast_value.py index c2e78770d9..3202121486 100644 --- a/tests/test_codecs/test_cast_value.py +++ b/tests/test_codecs/test_cast_value.py @@ -8,8 +8,12 @@ import zarr from tests.conftest import Expect, ExpectFail -from zarr.codecs import BytesCodec, TransposeCodec +from zarr.codecs import BytesCodec, ScaleOffset, TransposeCodec from zarr.codecs.cast_value import CastValue +from zarr.core.array_spec import ArrayConfig, ArraySpec +from zarr.core.buffer.core import default_buffer_prototype +from zarr.core.dtype import get_data_type_from_json +from zarr.errors import ZarrUserWarning from zarr.storage import MemoryStore try: @@ -611,3 +615,222 @@ def make_array(filters: list[Any]) -> Any: array[:] = data assert_array_equal(array[:], data) + + +# --------------------------------------------------------------------------- +# Fill value round trip +# --------------------------------------------------------------------------- + + +def _create_with_fill( + dtype: str, + fill_value: Any, + *filters: Any, + store: MemoryStore | None = None, + shards: tuple[int, ...] | None = None, +) -> Any: + """Create a 4-element array with 2-element chunks, filtered through `filters`, + in `shards` if given.""" + return zarr.create_array( + store=MemoryStore() if store is None else store, + shape=(4,), + chunks=(2,), + shards=shards, + dtype=dtype, + fill_value=fill_value, + filters=list(filters), + serializer=BytesCodec(endian="little"), + compressors=None, + ) + + +@requires_cast_value_rs +@pytest.mark.parametrize( + "case", + [ + Expect( + input=( + "float32", + np.nan, + CastValue( + data_type="int8", scalar_map={"encode": [("NaN", 0)], "decode": [(0, "NaN")]} + ), + ), + output=np.int8(0), + id="nan-through-scalar-map", + ), + Expect( + input=("float64", 5.0, CastValue(data_type="int8", out_of_range="clamp")), + output=np.int8(5), + id="clamp-in-range", + ), + Expect( + input=("float64", -3.0, CastValue(data_type="int8", out_of_range="wrap")), + output=np.int8(-3), + id="wrap-in-range", + ), + Expect( + input=("int16", 300, CastValue(data_type="float32")), + output=np.float32(300.0), + id="int-to-float", + ), + Expect( + input=("float64", 2.0, CastValue(data_type="int8", rounding="towards-zero")), + output=np.int8(2), + id="integral-float-to-int", + ), + Expect( + input=("float64", 1.5, CastValue(data_type="float32")), + output=np.float32(1.5), + id="exact-float-narrowing", + ), + Expect( + input=("float32", -0.0, CastValue(data_type="float16")), + output=np.float16(-0.0), + id="signed-zero-float-to-float", + ), + ], + ids=lambda c: c.id, +) +def test_fill_value_roundtrip_accepted(case: Expect[tuple[str, Any, CastValue], Any]) -> None: + """A fill value that survives encode then decode is accepted. The encoded + fill is stored for the unwritten part of a partly written chunk, and that + part reads back as the original fill value, bit for bit.""" + dtype, fill_value, codec = case.input + arr = _create_with_fill(dtype, fill_value, codec) + arr[0] = 1 + + stored = arr.store._store_dict["c/0"].to_bytes() + expected_encoded = np.array([case.output], dtype=case.output.dtype) + assert stored[expected_encoded.itemsize :] == expected_encoded.tobytes() + + read = arr[1:2] + assert read.tobytes() == np.array([fill_value], dtype=dtype).tobytes() + + +@requires_cast_value_rs +def test_fill_value_partial_chunk_without_check_reads_wrong_value( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression: a fill value that does not round-trip corrupts partly written chunks. + + With the round-trip check disabled, float32 fill 1000.0 encodes to int8 127 + under `out_of_range="clamp"`. The unwritten element of a partly written + chunk is stored as 127 and reads back as 127.0, not as the fill value. + Array creation refuses this metadata. The array already stored with it + still opens, with its fill value and a warning that says what its chunks + hold, so its data stays readable. + """ + import zarr.codecs.cast_value as cast_value_module + + codec = CastValue(data_type="int8", out_of_range="clamp") + store = MemoryStore() + + with monkeypatch.context() as m: + m.setattr(cast_value_module, "_same_scalar", lambda a, b: True) + arr = _create_with_fill("float32", 1000.0, codec, store=store) + arr[0] = 5.0 + assert_array_equal(arr[:], np.array([5.0, 127.0, 1000.0, 1000.0], dtype="float32")) + + with pytest.raises(ValueError, match=r"decodes back to 127\.0 in float32"): + _create_with_fill("float32", 1000.0, codec) + + with pytest.warns( + ZarrUserWarning, + match=r"decodes back to 127\.0 in float32.*Partly written chunks of this array", + ): + opened = zarr.open_array(store, mode="r") + assert opened.fill_value == np.float32(1000.0) + assert_array_equal(opened[:], np.array([5.0, 127.0, 1000.0, 1000.0], dtype="float32")) + + +@requires_cast_value_rs +def test_fill_value_checked_inside_shards() -> None: + """The codecs inside a shard are checked with the array's fill value, not with + the default scalar of the data type: a fill value the chain preserves is + accepted although the default scalar would be lost to clamping after scaling, + and a fill value the chain loses is refused.""" + scale = ScaleOffset(scale=655.35, offset=100.0) + clamp = CastValue(data_type="uint16", out_of_range="clamp") + arr = _create_with_fill("float32", 100.0, scale, clamp, shards=(4,)) + arr[0] = 150.0 + assert_array_equal(arr[1:], np.full(3, 100.0, dtype="float32")) + + with pytest.raises(ValueError, match=r"decodes back to 127\.0 in float32"): + _create_with_fill( + "float32", 1000.0, CastValue(data_type="int8", out_of_range="clamp"), shards=(4,) + ) + + +@requires_cast_value_rs +def test_resolve_metadata_only_casts_forward() -> None: + """`resolve_metadata` runs for every chunk, so it casts the fill value forward + without checking the round trip, which metadata creation checks once.""" + spec = ArraySpec( + shape=(2,), + dtype=get_data_type_from_json("float32", zarr_format=3), + fill_value=np.float32(1000.0), + config=ArrayConfig.from_dict({}), + prototype=default_buffer_prototype(), + ) + resolved = CastValue(data_type="int8", out_of_range="clamp").resolve_metadata(spec) + assert resolved.fill_value == np.int8(127) + + +@requires_cast_value_rs +def test_fill_value_clamp_out_of_range_rejected() -> None: + with pytest.raises( + ValueError, + match=r"fill value 1000\.0 .* from float32 to 127 in int8, which decodes back to 127\.0", + ): + _create_with_fill("float32", 1000.0, CastValue(data_type="int8", out_of_range="clamp")) + + +@requires_cast_value_rs +def test_fill_value_wrap_out_of_range_rejected() -> None: + with pytest.raises( + ValueError, + match=r"fill value 300 .* from int16 to 44 in int8, which decodes back to 44 in int16", + ): + _create_with_fill("int16", 300, CastValue(data_type="int8", out_of_range="wrap")) + + +@requires_cast_value_rs +def test_fill_value_lost_to_rounding_rejected() -> None: + with pytest.raises( + ValueError, + match=r"fill value 1\.5 .* from float64 to 2 in int8, which decodes back to 2\.0", + ): + _create_with_fill("float64", 1.5, CastValue(data_type="int8")) + + +@requires_cast_value_rs +def test_fill_value_signed_zero_lost_rejected() -> None: + """The spec requires signed zero to be preserved; an integer target cannot.""" + with pytest.raises( + ValueError, + match=r"fill value -0\.0 .* from float32 to 0 in int8, which decodes back to 0\.0", + ): + _create_with_fill("float32", -0.0, CastValue(data_type="int8")) + + +@requires_cast_value_rs +def test_fill_value_nan_to_int_without_scalar_map_rejected() -> None: + with pytest.raises( + ValueError, + match=r"fill value nan cannot be cast from float32 to int8 .*without scalar_map", + ): + _create_with_fill("float32", np.nan, CastValue(data_type="int8")) + + +@requires_cast_value_rs +def test_fill_value_cannot_decode_rejected() -> None: + """The fill encodes to NaN through the encode map, but with no decode map + NaN cannot be cast back to the integer source.""" + with pytest.raises( + ValueError, + match=r"encoded fill value nan cannot be cast from float64 to int32", + ): + _create_with_fill( + "int32", 0, CastValue(data_type="float64", scalar_map={"encode": [(0, "NaN")]}) + )