fix(sharding): flatten any multi-dim value for coordinate selections in partial writes (zarr-developers#4316)

* fix(sharding): flatten any multi-dim value for coordinate selections in partial writes

The guard added in zarr-developers#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#4315

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* docs: add changelog entry for #320

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Rename 320.bugfix.md to 4316.bugfix.md

* 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 <noreply@anthropic.com>

* 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 zarr-developers#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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Author: d-v-b

changes/4316.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 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.

src/zarr/codecs/sharding.py

@@ -57,6 +57,7 @@
 from zarr.core.indexing import (
     BasicIndexer,
     ChunkProjection,
+    CoordinateIndexer,
     SelectorTuple,
     SliceDimIndexer,
     _lexicographic_order,
@@ -399,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
@@ -791,23 +811,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 +1367,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 +1428,45 @@ 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. 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.
+        """
+        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
+            and _drops_only_unit_axes(value.shape, indexer.sel_shape)
+        ):
+            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:

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

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,58 @@ 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"
+        ),
+        pytest.param((np.array([3]), np.array([0, 2]), 1), id="3d-arr1-arr-int"),
+    ],
+)
+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 +1334,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)

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():