fix: apply drop_axes squeeze in partial decode path for sharding (#3691)

When reading sharded arrays with mixed integer/list indexing (e.g.
arr[0:10, 0, [0, 1]]), the outer OrthogonalIndexer produces chunk
selections that have been ix_()-transformed for orthogonal advanced
indexing. Integer indices become single-element ranges (size-1 dims)
via ix_() to enable NumPy orthogonal indexing.

In CodecPipeline.read_batch(), the non-partial path correctly applies
drop_axes.squeeze() to remove those size-1 integer dimensions before
writing to the output buffer. However, the partial decode path (used
by ShardingCodec) was missing this squeeze step.

The ShardingCodec._decode_partial_single() receives the ix_-transformed
chunk selection, interprets it as coordinate (fancy) indexing, and
returns an array with shape (10, 1, 2) instead of (10, 2). The
subsequent assignment to out[out_selection] then fails with:
  ValueError: could not broadcast input array from shape (10,1,2)
  into shape (10,2)

Fix: apply drop_axes squeeze to chunk_array in the partial decode
branch of read_batch(), matching the behaviour of the non-partial path.

Fixes #3691

Author: Cipher

src/zarr/core/codec_pipeline.py

@@ -263,6 +263,8 @@ async def read_batch(
                 chunk_array_batch, batch_info, strict=False
             ):
                 if chunk_array is not None:
+                    if drop_axes != ():
+                        chunk_array = chunk_array.squeeze(axis=drop_axes)
                     out[out_selection] = chunk_array
                 else:
                     out[out_selection] = fill_value_or_default(chunk_spec)

tests/test_codecs/test_sharding.py

@@ -501,3 +501,54 @@ def test_invalid_shard_shape() -> None:
             dtype=np.dtype("uint8"),
             fill_value=0,
         )
+
+
+@pytest.mark.parametrize("store", ["local"], indirect=["store"])
+def test_sharding_mixed_integer_list_indexing(store: Store) -> None:
+    """Regression test for https://github.com/zarr-developers/zarr-python/issues/3691.
+
+    Mixed integer/list indexing on sharded arrays should return the same
+    shape and data as on equivalent chunked arrays.
+    """
+    import numpy as np
+
+    data = np.arange(200 * 100 * 10, dtype=np.uint8).reshape(200, 100, 10)
+
+    chunked = zarr.create_array(
+        store,
+        name="chunked",
+        shape=(200, 100, 10),
+        dtype=np.uint8,
+        chunks=(200, 100, 1),
+        overwrite=True,
+    )
+    chunked[:, :, :] = data
+
+    sharded = zarr.create_array(
+        store,
+        name="sharded",
+        shape=(200, 100, 10),
+        dtype=np.uint8,
+        chunks=(200, 100, 1),
+        shards=(200, 100, 10),
+        overwrite=True,
+    )
+    sharded[:, :, :] = data
+
+    # Mixed integer + list indexing
+    c = chunked[0:10, 0, [0, 1]]
+    s = sharded[0:10, 0, [0, 1]]
+    assert c.shape == s.shape == (10, 2), f"Expected (10, 2), got chunked={c.shape}, sharded={s.shape}"
+    np.testing.assert_array_equal(c, s)
+
+    # Multiple integer axes
+    c2 = chunked[0, 0, [0, 1, 2]]
+    s2 = sharded[0, 0, [0, 1, 2]]
+    assert c2.shape == s2.shape == (3,)
+    np.testing.assert_array_equal(c2, s2)
+
+    # Slice + integer + slice
+    c3 = chunked[0:5, 1, 0:3]
+    s3 = sharded[0:5, 1, 0:3]
+    assert c3.shape == s3.shape == (5, 3)
+    np.testing.assert_array_equal(c3, s3)