fix: shape the value buffer for coordinate selections in sharded writes

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. Both partial-encode paths are fixed, so the write works under either codec pipeline.

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)
 
@@ -1359,13 +1364,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))

tests/test_codecs/test_sharding.py

@@ -1262,3 +1262,50 @@ 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(
+    "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, pipeline_path: str) -> 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.
+
+    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(),)))
+    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
+
+    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)