Hoist per-chunk config lookups and parallelize deep group traversal

src/zarr/abc/codec.py

@@ -9,7 +9,7 @@
 from zarr.abc.metadata import Metadata
 from zarr.core.buffer import Buffer, NDBuffer
 from zarr.core.common import NamedConfig, concurrent_map
-from zarr.core.config import config
+from zarr.core.config import get_async_concurrency
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable, Iterable
@@ -249,7 +249,7 @@ async def decode_partial(
         return await concurrent_map(
             list(batch_info),
             self._decode_partial_single,
-            config.get("async.concurrency"),
+            get_async_concurrency(),
         )
 
 
@@ -286,7 +286,7 @@ async def encode_partial(
         await concurrent_map(
             list(batch_info),
             self._encode_partial_single,
-            config.get("async.concurrency"),
+            get_async_concurrency(),
         )
 
 
@@ -493,7 +493,7 @@ async def _batching_helper[CI: CodecInput, CO: CodecOutput](
     return await concurrent_map(
         list(batch_info),
         _noop_for_none(func),
-        config.get("async.concurrency"),
+        get_async_concurrency(),
     )
 
 

src/zarr/core/codec_pipeline.py

@@ -17,7 +17,7 @@
     SupportsSyncCodec,
 )
 from zarr.core.common import concurrent_map
-from zarr.core.config import config
+from zarr.core.config import async_concurrency_scope, config, get_async_concurrency
 from zarr.core.indexing import SelectorTuple, is_scalar
 from zarr.errors import ZarrUserWarning
 from zarr.registry import register_pipeline
@@ -395,7 +395,7 @@ async def read_batch(
                     for byte_getter, array_spec, *_ in batch_info_list
                 ],
                 lambda byte_getter, prototype: byte_getter.get(prototype),
-                config.get("async.concurrency"),
+                get_async_concurrency(),
             )
             chunk_array_batch = await self.decode_batch(
                 [
@@ -505,7 +505,7 @@ async def _read_key(
                     for byte_setter, chunk_spec, chunk_selection, _, is_complete_chunk in batch_info
                 ],
                 _read_key,
-                config.get("async.concurrency"),
+                get_async_concurrency(),
             )
             chunk_array_decoded = await self.decode_batch(
                 [
@@ -571,25 +571,27 @@ async def _write_key(byte_setter: ByteSetter, chunk_bytes: Buffer | None) -> Non
                     )
                 ],
                 _write_key,
-                config.get("async.concurrency"),
+                get_async_concurrency(),
             )
 
     async def decode(
         self,
         chunk_bytes_and_specs: Iterable[tuple[Buffer | None, ArraySpec]],
     ) -> Iterable[NDBuffer | None]:
         output: list[NDBuffer | None] = []
-        for batch_info in batched(chunk_bytes_and_specs, self.batch_size):
-            output.extend(await self.decode_batch(batch_info))
+        with async_concurrency_scope():
+            for batch_info in batched(chunk_bytes_and_specs, self.batch_size):
+                output.extend(await self.decode_batch(batch_info))
         return output
 
     async def encode(
         self,
         chunk_arrays_and_specs: Iterable[tuple[NDBuffer | None, ArraySpec]],
     ) -> Iterable[Buffer | None]:
         output: list[Buffer | None] = []
-        for single_batch_info in batched(chunk_arrays_and_specs, self.batch_size):
-            output.extend(await self.encode_batch(single_batch_info))
+        with async_concurrency_scope():
+            for single_batch_info in batched(chunk_arrays_and_specs, self.batch_size):
+                output.extend(await self.encode_batch(single_batch_info))
         return output
 
     async def read(
@@ -598,14 +600,15 @@ async def read(
         out: NDBuffer,
         drop_axes: tuple[int, ...] = (),
     ) -> tuple[GetResult, ...]:
-        batch_results = await concurrent_map(
-            [
-                (single_batch_info, out, drop_axes)
-                for single_batch_info in batched(batch_info, self.batch_size)
-            ],
-            self.read_batch,
-            config.get("async.concurrency"),
-        )
+        with async_concurrency_scope():
+            batch_results = await concurrent_map(
+                [
+                    (single_batch_info, out, drop_axes)
+                    for single_batch_info in batched(batch_info, self.batch_size)
+                ],
+                self.read_batch,
+                get_async_concurrency(),
+            )
         results: list[GetResult] = []
         for batch in batch_results:
             results.extend(batch)
@@ -617,14 +620,15 @@ async def write(
         value: NDBuffer,
         drop_axes: tuple[int, ...] = (),
     ) -> None:
-        await concurrent_map(
-            [
-                (single_batch_info, value, drop_axes)
-                for single_batch_info in batched(batch_info, self.batch_size)
-            ],
-            self.write_batch,
-            config.get("async.concurrency"),
-        )
+        with async_concurrency_scope():
+            await concurrent_map(
+                [
+                    (single_batch_info, value, drop_axes)
+                    for single_batch_info in batched(batch_info, self.batch_size)
+                ],
+                self.write_batch,
+                get_async_concurrency(),
+            )
 
 
 def codecs_from_list(

src/zarr/core/config.py

@@ -29,11 +29,15 @@
 
 from __future__ import annotations
 
+import contextlib
+import contextvars
 from typing import TYPE_CHECKING, Any, Literal, cast
 
 from donfig import Config as DConfig
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
+
     from donfig.config_obj import ConfigSet
 
 
@@ -148,6 +152,47 @@ def enable_gpu(self) -> ConfigSet:
 )
 
 
+# Operation-scoped cache for ``async.concurrency``.
+#
+# The per-chunk codec read/write path looks up ``async.concurrency`` many times
+# (once per chunk, plus once per codec per chunk via the batching helper). Each
+# lookup walks Donfig's nested config. To avoid that, a single read/write/
+# encode/decode operation captures the value once via ``async_concurrency_scope``
+# and the hot paths read it through ``get_async_concurrency``. Outside of a
+# scope we fall back to a live ``config.get``, so configuration changes always
+# take effect between operations.
+_async_concurrency_var: contextvars.ContextVar[int | None] = contextvars.ContextVar(
+    "zarr_async_concurrency"
+)
+
+
+def get_async_concurrency() -> int | None:
+    """Return the ``async.concurrency`` limit for the current operation.
+
+    Within an :func:`async_concurrency_scope`, returns the value captured once at
+    the start of the operation. Outside of a scope, falls back to a live
+    ``config.get`` lookup.
+    """
+    try:
+        return _async_concurrency_var.get()
+    except LookupError:
+        return cast("int | None", config.get("async.concurrency"))
+
+
+@contextlib.contextmanager
+def async_concurrency_scope() -> Iterator[None]:
+    """Capture ``async.concurrency`` once for the duration of an operation.
+
+    Reads of :func:`get_async_concurrency` inside this context return the
+    captured value rather than performing a fresh ``config.get`` lookup.
+    """
+    token = _async_concurrency_var.set(cast("int | None", config.get("async.concurrency")))
+    try:
+        yield
+    finally:
+        _async_concurrency_var.reset(token)
+
+
 def parse_indexing_order(data: Any) -> Literal["C", "F"]:
     if data in ("C", "F"):
         return cast("Literal['C', 'F']", data)

src/zarr/core/group.py

@@ -3286,6 +3286,13 @@ async def _iter_members(
                 raise ValueError(f"Unexpected type: {type(fetched_node)}")
 
 
+class _Drained:
+    """Sentinel signalling that a subgroup iterator has been fully drained."""
+
+
+_DRAINED = _Drained()
+
+
 async def _iter_members_deep(
     group: AsyncGroup,
     *,
@@ -3344,10 +3351,48 @@ async def _iter_members_deep(
                 node, max_depth=new_depth, skip_keys=skip_keys, semaphore=semaphore
             )
 
-    for prefix, subgroup_iter in to_recurse.items():
-        async for name, node in subgroup_iter:
-            key = f"{prefix}/{name}".lstrip("/")
-            yield key, node
+    # Recurse into sibling subgroups concurrently rather than draining each
+    # subgroup's iterator in sequence, so that sibling subtrees overlap their
+    # latency. The shared ``semaphore`` (threaded into every recursive call)
+    # still bounds the total number of concurrent ``getitem`` operations.
+    # Members are not guaranteed to be ordered, so we yield them as they arrive.
+    queue: asyncio.Queue[tuple[str, AnyAsyncArray | AsyncGroup] | Exception | _Drained] = (
+        asyncio.Queue()
+    )
+
+    async def _drain(
+        prefix: str,
+        subgroup_iter: AsyncGenerator[tuple[str, AnyAsyncArray | AsyncGroup], None],
+    ) -> None:
+        try:
+            async for name, node in subgroup_iter:
+                key = f"{prefix}/{name}".lstrip("/")
+                await queue.put((key, node))
+        except Exception as exc:
+            # Forward the error to the consumer so it surfaces from this
+            # generator just as it would from sequential iteration.
+            await queue.put(exc)
+        finally:
+            await queue.put(_DRAINED)
+
+    tasks = [
+        asyncio.create_task(_drain(prefix, subgroup_iter))
+        for prefix, subgroup_iter in to_recurse.items()
+    ]
+    remaining = len(tasks)
+    try:
+        while remaining:
+            item = await queue.get()
+            if isinstance(item, _Drained):
+                remaining -= 1
+            elif isinstance(item, Exception):
+                raise item
+            else:
+                yield item
+    finally:
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
 
 
 async def _read_metadata_v3(store: Store, path: str) -> ArrayV3Metadata | GroupMetadata: