From 926b5c0b0607278502e7b3a8db9dce30c8f37c6e Mon Sep 17 00:00:00 2001 From: Dylan Pulver Date: Tue, 25 Aug 2026 16:49:23 +0300 Subject: [PATCH 1/3] fix: reshape the value buffer for coordinate selections in sharded writes --- src/zarr/codecs/sharding.py | 17 +++++++++------ tests/test_codecs/test_sharding.py | 34 ++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 41780e45b4..9eb90b4ffc 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -1359,13 +1359,18 @@ async def _encode_partial_single( chunks_per_shard = self._get_chunks_per_shard(shard_spec) chunk_spec = self._get_chunk_spec(shard_spec) - indexer = list( - get_indexer( - selection, - shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape), - ) + shard_indexer = get_indexer( + selection, + shape=shard_shape, + chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape), ) + # A coordinate indexer flattens the selection, so its projections address + # `shard_array` as 1-D while the caller shaped it like `sel_shape`. This + # mirrors the reshape `_decode_partial_single` applies on the way out. + sel_shape = getattr(shard_indexer, "sel_shape", None) + if sel_shape is not None and shard_array.shape == sel_shape: + shard_array = shard_array.reshape(shard_indexer.shape) + indexer = list(shard_indexer) if self._is_complete_shard_write(indexer, chunks_per_shard): shard_dict = dict.fromkeys(lexicographic_order_coords(chunks_per_shard)) diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index de576dbef5..4cc80c912f 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -1262,3 +1262,37 @@ def test_shard_reader_to_dict_vectorized(chunks_per_shard: tuple[int, ...]) -> N assert buf.to_bytes() == present[coords] else: assert buf is None + + +@pytest.mark.parametrize("nested", [False, True], ids=["single", "nested"]) +def test_sharding_orthogonal_set_multiple_array_dims(nested: bool) -> None: + """Orthogonal set with more than one array-indexed dimension. + + ``OrthogonalIndexer`` converts such a chunk selection to an ``np.ix_`` pair + of broadcastable arrays before handing it to the codec pipeline. The + sharding codec re-derives an indexer from that selection and gets a + ``CoordinateIndexer``, whose projections address the value buffer flat. + Regression test for the resulting shape mismatch on write. + """ + inner = ShardingCodec(chunk_shape=(1, 1), codecs=(BytesCodec(),)) + serializer = ShardingCodec(chunk_shape=(2, 2), codecs=((inner,) if nested else (BytesCodec(),))) + base = np.arange(16, dtype="int32").reshape(4, 4) + selection = (np.array([3, 1, 2]), np.array([0, 2])) + value = np.arange(6, dtype="int32").reshape(3, 2) + 100 + + a = zarr.create_array( + MemoryStore(), + shape=base.shape, + chunks=(2, 4), + dtype=base.dtype, + serializer=serializer, + compressors=None, + fill_value=0, + ) + a[:] = base + a.oindex[selection] = value + + expected = base.copy() + expected[np.ix_(*selection)] = value + assert np.array_equal(a[:], expected) + assert np.array_equal(a.oindex[selection], value) From 01aac1fc8a266ad37578c590e2e7dd17357a9967 Mon Sep 17 00:00:00 2001 From: Dylan Pulver Date: Tue, 25 Aug 2026 16:52:13 +0300 Subject: [PATCH 2/3] docs: add changelog entry --- changes/4284.bugfix.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/4284.bugfix.md diff --git a/changes/4284.bugfix.md b/changes/4284.bugfix.md new file mode 100644 index 0000000000..bee3c7d258 --- /dev/null +++ b/changes/4284.bugfix.md @@ -0,0 +1 @@ +Fixed a `ValueError` when setting an orthogonal selection on a sharded array where more than one dimension is indexed by an array. The sharding codec re-derives an indexer from the chunk selection it is handed, which turns such a selection into a coordinate selection addressing the value buffer flat, so the write failed on a shape mismatch. From 3ebedebcc34dae06dc9013de9330bf24d0e0d37d Mon Sep 17 00:00:00 2001 From: Dylan Pulver Date: Wed, 26 Aug 2026 12:49:15 +0300 Subject: [PATCH 3/3] fix: reshape the value buffer in the sync partial-encode twin too _encode_partial_sync derives its indexer the same way as _encode_partial_single and so hits the same coordinate-selection shape mismatch under FusedCodecPipeline. The regression test is parametrized over both pipelines. --- changes/4284.bugfix.md | 2 +- src/zarr/codecs/sharding.py | 17 +++++++---- tests/test_codecs/test_sharding.py | 45 +++++++++++++++++++----------- 3 files changed, 41 insertions(+), 23 deletions(-) diff --git a/changes/4284.bugfix.md b/changes/4284.bugfix.md index bee3c7d258..4ad1c456b4 100644 --- a/changes/4284.bugfix.md +++ b/changes/4284.bugfix.md @@ -1 +1 @@ -Fixed a `ValueError` when setting an orthogonal selection on a sharded array where more than one dimension is indexed by an array. The sharding codec re-derives an indexer from the chunk selection it is handed, which turns such a selection into a coordinate selection addressing the value buffer flat, so the write failed on a shape mismatch. +Fixed a `ValueError` when setting an orthogonal selection on a sharded array where more than one dimension is indexed by an array. The sharding codec re-derives an indexer from the chunk selection it is handed, which turns such a selection into a coordinate selection addressing the value buffer flat, so the write failed on a shape mismatch. Both partial-encode paths are fixed, so the write works under either codec pipeline. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 9eb90b4ffc..7a4a339f13 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -796,13 +796,18 @@ def _encode_partial_sync( chunk_spec = self._get_chunk_spec(shard_spec) inner_transform = self._get_inner_chunk_transform(shard_spec) - indexer = list( - get_indexer( - selection, - shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, self.chunk_shape), - ) + shard_indexer = get_indexer( + selection, + shape=shard_shape, + chunk_grid=ChunkGrid.from_sizes(shard_shape, self.chunk_shape), ) + # A coordinate indexer flattens the selection, so its projections address + # `value` as 1-D while the caller shaped it like `sel_shape`. Mirrors the + # reshape `_encode_partial_single` applies on the async path. + sel_shape = getattr(shard_indexer, "sel_shape", None) + if sel_shape is not None and value.shape == sel_shape: + value = value.reshape(shard_indexer.shape) + indexer = list(shard_indexer) is_complete = self._is_complete_shard_write(indexer, chunks_per_shard) diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index 4cc80c912f..80fea63760 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -1264,8 +1264,15 @@ def test_shard_reader_to_dict_vectorized(chunks_per_shard: tuple[int, ...]) -> N assert buf is None +@pytest.mark.parametrize( + "pipeline_path", + [ + "zarr.core.codec_pipeline.FusedCodecPipeline", + "zarr.core.codec_pipeline.BatchedCodecPipeline", + ], +) @pytest.mark.parametrize("nested", [False, True], ids=["single", "nested"]) -def test_sharding_orthogonal_set_multiple_array_dims(nested: bool) -> None: +def test_sharding_orthogonal_set_multiple_array_dims(nested: bool, pipeline_path: str) -> None: """Orthogonal set with more than one array-indexed dimension. ``OrthogonalIndexer`` converts such a chunk selection to an ``np.ix_`` pair @@ -1273,6 +1280,11 @@ def test_sharding_orthogonal_set_multiple_array_dims(nested: bool) -> None: sharding codec re-derives an indexer from that selection and gets a ``CoordinateIndexer``, whose projections address the value buffer flat. Regression test for the resulting shape mismatch on write. + + Parametrized over both pipelines because the partial-encode path is + written twice -- ``_encode_partial_single`` for ``BatchedCodecPipeline`` + and ``_encode_partial_sync`` for ``FusedCodecPipeline`` -- and each + derives its own indexer. """ inner = ShardingCodec(chunk_shape=(1, 1), codecs=(BytesCodec(),)) serializer = ShardingCodec(chunk_shape=(2, 2), codecs=((inner,) if nested else (BytesCodec(),))) @@ -1280,19 +1292,20 @@ def test_sharding_orthogonal_set_multiple_array_dims(nested: bool) -> None: selection = (np.array([3, 1, 2]), np.array([0, 2])) value = np.arange(6, dtype="int32").reshape(3, 2) + 100 - a = zarr.create_array( - MemoryStore(), - shape=base.shape, - chunks=(2, 4), - dtype=base.dtype, - serializer=serializer, - compressors=None, - fill_value=0, - ) - a[:] = base - a.oindex[selection] = value + with zarr.config.set({"codec_pipeline.path": pipeline_path}): + a = zarr.create_array( + MemoryStore(), + shape=base.shape, + chunks=(2, 4), + dtype=base.dtype, + serializer=serializer, + compressors=None, + fill_value=0, + ) + a[:] = base + a.oindex[selection] = value - expected = base.copy() - expected[np.ix_(*selection)] = value - assert np.array_equal(a[:], expected) - assert np.array_equal(a.oindex[selection], value) + expected = base.copy() + expected[np.ix_(*selection)] = value + assert np.array_equal(a[:], expected) + assert np.array_equal(a.oindex[selection], value)