fix: sharding codec with fancy indexing

Author: LDeakin

src/zarr/codecs/sharding.py

@@ -481,8 +481,12 @@ async def _decode_partial_single(
         )
 
         # setup output array
+        if hasattr(indexer, "sel_shape"):
+            out_shape = indexer.sel_shape
+        else:
+            out_shape = indexer.shape
         out = shard_spec.prototype.nd_buffer.create(
-            shape=indexer.shape, dtype=shard_spec.dtype, order=shard_spec.order, fill_value=0
+            shape=out_shape, dtype=shard_spec.dtype, order=shard_spec.order, fill_value=0
         )
 
         indexed_chunks = list(indexer)

tests/test_array.py

@@ -1420,3 +1420,18 @@ def test_multiprocessing(store: Store, method: Literal["fork", "spawn", "forkser
 
     results = pool.starmap(_index_array, [(arr, slice(len(data)))])
     assert all(np.array_equal(r, data) for r in results)
+
+
+async def test_sharding_coordinate_selection() -> None:
+    store = MemoryStore()
+    g = zarr.open_group(store, mode="w")
+    arr = g.create_array(
+        name="a",
+        shape=(10, 20, 30),
+        chunks=(5, 1, 30),
+        overwrite=True,
+        dtype=np.float32,
+        shards=(5, 20, 30),
+    )
+    arr[:] = 1
+    assert (arr[5, [0, 1]] == 1).all()