Shape mismatch with mixed integer/list indexing on Sharded arrays

[jn-cordulus]

Zarr version

v3.1.5

Numcodecs version

v0.16.5

Python Version

3.11.13

Operating System

Mac

Installation

uv

Description

I found a bug when using mixed indexing (combining slicing, integer indexing, lists indexing) on a sharded array. The indexing works fine on a chunked-only array.

Steps to reproduce

Minimal example:

import numpy as np
import zarr

chunked = zarr.create_array(
    "chunked.zarr",
    shape=(200, 100, 10),
    dtype=np.uint8,
    overwrite=True,
    chunks=(200, 100, 1),
)
chunked[:, :, :] = np.random.randint(0, 255, (200, 100, 10))

sharded = zarr.create_array(
    "sharded.zarr",
    shape=(200, 100, 10),
    dtype=np.uint8,
    overwrite=True,
    chunks=(200, 100, 1),
    shards=(200, 100, 10),
)
sharded[:, :, :] = np.random.randint(0, 255, (200, 100, 10))

print("chunked", chunked[0:10, 0, [0, 1]].shape)
print("sharded", sharded[0:10, 0, [0, 1]].shape)

Additional output

Outputs

chunked (10, 2)
...
ValueError: could not broadcast input array from shape (10,1,2) into shape (10,2)

Workaround: replacing 0 with a slice 0:1 works:

print(sharded[0:10, 0:1, [0, 1]].shape)

although the shape is then (10, 1, 2), not (10, 2).