added a BatchedCodecPipeline

Author: normanrz

src/zarr/v3/abc/codec.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Awaitable, Callable, Iterable, Optional, Tuple, TypeVar
 
 import numpy as np
 from zarr.v3.abc.metadata import Metadata
 
-from zarr.v3.common import ArraySpec
+from zarr.v3.common import ArraySpec, concurrent_map
 from zarr.v3.store import StorePath
 
 
@@ -18,6 +18,22 @@
         RuntimeConfiguration,
     )
 
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+def noop_for_none(
+    func: Callable[[Optional[T], ArraySpec, RuntimeConfiguration], Awaitable[U]]
+) -> Callable[[T, ArraySpec, RuntimeConfiguration], Awaitable[U]]:
+    async def wrap(
+        chunk: Optional[T], chunk_spec: ArraySpec, runtime_configuration: RuntimeConfiguration
+    ) -> U:
+        if chunk is None:
+            return None
+        return await func(chunk, chunk_spec, runtime_configuration)
+
+    return wrap
+
 
 class Codec(Metadata):
     is_fixed_size: bool
@@ -46,6 +62,20 @@ async def decode(
     ) -> np.ndarray:
         pass
 
+    async def decode_batch(
+        self,
+        chunk_arrays_and_specs: Iterable[Tuple[np.ndarray, ArraySpec]],
+        runtime_configuration: RuntimeConfiguration,
+    ) -> Iterable[np.ndarray]:
+        return await concurrent_map(
+            [
+                (chunk_array, chunk_spec, runtime_configuration)
+                for chunk_array, chunk_spec in chunk_arrays_and_specs
+            ],
+            noop_for_none(self.decode),
+            runtime_configuration.concurrency,
+        )
+
     @abstractmethod
     async def encode(
         self,
@@ -55,17 +85,45 @@ async def encode(
     ) -> Optional[np.ndarray]:
         pass
 
+    async def encode_batch(
+        self,
+        chunk_arrays_and_specs: Iterable[Tuple[Optional[np.ndarray], ArraySpec]],
+        runtime_configuration: RuntimeConfiguration,
+    ) -> Iterable[Optional[np.ndarray]]:
+        return await concurrent_map(
+            [
+                (chunk_array, chunk_spec, runtime_configuration)
+                for chunk_array, chunk_spec in chunk_arrays_and_specs
+            ],
+            noop_for_none(self.encode),
+            runtime_configuration.concurrency,
+        )
+
 
 class ArrayBytesCodec(Codec):
     @abstractmethod
     async def decode(
         self,
-        chunk_array: BytesLike,
+        chunk_bytes: BytesLike,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> np.ndarray:
         pass
 
+    async def decode_batch(
+        self,
+        chunk_bytes_and_specs: Iterable[Tuple[BytesLike, ArraySpec]],
+        runtime_configuration: RuntimeConfiguration,
+    ) -> Iterable[np.ndarray]:
+        return await concurrent_map(
+            [
+                (chunk_bytes, chunk_spec, runtime_configuration)
+                for chunk_bytes, chunk_spec in chunk_bytes_and_specs
+            ],
+            noop_for_none(self.decode),
+            runtime_configuration.concurrency,
+        )
+
     @abstractmethod
     async def encode(
         self,
@@ -75,6 +133,20 @@ async def encode(
     ) -> Optional[BytesLike]:
         pass
 
+    async def encode_batch(
+        self,
+        chunk_arrays_and_specs: Iterable[Tuple[Optional[np.ndarray], ArraySpec]],
+        runtime_configuration: RuntimeConfiguration,
+    ) -> Iterable[Optional[BytesLike]]:
+        return await concurrent_map(
+            [
+                (chunk_array, chunk_spec, runtime_configuration)
+                for chunk_array, chunk_spec in chunk_arrays_and_specs
+            ],
+            noop_for_none(self.encode),
+            runtime_configuration.concurrency,
+        )
+
 
 class ArrayBytesCodecPartialDecodeMixin:
     @abstractmethod
@@ -87,6 +159,20 @@ async def decode_partial(
     ) -> Optional[np.ndarray]:
         pass
 
+    async def decode_partial_batched(
+        self,
+        batch_info: Iterable[Tuple[StorePath, SliceSelection, ArraySpec]],
+        runtime_configuration: RuntimeConfiguration,
+    ) -> Iterable[Optional[np.ndarray]]:
+        return await concurrent_map(
+            [
+                (store_path, selection, chunk_spec, runtime_configuration)
+                for store_path, selection, chunk_spec in batch_info
+            ],
+            self.decode_partial,
+            runtime_configuration.concurrency,
+        )
+
 
 class ArrayBytesCodecPartialEncodeMixin:
     @abstractmethod
@@ -100,17 +186,45 @@ async def encode_partial(
     ) -> None:
         pass
 
+    async def encode_partial_batched(
+        self,
+        batch_info: Iterable[Tuple[StorePath, np.ndarray, SliceSelection, ArraySpec]],
+        runtime_configuration: RuntimeConfiguration,
+    ) -> None:
+        await concurrent_map(
+            [
+                (store_path, chunk_array, selection, chunk_spec, runtime_configuration)
+                for store_path, chunk_array, selection, chunk_spec in batch_info
+            ],
+            self.encode_partial,
+            runtime_configuration.concurrency,
+        )
+
 
 class BytesBytesCodec(Codec):
     @abstractmethod
     async def decode(
         self,
-        chunk_array: BytesLike,
+        chunk_bytes: BytesLike,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> BytesLike:
         pass
 
+    async def decode_batch(
+        self,
+        chunk_bytes_and_specs: Iterable[Tuple[BytesLike, ArraySpec]],
+        runtime_configuration: RuntimeConfiguration,
+    ) -> Iterable[BytesLike]:
+        return await concurrent_map(
+            [
+                (chunk_bytes, chunk_spec, runtime_configuration)
+                for chunk_bytes, chunk_spec in chunk_bytes_and_specs
+            ],
+            noop_for_none(self.decode),
+            runtime_configuration.concurrency,
+        )
+
     @abstractmethod
     async def encode(
         self,
@@ -119,3 +233,17 @@ async def encode(
         runtime_configuration: RuntimeConfiguration,
     ) -> Optional[BytesLike]:
         pass
+
+    async def encode_batch(
+        self,
+        chunk_bytes_and_specs: Iterable[Tuple[Optional[BytesLike], ArraySpec]],
+        runtime_configuration: RuntimeConfiguration,
+    ) -> Iterable[Optional[BytesLike]]:
+        return await concurrent_map(
+            [
+                (chunk_bytes, chunk_spec, runtime_configuration)
+                for chunk_bytes, chunk_spec in chunk_bytes_and_specs
+            ],
+            noop_for_none(self.encode),
+            runtime_configuration.concurrency,
+        )