feat: async pipeline

Cargo.toml

@@ -11,6 +11,10 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 pyo3 = { version = "0.27.1", features = ["abi3-py311"] }
 zarrs = { version = "0.23.6", features = ["async", "zlib", "pcodec", "bz2"] }
+# `AsyncStoragePartialDecoder` is `pub` in `zarrs_codec` but not re-exported by
+# `zarrs`; depend on it directly. Cargo unifies it to the same version `zarrs`
+# uses, so the codec types are identical to those re-exported via `zarrs`.
+zarrs_codec = { version = "0.2", features = ["async"] }
 rayon_iter_concurrent_limit = "0.2.0"
 rayon = "1.10.0"
 # fix for https://stackoverflow.com/questions/76593417/package-openssl-was-not-found-in-the-pkg-config-search-path
@@ -20,9 +24,11 @@ unsafe_cell_slice = "0.2.0"
 serde_json = "1.0.128"
 pyo3-stub-gen = "0.17.1"
 opendal = { version = "0.55.0", features = ["services-http"] }
-tokio = { version = "1.41.1", features = ["rt-multi-thread"] }
+tokio = { version = "1.41.1", features = ["rt-multi-thread", "sync"] }
+pyo3-async-runtimes = { version = "0.27.0", features = ["tokio-runtime"] }
 zarrs_opendal = "0.10.0"
 itertools = "0.14.0"
+futures = "0.3.31"
 bytemuck = { version = "1.24.0", features = ["must_cast"] }
 pyo3-object_store = "0.7.0" # object_store 0.12
 zarrs_object_store = "0.5.0" # object_store 0.12

python/zarrs/_internal.pyi

@@ -2,13 +2,53 @@
 # ruff: noqa: E501, F401
 
 import builtins
+import collections.abc
 import typing
 
+import numpy
 import numpy.typing
 import zarr.abc.store
 
+@typing.final
+class AsyncCodecPipelineImpl:
+    def __new__(
+        cls,
+        array_metadata: builtins.str,
+        store_config: zarr.abc.store.Store,
+        *,
+        validate_checksums: builtins.bool = False,
+        chunk_concurrent_minimum: builtins.int | None = None,
+        chunk_concurrent_maximum: builtins.int | None = None,
+        num_threads: builtins.int | None = None,
+        direct_io: builtins.bool = False,
+        file_handle_cache_size: builtins.int = 0,
+    ) -> AsyncCodecPipelineImpl: ...
+    def retrieve_chunks(
+        self, chunk_descriptions: typing.Sequence[ChunkItem]
+    ) -> builtins.list[collections.abc.Awaitable[numpy.typing.NDArray[numpy.uint8]]]:
+        r"""
+        One awaitable per chunk, each resolving to that chunk's decoded subset as a
+        flat `uint8` array; the caller copies it into the output.
+        """
+    def store_chunks_with_indices(
+        self,
+        chunk_descriptions: typing.Sequence[ChunkItem],
+        value: numpy.typing.NDArray[typing.Any],
+        write_empty_chunks: builtins.bool,
+    ) -> collections.abc.Awaitable[None]: ...
+
 @typing.final
 class ChunkItem:
+    @property
+    def io_shape(self) -> builtins.list[builtins.int]:
+        r"""
+        The shape of the input/output array, with dropped axes reinstated as 1.
+        """
+    @property
+    def io_selection(self) -> builtins.list[slice]:
+        r"""
+        The selection of this chunk's subset within an array of [`Self::io_shape`].
+        """
     def __new__(
         cls,
         key: builtins.str,

python/zarrs/pipeline.py

@@ -12,9 +12,17 @@
 from zarr.core import BatchedCodecPipeline
 from zarr.core.config import config
 from zarr.core.metadata import ArrayMetadata, ArrayV2Metadata, ArrayV3Metadata
+from zarr.storage import FsspecStore
+
+try:
+    from zarr.storage import ObjectStore
+
+    _ASYNC_STORES: tuple[type, ...] = (FsspecStore, ObjectStore)
+except ImportError:
+    _ASYNC_STORES = (FsspecStore,)
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable, Iterator
+    from collections.abc import Awaitable, Iterable, Iterator
     from typing import Self
 
     from zarr.abc.store import ByteGetter, ByteSetter, Store
@@ -24,7 +32,7 @@
     from zarr.core.indexing import SelectorTuple
     from zarr.dtype import ZDType
 
-from ._internal import CodecPipelineImpl
+from ._internal import AsyncCodecPipelineImpl, ChunkItem, CodecPipelineImpl
 from .utils import (
     DiscontiguousArrayError,
     FillValueNoneError,
@@ -41,16 +49,27 @@ class UnsupportedMetadataError(Exception):
     pass
 
 
+def _supports_async_pipeline(store: Store) -> bool:
+    # Async-backed stores issue many requests concurrently instead of blocking a
+    # worker thread per request; everything else (e.g. ``LocalStore``) uses the
+    # synchronous pipeline. Mirrors the stores accepted by the Rust
+    # ``AsyncReadableWritableListableStorage`` conversion (see ``src/store.rs``).
+    return isinstance(store, _ASYNC_STORES)
+
+
 def get_codec_pipeline_impl(
     metadata: ArrayMetadata, store: Store, *, strict: bool
-) -> CodecPipelineImpl | None:
+) -> CodecPipelineImpl | AsyncCodecPipelineImpl | None:
+    pipeline_class = (
+        AsyncCodecPipelineImpl if _supports_async_pipeline(store) else CodecPipelineImpl
+    )
     try:
         array_metadata_json = json.dumps(metadata.to_dict())
         # Maintain old behavior: https://github.com/zarrs/zarrs-python/tree/b36ba797cafec77f5f41a25316be02c718a2b4f8?tab=readme-ov-file#configuration
         validate_checksums = config.get("codec_pipeline.validate_checksums", True)
         if validate_checksums is None:
             validate_checksums = True
-        return CodecPipelineImpl(
+        return pipeline_class(
             array_metadata_json,
             store_config=store,
             validate_checksums=validate_checksums,
@@ -104,7 +123,7 @@ def array_metadata_to_codecs(metadata: ArrayMetadata) -> list[Codec]:
 class ZarrsCodecPipeline(CodecPipeline):
     metadata: ArrayMetadata
     store: Store
-    impl: CodecPipelineImpl | None
+    impl: CodecPipelineImpl | AsyncCodecPipelineImpl | None
     python_impl: BatchedCodecPipeline | None
 
     def __getstate__(self) -> ZarrsCodecPipelineState:
@@ -167,6 +186,22 @@ async def encode(
     ) -> Iterable[Buffer | None]:
         raise NotImplementedError("encode")
 
+    async def _retrieve_async(
+        self, impl: AsyncCodecPipelineImpl, chunks: list[ChunkItem], out: np.ndarray
+    ) -> None:
+        # Rust decodes each chunk on its tokio runtime; the copy into ``out``
+        # happens here as each one lands, so decoded chunks are freed eagerly.
+        if not out.flags.c_contiguous:
+            raise ValueError("output array must be a C contiguous array")
+
+        async def fill(chunk: ChunkItem, decoded: Awaitable[np.ndarray]) -> None:
+            dst = out.reshape(chunk.io_shape)[tuple(chunk.io_selection)]
+            dst[...] = (await decoded).view(out.dtype).reshape(dst.shape)
+
+        await asyncio.gather(
+            *map(fill, chunks, impl.retrieve_chunks(chunks)),
+        )
+
     async def read(
         self,
         batch_info: Iterable[
@@ -198,11 +233,16 @@ async def read(
             return None
         else:
             out: NDArrayLike = out.as_ndarray_like()
-            await asyncio.to_thread(
-                self.impl.retrieve_chunks_and_apply_index,
-                chunks_desc.chunk_info_with_indices,
-                out,
-            )
+            if isinstance(self.impl, AsyncCodecPipelineImpl):
+                await self._retrieve_async(
+                    self.impl, chunks_desc.chunk_info_with_indices, out
+                )
+            else:
+                await asyncio.to_thread(
+                    self.impl.retrieve_chunks_and_apply_index,
+                    chunks_desc.chunk_info_with_indices,
+                    out,
+                )
             return None
 
     async def write(
@@ -240,12 +280,15 @@ async def write(
                 )
             elif not value_np.flags.c_contiguous:
                 value_np = np.ascontiguousarray(value_np)
-            await asyncio.to_thread(
-                self.impl.store_chunks_with_indices,
+            args = (
                 chunks_desc.chunk_info_with_indices,
                 value_np,
                 chunks_desc.write_empty_chunks,
             )
+            if isinstance(self.impl, AsyncCodecPipelineImpl):
+                await self.impl.store_chunks_with_indices(*args)
+            else:
+                await asyncio.to_thread(self.impl.store_chunks_with_indices, *args)
             return None
 
     def _raise_error_on_unsupported_batch_dtype(