diff --git a/changes/4560.bugfix.md b/changes/4560.bugfix.md new file mode 100644 index 0000000000..a958b483e1 --- /dev/null +++ b/changes/4560.bugfix.md @@ -0,0 +1,8 @@ +Encoding a chunk through a numcodecs filter no longer passes the +filter's flattened output to the next codec: `_encode_sync` reshapes the +flattened array back to the chunk shape, so downstream codecs +(e.g. `TransposeCodec` or an `imagecodecs` serializer chained after a +`Delta` or `FixedScaleOffset` filter) receive a chunk of the expected +shape instead of a flat 1-D array that could make writing fail or +silently corrupt data. Shape-changing codecs such as `PackBits` are +unaffected. diff --git a/src/zarr/codecs/numcodecs/_codecs.py b/src/zarr/codecs/numcodecs/_codecs.py index c980743166..52e2dfdbaf 100644 --- a/src/zarr/codecs/numcodecs/_codecs.py +++ b/src/zarr/codecs/numcodecs/_codecs.py @@ -164,6 +164,11 @@ def _encode_sync(self, chunk_data: NDBuffer, chunk_spec: ArraySpec) -> NDBuffer: # would be encoded in transposed element order (gh-3558) chunk_ndarray = np.ascontiguousarray(chunk_data.as_ndarray_like()) out = self._codec.encode(chunk_ndarray) + if out.ndim == 1 and out.size == np.prod(chunk_spec.shape): + # most numcodecs filters return a flattened array; reshape it so + # downstream codecs receive a chunk of the shape the pipeline + # promised them. Shape-changing codecs (e.g. PackBits) keep out. + out = out.reshape(chunk_spec.shape) return chunk_spec.prototype.nd_buffer.from_ndarray_like(out) async def _encode_single(self, chunk_data: NDBuffer, chunk_spec: ArraySpec) -> NDBuffer: diff --git a/tests/test_codecs/test_numcodecs.py b/tests/test_codecs/test_numcodecs.py index d78b73f34c..d189aef2dd 100644 --- a/tests/test_codecs/test_numcodecs.py +++ b/tests/test_codecs/test_numcodecs.py @@ -11,6 +11,7 @@ from zarr import config, create_array, open_array from zarr.abc.numcodec import _is_numcodec, _is_numcodec_cls from zarr.codecs import numcodecs as _numcodecs +from zarr.codecs.transpose import TransposeCodec from zarr.errors import UnknownCodecError from zarr.registry import get_codec_class, get_numcodec @@ -192,6 +193,45 @@ def test_generic_filter_f_contiguous( np.testing.assert_array_equal(data, a[:, :]) +def test_generic_filter_flattened_output_reshaped() -> None: + # gh-4553: numcodecs filters return a flattened array; `_encode_sync` must + # reshape it back to the chunk shape, so downstream codecs receive chunks + # with the shape the pipeline promised them + data = np.arange(256, dtype="int64").reshape(16, 16) + + a = create_array( + {}, + shape=data.shape, + chunks=(16, 16), + dtype=data.dtype, + fill_value=0, + filters=[_numcodecs.Delta(dtype="int64"), TransposeCodec(order=(1, 0))], + ) + + a[:, :] = data + with codec_conf(): + b = open_array(a.store, mode="r") + np.testing.assert_array_equal(data, b[:, :]) + + +def test_generic_filter_packedbits_output_not_reshaped() -> None: + # PackBits drives the size of the chunk, so `_encode_sync` must not reshape + # its output + data = np.tril(np.ones((16, 16), dtype=bool)) + + a = create_array( + {}, + shape=data.shape, + chunks=(16, 16), + dtype=data.dtype, + fill_value=0, + filters=[_numcodecs.PackBits()], + ) + + a[:, :] = data + np.testing.assert_array_equal(data, a[:, :]) + + def test_generic_filter_bitround() -> None: data = np.linspace(0, 1, 256, dtype="float32").reshape((16, 16))