Support combo of advanced and basic indexing

[ilan-gold]

Discussed in #2485

^{Originally posted by ilan-gold November 12, 2024}
Hello,

I am not entirely clear about what is supported and what isn't in zarr indexing, especially in v3. I understand https://zarr.readthedocs.io/en/stable/tutorial.html#advanced-indexing as "the API is different because we explicitly implement oindex but otherwise the behavior should be the same as numpy."

In other words, it seems "coordinate indexing" -> vindex (or just normal bracket shortcutting on the array directly) and "orthogonal indexing" -> oindex.

But other than these things, I'm not sure I understand what is and isn't supported. For example

import tempfile
import numpy as np
import zarr
from zarr.storage import LocalStore

shape = (4, 4, 4)
chunks = (2, 2, 2)
fill_value = 32767

tmp = tempfile.TemporaryDirectory()
arr = zarr.create(
    shape,
    store=LocalStore(root=tmp.name, mode="w"),
    chunks=chunks,
    dtype=np.int16,
    fill_value=fill_value,
    codecs=[zarr.codecs.BytesCodec(), zarr.codecs.BloscCodec()],
)
index_tuple = (np.array([1, 2]), np.array([1, 2]), slice(None))
# works
np.arange(64).reshape(shape)[index_tuple]
# breaks
arr[index_tuple]

This is on zarr v3 3.0.0b1. Is this a bug? Is there any good documentation on this?

[dcherian]

This is a combination of "advanced indexing" or "vectorized indexing" along axes 0, 1 and basic indexing on axis 2. Looking at the code this is not supported (yet).

You're going to have expand the slice manually:

> arr[np.array([[1], [2]]), np.array([[1], [2]]), np.array([[0,1,2,3]])]
array([[20, 21, 22, 23],
       [40, 41, 42, 43]], dtype=int16)

[ilan-gold]

@dcherian So I guess I don't really mind one way or another whether it's supported, it's more that I would like to know what is supported and is not somehow. Is there some document somewhere or the like? I would expect "we replicate numpy indexing" to be the answer but it apparently is not.

https://github.com/ilan-gold/zarrs-python/blob/main/tests/test_pipeline.py has a nice suite of test - by expanding the >2d axis index types, there are a number of failures, and I don't know if they're bugs or not (not including some Ellipsis indexing issues in 2D I found as well: https://github.com/ilan-gold/zarrs-python/blob/bde5e3b9c800d21a2db2dd3d00430ee2f8cbfde3/tests/conftest.py#L69-L94?). I didn't want to post every failure because it's a lot and I don't know if they are bugs or not.