split codecs in 3 groups in CodecPipeline

Author: normanrz

zarr/v3/array.py

@@ -18,7 +18,6 @@
 from attr import evolve, frozen
 
 from zarr.v3.abc.array import SynchronousArray, AsynchronousArray
-from zarr.v3.abc.codec import ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin
 
 # from zarr.v3.array_v2 import ArrayV2
 from zarr.v3.codecs import CodecMetadata, CodecPipeline, bytes_codec
@@ -120,7 +119,7 @@ async def create(
             metadata=metadata,
             store_path=store_path,
             runtime_configuration=runtime_configuration,
-            codec_pipeline=CodecPipeline(
+            codec_pipeline=CodecPipeline.create(
                 [
                     get_codec_from_metadata(codec).evolve(ndim=len(shape), data_type=data_type)
                     for codec in codecs
@@ -149,7 +148,7 @@ def from_json(
             metadata=metadata,
             store_path=store_path,
             runtime_configuration=runtime_configuration,
-            codec_pipeline=CodecPipeline(codecs),
+            codec_pipeline=CodecPipeline.create(codecs),
         )
         async_array._validate_metadata()
         return async_array
@@ -263,10 +262,8 @@ async def _read_chunk(
         chunk_key = chunk_key_encoding.encode_chunk_key(chunk_coords)
         store_path = self.store_path / chunk_key
 
-        if len(self.codec_pipeline.codecs) == 1 and isinstance(
-            self.codec_pipeline.codecs[0], ArrayBytesCodecPartialDecodeMixin
-        ):
-            chunk_array = await self.codec_pipeline.codecs[0].decode_partial(
+        if self.codec_pipeline.supports_partial_decode:
+            chunk_array = await self.codec_pipeline.decode_partial(
                 store_path, chunk_selection, chunk_metadata, self.runtime_configuration
             )
             if chunk_array is not None:
@@ -346,12 +343,9 @@ async def _write_chunk(
                 chunk_array = value[out_selection]
             await self._write_chunk_to_store(store_path, chunk_array, chunk_metadata)
 
-        elif len(self.codec_pipeline.codecs) == 1 and isinstance(
-            self.codec_pipeline.codecs[0], ArrayBytesCodecPartialEncodeMixin
-        ):
-            codec_with_partial_encode = self.codec_pipeline.codecs[0]
+        elif self.codec_pipeline.supports_partial_encode:
             # print("encode_partial", chunk_coords, chunk_selection, repr(self))
-            await codec_with_partial_encode.encode_partial(
+            await self.codec_pipeline.encode_partial(
                 store_path,
                 value[out_selection],
                 chunk_selection,

zarr/v3/codecs/__init__.py

@@ -15,9 +15,17 @@
 
 import numpy as np
 
-from zarr.v3.abc.codec import Codec, ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec
-from zarr.v3.common import BytesLike
+from zarr.v3.abc.codec import (
+    ArrayBytesCodecPartialDecodeMixin,
+    ArrayBytesCodecPartialEncodeMixin,
+    Codec,
+    ArrayArrayCodec,
+    ArrayBytesCodec,
+    BytesBytesCodec,
+)
+from zarr.v3.common import BytesLike, SliceSelection
 from zarr.v3.metadata import CodecMetadata, ShardingCodecIndexLocation, RuntimeConfiguration
+from zarr.v3.store import StorePath
 
 if TYPE_CHECKING:
     from zarr.v3.metadata import ArrayMetadata, ArraySpec
@@ -33,26 +41,28 @@
 def _find_array_bytes_codec(
     codecs: Iterable[Tuple[Codec, ArraySpec]]
 ) -> Tuple[ArrayBytesCodec, ArraySpec]:
-    for codec, chunk_metadata in codecs:
+    for codec, array_spec in codecs:
         if isinstance(codec, ArrayBytesCodec):
-            return (codec, chunk_metadata)
+            return (codec, array_spec)
     raise KeyError
 
 
 @frozen
 class CodecPipeline:
-    codecs: List[Codec]
+    array_array_codecs: List[ArrayArrayCodec]
+    array_bytes_codec: ArrayBytesCodec
+    bytes_bytes_codecs: List[BytesBytesCodec]
 
-    def validate(self, array_metadata: ArrayMetadata) -> None:
+    @classmethod
+    def create(cls, codecs: List[Codec]) -> CodecPipeline:
         from zarr.v3.codecs.sharding import ShardingCodec
 
         assert any(
-            isinstance(codec, ArrayBytesCodec) for codec in self.codecs
+            isinstance(codec, ArrayBytesCodec) for codec in codecs
         ), "Exactly one array-to-bytes codec is required."
 
         prev_codec: Optional[Codec] = None
-        for codec in self.codecs:
-            codec.validate(array_metadata)
+        for codec in codecs:
             if prev_codec is not None:
                 assert not isinstance(codec, ArrayBytesCodec) or not isinstance(
                     prev_codec, ArrayBytesCodec
@@ -81,84 +91,155 @@ def validate(self, array_metadata: ArrayMetadata) -> None:
                 )
             prev_codec = codec
 
-        if any(isinstance(codec, ShardingCodec) for codec in self.codecs) and len(self.codecs) > 1:
+        if any(isinstance(codec, ShardingCodec) for codec in codecs) and len(codecs) > 1:
             warn(
                 "Combining a `sharding_indexed` codec disables partial reads and "
                 + "writes, which may lead to inefficient performance."
             )
 
+        return CodecPipeline(
+            array_array_codecs=[codec for codec in codecs if isinstance(codec, ArrayArrayCodec)],
+            array_bytes_codec=[codec for codec in codecs if isinstance(codec, ArrayBytesCodec)][0],
+            bytes_bytes_codecs=[codec for codec in codecs if isinstance(codec, BytesBytesCodec)],
+        )
+
+    @property
+    def supports_partial_decode(self) -> bool:
+        return (len(self.array_array_codecs) + len(self.bytes_bytes_codecs)) == 0 and isinstance(
+            self.array_bytes_codec, ArrayBytesCodecPartialDecodeMixin
+        )
+
+    @property
+    def supports_partial_encode(self) -> bool:
+        return (len(self.array_array_codecs) + len(self.bytes_bytes_codecs)) == 0 and isinstance(
+            self.array_bytes_codec, ArrayBytesCodecPartialEncodeMixin
+        )
+
+    def all_codecs(self) -> Iterator[Codec]:
+        for aa_codec in self.array_array_codecs:
+            yield aa_codec
+
+        yield self.array_bytes_codec
+
+        for bb_codec in self.bytes_bytes_codecs:
+            yield bb_codec
+
+    def validate(self, array_metadata: ArrayMetadata) -> None:
+        for codec in self.all_codecs():
+            codec.validate(array_metadata)
+
     def _codecs_with_resolved_metadata(
-        self, chunk_metadata: ArraySpec
-    ) -> Iterator[Tuple[Codec, ArraySpec]]:
-        for codec in self.codecs:
-            yield (codec, chunk_metadata)
-            chunk_metadata = codec.resolve_metadata(chunk_metadata)
+        self, array_spec: ArraySpec
+    ) -> Tuple[
+        List[Tuple[ArrayArrayCodec, ArraySpec]],
+        Tuple[ArrayBytesCodec, ArraySpec],
+        List[Tuple[BytesBytesCodec, ArraySpec]],
+    ]:
+        aa_codecs_with_spec: List[Tuple[ArrayArrayCodec, ArraySpec]] = []
+        for aa_codec in self.array_array_codecs:
+            aa_codecs_with_spec.append((aa_codec, array_spec))
+            array_spec = aa_codec.resolve_metadata(array_spec)
+
+        ab_codec_with_spec = (self.array_bytes_codec, array_spec)
+        array_spec = self.array_bytes_codec.resolve_metadata(array_spec)
+
+        bb_codecs_with_spec: List[Tuple[BytesBytesCodec, ArraySpec]] = []
+        for bb_codec in self.bytes_bytes_codecs:
+            bb_codecs_with_spec.append((bb_codec, array_spec))
+            array_spec = bb_codec.resolve_metadata(array_spec)
+
+        return (aa_codecs_with_spec, ab_codec_with_spec, bb_codecs_with_spec)
 
     async def decode(
         self,
         chunk_bytes: BytesLike,
-        chunk_metadata: ArraySpec,
+        array_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> np.ndarray:
-        codecs = list(self._codecs_with_resolved_metadata(chunk_metadata))[::-1]
+        (
+            aa_codecs_with_spec,
+            ab_codec_with_spec,
+            bb_codecs_with_spec,
+        ) = self._codecs_with_resolved_metadata(array_spec)
 
-        for bb_codec, chunk_metadata in codecs:
-            if isinstance(bb_codec, BytesBytesCodec):
-                chunk_bytes = await bb_codec.decode(
-                    chunk_bytes, chunk_metadata, runtime_configuration
-                )
+        for bb_codec, array_spec in bb_codecs_with_spec[::-1]:
+            chunk_bytes = await bb_codec.decode(chunk_bytes, array_spec, runtime_configuration)
 
-        ab_codec, chunk_metadata = _find_array_bytes_codec(codecs)
-        chunk_array = await ab_codec.decode(chunk_bytes, chunk_metadata, runtime_configuration)
+        ab_codec, array_spec = ab_codec_with_spec
+        chunk_array = await ab_codec.decode(chunk_bytes, array_spec, runtime_configuration)
 
-        for aa_codec, chunk_metadata in codecs:
-            if isinstance(aa_codec, ArrayArrayCodec):
-                chunk_array = await aa_codec.decode(
-                    chunk_array, chunk_metadata, runtime_configuration
-                )
+        for aa_codec, array_spec in aa_codecs_with_spec[::-1]:
+            chunk_array = await aa_codec.decode(chunk_array, array_spec, runtime_configuration)
 
         return chunk_array
 
+    async def decode_partial(
+        self,
+        store_path: StorePath,
+        selection: SliceSelection,
+        chunk_metadata: ArraySpec,
+        runtime_configuration: RuntimeConfiguration,
+    ) -> Optional[np.ndarray]:
+        assert self.supports_partial_decode
+        assert isinstance(self.array_bytes_codec, ArrayBytesCodecPartialDecodeMixin)
+        return await self.array_bytes_codec.decode_partial(
+            store_path, selection, chunk_metadata, runtime_configuration
+        )
+
     async def encode(
         self,
         chunk_array: np.ndarray,
-        chunk_metadata: ArraySpec,
+        array_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> Optional[BytesLike]:
-        codecs = list(self._codecs_with_resolved_metadata(chunk_metadata))
-
-        for aa_codec, chunk_metadata in codecs:
-            if isinstance(aa_codec, ArrayArrayCodec):
-                chunk_array_maybe = await aa_codec.encode(
-                    chunk_array, chunk_metadata, runtime_configuration
-                )
-                if chunk_array_maybe is None:
-                    return None
-                chunk_array = chunk_array_maybe
+        (
+            aa_codecs_with_spec,
+            ab_codec_with_spec,
+            bb_codecs_with_spec,
+        ) = self._codecs_with_resolved_metadata(array_spec)
+
+        for aa_codec, array_spec in aa_codecs_with_spec:
+            chunk_array_maybe = await aa_codec.encode(
+                chunk_array, array_spec, runtime_configuration
+            )
+            if chunk_array_maybe is None:
+                return None
+            chunk_array = chunk_array_maybe
 
-        ab_codec, chunk_metadata = _find_array_bytes_codec(codecs)
-        chunk_bytes_maybe = await ab_codec.encode(
-            chunk_array, chunk_metadata, runtime_configuration
-        )
+        ab_codec, array_spec = ab_codec_with_spec
+        chunk_bytes_maybe = await ab_codec.encode(chunk_array, array_spec, runtime_configuration)
         if chunk_bytes_maybe is None:
             return None
         chunk_bytes = chunk_bytes_maybe
 
-        for bb_codec, chunk_metadata in codecs:
-            if isinstance(bb_codec, BytesBytesCodec):
-                chunk_bytes_maybe = await bb_codec.encode(
-                    chunk_bytes, chunk_metadata, runtime_configuration
-                )
-                if chunk_bytes_maybe is None:
-                    return None
-                chunk_bytes = chunk_bytes_maybe
+        for bb_codec, array_spec in bb_codecs_with_spec:
+            chunk_bytes_maybe = await bb_codec.encode(
+                chunk_bytes, array_spec, runtime_configuration
+            )
+            if chunk_bytes_maybe is None:
+                return None
+            chunk_bytes = chunk_bytes_maybe
 
         return chunk_bytes
 
-    def compute_encoded_size(self, byte_length: int, chunk_metadata: ArraySpec) -> int:
-        for codec in self.codecs:
-            byte_length = codec.compute_encoded_size(byte_length, chunk_metadata)
-            chunk_metadata = codec.resolve_metadata(chunk_metadata)
+    async def encode_partial(
+        self,
+        store_path: StorePath,
+        chunk_array: np.ndarray,
+        selection: SliceSelection,
+        chunk_metadata: ArraySpec,
+        runtime_configuration: RuntimeConfiguration,
+    ) -> None:
+        assert self.supports_partial_encode
+        assert isinstance(self.array_bytes_codec, ArrayBytesCodecPartialEncodeMixin)
+        await self.array_bytes_codec.encode_partial(
+            store_path, chunk_array, selection, chunk_metadata, runtime_configuration
+        )
+
+    def compute_encoded_size(self, byte_length: int, array_spec: ArraySpec) -> int:
+        for codec in self.all_codecs():
+            byte_length = codec.compute_encoded_size(byte_length, array_spec)
+            array_spec = codec.resolve_metadata(array_spec)
         return byte_length
 
 

zarr/v3/codecs/sharding.py

@@ -624,11 +624,13 @@ def _get_chunks_per_shard(self, shard_metadata: ArraySpec) -> ChunkCoords:
 
     @cached_property
     def _index_codec_pipeline(self) -> CodecPipeline:
-        return CodecPipeline([get_codec_from_metadata(c) for c in self.configuration.index_codecs])
+        return CodecPipeline.create(
+            [get_codec_from_metadata(c) for c in self.configuration.index_codecs]
+        )
 
     @cached_property
     def _codec_pipeline(self) -> CodecPipeline:
-        return CodecPipeline([get_codec_from_metadata(c) for c in self.configuration.codecs])
+        return CodecPipeline.create([get_codec_from_metadata(c) for c in self.configuration.codecs])
 
     async def _load_shard_index_maybe(
         self, store_path: StorePath, chunks_per_shard: ChunkCoords