diff --git a/src/zarr/v3/__init__.py b/src/zarr/v3/__init__.py index 7ae49db697..776ff8945d 100644 --- a/src/zarr/v3/__init__.py +++ b/src/zarr/v3/__init__.py @@ -7,7 +7,7 @@ from zarr.v3.array_v2 import ArrayV2 from zarr.v3.common import RuntimeConfiguration # noqa: F401 from zarr.v3.group import Group # noqa: F401 -from zarr.v3.metadata import RuntimeConfiguration, runtime_configuration # noqa: F401 +from zarr.v3.metadata import runtime_configuration # noqa: F401 from zarr.v3.store import ( # noqa: F401 StoreLike, make_store_path, diff --git a/src/zarr/v3/abc/codec.py b/src/zarr/v3/abc/codec.py index a5be720d6b..598b3a15ac 100644 --- a/src/zarr/v3/abc/codec.py +++ b/src/zarr/v3/abc/codec.py @@ -1,35 +1,27 @@ from __future__ import annotations -from abc import abstractmethod, ABC -from typing import TYPE_CHECKING, Optional, Type +from abc import abstractmethod +from typing import TYPE_CHECKING, Optional import numpy as np +from zarr.v3.abc.metadata import Metadata from zarr.v3.common import ArraySpec from zarr.v3.store import StorePath if TYPE_CHECKING: - from zarr.v3.common import BytesLike, SliceSelection, NamedConfig + from typing_extensions import Self + from zarr.v3.common import BytesLike, SliceSelection from zarr.v3.metadata import ( ArrayMetadata, RuntimeConfiguration, ) -class Codec(ABC): - metadata: NamedConfig +class Codec(Metadata): is_fixed_size: bool - @classmethod - @abstractmethod - def from_metadata(cls, codec_metadata: "NamedConfig") -> Codec: - pass - - @classmethod - def get_metadata_class(cls) -> Type[NamedConfig]: - pass - @abstractmethod def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int: pass @@ -37,7 +29,7 @@ def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: return chunk_spec - def evolve(self, *, ndim: int, data_type: np.dtype) -> Codec: + def evolve(self, array_spec: ArraySpec) -> Self: return self def validate(self, array_metadata: ArrayMetadata) -> None: diff --git a/src/zarr/v3/abc/metadata.py b/src/zarr/v3/abc/metadata.py index 83a061f10e..eab61f7730 100644 --- a/src/zarr/v3/abc/metadata.py +++ b/src/zarr/v3/abc/metadata.py @@ -1,12 +1,14 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Sequence if TYPE_CHECKING: - from typing import Dict, Any, Sequence + from typing import Dict, Any from typing_extensions import Self from dataclasses import fields +from zarr.v3.common import JSON + class Metadata: def to_dict(self) -> Dict[str, Any]: @@ -23,6 +25,8 @@ def to_dict(self) -> Dict[str, Any]: value = getattr(self, key) if isinstance(value, Metadata): out_dict[field.name] = getattr(self, field.name).to_dict() + elif isinstance(value, str): + out_dict[key] = value elif isinstance(value, Sequence): out_dict[key] = [v.to_dict() if isinstance(v, Metadata) else v for v in value] else: @@ -31,7 +35,7 @@ def to_dict(self) -> Dict[str, Any]: return out_dict @classmethod - def from_dict(cls, data: Dict[str, Any]) -> Self: + def from_dict(cls, data: Dict[str, JSON]) -> Self: """ Create an instance of the model from a dictionary """ diff --git a/src/zarr/v3/array.py b/src/zarr/v3/array.py index 9cd47058d6..4b7a3bb0ad 100644 --- a/src/zarr/v3/array.py +++ b/src/zarr/v3/array.py @@ -13,17 +13,13 @@ from dataclasses import dataclass, replace import json -from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple, Union +from typing import Any, Dict, Iterable, Literal, Optional, Tuple, Union import numpy as np -from zarr.v3.abc.codec import Codec # from zarr.v3.array_v2 import ArrayV2 -from zarr.v3.codecs.common import decode, encode - -# from zarr.v3.array_v2 import ArrayV2 -from zarr.v3.codecs import bytes_codec +from zarr.v3.codecs import BytesCodec from zarr.v3.common import ( ZARR_JSON, ArraySpec, @@ -35,15 +31,9 @@ concurrent_map, ) from zarr.v3.indexing import BasicIndexer, all_chunk_coords, is_total_slice -from zarr.v3.metadata import ( - ArrayMetadata, - DefaultChunkKeyEncodingConfigurationMetadata, - DefaultChunkKeyEncodingMetadata, - RegularChunkGridConfigurationMetadata, - RegularChunkGridMetadata, - V2ChunkKeyEncodingConfigurationMetadata, - V2ChunkKeyEncodingMetadata, -) +from zarr.v3.chunk_grids import RegularChunkGrid +from zarr.v3.chunk_key_encodings import DefaultChunkKeyEncoding, V2ChunkKeyEncoding +from zarr.v3.metadata import ArrayMetadata from zarr.v3.store import StoreLike, StorePath, make_store_path from zarr.v3.sync import sync @@ -57,38 +47,27 @@ def parse_array_metadata(data: Any): raise TypeError +@dataclass(frozen=True) class AsyncArray: metadata: ArrayMetadata store_path: StorePath runtime_configuration: RuntimeConfiguration - codecs: List[Codec] @property def codecs(self): return self.metadata.codecs - @property - def store_path(self): - return self._store_path - def __init__( self, metadata: ArrayMetadata, store_path: StorePath, runtime_configuration: RuntimeConfiguration, ): - self.metadata = parse_array_metadata(metadata) - self._store_path = store_path - self.runtime_configuration = runtime_configuration + metadata_parsed = parse_array_metadata(metadata) - async def encode_chunk(self, data: np.ndarray): - """ - Encode a numpy array using the codec pipeline - """ - return await encode(self.codecs, data, self.runtime_configuration) - - async def decode_chunk(self, data: bytes): - return await decode(self.codecs, data, self.runtime_configuration) + object.__setattr__(self, "metadata", metadata_parsed) + object.__setattr__(self, "store_path", store_path) + object.__setattr__(self, "runtime_configuration", runtime_configuration) @classmethod async def create( @@ -113,7 +92,7 @@ async def create( if not exists_ok: assert not await (store_path / ZARR_JSON).exists() - codecs = list(codecs) if codecs is not None else [bytes_codec()] + codecs = list(codecs) if codecs is not None else [BytesCodec()] if fill_value is None: if dtype == np.dtype("bool"): @@ -124,21 +103,11 @@ async def create( metadata = ArrayMetadata( shape=shape, data_type=dtype, - chunk_grid=RegularChunkGridMetadata( - configuration=RegularChunkGridConfigurationMetadata(chunk_shape=chunk_shape) - ), + chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape), chunk_key_encoding=( - V2ChunkKeyEncodingMetadata( - configuration=V2ChunkKeyEncodingConfigurationMetadata( - separator=chunk_key_encoding[1] - ) - ) + V2ChunkKeyEncoding(separator=chunk_key_encoding[1]) if chunk_key_encoding[0] == "v2" - else DefaultChunkKeyEncodingMetadata( - configuration=DefaultChunkKeyEncodingConfigurationMetadata( - separator=chunk_key_encoding[1] - ) - ) + else DefaultChunkKeyEncoding(separator=chunk_key_encoding[1]) ), fill_value=fill_value, codecs=codecs, @@ -228,7 +197,7 @@ async def getitem(self, selection: Selection): indexer = BasicIndexer( selection, shape=self.metadata.shape, - chunk_shape=self.metadata.chunk_grid.configuration.chunk_shape, + chunk_shape=self.metadata.chunk_grid.chunk_shape, ) # setup output array @@ -260,13 +229,13 @@ async def _save_metadata(self) -> None: def _validate_metadata(self) -> None: assert len(self.metadata.shape) == len( - self.metadata.chunk_grid.configuration.chunk_shape + self.metadata.chunk_grid.chunk_shape ), "`chunk_shape` and `shape` need to have the same number of dimensions." assert self.metadata.dimension_names is None or len(self.metadata.shape) == len( self.metadata.dimension_names ), "`dimension_names` and `shape` need to have the same number of dimensions." assert self.metadata.fill_value is not None, "`fill_value` is required." - # self.codecs.validate(self.metadata) + self.codecs.validate(self.metadata) async def _read_chunk( self, @@ -280,8 +249,8 @@ async def _read_chunk( chunk_key = chunk_key_encoding.encode_chunk_key(chunk_coords) store_path = self.store_path / chunk_key - if self.codec_pipeline.supports_partial_decode: - chunk_array = await self.codec_pipeline.decode_partial( + if self.codecs.supports_partial_decode: + chunk_array = await self.codecs.decode_partial( store_path, chunk_selection, chunk_spec, self.runtime_configuration ) if chunk_array is not None: @@ -291,7 +260,7 @@ async def _read_chunk( else: chunk_bytes = await store_path.get() if chunk_bytes is not None: - chunk_array = await self.codec_pipeline.decode( + chunk_array = await self.codecs.decode( chunk_bytes, chunk_spec, self.runtime_configuration ) tmp = chunk_array[chunk_selection] @@ -300,7 +269,7 @@ async def _read_chunk( out[out_selection] = self.metadata.fill_value async def setitem(self, selection: Selection, value: np.ndarray) -> None: - chunk_shape = self.metadata.chunk_grid.configuration.chunk_shape + chunk_shape = self.metadata.chunk_grid.chunk_shape indexer = BasicIndexer( selection, shape=self.metadata.shape, @@ -361,9 +330,9 @@ async def _write_chunk( chunk_array = value[out_selection] await self._write_chunk_to_store(store_path, chunk_array, chunk_spec) - elif self.codec_pipeline.supports_partial_encode: + elif self.codecs.supports_partial_encode: # print("encode_partial", chunk_coords, chunk_selection, repr(self)) - await self.codec_pipeline.encode_partial( + await self.codecs.encode_partial( store_path, value[out_selection], chunk_selection, @@ -384,9 +353,7 @@ async def _write_chunk( chunk_array.fill(self.metadata.fill_value) else: chunk_array = ( - await self.codec_pipeline.decode( - chunk_bytes, chunk_spec, self.runtime_configuration - ) + await self.codecs.decode(chunk_bytes, chunk_spec, self.runtime_configuration) ).copy() # make a writable copy chunk_array[chunk_selection] = value[out_selection] @@ -399,7 +366,9 @@ async def _write_chunk_to_store( # chunks that only contain fill_value will be removed await store_path.delete() else: - chunk_bytes = await encode(self.codecs, chunk_array, self.runtime_configuration) + chunk_bytes = await self.codecs.encode( + chunk_array, chunk_spec, self.runtime_configuration + ) if chunk_bytes is None: await store_path.delete() else: @@ -410,7 +379,7 @@ async def resize(self, new_shape: ChunkCoords) -> AsyncArray: new_metadata = replace(self.metadata, shape=new_shape) # Remove all chunks outside of the new shape - chunk_shape = self.metadata.chunk_grid.configuration.chunk_shape + chunk_shape = self.metadata.chunk_grid.chunk_shape chunk_key_encoding = self.metadata.chunk_key_encoding old_chunk_coords = set(all_chunk_coords(self.metadata.shape, chunk_shape)) new_chunk_coords = set(all_chunk_coords(new_shape, chunk_shape)) @@ -429,14 +398,14 @@ async def _delete_key(key: str) -> None: # Write new metadata await (self.store_path / ZARR_JSON).set(new_metadata.to_bytes()) - return evolve(self, metadata=new_metadata) + return replace(self, metadata=new_metadata) async def update_attributes(self, new_attributes: Dict[str, Any]) -> Array: new_metadata = replace(self.metadata, attributes=new_attributes) # Write new metadata await (self.store_path / ZARR_JSON).set(new_metadata.to_bytes()) - return evolve(self, metadata=new_metadata) + return replace(self, metadata=new_metadata) def __repr__(self): return f"" diff --git a/src/zarr/v3/array_v2.py b/src/zarr/v3/array_v2.py index d788b6fa5c..860deb8ad9 100644 --- a/src/zarr/v3/array_v2.py +++ b/src/zarr/v3/array_v2.py @@ -23,7 +23,7 @@ to_thread, ) from zarr.v3.indexing import BasicIndexer, all_chunk_coords, is_total_slice -from zarr.v3.metadata import ArrayV2Metadata, RuntimeConfiguration +from zarr.v3.metadata import ArrayV2Metadata from zarr.v3.store import StoreLike, StorePath, make_store_path from zarr.v3.sync import sync @@ -444,34 +444,22 @@ async def convert_to_v3_async(self) -> Array: from zarr.v3.array import Array from zarr.v3.common import ZARR_JSON + from zarr.v3.chunk_grids import RegularChunkGrid from zarr.v3.metadata import ( ArrayMetadata, DataType, - RegularChunkGridConfigurationMetadata, - RegularChunkGridMetadata, - V2ChunkKeyEncodingConfigurationMetadata, - V2ChunkKeyEncodingMetadata, - dtype_to_data_type, + V2ChunkKeyEncoding, ) - from zarr.v3.codecs.blosc import ( - BloscCodecConfigurationMetadata, - BloscCodecMetadata, - blosc_shuffle_int_to_str, - ) - from zarr.v3.codecs.bytes import ( - BytesCodecConfigurationMetadata, - BytesCodecMetadata, - ) - from zarr.v3.codecs.gzip import ( - GzipCodecConfigurationMetadata, - GzipCodecMetadata, - ) - from zarr.v3.codecs.transpose import ( - TransposeCodecConfigurationMetadata, - TransposeCodecMetadata, + + from zarr.v3.codecs import ( + BloscCodec, + BloscShuffle, + BytesCodec, + GzipCodec, + TransposeCodec, ) - data_type = DataType[dtype_to_data_type[self.metadata.dtype.str]] + data_type = DataType.from_dtype(self.metadata.dtype) endian: Literal["little", "big"] if self.metadata.dtype.byteorder == "=": endian = sys_byteorder @@ -487,16 +475,8 @@ async def convert_to_v3_async(self) -> Array: codecs: List[NamedConfig] = [] if self.metadata.order == "F": - codecs.append( - TransposeCodecMetadata( - configuration=TransposeCodecConfigurationMetadata( - order=tuple(reversed(range(self.metadata.ndim))) - ) - ) - ) - codecs.append( - BytesCodecMetadata(configuration=BytesCodecConfigurationMetadata(endian=endian)) - ) + codecs.append(TransposeCodec(order=tuple(reversed(range(self.metadata.ndim))))) + codecs.append(BytesCodec(endian=endian)) if self.metadata.compressor is not None: v2_codec = numcodecs.get_codec(self.metadata.compressor).get_config() @@ -505,39 +485,24 @@ async def convert_to_v3_async(self) -> Array: "gzip", ), "Only blosc and gzip are supported by v3." if v2_codec["id"] == "blosc": - shuffle = blosc_shuffle_int_to_str[v2_codec.get("shuffle", 0)] codecs.append( - BloscCodecMetadata( - configuration=BloscCodecConfigurationMetadata( - typesize=data_type.byte_count, - cname=v2_codec["cname"], - clevel=v2_codec["clevel"], - shuffle=shuffle, - blocksize=v2_codec.get("blocksize", 0), - ) + BloscCodec( + typesize=data_type.byte_count, + cname=v2_codec["cname"], + clevel=v2_codec["clevel"], + shuffle=BloscShuffle.from_int(v2_codec.get("shuffle", 0)), + blocksize=v2_codec.get("blocksize", 0), ) ) elif v2_codec["id"] == "gzip": - codecs.append( - GzipCodecMetadata( - configuration=GzipCodecConfigurationMetadata(level=v2_codec.get("level", 5)) - ) - ) + codecs.append(GzipCodec(v2_codec.get("level", 5))) new_metadata = ArrayMetadata( shape=self.metadata.shape, - chunk_grid=RegularChunkGridMetadata( - configuration=RegularChunkGridConfigurationMetadata( - chunk_shape=self.metadata.chunks - ) - ), + chunk_grid=RegularChunkGrid(chunk_shape=self.metadata.chunks), data_type=data_type, fill_value=0 if self.metadata.fill_value is None else self.metadata.fill_value, - chunk_key_encoding=V2ChunkKeyEncodingMetadata( - configuration=V2ChunkKeyEncodingConfigurationMetadata( - separator=self.metadata.dimension_separator - ) - ), + chunk_key_encoding=V2ChunkKeyEncoding(separator=self.metadata.dimension_separator), codecs=codecs, attributes=self.attributes or {}, ) diff --git a/src/zarr/v3/chunk_grids.py b/src/zarr/v3/chunk_grids.py new file mode 100644 index 0000000000..470574cd2b --- /dev/null +++ b/src/zarr/v3/chunk_grids.py @@ -0,0 +1,38 @@ +from __future__ import annotations +from typing import TYPE_CHECKING, Any, Dict +from dataclasses import dataclass +from zarr.v3.abc.metadata import Metadata + +from zarr.v3.common import JSON, ChunkCoords, parse_name, parse_shapelike + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(frozen=True) +class ChunkGrid(Metadata): + @classmethod + def from_dict(cls, data: Dict[str, JSON]) -> Self: + if isinstance(data, ChunkGrid): + return data + if data["name"] == "regular": + return RegularChunkGrid.from_dict(data) + raise ValueError(f"Unknown chunk grid, got {data['name']}") + + +@dataclass(frozen=True) +class RegularChunkGrid(ChunkGrid): + chunk_shape: ChunkCoords + + def __init__(self, *, chunk_shape) -> None: + chunk_shape_parsed = parse_shapelike(chunk_shape) + + object.__setattr__(self, "chunk_shape", chunk_shape_parsed) + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> Self: + parse_name(data["name"], "regular") + return cls(**data["configuration"]) + + def to_dict(self) -> Dict[str, JSON]: + return {"name": "regular", "configuration": {"chunk_shape": self.chunk_shape}} diff --git a/src/zarr/v3/chunk_key_encodings.py b/src/zarr/v3/chunk_key_encodings.py new file mode 100644 index 0000000000..a0dc159519 --- /dev/null +++ b/src/zarr/v3/chunk_key_encodings.py @@ -0,0 +1,66 @@ +from __future__ import annotations +from typing import TYPE_CHECKING, Dict, Literal +from dataclasses import dataclass +from zarr.v3.abc.metadata import Metadata + +from zarr.v3.common import JSON, ChunkCoords + +if TYPE_CHECKING: + from typing_extensions import Self + +SeparatorLiteral = Literal[".", "/"] + + +def parse_separator(data: JSON) -> ChunkCoords: + if data not in (".", "/"): + raise ValueError(f"Expected an '.' or '/' separator. Got {data} instead.") + return data + + +@dataclass(frozen=True) +class ChunkKeyEncoding(Metadata): + name: str + separator: SeparatorLiteral = "." + + def __init__(self, *, separator) -> None: + separator_parsed = parse_separator(separator) + + object.__setattr__(self, "separator", separator_parsed) + + @classmethod + def from_dict(cls, data: Dict[str, JSON]) -> Self: + if isinstance(data, ChunkKeyEncoding): + return data + if data["name"] == "default": + return DefaultChunkKeyEncoding(**data["configuration"]) + if data["name"] == "v2": + return V2ChunkKeyEncoding(**data["configuration"]) + raise ValueError(f"Unknown chunk key encoding, got {data['name']}") + + def to_dict(self) -> Dict[str, JSON]: + return {"name": self.name, "configuration": {"separator": self.separator}} + + +@dataclass(frozen=True) +class DefaultChunkKeyEncoding(ChunkKeyEncoding): + name: Literal["default"] = "default" + + def decode_chunk_key(self, chunk_key: str) -> ChunkCoords: + if chunk_key == "c": + return () + return tuple(map(int, chunk_key[1:].split(self.separator))) + + def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str: + return self.separator.join(map(str, ("c",) + chunk_coords)) + + +@dataclass(frozen=True) +class V2ChunkKeyEncoding(ChunkKeyEncoding): + name: Literal["v2"] = "v2" + + def decode_chunk_key(self, chunk_key: str) -> ChunkCoords: + return tuple(map(int, chunk_key.split(self.separator))) + + def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str: + chunk_identifier = self.separator.join(map(str, chunk_coords)) + return "0" if chunk_identifier == "" else chunk_identifier diff --git a/src/zarr/v3/codecs/__init__.py b/src/zarr/v3/codecs/__init__.py index ba3ebe4140..474344ec25 100644 --- a/src/zarr/v3/codecs/__init__.py +++ b/src/zarr/v3/codecs/__init__.py @@ -1,326 +1,9 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import ( - TYPE_CHECKING, - Iterable, - Iterator, - List, - Literal, - Optional, - Tuple, - Union, -) -from warnings import warn - -import numpy as np - -from zarr.v3.abc.codec import ( - ArrayBytesCodecPartialDecodeMixin, - ArrayBytesCodecPartialEncodeMixin, - Codec, - ArrayArrayCodec, - ArrayBytesCodec, - BytesBytesCodec, -) -from zarr.v3.common import ArraySpec, BytesLike, NamedConfig, SliceSelection -from zarr.v3.metadata import RuntimeConfiguration -from zarr.v3.store import StorePath - -if TYPE_CHECKING: - from zarr.v3.metadata import ArrayMetadata - from zarr.v3.codecs.sharding import ShardingCodecMetadata - from zarr.v3.codecs.blosc import BloscCodecMetadata - from zarr.v3.codecs.bytes import BytesCodecMetadata - from zarr.v3.codecs.transpose import TransposeCodecMetadata - from zarr.v3.codecs.gzip import GzipCodecMetadata - from zarr.v3.codecs.zstd import ZstdCodecMetadata - from zarr.v3.codecs.crc32c_ import Crc32cCodecMetadata - - -ShardingCodecIndexLocation = Literal["start", "end"] - - -def _find_array_bytes_codec( - codecs: Iterable[Tuple[Codec, ArraySpec]] -) -> Tuple[ArrayBytesCodec, ArraySpec]: - for codec, array_spec in codecs: - if isinstance(codec, ArrayBytesCodec): - return (codec, array_spec) - raise KeyError - - -@dataclass(frozen=True) -class CodecPipeline: - array_array_codecs: List[ArrayArrayCodec] - array_bytes_codec: ArrayBytesCodec - bytes_bytes_codecs: List[BytesBytesCodec] - - @classmethod - def create(cls, codecs: List[Codec]) -> CodecPipeline: - from zarr.v3.codecs.sharding import ShardingCodec - - assert any( - isinstance(codec, ArrayBytesCodec) for codec in codecs - ), "Exactly one array-to-bytes codec is required." - - prev_codec: Optional[Codec] = None - for codec in codecs: - if prev_codec is not None: - assert not isinstance(codec, ArrayBytesCodec) or not isinstance( - prev_codec, ArrayBytesCodec - ), ( - f"ArrayBytesCodec '{type(codec)}' cannot follow after " - + f"ArrayBytesCodec '{type(prev_codec)}' because exactly " - + "1 ArrayBytesCodec is allowed." - ) - assert not isinstance(codec, ArrayBytesCodec) or not isinstance( - prev_codec, BytesBytesCodec - ), ( - f"ArrayBytesCodec '{type(codec)}' cannot follow after " - + f"BytesBytesCodec '{type(prev_codec)}'." - ) - assert not isinstance(codec, ArrayArrayCodec) or not isinstance( - prev_codec, ArrayBytesCodec - ), ( - f"ArrayArrayCodec '{type(codec)}' cannot follow after " - + f"ArrayBytesCodec '{type(prev_codec)}'." - ) - assert not isinstance(codec, ArrayArrayCodec) or not isinstance( - prev_codec, BytesBytesCodec - ), ( - f"ArrayArrayCodec '{type(codec)}' cannot follow after " - + f"BytesBytesCodec '{type(prev_codec)}'." - ) - prev_codec = codec - - 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 __iter__(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: - codec.validate(array_metadata) - - def _codecs_with_resolved_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, - array_spec: ArraySpec, - runtime_configuration: RuntimeConfiguration, - ) -> np.ndarray: - ( - aa_codecs_with_spec, - ab_codec_with_spec, - bb_codecs_with_spec, - ) = self._codecs_with_resolved_metadata(array_spec) - - 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, array_spec = ab_codec_with_spec - chunk_array = await ab_codec.decode(chunk_bytes, array_spec, 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_spec: 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_spec, runtime_configuration - ) - - async def encode( - self, - chunk_array: np.ndarray, - array_spec: ArraySpec, - runtime_configuration: RuntimeConfiguration, - ) -> Optional[BytesLike]: - ( - 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, 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, 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 - - async def encode_partial( - self, - store_path: StorePath, - chunk_array: np.ndarray, - selection: SliceSelection, - chunk_spec: 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_spec, runtime_configuration - ) - - def compute_encoded_size(self, byte_length: int, array_spec: ArraySpec) -> int: - for codec in self: - byte_length = codec.compute_encoded_size(byte_length, array_spec) - array_spec = codec.resolve_metadata(array_spec) - return byte_length - - -def blosc_codec( - typesize: int, - cname: Literal["lz4", "lz4hc", "blosclz", "zstd", "snappy", "zlib"] = "zstd", - clevel: int = 5, - shuffle: Literal["noshuffle", "shuffle", "bitshuffle"] = "noshuffle", - blocksize: int = 0, -) -> "BloscCodecMetadata": - from zarr.v3.codecs.blosc import BloscCodecMetadata, BloscCodecConfigurationMetadata - - return BloscCodecMetadata( - configuration=BloscCodecConfigurationMetadata( - cname=cname, - clevel=clevel, - shuffle=shuffle, - blocksize=blocksize, - typesize=typesize, - ) - ) - - -def bytes_codec(endian: Optional[Literal["big", "little"]] = "little") -> "BytesCodecMetadata": - from zarr.v3.codecs.bytes import BytesCodecMetadata, BytesCodecConfigurationMetadata - - return BytesCodecMetadata(configuration=BytesCodecConfigurationMetadata(endian)) - - -def transpose_codec( - order: Union[Tuple[int, ...], Literal["C", "F"]], ndim: Optional[int] = None -) -> "TransposeCodecMetadata": - from zarr.v3.codecs.transpose import TransposeCodecMetadata, TransposeCodecConfigurationMetadata - - if order == "C" or order == "F": - assert ( - isinstance(ndim, int) and ndim > 0 - ), 'When using "C" or "F" the `ndim` argument needs to be provided.' - if order == "C": - order = tuple(range(ndim)) - if order == "F": - order = tuple(ndim - i - 1 for i in range(ndim)) - - return TransposeCodecMetadata(configuration=TransposeCodecConfigurationMetadata(order)) - - -def gzip_codec(level: int = 5) -> "GzipCodecMetadata": - from zarr.v3.codecs.gzip import GzipCodecMetadata, GzipCodecConfigurationMetadata - - return GzipCodecMetadata(configuration=GzipCodecConfigurationMetadata(level)) - - -def zstd_codec(level: int = 0, checksum: bool = False) -> "ZstdCodecMetadata": - from zarr.v3.codecs.zstd import ZstdCodecMetadata, ZstdCodecConfigurationMetadata - - return ZstdCodecMetadata(configuration=ZstdCodecConfigurationMetadata(level, checksum)) - - -def crc32c_codec() -> "Crc32cCodecMetadata": - from zarr.v3.codecs.crc32c_ import Crc32cCodecMetadata - - return Crc32cCodecMetadata() - - -def sharding_codec( - chunk_shape: Tuple[int, ...], - codecs: Optional[List[NamedConfig]] = None, - index_codecs: Optional[List[NamedConfig]] = None, - index_location: ShardingCodecIndexLocation = "end", -) -> "ShardingCodecMetadata": - from zarr.v3.codecs.sharding import ShardingCodecMetadata, ShardingCodecConfigurationMetadata - - codecs = tuple(codecs) if codecs is not None else (bytes_codec(),) - index_codecs = ( - tuple(index_codecs) if index_codecs is not None else (bytes_codec(), crc32c_codec()) - ) - return ShardingCodecMetadata( - configuration=ShardingCodecConfigurationMetadata( - chunk_shape, codecs, index_codecs, index_location - ) - ) +from zarr.v3.codecs.blosc import BloscCodec, BloscCname, BloscShuffle # noqa: F401 +from zarr.v3.codecs.bytes import BytesCodec, Endian # noqa: F401 +from zarr.v3.codecs.crc32c_ import Crc32cCodec # noqa: F401 +from zarr.v3.codecs.gzip import GzipCodec # noqa: F401 +from zarr.v3.codecs.sharding import ShardingCodec, ShardingCodecIndexLocation # noqa: F401 +from zarr.v3.codecs.transpose import TransposeCodec # noqa: F401 +from zarr.v3.codecs.zstd import ZstdCodec # noqa: F401 diff --git a/src/zarr/v3/codecs/blosc.py b/src/zarr/v3/codecs/blosc.py index 2800e65a79..caaeb18e8b 100644 --- a/src/zarr/v3/codecs/blosc.py +++ b/src/zarr/v3/codecs/blosc.py @@ -1,39 +1,55 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass, replace +from enum import Enum from functools import lru_cache -from typing import ( - TYPE_CHECKING, - Any, - Dict, - Literal, - Optional, - Type, -) +from typing import TYPE_CHECKING import numcodecs import numpy as np from numcodecs.blosc import Blosc from zarr.v3.abc.codec import BytesBytesCodec -from zarr.v3.abc.metadata import Metadata from zarr.v3.codecs.registry import register_codec -from zarr.v3.common import to_thread -from zarr.v3.common import NamedConfig +from zarr.v3.common import parse_enum, parse_name, to_thread if TYPE_CHECKING: - from zarr.v3.common import ArraySpec + from typing import Dict, Optional from typing_extensions import Self - from zarr.v3.common import BytesLike, RuntimeConfiguration + from zarr.v3.common import JSON, ArraySpec, BytesLike, RuntimeConfiguration + + +class BloscShuffle(Enum): + noshuffle = "noshuffle" + shuffle = "shuffle" + bitshuffle = "bitshuffle" + + @classmethod + def from_int(cls, num: int) -> Self: + blosc_shuffle_int_to_str = { + 0: "noshuffle", + 1: "shuffle", + 2: "bitshuffle", + } + if num not in blosc_shuffle_int_to_str: + raise ValueError(f"Value must be between 0 and 2. Got {num}.") + return BloscShuffle[blosc_shuffle_int_to_str[num]] + + +class BloscCname(Enum): + lz4 = "lz4" + lz4hc = "lz4hc" + blosclz = "blosclz" + zstd = "zstd" + snappy = "snappy" + zlib = "zlib" -BloscShuffle = Literal["noshuffle", "shuffle", "bitshuffle"] -BloscCname = Literal["lz4", "lz4hc", "blosclz", "zstd", "snappy", "zlib"] # See https://zarr.readthedocs.io/en/stable/tutorial.html#configuring-blosc numcodecs.blosc.use_threads = False -def parse_typesize(data: Any) -> int: +def parse_typesize(data: JSON) -> int: if isinstance(data, int): if data >= 0: return data @@ -44,66 +60,44 @@ def parse_typesize(data: Any) -> int: raise TypeError(msg) -def parse_cname(data: Any) -> BloscCname: - if data in ["lz4", "lz4hc", "blosclz", "zstd", "snappy", "zlib"]: - return data - msg = ( - "Value must be one of ", - '["lz4", "lz4hc", "blosclz", "zstd", "snappy", "zlib"], ', - f"got {data} instead.", - ) - raise ValueError(msg) - - # todo: real validation -def parse_clevel(data: Any) -> int: +def parse_clevel(data: JSON) -> int: if isinstance(data, int): return data msg = f"Value should be an int, got {type(data)} instead" raise TypeError(msg) -def parse_shuffle(data: Any) -> BloscShuffle: - if data in ["noshuffle", "shuffle", "bitshuffle"]: - return data - msg = f'Value must be one of ["noshuffle", "shuffle", "bitshuffle"], got {data} instead.' - raise ValueError(msg) - - -def parse_blocksize(data: Any) -> int: +def parse_blocksize(data: JSON) -> int: if isinstance(data, int): return data msg = f"Value should be an int, got {type(data)} instead" raise TypeError(msg) -def parse_name(data: Any) -> Literal["blosc"]: - if data == "blosc": - return data - msg = f"Expected 'blosc', got {data} instead." - raise ValueError(msg) - - @dataclass(frozen=True) -class BloscCodecConfigurationMetadata(Metadata): +class BloscCodec(BytesBytesCodec): + is_fixed_size = False + typesize: int - cname: BloscCname = "zstd" + cname: BloscCname = BloscCname.zstd clevel: int = 5 - shuffle: BloscShuffle = "noshuffle" + shuffle: BloscShuffle = BloscShuffle.noshuffle blocksize: int = 0 def __init__( self, - typesize: int, - cname: BloscCname = "zstd", - clevel: int = 5, - shuffle: BloscShuffle = "noshuffle", - blocksize: int = 0, - ): + *, + typesize, + cname=BloscCname.zstd, + clevel=5, + shuffle=BloscShuffle.noshuffle, + blocksize=0, + ) -> None: typesize_parsed = parse_typesize(typesize) - cname_parsed = parse_cname(cname) + cname_parsed = parse_enum(cname, BloscCname) clevel_parsed = parse_clevel(clevel) - shuffle_parsed = parse_shuffle(shuffle) + shuffle_parsed = parse_enum(shuffle, BloscShuffle) blocksize_parsed = parse_blocksize(blocksize) object.__setattr__(self, "typesize", typesize_parsed) @@ -112,55 +106,42 @@ def __init__( object.__setattr__(self, "shuffle", shuffle_parsed) object.__setattr__(self, "blocksize", blocksize_parsed) - -blosc_shuffle_int_to_str: Dict[int, BloscShuffle] = { - 0: "noshuffle", - 1: "shuffle", - 2: "bitshuffle", -} - - -@dataclass(frozen=True) -class BloscCodecMetadata(Metadata): - configuration: BloscCodecConfigurationMetadata - name: Literal["blosc"] = field(default="blosc", init=False) - @classmethod - def from_dict(cls, data: Dict[str, Any]) -> Self: - _ = parse_name(data.pop("name")) - return cls(**data) - - -@dataclass(frozen=True) -class BloscCodec(BytesBytesCodec): - configuration: BloscCodecConfigurationMetadata - is_fixed_size = False - - @classmethod - def from_metadata(cls, codec_metadata: NamedConfig) -> BloscCodec: - assert isinstance(codec_metadata, BloscCodecMetadata) - return cls(configuration=codec_metadata.configuration) - - @classmethod - def get_metadata_class(cls) -> Type[BloscCodecMetadata]: - return BloscCodecMetadata + def from_dict(cls, data: Dict[str, JSON]) -> Self: + parse_name(data["name"], "blosc") + return cls(**data["configuration"]) + + def to_dict(self) -> Dict[str, JSON]: + return { + "name": "blosc", + "configuration": { + "typesize": self.typesize, + "cname": self.cname, + "clevel": self.clevel, + "shuffle": self.shuffle, + "blocksize": self.blocksize, + }, + } - def evolve(self, *, data_type: np.dtype, **_kwargs) -> BloscCodec: + def evolve(self, array_spec: ArraySpec) -> Self: new_codec = self - if new_codec.configuration.typesize == 0: - new_configuration = evolve(new_codec.configuration, typesize=data_type.byte_count) - new_codec = evolve(new_codec, configuration=new_configuration) + if new_codec.typesize == 0: + new_codec = replace(new_codec, typesize=array_spec.dtype.itemsize) return new_codec @lru_cache def get_blosc_codec(self) -> Blosc: - map_shuffle_str_to_int = {"noshuffle": 0, "shuffle": 1, "bitshuffle": 2} + map_shuffle_str_to_int = { + BloscShuffle.noshuffle: 0, + BloscShuffle.shuffle: 1, + BloscShuffle.bitshuffle: 2, + } config_dict = { - "cname": self.configuration.cname, - "clevel": self.configuration.clevel, - "shuffle": map_shuffle_str_to_int[self.configuration.shuffle], - "blocksize": self.configuration.blocksize, + "cname": self.cname.name, + "clevel": self.clevel, + "shuffle": map_shuffle_str_to_int[self.shuffle], + "blocksize": self.blocksize, } return Blosc.from_config(config_dict) @@ -184,12 +165,5 @@ async def encode( def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int: raise NotImplementedError - def to_dict(self) -> Dict[str, Any]: - return BloscCodecMetadata(configuration=self.configuration) - - @classmethod - def from_dict(cls, data: Dict[str, Any]): - return cls(configuration=data["configuration"]) - register_codec("blosc", BloscCodec) diff --git a/src/zarr/v3/codecs/bytes.py b/src/zarr/v3/codecs/bytes.py index 95fbe7a2ef..705a63c02e 100644 --- a/src/zarr/v3/codecs/bytes.py +++ b/src/zarr/v3/codecs/bytes.py @@ -1,81 +1,62 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass, replace +from enum import Enum -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Dict, Optional import numpy as np from zarr.v3.abc.codec import ArrayBytesCodec -from zarr.v3.abc.metadata import Metadata from zarr.v3.codecs.registry import register_codec +from zarr.v3.common import parse_enum, parse_name if TYPE_CHECKING: - from zarr.v3.common import ArraySpec, NamedConfig, BytesLike, RuntimeConfiguration + from zarr.v3.common import JSON, ArraySpec, BytesLike, RuntimeConfiguration from typing_extensions import Self - from typing import Any, Dict, Optional, Type -def parse_endian(data: Any) -> Literal["big", "little"]: - if data in ("big", "little"): - return data - msg = f"Expected on of ('big', 'little'), got {data} instead." - raise ValueError(msg) - - -def parse_name(data: Any) -> Literal["bytes"]: - if data == "bytes": - return data - msg = f"Expected 'bytes', got {data} instead." - raise ValueError(msg) - - -Endian = Literal["big", "little"] - - -@dataclass(frozen=True) -class BytesCodecConfigurationMetadata(Metadata): - endian: Optional[Endian] = "little" - - def __init__(self, endian: Literal["big", "little"]): - endian_parsed = parse_endian(endian) - object.__setattr__(self, "endian", endian_parsed) - - -@dataclass(frozen=True) -class BytesCodecMetadata(Metadata): - configuration: BytesCodecConfigurationMetadata - name: Literal["bytes"] = field(default="bytes", init=False) - - @classmethod - def from_dict(cls, data: Dict[str, Any]) -> Self: - _ = parse_name(data.pop("name")) - return cls(**data) +class Endian(Enum): + big = "big" + little = "little" @dataclass(frozen=True) class BytesCodec(ArrayBytesCodec): - configuration: BytesCodecConfigurationMetadata is_fixed_size = True - @classmethod - def from_metadata(cls, codec_metadata: NamedConfig) -> BytesCodec: - assert isinstance(codec_metadata, BytesCodecMetadata) - return cls(configuration=codec_metadata.configuration) + endian: Optional[Endian] + + def __init__(self, *, endian=Endian.little) -> None: + endian_parsed = None if endian is None else parse_enum(endian, Endian) + + object.__setattr__(self, "endian", endian_parsed) @classmethod - def get_metadata_class(cls) -> Type[BytesCodecMetadata]: - return BytesCodecMetadata + def from_dict(cls, data: Dict[str, JSON]) -> Self: + parse_name(data["name"], "bytes") + return cls(**data.get("configuration", {})) - def validate(self, array_metadata: ArraySpec) -> None: - assert ( - not array_metadata.data_type.has_endianness or self.configuration.endian is not None - ), "The `endian` configuration needs to be specified for multi-byte data types." + def to_dict(self) -> Dict[str, JSON]: + if self.endian is None: + return {"name": "bytes"} + else: + return {"name": "bytes", "configuration": {"endian": self.endian}} + + def evolve(self, array_spec: ArraySpec) -> Self: + if array_spec.dtype.itemsize == 0: + if self.endian is not None: + return replace(self, endian=None) + elif self.endian is None: + raise ValueError( + "The `endian` configuration needs to be specified for multi-byte data types." + ) + return self def _get_byteorder(self, array: np.ndarray) -> Endian: if array.dtype.byteorder == "<": - return "little" + return Endian.little elif array.dtype.byteorder == ">": - return "big" + return Endian.big else: import sys @@ -88,13 +69,14 @@ async def decode( _runtime_configuration: RuntimeConfiguration, ) -> np.ndarray: if chunk_spec.dtype.itemsize > 0: - if self.configuration.endian == "little": + if self.endian == Endian.little: prefix = "<" else: prefix = ">" - dtype = np.dtype(f"{prefix}{self.array_metadata.dtype.str[1:]}") + dtype = np.dtype(f"{prefix}{chunk_spec.dtype.str[1:]}") else: - dtype = np.dtype(f"|{self.array_metadata.dtype.str[1:]}") + dtype = np.dtype(f"|{chunk_spec.dtype.str[1:]}") + print(dtype) chunk_array = np.frombuffer(chunk_bytes, dtype) # ensure correct chunk shape @@ -112,21 +94,14 @@ async def encode( ) -> Optional[BytesLike]: if chunk_array.dtype.itemsize > 1: byteorder = self._get_byteorder(chunk_array) - if self.configuration.endian != byteorder: - new_dtype = chunk_array.dtype.newbyteorder(self.configuration.endian) + if self.endian != byteorder: + new_dtype = chunk_array.dtype.newbyteorder(self.endian.name) chunk_array = chunk_array.astype(new_dtype) return chunk_array.tobytes() def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int: return input_byte_length - def to_dict(self) -> Dict[str, Any]: - return BytesCodecMetadata(configuration=self.configuration) - - @classmethod - def from_dict(cls, data: Dict[str, Any]): - return cls(configuration=data["configuration"]) - register_codec("bytes", BytesCodec) diff --git a/src/zarr/v3/codecs/common.py b/src/zarr/v3/codecs/common.py deleted file mode 100644 index 0221bb7b9f..0000000000 --- a/src/zarr/v3/codecs/common.py +++ /dev/null @@ -1,198 +0,0 @@ -from __future__ import annotations -from zarr.v3.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec -from zarr.v3.codecs.registry import get_codec_class -from zarr.v3.common import BytesLike, NamedConfig, RuntimeConfiguration -from zarr.v3.common import ArraySpec - - -import numpy as np - - -from dataclasses import dataclass -from functools import reduce -from typing import Iterable, List, Optional -from warnings import warn - - -@dataclass(frozen=True) -class CodecPipeline: - codecs: List[Codec] - - @classmethod - def from_metadata( - cls, - codecs_metadata: Iterable[NamedConfig], - array_metadata: ArraySpec, - ) -> CodecPipeline: - out: List[Codec] = [] - for codec_metadata in codecs_metadata or []: - codec_cls = get_codec_class(codec_metadata.name) - codec = codec_cls.from_metadata(codec_metadata, array_metadata) - out.append(codec) - array_metadata = codec.resolve_metadata() - CodecPipeline._validate_codecs(out, array_metadata) - return cls(out) - - @staticmethod - def _validate_codecs(codecs: List[Codec], array_metadata: ArraySpec) -> None: - from zarr.v3.codecs.sharding import ShardingCodec - - assert any( - isinstance(codec, ArrayBytesCodec) for codec in codecs - ), "Exactly one array-to-bytes codec is required." - - prev_codec: Optional[Codec] = None - for codec in codecs: - if prev_codec is not None: - assert not isinstance(codec, ArrayBytesCodec) or not isinstance( - prev_codec, ArrayBytesCodec - ), ( - f"ArrayBytesCodec '{type(codec)}' cannot follow after " - + f"ArrayBytesCodec '{type(prev_codec)}' because exactly " - + "1 ArrayBytesCodec is allowed." - ) - assert not isinstance(codec, ArrayBytesCodec) or not isinstance( - prev_codec, BytesBytesCodec - ), ( - f"ArrayBytesCodec '{type(codec)}' cannot follow after " - + f"BytesBytesCodec '{type(prev_codec)}'." - ) - assert not isinstance(codec, ArrayArrayCodec) or not isinstance( - prev_codec, ArrayBytesCodec - ), ( - f"ArrayArrayCodec '{type(codec)}' cannot follow after " - + f"ArrayBytesCodec '{type(prev_codec)}'." - ) - assert not isinstance(codec, ArrayArrayCodec) or not isinstance( - prev_codec, BytesBytesCodec - ), ( - f"ArrayArrayCodec '{type(codec)}' cannot follow after " - + f"BytesBytesCodec '{type(prev_codec)}'." - ) - - if isinstance(codec, ShardingCodec): - assert len(codec.configuration.chunk_shape) == len(array_metadata.shape), ( - "The shard's `chunk_shape` and array's `shape` need to have the " - + "same number of dimensions." - ) - assert all( - s % c == 0 - for s, c in zip( - array_metadata.chunk_shape, - codec.configuration.chunk_shape, - ) - ), ( - "The array's `chunk_shape` needs to be divisible by the " - + "shard's inner `chunk_shape`." - ) - prev_codec = codec - - 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." - ) - - def _array_array_codecs(self) -> List[ArrayArrayCodec]: - return [codec for codec in self.codecs if isinstance(codec, ArrayArrayCodec)] - - def _array_bytes_codec(self) -> ArrayBytesCodec: - return next(codec for codec in self.codecs if isinstance(codec, ArrayBytesCodec)) - - def _bytes_bytes_codecs(self) -> List[BytesBytesCodec]: - return [codec for codec in self.codecs if isinstance(codec, BytesBytesCodec)] - - async def decode(self, chunk_bytes: BytesLike) -> np.ndarray: - for bb_codec in self._bytes_bytes_codecs()[::-1]: - chunk_bytes = await bb_codec.decode(chunk_bytes) - - chunk_array = await self._array_bytes_codec().decode(chunk_bytes) - - for aa_codec in self._array_array_codecs()[::-1]: - chunk_array = await aa_codec.decode(chunk_array) - - return chunk_array - - async def encode(self, chunk_array: np.ndarray) -> Optional[BytesLike]: - for aa_codec in self._array_array_codecs(): - chunk_array_maybe = await aa_codec.encode(chunk_array) - if chunk_array_maybe is None: - return None - chunk_array = chunk_array_maybe - - chunk_bytes_maybe = await self._array_bytes_codec().encode(chunk_array) - if chunk_bytes_maybe is None: - return None - chunk_bytes = chunk_bytes_maybe - - for bb_codec in self._bytes_bytes_codecs(): - chunk_bytes_maybe = await bb_codec.encode(chunk_bytes) - if chunk_bytes_maybe is None: - return None - chunk_bytes = chunk_bytes_maybe - - return chunk_bytes - - def compute_encoded_size(self, byte_length: int) -> int: - return reduce(lambda acc, codec: codec.compute_encoded_size(acc), self.codecs, byte_length) - - -async def decode( - codecs: List[Codec], chunk_bytes: BytesLike, runtime_configuration: RuntimeConfiguration -) -> np.ndarray: - # todo: increase the arity of the function signature with - # positions for array_array, array_bytes, and bytes_bytes - _array_array_codecs = [codec for codec in codecs if isinstance(codec, ArrayArrayCodec)] - _array_bytes_codec = next(codec for codec in codecs if isinstance(codec, ArrayBytesCodec)) - _bytes_bytes_codecs = [codec for codec in codecs if isinstance(codec, BytesBytesCodec)] - - for bb_codec in _bytes_bytes_codecs[::-1]: - chunk_bytes = await bb_codec.decode( - chunk_bytes, runtime_configuration=runtime_configuration - ) - - chunk_array = await _array_bytes_codec.decode( - chunk_bytes, runtime_configuration=runtime_configuration - ) - - for aa_codec in _array_array_codecs[::-1]: - chunk_array = await aa_codec.decode( - chunk_array, runtime_configuration=runtime_configuration - ) - - return chunk_array - - -async def encode( - codecs: List[Codec], chunk_array: np.ndarray, runtime_configuration: RuntimeConfiguration -) -> Optional[BytesLike]: - # todo: increase the arity of the function signature - # with positions for array_array, array_bytes, and bytes_bytes - _array_array_codecs = [codec for codec in codecs if isinstance(codec, ArrayArrayCodec)] - _array_bytes_codec = next(codec for codec in codecs if isinstance(codec, ArrayBytesCodec)) - _bytes_bytes_codecs = [codec for codec in codecs if isinstance(codec, BytesBytesCodec)] - - for aa_codec in _array_array_codecs: - chunk_array_maybe = await aa_codec.encode( - chunk_array, runtime_configuration=runtime_configuration - ) - if chunk_array_maybe is None: - return None - chunk_array = chunk_array_maybe - - chunk_bytes_maybe = await _array_bytes_codec.encode( - chunk_array, runtime_configuration=runtime_configuration - ) - if chunk_bytes_maybe is None: - return None - chunk_bytes = chunk_bytes_maybe - - for bb_codec in _bytes_bytes_codecs: - chunk_bytes_maybe = await bb_codec.encode( - chunk_bytes, runtime_configuration=runtime_configuration - ) - if chunk_bytes_maybe is None: - return None - chunk_bytes = chunk_bytes_maybe - - return chunk_bytes diff --git a/src/zarr/v3/codecs/crc32c_.py b/src/zarr/v3/codecs/crc32c_.py index a3e7602a8c..a5b9eec3bc 100644 --- a/src/zarr/v3/codecs/crc32c_.py +++ b/src/zarr/v3/codecs/crc32c_.py @@ -1,43 +1,20 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass -from typing import ( - TYPE_CHECKING, - Any, - Dict, - Literal, - Optional, - Type, -) +from typing import TYPE_CHECKING import numpy as np from crc32c import crc32c from zarr.v3.abc.codec import BytesBytesCodec -from zarr.v3.abc.metadata import Metadata from zarr.v3.codecs.registry import register_codec -from zarr.v3.common import NamedConfig +from zarr.v3.common import parse_name if TYPE_CHECKING: - from zarr.v3.common import BytesLike, RuntimeConfiguration, ArraySpec - - -def parse_name(data: Any) -> Literal["crc32c"]: - if data == "crc32c": - return data - msg = f"Expected 'crc32c', got {data} instead." - raise ValueError(msg) - - -@dataclass(frozen=True) -class Crc32cCodecMetadata(Metadata): - name: Literal["crc32c"] = field(default="crc32c", init=False) - - @classmethod - def from_dict(cls, data: Any): - _ = parse_name(data.pop("name")) - return cls(**data) + from typing import Dict, Optional + from typing_extensions import Self + from zarr.v3.common import JSON, BytesLike, RuntimeConfiguration, ArraySpec @dataclass(frozen=True) @@ -45,13 +22,12 @@ class Crc32cCodec(BytesBytesCodec): is_fixed_size = True @classmethod - def from_metadata(cls, codec_metadata: NamedConfig) -> Crc32cCodec: - assert isinstance(codec_metadata, Crc32cCodecMetadata) + def from_dict(cls, data: Dict[str, JSON]) -> Self: + parse_name(data["name"], "crc32c") return cls() - @classmethod - def get_metadata_class(cls) -> Type[Crc32cCodecMetadata]: - return Crc32cCodecMetadata + def to_dict(self) -> Dict[str, JSON]: + return {"name": "crc32c"} async def decode( self, @@ -76,12 +52,5 @@ async def encode( def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int: return input_byte_length + 4 - def to_dict(self) -> Dict[str, Any]: - return Crc32cCodecMetadata() - - @classmethod - def from_dict(cls, data: Dict[str, Any]): - return cls(configuration=data["configuration"]) - register_codec("crc32c", Crc32cCodec) diff --git a/src/zarr/v3/codecs/gzip.py b/src/zarr/v3/codecs/gzip.py index 1b1d49f94e..1c6c34270d 100644 --- a/src/zarr/v3/codecs/gzip.py +++ b/src/zarr/v3/codecs/gzip.py @@ -1,61 +1,46 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING -from zarr.v3.abc.metadata import Metadata from numcodecs.gzip import GZip from zarr.v3.abc.codec import BytesBytesCodec from zarr.v3.codecs.registry import register_codec -from zarr.v3.common import ArraySpec, to_thread +from zarr.v3.common import parse_name, to_thread if TYPE_CHECKING: - from zarr.v3.metadata import RuntimeConfiguration - from zarr.v3.common import BytesLike, NamedConfig + from typing import Optional, Dict from typing_extensions import Self - from typing import Any, Optional, Dict, Literal, Type + from zarr.v3.metadata import RuntimeConfiguration + from zarr.v3.common import JSON, ArraySpec, BytesLike -def parse_gzip_level(data: Any) -> int: +def parse_gzip_level(data: JSON) -> int: if data not in range(0, 10): - msg = f"Expected an integer from the inclusive range (0, 9). Got {data} instead." - raise ValueError(msg) + raise ValueError( + f"Expected an integer from the inclusive range (0, 9). Got {data} instead." + ) return data @dataclass(frozen=True) -class GzipCodecConfigurationMetadata(Metadata): +class GzipCodec(BytesBytesCodec): + is_fixed_size = False + level: int = 5 - def __init__(self, level: int): + def __init__(self, *, level=5) -> None: level_parsed = parse_gzip_level(level) - object.__setattr__(self, "level", level_parsed) - - -@dataclass(frozen=True) -class GzipCodecMetadata(Metadata): - configuration: GzipCodecConfigurationMetadata - name: Literal["gzip"] = field(default="gzip", init=False) - - @classmethod - def from_dict(cls, data: Dict[str, Any]) -> Self: - return cls(configuration=GzipCodecConfigurationMetadata.from_dict(data["configuration"])) - -@dataclass(frozen=True) -class GzipCodec(BytesBytesCodec): - configuration: GzipCodecConfigurationMetadata - is_fixed_size: Literal[True] = field(default=True, init=False) + object.__setattr__(self, "level", level_parsed) @classmethod - def from_metadata(cls, codec_metadata: NamedConfig) -> GzipCodec: - assert isinstance(codec_metadata, GzipCodecMetadata) + def from_dict(cls, data: Dict[str, JSON]) -> Self: + parse_name(data["name"], "gzip") + return cls(**data["configuration"]) - return cls(configuration=codec_metadata.configuration) - - @classmethod - def get_metadata_class(cls) -> Type[GzipCodecMetadata]: - return GzipCodecMetadata + def to_dict(self) -> Dict[str, JSON]: + return {"name": "gzip", "configuration": {"level": self.level}} async def decode( self, @@ -63,7 +48,7 @@ async def decode( _chunk_spec: ArraySpec, _runtime_configuration: RuntimeConfiguration, ) -> BytesLike: - return await to_thread(GZip(self.configuration.level).decode, chunk_bytes) + return await to_thread(GZip(self.level).decode, chunk_bytes) async def encode( self, @@ -71,7 +56,7 @@ async def encode( _chunk_spec: ArraySpec, _runtime_configuration: RuntimeConfiguration, ) -> Optional[BytesLike]: - return await to_thread(GZip(self.configuration.level).encode, chunk_bytes) + return await to_thread(GZip(self.level).encode, chunk_bytes) def compute_encoded_size( self, @@ -80,12 +65,5 @@ def compute_encoded_size( ) -> int: raise NotImplementedError - def to_dict(self) -> Dict[str, Any]: - return GzipCodecMetadata(configuration=self.configuration) - - @classmethod - def from_dict(cls, data: Dict[str, Any]): - return cls(configuration=data["configuration"]) - register_codec("gzip", GzipCodec) diff --git a/src/zarr/v3/codecs/pipeline.py b/src/zarr/v3/codecs/pipeline.py new file mode 100644 index 0000000000..50c69879c8 --- /dev/null +++ b/src/zarr/v3/codecs/pipeline.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +import numpy as np +from dataclasses import dataclass +from warnings import warn + +from zarr.v3.abc.codec import ( + ArrayArrayCodec, + ArrayBytesCodec, + ArrayBytesCodecPartialDecodeMixin, + ArrayBytesCodecPartialEncodeMixin, + BytesBytesCodec, + Codec, +) +from zarr.v3.abc.metadata import Metadata +from zarr.v3.codecs.registry import get_codec_class + +if TYPE_CHECKING: + from typing import Dict, Iterator, List, Optional, Tuple, Union + from typing_extensions import Self + from zarr.v3.store import StorePath + from zarr.v3.metadata import ArrayMetadata, RuntimeConfiguration + from zarr.v3.common import JSON, ArraySpec, BytesLike, SliceSelection + + +@dataclass(frozen=True) +class CodecPipeline(Metadata): + array_array_codecs: Tuple[ArrayArrayCodec, ...] + array_bytes_codec: ArrayBytesCodec + bytes_bytes_codecs: Tuple[BytesBytesCodec, ...] + + @classmethod + def from_dict(cls, data: Union[JSON, List[Codec]]) -> Self: + out: List[Codec] = [] + for c in data: + if isinstance(c, Codec): + out.append(c) + else: + out.append(get_codec_class(c["name"]).from_dict(c)) + return CodecPipeline.from_list(out) + + def to_dict(self) -> Dict[str, JSON]: + return [c.to_dict() for c in self] + + def evolve(self, array_spec: ArraySpec) -> Self: + return CodecPipeline.from_list([c.evolve(array_spec) for c in self]) + + @classmethod + def from_list(cls, codecs: List[Codec]) -> CodecPipeline: + from zarr.v3.codecs.sharding import ShardingCodec + + assert any( + isinstance(codec, ArrayBytesCodec) for codec in codecs + ), "Exactly one array-to-bytes codec is required." + + prev_codec: Optional[Codec] = None + for codec in codecs: + if prev_codec is not None: + assert not isinstance(codec, ArrayBytesCodec) or not isinstance( + prev_codec, ArrayBytesCodec + ), ( + f"ArrayBytesCodec '{type(codec)}' cannot follow after " + + f"ArrayBytesCodec '{type(prev_codec)}' because exactly " + + "1 ArrayBytesCodec is allowed." + ) + assert not isinstance(codec, ArrayBytesCodec) or not isinstance( + prev_codec, BytesBytesCodec + ), ( + f"ArrayBytesCodec '{type(codec)}' cannot follow after " + + f"BytesBytesCodec '{type(prev_codec)}'." + ) + assert not isinstance(codec, ArrayArrayCodec) or not isinstance( + prev_codec, ArrayBytesCodec + ), ( + f"ArrayArrayCodec '{type(codec)}' cannot follow after " + + f"ArrayBytesCodec '{type(prev_codec)}'." + ) + assert not isinstance(codec, ArrayArrayCodec) or not isinstance( + prev_codec, BytesBytesCodec + ), ( + f"ArrayArrayCodec '{type(codec)}' cannot follow after " + + f"BytesBytesCodec '{type(prev_codec)}'." + ) + prev_codec = codec + + 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=tuple( + 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=tuple( + 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 __iter__(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: + codec.validate(array_metadata) + + def _codecs_with_resolved_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, + array_spec: ArraySpec, + runtime_configuration: RuntimeConfiguration, + ) -> np.ndarray: + ( + aa_codecs_with_spec, + ab_codec_with_spec, + bb_codecs_with_spec, + ) = self._codecs_with_resolved_metadata(array_spec) + + 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, array_spec = ab_codec_with_spec + chunk_array = await ab_codec.decode(chunk_bytes, array_spec, 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_spec: 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_spec, runtime_configuration + ) + + async def encode( + self, + chunk_array: np.ndarray, + array_spec: ArraySpec, + runtime_configuration: RuntimeConfiguration, + ) -> Optional[BytesLike]: + ( + 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, 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, 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 + + async def encode_partial( + self, + store_path: StorePath, + chunk_array: np.ndarray, + selection: SliceSelection, + chunk_spec: 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_spec, runtime_configuration + ) + + def compute_encoded_size(self, byte_length: int, array_spec: ArraySpec) -> int: + for codec in self: + byte_length = codec.compute_encoded_size(byte_length, array_spec) + array_spec = codec.resolve_metadata(array_spec) + return byte_length diff --git a/src/zarr/v3/codecs/registry.py b/src/zarr/v3/codecs/registry.py index 9aae3648db..4cf2736685 100644 --- a/src/zarr/v3/codecs/registry.py +++ b/src/zarr/v3/codecs/registry.py @@ -1,20 +1,14 @@ from __future__ import annotations -from typing import TYPE_CHECKING, NamedTuple +from typing import TYPE_CHECKING if TYPE_CHECKING: from typing import Dict, Type from zarr.v3.abc.codec import Codec - from zarr.v3.common import NamedConfig from importlib.metadata import EntryPoint, entry_points as get_entry_points -class CodecRegistryItem(NamedTuple): - codec_cls: Type[Codec] - codec_metadata_cls: Type[NamedConfig] - - -__codec_registry: Dict[str, CodecRegistryItem] = {} +__codec_registry: Dict[str, Type[Codec]] = {} __lazy_load_codecs: Dict[str, EntryPoint] = {} @@ -31,10 +25,10 @@ def _collect_entrypoints() -> None: def register_codec(key: str, codec_cls: Type[Codec]) -> None: - __codec_registry[key] = CodecRegistryItem(codec_cls, codec_cls.get_metadata_class()) + __codec_registry[key] = codec_cls -def _get_codec_item(key: str) -> CodecRegistryItem: +def get_codec_class(key: str) -> Type[Codec]: item = __codec_registry.get(key) if item is None: if key in __lazy_load_codecs: @@ -47,17 +41,4 @@ def _get_codec_item(key: str) -> CodecRegistryItem: raise KeyError(key) -def get_codec_from_metadata(val: NamedConfig) -> Codec: - key = val.name - return _get_codec_item(key).codec_cls.from_metadata(val) - - -def get_codec_metadata_class(key: str) -> Type[NamedConfig]: - return _get_codec_item(key).codec_metadata_cls - - -def get_codec_class(key: str) -> Type[Codec]: - return _get_codec_item(key).codec_cls - - _collect_entrypoints() diff --git a/src/zarr/v3/codecs/sharding.py b/src/zarr/v3/codecs/sharding.py index 7573df586e..fb60dc4ab8 100644 --- a/src/zarr/v3/codecs/sharding.py +++ b/src/zarr/v3/codecs/sharding.py @@ -1,79 +1,73 @@ from __future__ import annotations -from typing import TYPE_CHECKING, NamedTuple, Mapping -from dataclasses import dataclass, field -from functools import cached_property, lru_cache +from enum import Enum +from typing import TYPE_CHECKING, Mapping, NamedTuple +from dataclasses import dataclass +from functools import lru_cache -if TYPE_CHECKING: - from typing import ( - Any, - Awaitable, - Callable, - Dict, - Iterator, - List, - Optional, - Set, - Tuple, - Type, - Literal, - ) - from zarr.v3.common import BytesLike, NamedConfig, RuntimeConfiguration, SliceSelection import numpy as np from zarr.v3.abc.codec import ( ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin, - Codec, ) -from zarr.v3.abc.metadata import Metadata - -from zarr.v3.codecs.common import CodecPipeline -from zarr.v3.codecs.registry import get_codec_from_metadata, register_codec +from zarr.v3.codecs.bytes import BytesCodec +from zarr.v3.codecs.crc32c_ import Crc32cCodec +from zarr.v3.codecs.pipeline import CodecPipeline +from zarr.v3.codecs.registry import register_codec from zarr.v3.common import ( ArraySpec, - ChunkCoords, concurrent_map, + parse_enum, + parse_name, + parse_shapelike, product, ) - +from zarr.v3.chunk_grids import RegularChunkGrid from zarr.v3.indexing import ( BasicIndexer, c_order_iter, is_total_slice, morton_order_iter, ) - from zarr.v3.metadata import ( ArrayMetadata, - RegularChunkGridMetadata, - ShardingCodecIndexLocation, runtime_configuration as make_runtime_configuration, + parse_codecs, ) -from zarr.v3.store import StorePath +if TYPE_CHECKING: + from typing import ( + Awaitable, + Callable, + Dict, + Iterator, + List, + Optional, + Set, + Tuple, + ) + from typing_extensions import Self + + from zarr.v3.store import StorePath + from zarr.v3.common import ( + JSON, + ChunkCoords, + BytesLike, + RuntimeConfiguration, + SliceSelection, + ) MAX_UINT_64 = 2**64 - 1 -@dataclass(frozen=True) -class ShardingCodecConfigurationMetadata(Metadata): - chunk_shape: ChunkCoords - codecs: List[Codec] - index_codecs: List["NamedConfig"] - index_location: ShardingCodecIndexLocation = "end" - +class ShardingCodecIndexLocation(Enum): + start = "start" + end = "end" -@dataclass(frozen=True) -class ShardingCodecMetadata(Metadata): - configuration: ShardingCodecConfigurationMetadata - name: Literal["sharding_indexed"] = field(default="sharding_indexed", init=False) - @classmethod - def from_dict(cls, data: Dict[str, Any]): - return cls( - configuration=ShardingCodecConfigurationMetadata.from_dict(data["configuration"]) - ) +def parse_index_location(data: JSON) -> ShardingCodecIndexLocation: + return parse_enum(data, ShardingCodecIndexLocation) class _ShardIndex(NamedTuple): @@ -150,7 +144,7 @@ async def from_bytes( shard_index_size = codec._shard_index_size(chunks_per_shard) obj = cls() obj.buf = memoryview(buf) - if codec.configuration.index_location == "start": + if codec.index_location == ShardingCodecIndexLocation.start: shard_index_bytes = obj.buf[:shard_index_size] else: shard_index_bytes = obj.buf[-shard_index_size:] @@ -218,14 +212,11 @@ async def finalize( self, index_location: ShardingCodecIndexLocation, index_encoder: Callable[[_ShardIndex], Awaitable[BytesLike]], - runtime_configuration: RuntimeConfiguration, ) -> BytesLike: - index_bytes = await index_encoder(self.index, runtime_configuration=runtime_configuration) - if index_location == "start": + index_bytes = await index_encoder(self.index) + if index_location == ShardingCodecIndexLocation.start: self.index.offsets_and_lengths[..., 0] += len(index_bytes) - index_bytes = await index_encoder( - self.index, runtime_configuration=runtime_configuration - ) # encode again with corrected offsets + index_bytes = await index_encoder(self.index) # encode again with corrected offsets out_buf = bytearray(index_bytes) out_buf.extend(self.buf) else: @@ -236,35 +227,70 @@ async def finalize( @dataclass(frozen=True) class ShardingCodec( - ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin, Metadata + ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin ): - configuration: ShardingCodecConfigurationMetadata + chunk_shape: ChunkCoords + codecs: CodecPipeline + index_codecs: CodecPipeline + index_location: ShardingCodecIndexLocation = "end" - @classmethod - def from_metadata( - cls, - codec_metadata: NamedConfig, - ) -> ShardingCodec: - assert isinstance(codec_metadata, ShardingCodecMetadata) - return cls(configuration=codec_metadata.configuration) + def __init__( + self, + *, + chunk_shape, + codecs=None, + index_codecs=None, + index_location=None, + ) -> None: + chunk_shape_parsed = parse_shapelike(chunk_shape) + codecs_parsed = ( + parse_codecs(codecs) if codecs is not None else CodecPipeline.from_list([BytesCodec()]) + ) + index_codecs_parsed = ( + parse_codecs(index_codecs) + if index_codecs is not None + else CodecPipeline.from_list([BytesCodec(), Crc32cCodec()]) + ) + index_location_parsed = ( + parse_index_location(index_location) + if index_location is not None + else ShardingCodecIndexLocation.end + ) + + object.__setattr__(self, "chunk_shape", chunk_shape_parsed) + object.__setattr__(self, "codecs", codecs_parsed) + object.__setattr__(self, "index_codecs", index_codecs_parsed) + object.__setattr__(self, "index_location", index_location_parsed) @classmethod - def get_metadata_class(cls) -> Type[ShardingCodecMetadata]: - return ShardingCodecMetadata + def from_dict(cls, data: Dict[str, JSON]) -> Self: + parse_name(data["name"], "sharding_indexed") + return cls(**data["configuration"]) + + def to_dict(self) -> Dict[str, JSON]: + return { + "name": "sharding_indexed", + "configuration": { + "chunk_shape": list(self.chunk_shape), + "codecs": self.codecs.to_dict(), + "index_codecs": self.index_codecs.to_dict(), + "index_location": self.index_location, + }, + } def validate(self, array_metadata: ArrayMetadata) -> None: - assert len(self.configuration.chunk_shape) == array_metadata.ndim, ( + assert len(self.chunk_shape) == array_metadata.ndim, ( "The shard's `chunk_shape` and array's `shape` need to have the " + "same number of dimensions." ) assert isinstance( - array_metadata.chunk_grid, RegularChunkGridMetadata + array_metadata.chunk_grid, RegularChunkGrid ), "Sharding is only compatible with regular chunk grids." assert all( s % c == 0 for s, c in zip( - array_metadata.chunk_grid.configuration.chunk_shape, - self.configuration.chunk_shape, + array_metadata.chunk_grid.chunk_shape, + self.chunk_shape, ) ), ( "The array's `chunk_shape` needs to be divisible by the " @@ -279,7 +305,7 @@ async def decode( ) -> np.ndarray: # print("decode") shard_shape = shard_spec.shape - chunk_shape = self.configuration.chunk_shape + chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) indexer = BasicIndexer( @@ -311,7 +337,6 @@ async def decode( shard_spec, runtime_configuration, out, - runtime_configuration, ) for chunk_coords, chunk_selection, out_selection in indexer ], @@ -329,7 +354,7 @@ async def decode_partial( runtime_configuration: RuntimeConfiguration, ) -> Optional[np.ndarray]: shard_shape = shard_spec.shape - chunk_shape = self.configuration.chunk_shape + chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) indexer = BasicIndexer( @@ -380,7 +405,6 @@ async def decode_partial( shard_spec, runtime_configuration, out, - runtime_configuration, ) for chunk_coords, chunk_selection, out_selection in indexed_chunks ], @@ -403,9 +427,7 @@ async def _read_chunk( chunk_spec = self._get_chunk_spec(shard_spec) chunk_bytes = shard_dict.get(chunk_coords, None) if chunk_bytes is not None: - chunk_array = await self._codec_pipeline.decode( - chunk_bytes, chunk_spec, runtime_configuration - ) + chunk_array = await self.codecs.decode(chunk_bytes, chunk_spec, runtime_configuration) tmp = chunk_array[chunk_selection] out[out_selection] = tmp else: @@ -418,7 +440,7 @@ async def encode( runtime_configuration: RuntimeConfiguration, ) -> Optional[BytesLike]: shard_shape = shard_spec.shape - chunk_shape = self.configuration.chunk_shape + chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) indexer = list( @@ -449,9 +471,7 @@ async def _write_chunk( chunk_spec = self._get_chunk_spec(shard_spec) return ( chunk_coords, - await self._codec_pipeline.encode( - chunk_array, chunk_spec, runtime_configuration - ), + await self.codecs.encode(chunk_array, chunk_spec, runtime_configuration), ) return (chunk_coords, None) @@ -472,11 +492,7 @@ async def _write_chunk( if chunk_bytes is not None: shard_builder.append(chunk_coords, chunk_bytes) - return await shard_builder.finalize( - self.configuration.index_location, - self._encode_shard_index, - runtime_configuration=runtime_configuration, - ) + return await shard_builder.finalize(self.index_location, self._encode_shard_index) async def encode_partial( self, @@ -488,7 +504,7 @@ async def encode_partial( ) -> None: # print("encode_partial") shard_shape = shard_spec.shape - chunk_shape = self.configuration.chunk_shape + chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) chunk_spec = self._get_chunk_spec(shard_spec) @@ -512,7 +528,7 @@ async def _write_chunk( out_selection: SliceSelection, ) -> Tuple[ChunkCoords, Optional[BytesLike]]: chunk_array = None - if is_total_slice(chunk_selection, self.configuration.chunk_shape): + if is_total_slice(chunk_selection, self.chunk_shape): chunk_array = shard_array[out_selection] else: # handling writing partial chunks @@ -522,24 +538,20 @@ async def _write_chunk( # merge new value if chunk_bytes is None: chunk_array = np.empty( - self.configuration.chunk_shape, + self.chunk_shape, dtype=shard_spec.dtype, ) chunk_array.fill(shard_spec.fill_value) else: chunk_array = ( - await self._codec_pipeline.decode( - chunk_bytes, chunk_spec, runtime_configuration - ) + await self.codecs.decode(chunk_bytes, chunk_spec, runtime_configuration) ).copy() # make a writable copy chunk_array[chunk_selection] = shard_array[out_selection] if not np.array_equiv(chunk_array, shard_spec.fill_value): return ( chunk_coords, - await self._codec_pipeline.encode( - chunk_array, chunk_spec, runtime_configuration - ), + await self.codecs.encode(chunk_array, chunk_spec, runtime_configuration), ) else: return (chunk_coords, None) @@ -575,9 +587,8 @@ async def _write_chunk( else: await store_path.set( await shard_builder.finalize( - self.configuration.index_location, + self.index_location, self._encode_shard_index, - runtime_configuration=runtime_configuration, ) ) @@ -592,7 +603,7 @@ async def _decode_shard_index( self, index_bytes: BytesLike, chunks_per_shard: ChunkCoords ) -> _ShardIndex: return _ShardIndex( - await self._index_codec_pipeline.decode( + await self.index_codecs.decode( index_bytes, self._get_index_chunk_spec(chunks_per_shard), make_runtime_configuration("C"), @@ -600,7 +611,7 @@ async def _decode_shard_index( ) async def _encode_shard_index(self, index: _ShardIndex) -> BytesLike: - index_bytes = await self._index_codec_pipeline.encode( + index_bytes = await self.index_codecs.encode( index.offsets_and_lengths, self._get_index_chunk_spec(index.chunks_per_shard), make_runtime_configuration("C"), @@ -609,7 +620,7 @@ async def _encode_shard_index(self, index: _ShardIndex) -> BytesLike: return index_bytes def _shard_index_size(self, chunks_per_shard: ChunkCoords) -> int: - return self._index_codec_pipeline.compute_encoded_size( + return self.index_codecs.compute_encoded_size( 16 * product(chunks_per_shard), self._get_index_chunk_spec(chunks_per_shard) ) @@ -617,14 +628,14 @@ def _shard_index_size(self, chunks_per_shard: ChunkCoords) -> int: def _get_index_chunk_spec(self, chunks_per_shard: ChunkCoords) -> ArraySpec: return ArraySpec( shape=chunks_per_shard + (2,), - data=np.uint64, + dtype=np.dtype(" ArraySpec: return ArraySpec( - shape=self.configuration.chunk_shape, + shape=self.chunk_shape, dtype=shard_spec.dtype, fill_value=shard_spec.fill_value, ) @@ -635,25 +646,15 @@ def _get_chunks_per_shard(self, shard_spec: ArraySpec) -> ChunkCoords: s // c for s, c in zip( shard_spec.shape, - self.configuration.chunk_shape, + self.chunk_shape, ) ) - @cached_property - def _index_codec_pipeline(self) -> CodecPipeline: - return CodecPipeline.create( - [get_codec_from_metadata(c) for c in self.configuration.index_codecs] - ) - - @cached_property - def _codec_pipeline(self) -> CodecPipeline: - 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 ) -> Optional[_ShardIndex]: shard_index_size = self._shard_index_size(chunks_per_shard) - if self.configuration.index_location == ShardingCodecIndexLocation.start: + if self.index_location == ShardingCodecIndexLocation.start: index_bytes = await store_path.get((0, shard_index_size)) else: index_bytes = await store_path.get((-shard_index_size, None)) @@ -683,12 +684,5 @@ def compute_encoded_size(self, input_byte_length: int, shard_spec: ArraySpec) -> chunks_per_shard = self._get_chunks_per_shard(shard_spec) return input_byte_length + self._shard_index_size(chunks_per_shard) - def to_dict(self) -> Dict[str, Any]: - return ShardingCodecMetadata(configuration=self.configuration).to_dict() - - @classmethod - def from_dict(cls, data: Dict[str, Any]): - return cls(configuration=data["configuration"]) - register_codec("sharding_indexed", ShardingCodec) diff --git a/src/zarr/v3/codecs/transpose.py b/src/zarr/v3/codecs/transpose.py index 9dc9753e1d..e00d70ec6f 100644 --- a/src/zarr/v3/codecs/transpose.py +++ b/src/zarr/v3/codecs/transpose.py @@ -1,14 +1,18 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, Iterable -from dataclasses import dataclass, field -from zarr.v3.abc.metadata import Metadata +from dataclasses import dataclass, replace -from zarr.v3.common import ArraySpec +from zarr.v3.common import JSON, ArraySpec, parse_name if TYPE_CHECKING: - from zarr.v3.common import NamedConfig, RuntimeConfiguration - from typing import TYPE_CHECKING, Literal, Optional, Tuple, Type, Dict, Any + from zarr.v3.common import RuntimeConfiguration + from typing import ( + TYPE_CHECKING, + Optional, + Tuple, + ) + from typing_extensions import Self import numpy as np @@ -16,57 +20,54 @@ from zarr.v3.codecs.registry import register_codec -@dataclass(frozen=True) -class TransposeCodecConfigurationMetadata(Metadata): - order: Tuple[int, ...] +def parse_transpose_order(data: JSON) -> Tuple[int]: + if not isinstance(data, Iterable): + raise TypeError(f"Expected an iterable. Got {data} instead.") + if not all(isinstance(a, int) for a in data): + raise TypeError(f"Expected an iterable of integers. Got {data} instead.") + return tuple(data) @dataclass(frozen=True) -class TransposeCodecMetadata(Metadata): - configuration: TransposeCodecConfigurationMetadata - name: Literal["transpose"] = field(default="transpose", init=False) - +class TransposeCodec(ArrayArrayCodec): + is_fixed_size = True -@dataclass(frozen=True) -class TransposeCodec(ArrayArrayCodec, Metadata): order: Tuple[int, ...] - is_fixed_size = True - @classmethod - def from_metadata(cls, codec_metadata: NamedConfig) -> TransposeCodec: - assert isinstance(codec_metadata, TransposeCodecMetadata) - return cls(order=codec_metadata.configuration.order) + def __init__(self, *, order) -> None: + order_parsed = parse_transpose_order(order) - def evolve(self, *, ndim: int, **_kwargs) -> TransposeCodec: - # Compatibility with older version of ZEP1 - if self.order == "F": # type: ignore - order = tuple(ndim - x - 1 for x in range(ndim)) + object.__setattr__(self, "order", order_parsed) - elif self.order == "C": # type: ignore - order = tuple(range(ndim)) + @classmethod + def from_dict(cls, data: Dict[str, JSON]) -> Self: + parse_name(data["name"], "transpose") + return cls(**data["configuration"]) + + def to_dict(self) -> Dict[str, JSON]: + return {"name": "transpose", "configuration": {"order": list(self.order)}} - else: - assert len(self.order) == ndim, ( + def evolve(self, array_spec: ArraySpec) -> Self: + if len(self.order) != array_spec.ndim: + raise ValueError( "The `order` tuple needs have as many entries as " + f"there are dimensions in the array. Got: {self.order}" ) - assert len(self.order) == len(set(self.order)), ( + if len(self.order) != len(set(self.order)): + raise ValueError( "There must not be duplicates in the `order` tuple. " + f"Got: {self.order}" ) - assert all(0 <= x < ndim for x in self.order), ( + if not all(0 <= x < array_spec.ndim for x in self.order): + raise ValueError( "All entries in the `order` tuple must be between 0 and " + f"the number of dimensions in the array. Got: {self.order}" ) - order = tuple(self.order) + order = tuple(self.order) if order != self.order: - return evolve(self, order=order) + return replace(self, order=order) return self - @classmethod - def get_metadata_class(cls) -> Type[TransposeCodecMetadata]: - return TransposeCodecMetadata - def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: from zarr.v3.common import ArraySpec @@ -100,12 +101,5 @@ async def encode( def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int: return input_byte_length - def to_dict(self) -> Dict[str, Any]: - return TransposeCodecMetadata(configuration=self.configuration).to_dict() - - @classmethod - def from_dict(cls, data: Dict[str, Any]): - return cls(configuration=data["configuration"]) - register_codec("transpose", TransposeCodec) diff --git a/src/zarr/v3/codecs/zstd.py b/src/zarr/v3/codecs/zstd.py index 2beb18767d..47fac495fb 100644 --- a/src/zarr/v3/codecs/zstd.py +++ b/src/zarr/v3/codecs/zstd.py @@ -1,23 +1,22 @@ from __future__ import annotations from typing import TYPE_CHECKING -from dataclasses import dataclass, field +from dataclasses import dataclass -from zarr.v3.abc.metadata import Metadata from zstandard import ZstdCompressor, ZstdDecompressor from zarr.v3.abc.codec import BytesBytesCodec from zarr.v3.codecs.registry import register_codec -from zarr.v3.common import to_thread, ArraySpec +from zarr.v3.common import parse_name, to_thread if TYPE_CHECKING: - from zarr.v3.metadata import RuntimeConfiguration - from typing import Any, Literal, Dict, Type, Optional + from typing import Dict, Optional from typing_extensions import Self - from zarr.v3.common import BytesLike, NamedConfig + from zarr.v3.metadata import RuntimeConfiguration + from zarr.v3.common import BytesLike, JSON, ArraySpec -def parse_zstd_level(data: Any) -> int: +def parse_zstd_level(data: JSON) -> int: if isinstance(data, int): if data >= 23: msg = f"Value must be less than or equal to 22. Got {data} instead." @@ -27,59 +26,37 @@ def parse_zstd_level(data: Any) -> int: raise TypeError(msg) -def parse_checksum(data: Any) -> bool: +def parse_checksum(data: JSON) -> bool: if isinstance(data, bool): return data msg = f"Expected bool, got {type(data)}" raise TypeError(msg) -def parse_name(data: Any) -> Literal["zstd"]: - if data == "zstd": - return data - msg = f"Expected 'zstd', got {data}" - raise ValueError(msg) - - @dataclass(frozen=True) -class ZstdCodecConfigurationMetadata(Metadata): +class ZstdCodec(BytesBytesCodec): + is_fixed_size = True + level: int = 0 checksum: bool = False - def __init__(self, level: int, checksum: bool): + def __init__(self, *, level, checksum) -> None: level_parsed = parse_zstd_level(level) checksum_parsed = parse_checksum(checksum) + object.__setattr__(self, "level", level_parsed) object.__setattr__(self, "checksum", checksum_parsed) - -@dataclass(frozen=True) -class ZstdCodecMetadata(Metadata): - configuration: ZstdCodecConfigurationMetadata - name: Literal["zstd"] = field(default="zstd", init=False) - @classmethod - def from_dict(cls, data: Dict[str, Any]) -> Self: - _ = parse_name(data.pop("name")) - return cls(**data) + def from_dict(cls, data: Dict[str, JSON]) -> Self: + parse_name(data["name"], "zstd") + return cls(**data["configuration"]) - -@dataclass(frozen=True) -class ZstdCodec(BytesBytesCodec, Metadata): - configuration: ZstdCodecConfigurationMetadata - is_fixed_size = True - - @classmethod - def from_metadata(cls, codec_metadata: NamedConfig) -> ZstdCodec: - assert isinstance(codec_metadata, ZstdCodecMetadata) - return cls(configuration=codec_metadata.configuration) - - @classmethod - def get_metadata_class(cls) -> Type[ZstdCodecMetadata]: - return ZstdCodecMetadata + def to_dict(self) -> Dict[str, JSON]: + return {"name": "zstd", "configuration": {"level": self.level, "checksum": self.checksum}} def _compress(self, data: bytes) -> bytes: - ctx = ZstdCompressor(level=self.metadata.level, write_checksum=self.metadata.checksum) + ctx = ZstdCompressor(level=self.level, write_checksum=self.checksum) return ctx.compress(data) def _decompress(self, data: bytes) -> bytes: @@ -105,12 +82,5 @@ async def encode( def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int: raise NotImplementedError - def to_dict(self) -> Dict[str, Any]: - return ZstdCodecMetadata(configuration=self.configuration).to_dict() - - @classmethod - def from_dict(cls, data: Dict[str, Any]): - return cls(configuration=data["configuration"]) - register_codec("zstd", ZstdCodec) diff --git a/src/zarr/v3/common.py b/src/zarr/v3/common.py index deac54ee7a..a3c218ee9b 100644 --- a/src/zarr/v3/common.py +++ b/src/zarr/v3/common.py @@ -4,17 +4,21 @@ from asyncio import AbstractEventLoop import contextvars from dataclasses import dataclass +from enum import Enum import functools from typing import ( Any, Awaitable, Callable, Dict, + Iterable, + Iterator, List, Literal, Optional, Protocol, Tuple, + Type, TypeVar, Union, ) @@ -30,6 +34,7 @@ ChunkCoords = Tuple[int, ...] SliceSelection = Tuple[slice, ...] Selection = Union[slice, SliceSelection] +JSON = Union[str, None, int, float, Dict[str, "JSON"], List["JSON"]] def product(tup: ChunkCoords) -> int: @@ -63,6 +68,23 @@ async def to_thread(func, /, *args, **kwargs): return await loop.run_in_executor(None, func_call) +def enum_names(enum: Enum) -> Iterator[str]: + for item in enum: + yield item.name + + +E = TypeVar("E", bound=Enum) + + +def parse_enum(data: JSON, cls: Type[E]) -> E: + if isinstance(data, cls): + return data + if data in enum_names(cls): + return cls(data) + msg = f"Value must be one of {repr(list(enum_names(cls)))}, got {data} instead." + raise ValueError(msg) + + class NamedConfig(Protocol): @property def name(self) -> str: @@ -83,18 +105,15 @@ class RuntimeConfiguration: @dataclass(frozen=True) class ArraySpec: shape: ChunkCoords - chunk_shape: ChunkCoords dtype: np.dtype fill_value: Any - def __init__(self, shape, chunk_shape, dtype, fill_value): + def __init__(self, shape, dtype, fill_value): shape_parsed = parse_shapelike(shape) dtype_parsed = parse_dtype(dtype) - chunk_shape_parsed = parse_shapelike(chunk_shape) fill_value_parsed = parse_fill_value(fill_value) object.__setattr__(self, "shape", shape_parsed) - object.__setattr__(self, "chunk_shape", chunk_shape_parsed) object.__setattr__(self, "dtype", dtype_parsed) object.__setattr__(self, "fill_value", fill_value_parsed) @@ -103,9 +122,18 @@ def ndim(self) -> int: return len(self.shape) +def parse_name(data: JSON, expected: str) -> str: + if data == expected: + return data + raise ValueError(f"Expected '{expected}' chunk, got {data} instead.") + + def parse_shapelike(data: Any) -> Tuple[int, ...]: - # todo: handle empty tuple - return tuple(int(x) for x in data) + if not isinstance(data, Iterable): + raise TypeError(f"Expected an iterable. Got {data} instead.") + if not all(isinstance(a, int) for a in data): + raise TypeError(f"Expected an iterable of integers. Got {data} instead.") + return tuple(data) def parse_dtype(data: Any) -> np.dtype: diff --git a/src/zarr/v3/group.py b/src/zarr/v3/group.py index 846a9abe2b..20f2c53772 100644 --- a/src/zarr/v3/group.py +++ b/src/zarr/v3/group.py @@ -50,10 +50,7 @@ def to_bytes(self) -> Dict[str, bytes]: ZATTRS_JSON: json.dumps(self.attributes).encode(), } - def __init__( - self, - attributes: Dict[str, Any] = None, - zarr_format: Literal[2, 3] = 3): + def __init__(self, attributes: Dict[str, Any] = None, zarr_format: Literal[2, 3] = 3): attributes_parsed = parse_attributes(attributes) zarr_format_parsed = parse_zarr_format(zarr_format) diff --git a/src/zarr/v3/metadata.py b/src/zarr/v3/metadata.py index 90d0d6d7c6..3b6c1346bb 100644 --- a/src/zarr/v3/metadata.py +++ b/src/zarr/v3/metadata.py @@ -1,19 +1,22 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Literal, Union +from enum import Enum +from typing import TYPE_CHECKING from dataclasses import dataclass, field - import json +import numpy as np + +from zarr.v3.chunk_grids import ChunkGrid, RegularChunkGrid +from zarr.v3.chunk_key_encodings import ChunkKeyEncoding + if TYPE_CHECKING: - from typing import Any, Dict, Iterable, List, Optional, Tuple + from typing import Any, Dict, Literal, Union, Iterable, List, Optional, Tuple from typing_extensions import Self + from zarr.v3.codecs.pipeline import CodecPipeline -from warnings import warn -import numpy as np -from zarr.v3.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec +from zarr.v3.abc.codec import Codec from zarr.v3.abc.metadata import Metadata -from zarr.v3.codecs.registry import get_codec_class from zarr.v3.common import ( ArraySpec, @@ -25,8 +28,6 @@ parse_shapelike, ) -ShardingCodecIndexLocation = Literal["start", "end"] - def runtime_configuration( order: Literal["C", "F"], concurrency: Optional[int] = None @@ -37,7 +38,7 @@ def runtime_configuration( # For type checking _bool = bool -""" + class DataType(Enum): bool = "bool" int8 = "int8" @@ -89,105 +90,33 @@ def to_numpy_shortname(self) -> str: } return data_type_to_numpy[self] - -dtype_to_data_type = { - "|b1": "bool", - "bool": "bool", - "|i1": "int8", - " Self: - return cls( - configuration=RegularChunkGridConfigurationMetadata.from_dict(data["configuration"]) - ) - - -@dataclass(frozen=True) -class DefaultChunkKeyEncodingConfigurationMetadata(Metadata): - separator: Literal[".", "/"] = "/" - - -@dataclass(frozen=True) -class DefaultChunkKeyEncodingMetadata(Metadata): - configuration: DefaultChunkKeyEncodingConfigurationMetadata = ( - DefaultChunkKeyEncodingConfigurationMetadata() - ) - name: Literal["default"] = field(default="default", init=False) - - def decode_chunk_key(self, chunk_key: str) -> ChunkCoords: - if chunk_key == "c": - return () - return tuple(map(int, chunk_key[1:].split(self.configuration.separator))) - - def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str: - return self.configuration.separator.join(map(str, ("c",) + chunk_coords)) - - @classmethod - def from_dict(cls, data: Dict[str, Any]): - return cls( - configuration=DefaultChunkKeyEncodingConfigurationMetadata.from_dict( - data["configuration"] - ) - ) - - -@dataclass(frozen=True) -class V2ChunkKeyEncodingConfigurationMetadata(Metadata): - separator: Literal[".", "/"] = "." - - -@dataclass(frozen=True) -class V2ChunkKeyEncodingMetadata(Metadata): - configuration: V2ChunkKeyEncodingConfigurationMetadata = ( - V2ChunkKeyEncodingConfigurationMetadata() - ) - name: Literal["v2"] = field(init=False, default="v2") - - def decode_chunk_key(self, chunk_key: str) -> ChunkCoords: - return tuple(map(int, chunk_key.split(self.configuration.separator))) - - def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str: - chunk_identifier = self.configuration.separator.join(map(str, chunk_coords)) - return "0" if chunk_identifier == "" else chunk_identifier - @classmethod - def from_dict(cls, data: Dict[str, Any]): - return cls(configuration=data["configuration"]) - - -ChunkKeyEncodingMetadata = Union[DefaultChunkKeyEncodingMetadata, V2ChunkKeyEncodingMetadata] + def from_dtype(cls, dtype: np.dtype) -> Self: + dtype_to_data_type = { + "|b1": "bool", + "bool": "bool", + "|i1": "int8", + " int: def get_chunk_spec(self, _chunk_coords: ChunkCoords) -> ArraySpec: assert isinstance( - self.chunk_grid, RegularChunkGridMetadata + self.chunk_grid, RegularChunkGrid ), "Currently, only regular chunk grid is supported" return ArraySpec( - shape=self.chunk_grid.configuration.chunk_shape, + shape=self.chunk_grid.chunk_shape, dtype=self.dtype, fill_value=self.fill_value, - chunk_shape=self.chunk_grid.configuration.chunk_shape, ) def to_bytes(self) -> bytes: def _json_convert(o): if isinstance(o, np.dtype): return str(o) + if isinstance(o, Enum): + return o.name # this serializes numcodecs compressors # todo: implement to_dict for codecs elif hasattr(o, "get_config"): @@ -347,33 +273,6 @@ def from_dict(cls, data: Dict[str, Any]) -> ArrayV2Metadata: return cls(**data) -def parse_chunk_grid(data: Any) -> RegularChunkGridMetadata: - if isinstance(data, dict): - return RegularChunkGridMetadata.from_dict(data) - if isinstance(data, RegularChunkGridMetadata): - return data - msg = f"Expected dict or instance of RegularChunkGridMetadata, got {type(data)}" - raise TypeError(msg) - - -def parse_chunk_key_encoding(data: Any) -> ChunkKeyEncodingMetadata: - if isinstance(data, dict): - # todo: consider handling keyerrors gracefully here - if data["name"] == "v2": - return V2ChunkKeyEncodingMetadata.from_dict(data) - elif data["name"] == "default": - return DefaultChunkKeyEncodingMetadata.from_dict(data) - msg = f'Invalid `name` attribute. Got {data["name"]}, expected one of ("v2", "default")' - raise ValueError(msg) - if isinstance(data, (V2ChunkKeyEncodingMetadata, DefaultChunkKeyEncodingMetadata)): - return data - msg = ( - f"Expected a dict or an instance of V2ChunkKeyEncodingMetadata " - f"or an instance of DefaultChunkKeyEncodingMetadata, got input with type={type(data)}" - ) - raise TypeError(msg) - - def parse_dimension_names(data: Any) -> Tuple[str, ...] | None: if data is None: return data @@ -419,11 +318,11 @@ def parse_compressor(data: Any) -> Codec: def parse_v3_metadata(data: ArrayMetadata) -> ArrayMetadata: - if (l_chunks := len(data.chunk_grid.configuration.chunk_shape)) != (l_shape := len(data.shape)): + if (l_chunks := len(data.chunk_grid.chunk_shape)) != (l_shape := len(data.shape)): msg = ( - f"The `shape` and `chunk_grid.configuration.chunk_shape` attributes " + f"The `shape` and `chunk_grid.chunk_shape` attributes " "must have the same length. " - f"`chunk_grid.configuration.chunk_shape` has length {l_chunks}, " + f"`chunk_grid.chunk_shape` has length {l_chunks}, " f"but `shape` has length {l_shape}" ) raise ValueError(msg) @@ -446,77 +345,7 @@ def parse_v2_metadata(data: ArrayV2Metadata) -> ArrayV2Metadata: return data -def parse_codecs( - data: Iterable[NamedConfig], - array_metadata: ArraySpec, -) -> List[Codec]: - out: List[Codec] = [] - - for codec_metadata in data or []: - codec_cls = get_codec_class(codec_metadata.name) - codec = codec_cls.from_metadata(codec_metadata) - out.append(codec) - array_metadata = codec.resolve_metadata(array_metadata) - out = _validate_codecs(out, array_metadata) - return out - - -def _validate_codecs(codecs: List[Codec], array_metadata: ArraySpec) -> None: - from zarr.v3.codecs.sharding import ShardingCodec - - assert any( - isinstance(codec, ArrayBytesCodec) for codec in codecs - ), "Exactly one array-to-bytes codec is required." - - prev_codec: Optional[Codec] = None - for codec in codecs: - if prev_codec is not None: - assert not isinstance(codec, ArrayBytesCodec) or not isinstance( - prev_codec, ArrayBytesCodec - ), ( - f"ArrayBytesCodec '{type(codec)}' cannot follow after " - + f"ArrayBytesCodec '{type(prev_codec)}' because exactly " - + "1 ArrayBytesCodec is allowed." - ) - assert not isinstance(codec, ArrayBytesCodec) or not isinstance( - prev_codec, BytesBytesCodec - ), ( - f"ArrayBytesCodec '{type(codec)}' cannot follow after " - + f"BytesBytesCodec '{type(prev_codec)}'." - ) - assert not isinstance(codec, ArrayArrayCodec) or not isinstance( - prev_codec, ArrayBytesCodec - ), ( - f"ArrayArrayCodec '{type(codec)}' cannot follow after " - + f"ArrayBytesCodec '{type(prev_codec)}'." - ) - assert not isinstance(codec, ArrayArrayCodec) or not isinstance( - prev_codec, BytesBytesCodec - ), ( - f"ArrayArrayCodec '{type(codec)}' cannot follow after " - + f"BytesBytesCodec '{type(prev_codec)}'." - ) - - if isinstance(codec, ShardingCodec): - assert len(codec.configuration.chunk_shape) == len(array_metadata.shape), ( - "The shard's `chunk_shape` and array's `shape` need to have the " - + "same number of dimensions." - ) - assert all( - s % c == 0 - for s, c in zip( - array_metadata.chunk_shape, - codec.configuration.chunk_shape, - ) - ), ( - "The array's `chunk_shape` needs to be divisible by the " - + "shard's inner `chunk_shape`." - ) - prev_codec = codec - - 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 codecs +def parse_codecs(data: Iterable[NamedConfig]) -> CodecPipeline: + from zarr.v3.codecs.pipeline import CodecPipeline + + return CodecPipeline.from_dict(data) diff --git a/tests/test_codecs_v3.py b/tests/test_codecs_v3.py index a14a63a28c..6b32924847 100644 --- a/tests/test_codecs_v3.py +++ b/tests/test_codecs_v3.py @@ -2,17 +2,25 @@ from dataclasses import dataclass import json -from typing import Iterator, List, Literal, Optional +from typing import Iterator, List, Literal, Optional, Tuple import numpy as np import pytest import zarr -from zarr.v3 import codecs from zarr.v3.array import Array, AsyncArray from zarr.v3.common import NamedConfig, Selection from zarr.v3.indexing import morton_order_iter -from zarr.v3.metadata import ShardingCodecIndexLocation, runtime_configuration +from zarr.v3.codecs import ( + ShardingCodec, + ShardingCodecIndexLocation, + BloscCodec, + BytesCodec, + GzipCodec, + TransposeCodec, + ZstdCodec, +) +from zarr.v3.metadata import runtime_configuration from zarr.v3.abc.store import Store from zarr.v3.store import MemoryStore, StorePath @@ -48,6 +56,13 @@ def sample_data() -> np.ndarray: return np.arange(0, 128 * 128 * 128, dtype="uint16").reshape((128, 128, 128), order="F") +def order_from_dim(order: Literal["F", "C"], ndim: int) -> Tuple[int, ...]: + if order == "F": + return tuple(ndim - x - 1 for x in range(ndim)) + else: + return tuple(range(ndim)) + + @pytest.mark.parametrize("index_location", ["start", "end"]) def test_sharding( store: Store, sample_data: np.ndarray, index_location: ShardingCodecIndexLocation @@ -59,12 +74,12 @@ def test_sharding( dtype=sample_data.dtype, fill_value=0, codecs=[ - codecs.sharding_codec( - (32, 32, 32), - [ - codecs.transpose_codec("F", sample_data.ndim), - codecs.bytes_codec(), - codecs.blosc_codec(typesize=sample_data.dtype.itemsize, cname="lz4"), + ShardingCodec( + chunk_shape=(32, 32, 32), + codecs=[ + TransposeCodec(order=order_from_dim("F", sample_data.ndim)), + BytesCodec(), + BloscCodec(typesize=sample_data.dtype.itemsize, cname="lz4"), ], index_location=index_location, ) @@ -89,12 +104,12 @@ def test_sharding_partial( dtype=sample_data.dtype, fill_value=0, codecs=[ - codecs.sharding_codec( - (32, 32, 32), - [ - codecs.transpose_codec("F", sample_data.ndim), - codecs.bytes_codec(), - codecs.blosc_codec(typesize=sample_data.dtype.itemsize, cname="lz4"), + ShardingCodec( + chunk_shape=(32, 32, 32), + codecs=[ + TransposeCodec(order=order_from_dim("F", sample_data.ndim)), + BytesCodec(), + BloscCodec(typesize=sample_data.dtype.itemsize, cname="lz4"), ], index_location=index_location, ) @@ -122,12 +137,12 @@ def test_sharding_partial_read( dtype=sample_data.dtype, fill_value=1, codecs=[ - codecs.sharding_codec( - (32, 32, 32), - [ - codecs.transpose_codec("F", sample_data.ndim), - codecs.bytes_codec(), - codecs.blosc_codec(typesize=sample_data.dtype.itemsize, cname="lz4"), + ShardingCodec( + chunk_shape=(32, 32, 32), + codecs=[ + TransposeCodec(order=order_from_dim("F", sample_data.ndim)), + BytesCodec(), + BloscCodec(typesize=sample_data.dtype.itemsize, cname="lz4"), ], index_location=index_location, ) @@ -151,12 +166,12 @@ def test_sharding_partial_overwrite( dtype=data.dtype, fill_value=1, codecs=[ - codecs.sharding_codec( - (32, 32, 32), - [ - codecs.transpose_codec("F", data.ndim), - codecs.bytes_codec(), - codecs.blosc_codec(typesize=data.dtype.itemsize, cname="lz4"), + ShardingCodec( + chunk_shape=(32, 32, 32), + codecs=[ + TransposeCodec(order=order_from_dim("F", data.ndim)), + BytesCodec(), + BloscCodec(typesize=data.dtype.itemsize, cname="lz4"), ], index_location=index_location, ) @@ -195,9 +210,11 @@ def test_nested_sharding( dtype=sample_data.dtype, fill_value=0, codecs=[ - codecs.sharding_codec( - (32, 32, 32), - [codecs.sharding_codec((16, 16, 16), index_location=inner_index_location)], + ShardingCodec( + chunk_shape=(32, 32, 32), + codecs=[ + ShardingCodec(chunk_shape=(16, 16, 16), index_location=inner_index_location) + ], index_location=outer_index_location, ) ], @@ -228,13 +245,13 @@ async def test_order( codecs_: List[NamedConfig] = ( [ - codecs.sharding_codec( - (16, 8), - codecs=[codecs.transpose_codec(store_order, data.ndim), codecs.bytes_codec()], + ShardingCodec( + chunk_shape=(16, 8), + codecs=[TransposeCodec(order=order_from_dim(store_order, data.ndim)), BytesCodec()], ) ] if with_sharding - else [codecs.transpose_codec(store_order, data.ndim), codecs.bytes_codec()] + else [TransposeCodec(order=order_from_dim(store_order, data.ndim)), BytesCodec()] ) a = await AsyncArray.create( @@ -294,7 +311,7 @@ def test_order_implicit( data = np.arange(0, 256, dtype="uint16").reshape((16, 16), order=input_order) codecs_: Optional[List[NamedConfig]] = ( - [codecs.sharding_codec((8, 8))] if with_sharding else None + [ShardingCodec(chunk_shape=(8, 8))] if with_sharding else None ) a = Array.create( @@ -340,13 +357,13 @@ async def test_transpose( codecs_: List[NamedConfig] = ( [ - codecs.sharding_codec( - (1, 16, 8), - codecs=[codecs.transpose_codec((2, 1, 0)), codecs.bytes_codec()], + ShardingCodec( + chunk_shape=(1, 16, 8), + codecs=[TransposeCodec(order=(2, 1, 0)), BytesCodec()], ) ] if with_sharding - else [codecs.transpose_codec((2, 1, 0)), codecs.bytes_codec()] + else [TransposeCodec(order=(2, 1, 0)), BytesCodec()] ) a = await AsyncArray.create( @@ -398,7 +415,7 @@ def test_transpose_invalid( data = np.arange(0, 256, dtype="uint16").reshape((1, 32, 8)) for order in [(1, 0), (3, 2, 1), (3, 3, 1)]: - with pytest.raises(AssertionError): + with pytest.raises(ValueError): Array.create( store / "transpose_invalid", shape=data.shape, @@ -406,7 +423,7 @@ def test_transpose_invalid( dtype=data.dtype, fill_value=0, chunk_key_encoding=("v2", "."), - codecs=[codecs.transpose_codec(order), codecs.bytes_codec()], + codecs=[TransposeCodec(order=order), BytesCodec()], ) @@ -430,12 +447,12 @@ def test_open_sharding(store: Store): dtype="int32", fill_value=0, codecs=[ - codecs.sharding_codec( - (8, 8), - [ - codecs.transpose_codec("F", 2), - codecs.bytes_codec(), - codecs.blosc_codec(typesize=4), + ShardingCodec( + chunk_shape=(8, 8), + codecs=[ + TransposeCodec(order=order_from_dim("F", 2)), + BytesCodec(), + BloscCodec(typesize=4), ], ) ], @@ -573,11 +590,11 @@ def test_write_partial_sharded_chunks(store: Store): dtype=data.dtype, fill_value=1, codecs=[ - codecs.sharding_codec( + ShardingCodec( chunk_shape=(10, 10), codecs=[ - codecs.bytes_codec(), - codecs.blosc_codec(typesize=data.dtype.itemsize), + BytesCodec(), + BloscCodec(typesize=data.dtype.itemsize), ], ) ], @@ -611,7 +628,7 @@ async def test_delete_empty_sharded_chunks(store: Store): chunk_shape=(8, 16), dtype="uint16", fill_value=1, - codecs=[codecs.sharding_codec(chunk_shape=(8, 8))], + codecs=[ShardingCodec(chunk_shape=(8, 8))], ) await _AsyncArrayProxy(a)[:, :].set(np.zeros((16, 16))) await _AsyncArrayProxy(a)[8:, :].set(np.ones((8, 16))) @@ -672,7 +689,7 @@ async def test_zarr_compat_F(store: Store): dtype=data.dtype, chunk_key_encoding=("v2", "."), fill_value=1, - codecs=[codecs.transpose_codec("F", data.ndim), codecs.bytes_codec()], + codecs=[TransposeCodec(order=order_from_dim("F", data.ndim)), BytesCodec()], ) z2 = zarr.create( @@ -736,7 +753,7 @@ def test_gzip(store: Store): chunk_shape=(16, 16), dtype=data.dtype, fill_value=0, - codecs=[codecs.bytes_codec(), codecs.gzip_codec()], + codecs=[BytesCodec(), GzipCodec()], ) a[:, :] = data @@ -753,7 +770,7 @@ def test_zstd(store: Store, checksum: bool): chunk_shape=(16, 16), dtype=data.dtype, fill_value=0, - codecs=[codecs.bytes_codec(), codecs.zstd_codec(level=0, checksum=checksum)], + codecs=[BytesCodec(), ZstdCodec(level=0, checksum=checksum)], ) a[:, :] = data @@ -772,7 +789,7 @@ async def test_endian(store: Store, endian: Literal["big", "little"]): dtype=data.dtype, fill_value=0, chunk_key_encoding=("v2", "."), - codecs=[codecs.bytes_codec(endian)], + codecs=[BytesCodec(endian=endian)], ) await _AsyncArrayProxy(a)[:, :].set(data) @@ -808,7 +825,7 @@ async def test_endian_write( dtype="uint16", fill_value=0, chunk_key_encoding=("v2", "."), - codecs=[codecs.bytes_codec(dtype_store_endian)], + codecs=[BytesCodec(endian=dtype_store_endian)], ) await _AsyncArrayProxy(a)[:, :].set(data) @@ -845,12 +862,12 @@ def test_invalid_metadata(store: Store): dtype=np.dtype("uint8"), fill_value=0, codecs=[ - codecs.bytes_codec("big"), - codecs.transpose_codec("F", 2), + BytesCodec(endian="big"), + TransposeCodec(order=order_from_dim("F", 2)), ], ) - with pytest.raises(AssertionError): + with pytest.raises(TypeError): Array.create( store / "invalid_order", shape=(16, 16), @@ -858,8 +875,8 @@ def test_invalid_metadata(store: Store): dtype=np.dtype("uint8"), fill_value=0, codecs=[ - codecs.bytes_codec(), - codecs.transpose_codec("F"), + BytesCodec(), + TransposeCodec(order="F"), ], ) @@ -871,7 +888,7 @@ def test_invalid_metadata(store: Store): dtype=np.dtype("uint8"), fill_value=0, codecs=[ - codecs.transpose_codec("F", 2), + TransposeCodec(order=order_from_dim("F", 2)), ], ) @@ -883,7 +900,7 @@ def test_invalid_metadata(store: Store): dtype=np.dtype("uint8"), fill_value=0, codecs=[ - codecs.sharding_codec(chunk_shape=(8,)), + ShardingCodec(chunk_shape=(8,)), ], ) with pytest.raises(AssertionError): @@ -894,7 +911,7 @@ def test_invalid_metadata(store: Store): dtype=np.dtype("uint8"), fill_value=0, codecs=[ - codecs.sharding_codec(chunk_shape=(8, 7)), + ShardingCodec(chunk_shape=(8, 7)), ], ) @@ -906,8 +923,8 @@ def test_invalid_metadata(store: Store): dtype=np.dtype("uint8"), fill_value=0, codecs=[ - codecs.sharding_codec(chunk_shape=(8, 8)), - codecs.gzip_codec(), + ShardingCodec(chunk_shape=(8, 8)), + GzipCodec(), ], )