From 4ccec6ae7d7d249d86c3423a7012b606d587e22c Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 4 Sep 2026 10:57:50 +0200 Subject: [PATCH 1/6] fix(sharding): flatten any multi-dim value for coordinate selections in partial writes The guard added in #4284 only reshaped the value when its shape equalled the re-derived CoordinateIndexer's sel_shape. An orthogonal selection that mixes an integer index with two or more array indices defeats that: OrthogonalIndexer drops the integer axis from the value but np.ix_ keeps it as a length-1 axis in the chunk selection, so the shapes differ in rank while agreeing in element count, the reshape was skipped, and the write still raised the shape-mismatch ValueError. The invariant is that a coordinate indexer addresses the value flat, so ravel any multi-dimensional value instead. Both partial-encode paths now share one helper for deriving the shard indexer and shaping the value, and the check is an isinstance on CoordinateIndexer so mypy types sel_shape. The regression test is parametrized over selections with an integer axis in each position, three array axes, and an unsorted selection spanning two shards. Closes zarr-developers/zarr-python#4315 Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 --- src/zarr/codecs/sharding.py | 57 ++++++++++++++++-------------- tests/test_codecs/test_sharding.py | 51 +++++++++++++++++++------- 2 files changed, 69 insertions(+), 39 deletions(-) diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 7a4a339f13..6a9c1da1f0 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -57,6 +57,7 @@ from zarr.core.indexing import ( BasicIndexer, ChunkProjection, + CoordinateIndexer, SelectorTuple, SliceDimIndexer, _lexicographic_order, @@ -791,23 +792,11 @@ def _encode_partial_sync( Loads the existing shard, merges the written region into the affected inner chunks, and rewrites the whole shard. """ - shard_shape = shard_spec.shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) chunk_spec = self._get_chunk_spec(shard_spec) inner_transform = self._get_inner_chunk_transform(shard_spec) - 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) + indexer, value = self._get_shard_indexer_and_value(selection, shard_spec, value) is_complete = self._is_complete_shard_write(indexer, chunks_per_shard) @@ -1359,23 +1348,10 @@ async def _encode_partial_single( selection: SelectorTuple, shard_spec: ArraySpec, ) -> None: - shard_shape = shard_spec.shape - chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) chunk_spec = self._get_chunk_spec(shard_spec) - 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) + indexer, shard_array = self._get_shard_indexer_and_value(selection, shard_spec, shard_array) if self._is_complete_shard_write(indexer, chunks_per_shard): shard_dict = dict.fromkeys(lexicographic_order_coords(chunks_per_shard)) @@ -1433,6 +1409,33 @@ async def _encode_shard_dict( index_bytes, buffers, buffer_prototype, chunks_per_shard=chunks_per_shard ) + def _get_shard_indexer_and_value( + self, selection: SelectorTuple, shard_spec: ArraySpec, value: NDBuffer + ) -> tuple[list[ChunkProjection], NDBuffer]: + """Index ``selection`` over the inner chunk grid, flattening ``value`` to match. + + ``get_indexer`` classifies a tuple of integer arrays as a coordinate + selection, and a ``CoordinateIndexer`` addresses the value buffer as + 1-D. The caller shaped ``value`` like the selection it came from: an + ``OrthogonalIndexer`` hands down an ``np.ix_`` tuple with a value of + the broadcast shape, minus any integer-indexed axes it dropped. So the + value can have any number of dimensions but the same element count and + C order as the coordinate indexer's flattened projections; ravel it. + Scalars pass through and are broadcast downstream. + + The partial-decode paths apply the inverse reshape, to + ``indexer.sel_shape``, on the way out. + """ + shard_shape = shard_spec.shape + indexer = get_indexer( + selection, + shape=shard_shape, + chunk_grid=ChunkGrid.from_sizes(shard_shape, self.chunk_shape), + ) + if isinstance(indexer, CoordinateIndexer) and len(value.shape) > 1: + value = value.reshape(indexer.shape) + return list(indexer), value + def _is_total_shard( self, all_chunk_coords: set[tuple[int, ...]], chunks_per_shard: tuple[int, ...] ) -> bool: diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index 80fea63760..0954b7209f 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -1,4 +1,5 @@ import enum +import math import pickle import warnings from typing import Any, cast, get_args @@ -1272,31 +1273,57 @@ def test_shard_reader_to_dict_vectorized(chunks_per_shard: tuple[int, ...]) -> N ], ) @pytest.mark.parametrize("nested", [False, True], ids=["single", "nested"]) -def test_sharding_orthogonal_set_multiple_array_dims(nested: bool, pipeline_path: str) -> None: +@pytest.mark.parametrize( + "selection", + [ + pytest.param((np.array([3, 1, 2]), np.array([0, 2])), id="2d-arr-arr"), + pytest.param((np.array([3, 0]), np.array([2, 0])), id="2d-arr-arr-unsorted-two-shards"), + pytest.param((np.array([3, 1, 2]), 1, np.array([0, 2])), id="3d-arr-int-arr"), + pytest.param((1, np.array([0, 2]), np.array([1, 3])), id="3d-int-arr-arr"), + pytest.param((np.array([3, 1]), np.array([0, 2]), 2), id="3d-arr-arr-int"), + pytest.param( + (np.array([3, 1, 2]), np.array([0, 2]), np.array([1, 3])), id="3d-arr-arr-arr" + ), + ], +) +def test_sharding_orthogonal_set_multiple_array_dims( + selection: tuple[int | npt.NDArray[np.intp], ...], 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 + ``OrthogonalIndexer`` converts such a chunk selection to an ``np.ix_`` tuple 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. + An integer index alongside the arrays is the case that a shape-equality + guard misses: ``OrthogonalIndexer`` drops that axis from the value but + ``np.ix_`` keeps it as a length-1 axis in the chunk selection, so the value + and the re-derived indexer's ``sel_shape`` differ in rank while agreeing in + element count. + 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. + and ``_encode_partial_sync`` for ``FusedCodecPipeline``. """ - 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 + ndim = len(selection) + shape = (4,) * ndim + inner = ShardingCodec(chunk_shape=(1,) * ndim, codecs=(BytesCodec(),)) + serializer = ShardingCodec( + chunk_shape=(2,) * ndim, codecs=((inner,) if nested else (BytesCodec(),)) + ) + base = np.arange(4**ndim, dtype="int32").reshape(shape) + ix = np.ix_(*(np.atleast_1d(s) for s in selection)) + # The value is shaped like the orthogonal result: integer axes dropped. + value_shape = tuple(len(s) for s in selection if not isinstance(s, int)) + value = np.arange(math.prod(value_shape), dtype="int32").reshape(value_shape) + 100 with zarr.config.set({"codec_pipeline.path": pipeline_path}): a = zarr.create_array( MemoryStore(), - shape=base.shape, - chunks=(2, 4), + shape=shape, + chunks=(2,) + (4,) * (ndim - 1), dtype=base.dtype, serializer=serializer, compressors=None, @@ -1306,6 +1333,6 @@ def test_sharding_orthogonal_set_multiple_array_dims(nested: bool, pipeline_path a.oindex[selection] = value expected = base.copy() - expected[np.ix_(*selection)] = value + expected[ix] = value.reshape(expected[ix].shape) assert np.array_equal(a[:], expected) assert np.array_equal(a.oindex[selection], value) From ee258844eb6e86ec94764bc442f4fd2e5b5c6b2d Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 4 Sep 2026 10:58:11 +0200 Subject: [PATCH 2/6] docs: add changelog entry for #320 Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 --- changes/320.bugfix.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/320.bugfix.md diff --git a/changes/320.bugfix.md b/changes/320.bugfix.md new file mode 100644 index 0000000000..a263bafce7 --- /dev/null +++ b/changes/320.bugfix.md @@ -0,0 +1 @@ +Fixed a `ValueError` when setting an orthogonal selection on a sharded array that mixes an integer index with two or more array indices, such as `a.oindex[[3, 1, 2], 1, [0, 2]] = value`. The fix for the pure array-only case in #4284 reshaped the value only when its shape matched the coordinate selection; the value now flattens whenever the sharding codec re-derives a coordinate indexer, which covers the dropped integer axis as well. Both partial-encode paths share one helper for this. From 8a61445bfccd1a6d10a00a6b974228f3f4f49d1e Mon Sep 17 00:00:00 2001 From: Davis Bennett Date: Fri, 4 Sep 2026 11:28:26 +0200 Subject: [PATCH 3/6] Rename 320.bugfix.md to 4316.bugfix.md --- changes/{320.bugfix.md => 4316.bugfix.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename changes/{320.bugfix.md => 4316.bugfix.md} (100%) diff --git a/changes/320.bugfix.md b/changes/4316.bugfix.md similarity index 100% rename from changes/320.bugfix.md rename to changes/4316.bugfix.md From 49f887221e252964173e3ef21f0025868e4bad9e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 4 Sep 2026 12:02:49 +0200 Subject: [PATCH 4/6] fix(sharding): ravel only a value shaped like the selection minus its unit axes Ravelling every multi-dimensional value for a coordinate selection was too lenient. A mask write with a (2, 2) value for four selected elements, or an orthogonal write with a spurious trailing axis, raises on an unsharded array but was silently accepted on a sharded one, because the element count matched and the shard-level selection cannot tell orthogonal from mask indexing. The value shape can. An np.ix_ selection has an N-D sel_shape and the caller's value is that shape minus the integer-indexed axes, which np.ix_ keeps as length-1 axes. Ravel exactly that shape and leave any other rank alone, so an invalid write fails the same way it does without sharding. Adds an error test for both leniencies and a positive case with a length-1 array axis next to an integer axis. Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 --- changes/4316.bugfix.md | 2 +- src/zarr/codecs/sharding.py | 45 +++++++++++++++++++++++++----- tests/test_codecs/test_sharding.py | 45 ++++++++++++++++++++++++++++++ 3 files changed, 84 insertions(+), 8 deletions(-) diff --git a/changes/4316.bugfix.md b/changes/4316.bugfix.md index a263bafce7..e4b9bd8d94 100644 --- a/changes/4316.bugfix.md +++ b/changes/4316.bugfix.md @@ -1 +1 @@ -Fixed a `ValueError` when setting an orthogonal selection on a sharded array that mixes an integer index with two or more array indices, such as `a.oindex[[3, 1, 2], 1, [0, 2]] = value`. The fix for the pure array-only case in #4284 reshaped the value only when its shape matched the coordinate selection; the value now flattens whenever the sharding codec re-derives a coordinate indexer, which covers the dropped integer axis as well. Both partial-encode paths share one helper for this. +Fixed a `ValueError` when setting an orthogonal selection on a sharded array that mixes an integer index with two or more array indices, such as `a.oindex[[3, 1, 2], 1, [0, 2]] = value`. The fix for the array-only case in #4284 reshaped the value only when its shape matched the coordinate selection exactly; the sharding codec now also ravels a value that is the selection's shape minus the integer-indexed axes. Values of any other rank are left alone, so a write that is invalid on an unsharded array fails the same way on a sharded one. Both partial-encode paths share one helper for this. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 6a9c1da1f0..bd0760f7e3 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -400,6 +400,25 @@ def to_dict_vectorized(self) -> dict[tuple[int, ...], Buffer | None]: return result +def _drops_only_unit_axes(shape: tuple[int, ...], full: tuple[int, ...]) -> bool: + """Return whether ``shape`` is ``full`` with zero or more length-1 axes removed. + + ``(2, 2)`` is ``(2, 1, 2)`` minus its unit axis, and ``(1, 2)`` is + ``(1, 2, 1)`` minus its last; ``(2, 2)`` is not ``(4,)``, and + ``(3, 2, 1)`` is not ``(3, 2)`` because it adds an axis. + """ + remaining = iter(full) + for size in shape: + for full_size in remaining: + if full_size == size: + break + if full_size != 1: + return False + else: + return False + return all(full_size == 1 for full_size in remaining) + + @dataclass(frozen=True) class ShardingCodec( ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin @@ -1416,12 +1435,20 @@ def _get_shard_indexer_and_value( ``get_indexer`` classifies a tuple of integer arrays as a coordinate selection, and a ``CoordinateIndexer`` addresses the value buffer as - 1-D. The caller shaped ``value`` like the selection it came from: an - ``OrthogonalIndexer`` hands down an ``np.ix_`` tuple with a value of - the broadcast shape, minus any integer-indexed axes it dropped. So the - value can have any number of dimensions but the same element count and - C order as the coordinate indexer's flattened projections; ravel it. - Scalars pass through and are broadcast downstream. + 1-D. An ``OrthogonalIndexer`` with two or more array-indexed axes hands + down an ``np.ix_`` tuple, so ``sel_shape`` is N-D, while the caller + shaped ``value`` like the orthogonal result: the broadcast shape minus + the integer-indexed axes, which ``np.ix_`` keeps as length-1 axes. That + value has the element count and C order of the flattened projections, + so ravel it. + + Only that value shape is ravelled. A mask or coordinate selection + arrives with a 1-D ``sel_shape`` and a value that must already be flat, + and an orthogonal value with an axis the selection does not have is + invalid. Both are left alone so the write fails the same way it does + on an unsharded array; the shard-level selection alone cannot tell + orthogonal from mask indexing, the value shape can. Scalars pass + through and are broadcast downstream. The partial-decode paths apply the inverse reshape, to ``indexer.sel_shape``, on the way out. @@ -1432,7 +1459,11 @@ def _get_shard_indexer_and_value( shape=shard_shape, chunk_grid=ChunkGrid.from_sizes(shard_shape, self.chunk_shape), ) - if isinstance(indexer, CoordinateIndexer) and len(value.shape) > 1: + if ( + isinstance(indexer, CoordinateIndexer) + and len(value.shape) > 1 + and _drops_only_unit_axes(value.shape, indexer.sel_shape) + ): value = value.reshape(indexer.shape) return list(indexer), value diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index 0954b7209f..56efa8841b 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -1284,6 +1284,7 @@ def test_shard_reader_to_dict_vectorized(chunks_per_shard: tuple[int, ...]) -> N pytest.param( (np.array([3, 1, 2]), np.array([0, 2]), np.array([1, 3])), id="3d-arr-arr-arr" ), + pytest.param((np.array([3]), np.array([0, 2]), 1), id="3d-arr1-arr-int"), ], ) def test_sharding_orthogonal_set_multiple_array_dims( @@ -1336,3 +1337,47 @@ def test_sharding_orthogonal_set_multiple_array_dims( expected[ix] = value.reshape(expected[ix].shape) assert np.array_equal(a[:], expected) assert np.array_equal(a.oindex[selection], value) + + +@pytest.mark.parametrize( + "pipeline_path", + [ + "zarr.core.codec_pipeline.FusedCodecPipeline", + "zarr.core.codec_pipeline.BatchedCodecPipeline", + ], +) +@pytest.mark.parametrize( + ("kind", "selection", "value_shape"), + [ + pytest.param("vindex", np.eye(4, dtype=bool), (2, 2), id="mask-2d-value"), + pytest.param( + "oindex", (np.array([3, 1, 2]), np.array([0, 2])), (3, 2, 1), id="oindex-extra-axis" + ), + ], +) +def test_sharding_set_rejects_value_with_wrong_rank( + kind: str, selection: Any, value_shape: tuple[int, ...], pipeline_path: str +) -> None: + """A value whose rank does not fit the selection fails on sharded arrays too. + + The sharding codec ravels a multi-dimensional value for a coordinate + selection only when that value is the selection's broadcast shape minus + integer-indexed axes. A 2-D value for a mask, or an orthogonal value with + an axis the selection lacks, is rejected on an unsharded array and must not + be silently accepted on a sharded one just because the element count + matches. + """ + value = np.arange(math.prod(value_shape), dtype="int32").reshape(value_shape) + with zarr.config.set({"codec_pipeline.path": pipeline_path}): + a = zarr.create_array( + MemoryStore(), + shape=(4, 4), + chunks=(4, 4), + dtype="int32", + serializer=ShardingCodec(chunk_shape=(2, 2), codecs=(BytesCodec(),)), + compressors=None, + fill_value=0, + ) + with pytest.raises(ValueError, match="shape mismatch"): + getattr(a, kind)[selection] = value + assert np.all(a[:] == 0) From aaf30c86901da0e845e35508bdb7d0534ffcf246 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 4 Sep 2026 12:50:16 +0200 Subject: [PATCH 5/6] test: exercise sharded writes and bare integer axes in the indexing property tests The property tests could not have found the sharded orthogonal-write bugs: - test_oindex, test_mask_indexing and test_block_indexing skipped their set half on sharded arrays with assume(zarray.shards is None), added in #2825 when the bug was first seen and never lifted. test_vindex had its set half commented out. - orthogonal_indices wrapped every bare integer as a one-element array, so zarr never received an integer index and OrthogonalIndexer's dropped-axis path was unreachable. basic_indices(min_dims=1) never yields an integer either, so that branch was dead. - arrays() only drew a shard shape when every axis had a chunk strictly between 1 and the axis length, on top of the v3 and regular-grid draws: 2 of 500 test_oindex examples were sharded. Lift the skips, draw integers explicitly and give the numpy indexer the same dropped-axis result, enable the vindex write with a duplicate-point filter, and let any chunk that fits the array be sharded (33 of 500 now). With these changes test_oindex fails against the code before this PR with the mixed-integer shape mismatch, and passes with it. Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 --- src/zarr/testing/strategies.py | 53 +++++++++++++++++++++++----------- tests/test_properties.py | 24 +++++++-------- 2 files changed, 47 insertions(+), 30 deletions(-) diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 6679dbcee4..f7ddca969e 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -328,11 +328,14 @@ def arrays( else: chunks_param = draw(chunk_shapes(shape=nparray.shape), label="chunk shape") - if all(s > c > 1 for s, c in zip(nparray.shape, chunks_param, strict=True)): + # Any chunk that fits the array can be sharded: shard_shapes draws a + # whole number of chunks per axis, one inner chunk included. + if all(s >= c >= 1 for s, c in zip(nparray.shape, chunks_param, strict=True)): shard_shape = draw( st.none() | shard_shapes(shape=nparray.shape, chunk_shape=chunks_param), label="shard shape", ) + event("sharded" if shard_shape is not None else "unsharded") if shard_shape is not None: subchunk_write_order = draw(subchunk_write_orders) inner_codecs = draw(sharding_inner_codecs, label="sharding inner codecs") @@ -572,20 +575,30 @@ def basic_indices( @st.composite def orthogonal_indices( draw: st.DrawFn, *, shape: tuple[int, ...] -) -> tuple[tuple[np.ndarray[Any, Any], ...], tuple[np.ndarray[Any, Any], ...]]: +) -> tuple[tuple[int | slice | np.ndarray[Any, Any], ...], tuple[np.ndarray[Any, Any], ...]]: """ Strategy that returns - (1) a tuple of integer arrays used for orthogonal indexing of Zarr arrays. - (2) a tuple of integer arrays that can be used for equivalent indexing of numpy arrays + (1) a tuple of per-axis selectors (integer array, slice, or bare integer) for + orthogonal indexing of Zarr arrays. + (2) a tuple of broadcast integer arrays that index a numpy array to the same + result. A bare integer drops its axis, as ``oindex`` does, so it is + given as a 0-d array and does not contribute a result dimension. """ - zindexer = [] - npindexer = [] - ndim = len(shape) + zindexer: list[int | slice | np.ndarray[Any, Any]] = [] + kept: list[tuple[int, np.ndarray[Any, Any]]] = [] + npindexer: dict[int, np.ndarray[Any, Any]] = {} for axis, size in enumerate(shape): if size != 0: - strategy = npst.integer_array_indices( - shape=(size,), result_shape=npst.array_shapes(min_side=1, max_side=size, max_dims=1) - ) | basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False) + strategy = ( + npst.integer_array_indices( + shape=(size,), + result_shape=npst.array_shapes(min_side=1, max_side=size, max_dims=1), + ) + | basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False) + # basic_indices(min_dims=1) never yields a bare integer, so draw + # one explicitly: it is the only selector that drops an axis. + | st.integers(min_value=-size, max_value=size - 1) + ) else: strategy = basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False) @@ -597,19 +610,25 @@ def orthogonal_indices( .filter(bool) ) (idxr,) = val - if isinstance(idxr, int): - idxr = np.array([idxr]) zindexer.append(idxr) + if isinstance(idxr, int): + npindexer[axis] = np.array(idxr) + continue if isinstance(idxr, slice): idxr = np.arange(*idxr.indices(size)) - elif isinstance(idxr, (tuple, int)): + elif isinstance(idxr, tuple): idxr = np.array(idxr) - newshape = [1] * ndim - newshape[axis] = idxr.size - npindexer.append(idxr.reshape(newshape)) + kept.append((axis, idxr)) + + for pos, (axis, idxr) in enumerate(kept): + newshape = [1] * len(kept) + newshape[pos] = idxr.size + npindexer[axis] = idxr.reshape(newshape) # casting the output of broadcast_arrays is needed for numpy < 2 - return tuple(zindexer), tuple(np.broadcast_arrays(*npindexer)) + return tuple(zindexer), tuple( + np.broadcast_arrays(*(npindexer[axis] for axis in range(len(shape)))) + ) @st.composite diff --git a/tests/test_properties.py b/tests/test_properties.py index 33888bfd4e..d57475df9e 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -176,9 +176,8 @@ async def test_oindex(data: st.DataObject) -> None: actual = await async_zarray.oindex.getitem(zindexer) assert_array_equal(nparray[npindexer], actual) - # sync get - assume(zarray.shards is None) # GH2834 - for idxr in npindexer: + # sync set + for idxr in zindexer: if isinstance(idxr, np.ndarray) and idxr.size != np.unique(idxr).size: # behaviour of setitem with repeated indices is not guaranteed in practice assume(False) @@ -217,13 +216,14 @@ async def test_vindex(data: st.DataObject) -> None: assert_array_equal(nparray[indexer], actual) # sync set - # FIXME! - # when the indexer is such that a value gets overwritten multiple times, - # I think the output depends on chunking. - # new_data = data.draw(npst.arrays(shape=st.just(actual.shape), dtype=nparray.dtype)) - # nparray[indexer] = new_data - # zarray.vindex[indexer] = new_data - # assert_array_equal(nparray, zarray[:]) + points = np.stack([idxr.ravel() for idxr in np.broadcast_arrays(*indexer)], axis=-1) + if len(np.unique(points, axis=0)) != len(points): + # behaviour of setitem with repeated coordinates is not guaranteed in practice + assume(False) + new_data = data.draw(numpy_arrays(shapes=st.just(actual.shape), dtype=nparray.dtype)) + nparray[indexer] = new_data + zarray.vindex[indexer] = new_data + assert_array_equal(nparray, zarray[:]) # note: async vindex setitem not yet implemented @@ -243,7 +243,6 @@ def test_mask_indexing(data: st.DataObject) -> None: assert_array_equal(expected, zarray.vindex[mask]) # sync set, via both interfaces - assume(zarray.shards is None) # GH2834 new_data = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=nparray.dtype)) nparray[mask] = new_data zarray.set_mask_selection(mask, new_data) @@ -270,8 +269,7 @@ def test_block_indexing(data: st.DataObject) -> None: assert_array_equal(expected, zarray.blocks[block_indexer]) assert_array_equal(expected, zarray.get_block_selection(block_indexer)) - # sync set, via both interfaces; sharded set is broken upstream (GH2834) - assume(zarray.shards is None) + # sync set, via both interfaces new_data = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=nparray.dtype)) nparray[array_indexer] = new_data zarray.blocks[block_indexer] = new_data From 23a0003aa97f4e9e5f11a894ae876a8ba928b563 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 4 Sep 2026 13:05:03 +0200 Subject: [PATCH 6/6] test: assert invalid value ranks are rejected on chunked and sharded arrays alike Replace the sharded-only wrong-rank test in test_sharding.py and the GH2469 one-off in test_indexing.py with one parametrized error test: a coordinate write with twice the elements, a mask write with a 2-D value, and an orthogonal write with an extra axis each raise ValueError on chunked and sharded arrays under both codec pipelines. The property under test is that storage layout does not change which writes are rejected, which a sharded-only test could not state. zarr_array_from_numpy_array grows a shards argument for it. Only the rejection is asserted; a write that fails inside the chunk merge may already have touched other chunks on a chunked array. Assisted-by: ClaudeCode:claude-fable-5-1 Co-Authored-By: Claude Fable 5.1 --- tests/test_codecs/test_sharding.py | 44 -------------------- tests/test_indexing.py | 64 ++++++++++++++++++++++-------- 2 files changed, 47 insertions(+), 61 deletions(-) diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index 56efa8841b..42a88ddc7a 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -1337,47 +1337,3 @@ def test_sharding_orthogonal_set_multiple_array_dims( expected[ix] = value.reshape(expected[ix].shape) assert np.array_equal(a[:], expected) assert np.array_equal(a.oindex[selection], value) - - -@pytest.mark.parametrize( - "pipeline_path", - [ - "zarr.core.codec_pipeline.FusedCodecPipeline", - "zarr.core.codec_pipeline.BatchedCodecPipeline", - ], -) -@pytest.mark.parametrize( - ("kind", "selection", "value_shape"), - [ - pytest.param("vindex", np.eye(4, dtype=bool), (2, 2), id="mask-2d-value"), - pytest.param( - "oindex", (np.array([3, 1, 2]), np.array([0, 2])), (3, 2, 1), id="oindex-extra-axis" - ), - ], -) -def test_sharding_set_rejects_value_with_wrong_rank( - kind: str, selection: Any, value_shape: tuple[int, ...], pipeline_path: str -) -> None: - """A value whose rank does not fit the selection fails on sharded arrays too. - - The sharding codec ravels a multi-dimensional value for a coordinate - selection only when that value is the selection's broadcast shape minus - integer-indexed axes. A 2-D value for a mask, or an orthogonal value with - an axis the selection lacks, is rejected on an unsharded array and must not - be silently accepted on a sharded one just because the element count - matches. - """ - value = np.arange(math.prod(value_shape), dtype="int32").reshape(value_shape) - with zarr.config.set({"codec_pipeline.path": pipeline_path}): - a = zarr.create_array( - MemoryStore(), - shape=(4, 4), - chunks=(4, 4), - dtype="int32", - serializer=ShardingCodec(chunk_shape=(2, 2), codecs=(BytesCodec(),)), - compressors=None, - fill_value=0, - ) - with pytest.raises(ValueError, match="shape mismatch"): - getattr(a, kind)[selection] = value - assert np.all(a[:] == 0) diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 04fbdad8c6..4a1a4172d2 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -51,12 +51,14 @@ def zarr_array_from_numpy_array( store: StorePath, a: npt.NDArray[Any], chunk_shape: tuple[int, ...] | None = None, + shards: tuple[int, ...] | None = None, ) -> zarr.Array: z = zarr.create_array( store=store / str(uuid4()), shape=a.shape, dtype=a.dtype, chunks=chunk_shape or a.shape, + shards=shards, chunk_key_encoding={"name": "v2", "separator": "."}, ) z[()] = a @@ -2103,23 +2105,51 @@ def test_zero_sized_chunks(store: StorePath, shape: list[int]) -> None: assert_array_equal(z[...], np.zeros(shape, dtype="f8")) -@pytest.mark.parametrize("store", ["memory"], indirect=["store"]) -def test_vectorized_indexing_incompatible_shape(store) -> None: - """Regression for GH2469: vectorized set-indexing raises ValueError when the value shape is incompatible with the indexer shape.""" - # GH2469 - shape = (4, 4) - chunks = (2, 2) - fill_value = 32767 - arr = zarr.create( - shape, - store=store, - chunks=chunks, - dtype=np.int16, - fill_value=fill_value, - codecs=[zarr.codecs.BytesCodec(), zarr.codecs.BloscCodec()], - ) - with pytest.raises(ValueError, match="Attempting to set"): - arr[np.array([1, 2]), np.array([1, 2])] = np.array([[-1, -2], [-3, -4]]) +@pytest.mark.parametrize( + "pipeline_path", + [ + "zarr.core.codec_pipeline.BatchedCodecPipeline", + "zarr.core.codec_pipeline.FusedCodecPipeline", + ], + ids=["batched", "fused"], +) +@pytest.mark.parametrize("shards", [None, (4, 4)], ids=["chunked", "sharded"]) +@pytest.mark.parametrize( + ("kind", "selection", "value_shape"), + [ + # GH2469: array-level check, the value has twice the selected elements. + pytest.param("vindex", (np.array([1, 2]), np.array([1, 2])), (2, 2), id="coord-2d-value"), + # Right element count, wrong rank: a mask takes a flat value. + pytest.param("vindex", np.eye(4, dtype=bool), (2, 2), id="mask-2d-value"), + # Right element count, an axis the selection does not have. + pytest.param( + "oindex", (np.array([3, 1, 2]), np.array([0, 2])), (3, 2, 1), id="oindex-extra-axis" + ), + ], +) +def test_set_selection_rejects_value_with_wrong_rank( + store: StorePath, + kind: str, + selection: Any, + value_shape: tuple[int, ...], + shards: tuple[int, ...] | None, + pipeline_path: str, +) -> None: + """A value whose rank does not fit the selection raises regardless of storage layout. + + The sharding codec re-derives an indexer from the selection it is handed + and ravels the value when it is the selection's broadcast shape minus + integer-indexed axes. Any other rank must fail on a sharded array exactly + as it does on a chunked one; an element count that happens to match is + not grounds to accept it. Only the rejection is asserted: a write that + fails inside the chunk merge may already have touched other chunks. + """ + a = np.zeros((4, 4), dtype=np.int32) + value = np.arange(np.prod(value_shape), dtype=np.int32).reshape(value_shape) + with zarr.config.set({"codec_pipeline.path": pipeline_path}): + z = zarr_array_from_numpy_array(store, a, chunk_shape=(2, 2), shards=shards) + with pytest.raises(ValueError, match="Attempting to set|shape mismatch"): + getattr(z, kind)[selection] = value def test_iter_chunk_regions():