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

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