Indexing sharded array with oindex is slow

[slevang]

Indexing a sharded array with oindex[[0], ...] is much slower than the equivalent slice 0:1.

import time
from tempfile import TemporaryDirectory

import numpy as np
import zarr

with TemporaryDirectory() as tmp:
    shape = (2, 1000, 100, 100)
    chunks = (1, 1000, 10, 10)
    shards = (1, 1000, 100, 100)

    z = zarr.create_array(
        tmp,
        shape=shape,
        chunks=chunks,
        shards=shards,
        dtype="f4",
    )
    z[:] = np.random.randn(*shape).astype(np.float32)

    t0 = time.perf_counter()
    a = z[0:1, :, :, :]
    t_slice = time.perf_counter() - t0

    t0 = time.perf_counter()
    b = z.oindex[[0], :, :, :]
    t_oindex = time.perf_counter() - t0

    print(f"slice   : {t_slice:.3f} s")
    print(f"oindex  : {t_oindex:.3f} s")

Gives:

slice   : 0.098 s
oindex  : 1.178 s