fix: sharding codec with fancy indexing (#2817)

* fix: sharding codec with fancy indexing * changelog * add a better test * proper fix * fix: ArrayOfIntOrBool typing * Revert "fix: ArrayOfIntOrBool typing" This reverts commit 1a30563. * ignore typing error in test --------- Co-authored-by: Deepak Cherian <[email protected]>
zarr-developers · Feb 13, 2025 · 870265a · 870265a
1 parent feeb08f
commit 870265a
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 1 deletion.
diff --git a/changes/2817.bugfix.rst b/changes/2817.bugfix.rst
@@ -0,0 +1 @@
+Fix fancy indexing (e.g. arr[5, [0, 1]]) with the sharding codec
diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py
@@ -531,7 +531,11 @@ async def _decode_partial_single(
             ],
             out,
         )
-        return out
+
+        if hasattr(indexer, "sel_shape"):
+            return out.reshape(indexer.sel_shape)
+        else:
+            return out
 
     async def _encode_single(
         self,

diff --git a/tests/test_array.py b/tests/test_array.py
@@ -1429,3 +1429,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=(2, 3, 4),
+        chunks=(1, 2, 2),
+        overwrite=True,
+        dtype=np.float32,
+        shards=(2, 4, 4),
+    )
+    arr[:] = np.arange(2 * 3 * 4).reshape((2, 3, 4))
+    assert (arr[1, [0, 1]] == np.array([[12, 13, 14, 15], [16, 17, 18, 19]])).all()  # type: ignore[index]
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Fix fancy indexing (e.g. arr[5, [0, 1]]) with the sharding codec