diff --git a/changes/4542.feature.md b/changes/4542.feature.md new file mode 100644 index 0000000000..8cba486083 --- /dev/null +++ b/changes/4542.feature.md @@ -0,0 +1 @@ +Implement the [reshape extension codec](https://github.com/zarr-developers/zarr-extensions/tree/main/codecs/reshape). diff --git a/src/zarr/codecs/__init__.py b/src/zarr/codecs/__init__.py index 9a1b47b351..ecd06ee0f5 100644 --- a/src/zarr/codecs/__init__.py +++ b/src/zarr/codecs/__init__.py @@ -28,6 +28,7 @@ Zlib, Zstd, ) +from zarr.codecs.reshape import ReshapeCodec from zarr.codecs.scale_offset import ScaleOffset from zarr.codecs.sharding import ShardingCodec, ShardingCodecIndexLocation, SubchunkWriteOrder from zarr.codecs.transpose import TransposeCodec @@ -44,6 +45,7 @@ "Crc32cCodec", "Endian", "GzipCodec", + "ReshapeCodec", "ScaleOffset", "ShardingCodec", "ShardingCodecIndexLocation", @@ -57,6 +59,7 @@ register_codec("blosc", BloscCodec) register_codec("cast_value", CastValue) register_codec("bytes", BytesCodec) +register_codec("reshape", ReshapeCodec) # compatibility with earlier versions of ZEP1 register_codec("endian", BytesCodec) diff --git a/src/zarr/codecs/reshape.py b/src/zarr/codecs/reshape.py new file mode 100644 index 0000000000..b52ab8a688 --- /dev/null +++ b/src/zarr/codecs/reshape.py @@ -0,0 +1,257 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from zarr.abc.codec import ArrayArrayCodec +from zarr.core.array_spec import ArraySpec +from zarr.core.common import JSON, parse_named_configuration + +if TYPE_CHECKING: + from collections.abc import Iterable, Sequence + from typing import Self + + from zarr.core.buffer import NDBuffer + from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType + from zarr.core.metadata.v3 import ChunkGridMetadata + + +# Type for shape specification: int, list of ints, or -1 +type ShapeElement = int | tuple[int, ...] +type ShapeSpec = tuple[ShapeElement, ...] + + +def prod(values: Iterable[int]) -> int: + p = 1 + for v in values: + p *= v + return p + + +@dataclass(frozen=True) +class ReshapeCodec(ArrayArrayCodec): + """Reshape codec. + + An array-to-array codec that performs a reshape operation, similar to numpy's + reshape function. The reshape operation preserves the lexicographical (C-order) + traversal of elements. + + The given `shape`'s elements must be one of: + - A positive integer specifying an explicit size. + - A list of integers specifying input dimensions whose product determines + the output dimension size. + - The special value -1 (at most once), which is inferred to satisfy + prod(output_shape) == prod(input_shape). + + Notes + ----- + This codec does NOT alter the lexicographical order of elements. The contents + of the output array B is related to the input array A by: + ``ravel(B) == ravel(A)``. + + When possible, implementations should construct a virtual view rather than + copy the array. + + References + ---------- + + - The `reshape` codec spec: https://github.com/zarr-developers/zarr-extensions/tree/main/codecs/reshape + """ + + is_fixed_size = True + + shape: ShapeSpec + + def __post_init__(self) -> None: + self._validate_shape_spec() + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> Self: + _, configuration_parsed = parse_named_configuration(data, "reshape") + return cls(**configuration_parsed) # type: ignore[arg-type] + + def to_dict(self) -> dict[str, JSON]: + return {"name": "reshape", "configuration": {"shape": self.shape}} + + def validate( + self, + shape: tuple[int, ...], + dtype: ZDType[TBaseDType, TBaseScalar], + chunk_grid: ChunkGridMetadata, + ) -> None: + """Validate the codec configuration against the array metadata.""" + self._compute_output_shape(shape) + + def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: + """Evolve the codec from an array specification.""" + self._compute_output_shape(array_spec.shape) + return self + + def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: + """Resolve the output metadata after applying this codec.""" + output_shape = self._compute_output_shape(chunk_spec.shape) + return ArraySpec( + shape=output_shape, + dtype=chunk_spec.dtype, + fill_value=chunk_spec.fill_value, + config=chunk_spec.config, + prototype=chunk_spec.prototype, + ) + + def _decode_sync( + self, + chunk_array: NDBuffer, + chunk_spec: ArraySpec, + ) -> NDBuffer: + """Decode: reshape back to the original shape.""" + # chunk_spec contains the *decoded* (original) shape + # The input chunk_array has the encoded (reshaped) shape + # We need to reshape it back to the original shape + inferred_shape = self._compute_output_shape(chunk_spec.shape) + if inferred_shape != chunk_array.shape: + raise RuntimeError( + "Shape of chunk to decode %s does not match shape of requested array %s encoded with reshape codec shape=%s", + chunk_array.shape, + chunk_spec.shape, + self.shape, + ) + return chunk_array.reshape(chunk_spec.shape) + + async def _decode_single( + self, + chunk_array: NDBuffer, + chunk_spec: ArraySpec, + ) -> NDBuffer: + return self._decode_sync(chunk_array, chunk_spec) + + def _encode_sync( + self, + chunk_array: NDBuffer, + chunk_spec: ArraySpec, + ) -> NDBuffer | None: + """Encode: reshape to the output shape.""" + output_shape = self._compute_output_shape(chunk_spec.shape) + return chunk_array.reshape(output_shape) + + async def _encode_single( + self, + chunk_array: NDBuffer, + chunk_spec: ArraySpec, + ) -> NDBuffer | None: + return self._encode_sync(chunk_array, chunk_spec) + + def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int: + """Compute the encoded size (same as input since reshape doesn't change data).""" + return input_byte_length + + def _compute_output_shape( + self, + input_shape: Sequence[int], + ) -> tuple[int, ...]: + """ + Compute the output shape from the shape specification and input shape. + + Parameters + ---------- + input_shape : tuple[int, ...] + The shape of the input array. + + Returns + ------- + tuple[int, ...] + The computed output shape. + + Raises + ------ + ValueError + If the shape specification is invalid or the invariant + prod(output_shape) == prod(input_shape) cannot be satisfied. + """ + output_shape: list[int] = [] + infer_idx: int | None = None + input_total = int(prod(input_shape)) + + # we can skip a lot of validity checks here because we have already used _validate_shape_spec + for i, elem in enumerate(self.shape): + if isinstance(elem, int): + if elem == -1: + infer_idx = i + output_shape.append(-1) # placeholder + else: + output_shape.append(elem) + continue + + os = 1 + + for d in elem: + if d >= len(input_shape): + raise ValueError(f"Dimension index {d} out of range for reshape spec") + os *= input_shape[d] + + # Product of input dimensions specified + output_shape.append(os) + + # Handle inference of -1 dimension + if infer_idx is not None: + known_product = prod(s for s in output_shape if s != -1) + if known_product == 0 or input_total % known_product != 0: + raise ValueError( + f"Cannot infer dimension size: input total {input_total} is not divisible " + f"by known product {known_product}." + ) + output_shape[infer_idx] = input_total // known_product + + output_total = prod(output_shape) + if output_total != input_total: + raise ValueError( + f"The invariant prod(output_shape) == prod(input_shape) is not satisfied. " + f"Got prod({output_shape}) = {output_total} != prod({input_shape}) = {input_total}." + ) + + return tuple(output_shape) + + def _validate_shape_spec(self) -> None: + """ + Validate the shape specification according to the constraints. + + Validates: + 1. Input dimensions are in monotonically increasing order. + 2. For each input_dims list with k > 0 dimensions, the constraints about + raveled indices are satisfiable. + + Raises + ------ + ValueError + If more than explicit length of -1 is found, + or any explicit length < -1 is found, + or any indices into the array shape are below 0, + or the indices into the array shape are not monotonically increasing. + """ + last_dim_idx = -1 + found_minus_one = False + + for elem in self.shape: + if isinstance(elem, int): + if elem == -1: + if found_minus_one: + raise ValueError("-1 should appear at most once in the reshape spec") + else: + found_minus_one = True + elif elem < 1: + raise ValueError("Explicit axis size must be >=1 in the reshape spec") + continue + + for d in elem: + if d < 0: + raise ValueError( + "Index of input dimension must not be negative in the reshape spec" + ) + if d <= last_dim_idx: + raise ValueError( + "Index of input dimension must be monotonically increasing in reshape spec" + ) + + last_dim_idx = d + + +__all__ = ["ReshapeCodec"] diff --git a/tests/test_codecs/test_reshape.py b/tests/test_codecs/test_reshape.py new file mode 100644 index 0000000000..2803c34c59 --- /dev/null +++ b/tests/test_codecs/test_reshape.py @@ -0,0 +1,176 @@ +import json + +import numpy as np +import pytest + +from zarr.codecs.reshape import ReshapeCodec, ShapeSpec +from zarr.core.array_spec import ArrayConfig, ArraySpec +from zarr.core.buffer import default_buffer_prototype +from zarr.core.dtype import UInt8 +from zarr.core.metadata.v3 import RegularChunkGridMetadata + + +def array_spec(shape: tuple[int, ...]) -> ArraySpec: + return ArraySpec(shape, UInt8(), 0, ArrayConfig("C", False), default_buffer_prototype()) + + +def test_deserialize() -> None: + s = '{"name": "reshape", "configuration": {"shape": [5, [1, 2], -1]}}' + d = json.loads(s) + ReshapeCodec.from_dict(d) + + +@pytest.mark.asyncio +async def test_encode() -> None: + arr_flat = np.arange(60, dtype="uint8") + arr_3x2 = np.reshape(arr_flat, (3, 4, 5)) + codec = ReshapeCodec(shape=((0, 1), 5)) + proto = default_buffer_prototype() + (encoded, *_) = await codec.encode( + [(proto.nd_buffer.from_numpy_array(arr_3x2), array_spec(arr_3x2.shape))] + ) + assert encoded is not None + assert encoded.shape == (12, 5) + assert np.all(np.equal(np.ravel(encoded.as_numpy_array()), arr_flat)) + + +@pytest.mark.asyncio +async def test_decode() -> None: + arr_flat = np.arange(60, dtype="uint8") + arr_12x5 = np.reshape(arr_flat, (12, 5)) + dt = UInt8() + codec = ReshapeCodec(shape=((0, 1), 5)) + proto = default_buffer_prototype() + (encoded, *_) = await codec.decode( + [ + ( + proto.nd_buffer.from_numpy_array(arr_12x5), + ArraySpec( + (3, 4, 5), + dt, + fill_value=0, + config=ArrayConfig("C", False), + prototype=proto, + ), + ) + ] + ) + assert encoded is not None + assert encoded.shape == (3, 4, 5) + assert np.all(np.equal(np.ravel(encoded.as_numpy_array()), arr_flat)) + + +@pytest.mark.parametrize( + "shape_spec", + [ + (-2,), # negative number isn't -1 + (-1, -1), # more than one -1 + ((-2, 1),), # negative index + ((1, 0),), # decreasing + # (((0, 4),), 3), # missing dim + ], +) +def test_invalid_shape_spec(shape_spec: ShapeSpec) -> None: + with pytest.raises(ValueError): + ReshapeCodec(shape=shape_spec) + + +def test_compute_encoded_size() -> None: + c = ReshapeCodec(shape=((0, 1), 5)) + + dec_shape = (64, 64) + dec_size = np.prod(dec_shape).item() + enc_size = c.compute_encoded_size( + dec_size, + # ignored + ArraySpec( + (64, 64), + UInt8(), + fill_value=0, + config=ArrayConfig("C", False), + prototype=default_buffer_prototype(), + ), + ) + assert dec_size == enc_size + + +@pytest.mark.parametrize( + ("shape", "shape_spec"), + [ + ((3, 4), ((0, 1, 2),)), # index out of range + ((3, 4), (5, -1)), # indivisible -1 + ((3, 4), (5, 2)), # wrong shape + ], +) +def test_validate_invalid(shape: tuple[int, ...], shape_spec: ShapeSpec) -> None: + """Stand-in for _compute_output_shape""" + codec = ReshapeCodec(shape=shape_spec) + with pytest.raises(ValueError): + codec.validate( + shape, + UInt8(), # ignored + RegularChunkGridMetadata(chunk_shape=shape), + ) + + +@pytest.mark.asyncio +async def test_decode_unexpected_shape() -> None: + codec = ReshapeCodec(shape=((0, 1),)) + buf = default_buffer_prototype().nd_buffer.from_numpy_array( + np.zeros(shape=(30, 40), dtype="uint8") + ) + spec = array_spec((3, 4)) + with pytest.raises(RuntimeError): + await codec.decode([(buf, spec)]) + + +@pytest.mark.parametrize( + ("input_shape", "shape_spec", "output_shape"), + [ + ((3, 4), (4, 3), (4, 3)), + ((3, 4), (6, -1), (6, 2)), + ((3, 4), ((0, 1),), (12,)), + ], +) +def test_valid( + input_shape: tuple[int, ...], shape_spec: ShapeSpec, output_shape: tuple[int, ...] +) -> None: + codec = ReshapeCodec(shape=shape_spec) + dtype = UInt8() + codec.validate(input_shape, dtype, RegularChunkGridMetadata(chunk_shape=input_shape)) + input_spec = array_spec(input_shape) + output_spec = codec.resolve_metadata(input_spec) + + assert output_spec.dtype == input_spec.dtype + assert output_spec.config == input_spec.config + assert output_spec.fill_value == input_spec.fill_value + assert output_spec.prototype == input_spec.prototype + + assert output_spec.shape == output_shape + + +@pytest.mark.asyncio +async def test_example() -> None: + import json + + import zarr + from zarr.storage import ManagedMemoryStore + + # Create an array with the reshape codec + # This reshapes chunks from (100, 50, 64, 3) to (5000, 64, 3) + store = ManagedMemoryStore() + _ = zarr.create_array( + store, + shape=(1000, 500, 64, 3), + chunks=(100, 50, 64, 3), + dtype="float32", + filters=[ReshapeCodec(shape=((0, 1), (2,), 3))], + ) + b = await store.get("zarr.json") + assert b is not None + s = b.to_bytes().decode("utf-8") + d = json.loads(s) + assert d["codecs"][0] == { + "name": "reshape", + "configuration": {"shape": [[0, 1], [2], 3]}, + }