[v3] Batch array / group access

[d-v-b]

In v3, since the storage API is asynchronous, we can open multiple array or groups concurrently. This would be broadly useful, but we don't have a good template from zarr-python v2 to extrapolate from, so we have to invent something new here (new, relative to zarr-python, that is).

Over in #1804 @martindurant brought this up, and I suggested something like this:

def open_nodes(store: Store, paths: tuple[str, ...], options: dict[Literal["array", "group"], dict[str, Any]]) -> Array | Group:
  ...
 
def open_arrays(store: Store, paths: tuple[str, ...], options: dict[str, Any]) -> Array:
  ...

def open_groups(store: Store, paths: tuple[str, ...], options: dict[str, Any]) -> Group:
  ...

I was imagining that the arguments to these functions would be the paths of arrays / groups anywhere in a Zarr hierarchy; we could also have a group.open_groups() method which can only "see" sub-groups, and similarly for group.open_arrays().

An alternative would be to use a more general transactional context manager:

with transaction(store) as tx:
     a1_maybe = tx.open_array(...)
     a2_maybe = tx.open_array(...)
    # IO gets run concurrently in `__aexit__`

a1 = a1_maybe.result()
a2 = a2_maybe.result()

I'm a lot less sure of this second design, since I have never implemented anything like it. For example, should we use futures for the results of tx.open_array()?

Are there other ideas, or examples from other implementations / domains we could draw from?

[martindurant]

I would also add data-getter, so in the first model something like

def get_data({arr_obj1: ((0, [1, 2, 3], slice(None, None, None)), arr_obj2: (Ellipsis, 0), ...})

or similar,

In addition to a convenience addition to the functions above

group.descent_to_path(path1, path2, ...

(where store and options are fixed).

I am not suggesting these are the right signatures, but this is the functionality I would want. After all "open many nodes" is sort of already covered in v2 for the special case of consolidated metadata (one call, no more latency, rather than many concurrent calls).

[d-v-b]

just noting that in the context of xarray, some function or an AsyncGroup method for concurrently creating sub-arrays / sub-groups would be extremely useful from a performance POV. In particular, it would be great to replace this "create arrays serially" code in xarray with a concurrent zarr method that can create a bunch of arrays in one go. I will make some progress on this shortly.

[oliverwm1]

Just dropping by to say that a feature to enable concurrently loading many (e.g. 50-100) small (~1MB) arrays would be very useful :) let me know if I can support in any way.

[martindurant]

You should already be able to use the async class to open them and then asyncio.gather to await a bunch of them. Providing this as a convenience method callable from sync code is a good idea.

xarray does something similar to read coordinates of a large hierarchy on open.

[oliverwm1]

You should already be able to use the async class to open them and then asyncio.gather to await a bunch of them

I am struggling with conflicting event loops when I try do to this, but this may just be me not understanding how to do this properly. E.g.

RuntimeError: Task <Task pending name='Task-500' coro=<StorePath.get() running at /home/OliWM/miniconda3/envs/fme/lib/python3.11/site-packages/zarr/storage/_common.py:124> cb=[gather.<locals>._done_callback() at /home/OliWM/miniconda3/envs/fme/lib/python3.11/asyncio/tasks.py:764]> got Future <Future pending> attached to a different loop

xarray does something similar to read coordinates of a large hierarchy on open.

Do you have a link to this code? would be a helpful pointer.

[d-v-b]

@oliverwm1 do you have a runnable code example?