Reworked codec pipelines

pyproject.toml

@@ -180,10 +180,6 @@ check_untyped_defs = false
 module = [
     "zarr.v2.*",
     "zarr.abc.codec",
-    "zarr.codecs.bytes",
-    "zarr.codecs.pipeline",
-    "zarr.codecs.sharding",
-    "zarr.codecs.transpose",
     "zarr.array_v2",
     "zarr.array",
     "zarr.sync",

src/zarr/abc/codec.py

@@ -1,13 +1,22 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Optional
+from typing import (
+    TYPE_CHECKING,
+    Awaitable,
+    Callable,
+    Iterable,
+    Optional,
+    Protocol,
+    Tuple,
+    TypeVar,
+    runtime_checkable,
+)
 
 import numpy as np
 from zarr.abc.metadata import Metadata
 
-from zarr.common import ArraySpec
-from zarr.store import StorePath
+from zarr.common import ArraySpec, concurrent_map
 
 
 if TYPE_CHECKING:
@@ -16,6 +25,40 @@
     from zarr.metadata import ArrayMetadata
     from zarr.config import RuntimeConfiguration
 
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+def noop_for_none(
+    func: Callable[[T, ArraySpec, RuntimeConfiguration], Awaitable[Optional[U]]],
+) -> Callable[[Optional[T], ArraySpec, RuntimeConfiguration], Awaitable[Optional[U]]]:
+    async def wrap(
+        chunk: Optional[T], chunk_spec: ArraySpec, runtime_configuration: RuntimeConfiguration
+    ) -> Optional[U]:
+        if chunk is None:
+            return None
+        return await func(chunk, chunk_spec, runtime_configuration)
+
+    return wrap
+
+
+@runtime_checkable
+class ByteGetter(Protocol):
+    async def get(
+        self, byte_range: Optional[Tuple[int, Optional[int]]] = None
+    ) -> Optional[BytesLike]: ...
+
+
+@runtime_checkable
+class ByteSetter(Protocol):
+    async def get(
+        self, byte_range: Optional[Tuple[int, Optional[int]]] = None
+    ) -> Optional[BytesLike]: ...
+
+    async def set(self, value: BytesLike, byte_range: Optional[Tuple[int, int]] = None) -> None: ...
+
+    async def delete(self) -> None: ...
+
 
 class Codec(Metadata):
     is_fixed_size: bool
@@ -44,6 +87,20 @@ async def decode(
     ) -> np.ndarray:
         pass
 
+    async def decode_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.decode),
+            runtime_configuration.concurrency,
+        )
+
     @abstractmethod
     async def encode(
         self,
@@ -53,17 +110,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[Optional[BytesLike], ArraySpec]],
+        runtime_configuration: RuntimeConfiguration,
+    ) -> Iterable[Optional[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,
@@ -73,42 +158,98 @@ 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
     async def decode_partial(
         self,
-        store_path: StorePath,
+        byte_getter: ByteGetter,
         selection: SliceSelection,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> Optional[np.ndarray]:
         pass
 
+    async def decode_partial_batch(
+        self,
+        batch_info: Iterable[Tuple[ByteGetter, SliceSelection, ArraySpec]],
+        runtime_configuration: RuntimeConfiguration,
+    ) -> Iterable[Optional[np.ndarray]]:
+        return await concurrent_map(
+            [
+                (byte_getter, selection, chunk_spec, runtime_configuration)
+                for byte_getter, selection, chunk_spec in batch_info
+            ],
+            self.decode_partial,
+            runtime_configuration.concurrency,
+        )
+
 
 class ArrayBytesCodecPartialEncodeMixin:
     @abstractmethod
     async def encode_partial(
         self,
-        store_path: StorePath,
+        byte_setter: ByteSetter,
         chunk_array: np.ndarray,
         selection: SliceSelection,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> None:
         pass
 
+    async def encode_partial_batch(
+        self,
+        batch_info: Iterable[Tuple[ByteSetter, np.ndarray, SliceSelection, ArraySpec]],
+        runtime_configuration: RuntimeConfiguration,
+    ) -> None:
+        await concurrent_map(
+            [
+                (byte_setter, chunk_array, selection, chunk_spec, runtime_configuration)
+                for byte_setter, 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[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.decode),
+            runtime_configuration.concurrency,
+        )
+
     @abstractmethod
     async def encode(
         self,
@@ -117,3 +258,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,
+        )