From 71196b627a282abc6fbc561bc555fe292183166e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 23 Sep 2026 23:58:25 +0200 Subject: [PATCH 1/2] refactor(array): share one immutable ArrayState between Array and AsyncArray Array and AsyncArray now hold the same `_StateCell`, a single mutable slot containing a frozen `ArrayState` (store path, metadata, runtime config, chunk grid, codec pipeline). Members that are derived from that state without IO are defined once on a shared `_ArrayBase`. `Array` does its own IO through the module-level functions instead of calling through its `AsyncArray`, which is kept only as the handle returned by `async_array`. resize, append and update_attributes are split into pure plan functions and a short IO step that swaps the cell's state, replacing in-place mutation of a frozen dataclass via `object.__setattr__` and of the attributes dict. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 2402 ++++++++++++++----------------- src/zarr/core/attributes.py | 3 +- src/zarr/metadata/migrate_v3.py | 2 +- tests/test_array_state.py | 162 +++ 4 files changed, 1213 insertions(+), 1356 deletions(-) create mode 100644 tests/test_array_state.py diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 5a8d6bf57e..04be4de4e2 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -11,6 +11,7 @@ from typing import ( TYPE_CHECKING, Any, + ClassVar, Literal, TypedDict, cast, @@ -351,482 +352,112 @@ async def _prepare_overwrite( @dataclass(frozen=True) -class AsyncArray[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: +class ArrayState[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: """ - An asynchronous array class representing a chunked array stored in a Zarr store. - - Parameters - ---------- - metadata : ArrayMetadata - The metadata of the array. - store_path : StorePath - The path to the Zarr store. - config : ArrayConfigLike, optional - The runtime configuration of the array, by default None. + The state of an open array: where it is stored, its metadata, and its runtime + configuration, together with the chunk grid and codec pipeline derived from them. - Attributes - ---------- - metadata : ArrayMetadata - The metadata of the array. - store_path : StorePath - The path to the Zarr store. - codec_pipeline : CodecPipeline - The codec pipeline used for encoding and decoding chunks. - config : ArrayConfig - The runtime configuration of the array. + `ArrayState` is immutable. An operation that changes the array, like a resize, + produces a new `ArrayState` instead of modifying this one. """ - metadata: T_ArrayMetadata store_path: StorePath - codec_pipeline: CodecPipeline = field(init=False) - _chunk_grid: ChunkGrid = field(init=False) + metadata: T_ArrayMetadata config: ArrayConfig + chunk_grid: ChunkGrid = field(compare=False, repr=False) + codec_pipeline: CodecPipeline = field(compare=False, repr=False) - @overload - def __init__( - self: AsyncArrayV2, - metadata: ArrayV2Metadata | ArrayV2MetadataDict, - store_path: StorePath, - config: ArrayConfigLike | None = None, - ) -> None: ... - - @overload - def __init__( - self: AsyncArrayV3, - metadata: ArrayV3Metadata | ArrayMetadataJSON_V3, - store_path: StorePath, - config: ArrayConfigLike | None = None, - ) -> None: ... - - def __init__( - self, - metadata: ArrayMetadata | ArrayMetadataDict, - store_path: StorePath, - config: ArrayConfigLike | None = None, - ) -> None: - metadata_parsed = parse_array_metadata(metadata) - config_parsed = parse_array_config(config) - - object.__setattr__(self, "metadata", metadata_parsed) - object.__setattr__(self, "store_path", store_path) - object.__setattr__(self, "config", config_parsed) - object.__setattr__(self, "_chunk_grid", ChunkGrid.from_metadata(metadata_parsed)) - object.__setattr__( - self, - "codec_pipeline", - create_codec_pipeline(metadata=metadata_parsed, store=store_path.store), + @classmethod + def from_parts( + cls, store_path: StorePath, metadata: T_ArrayMetadata, config: ArrayConfig + ) -> Self: + return cls( + store_path=store_path, + metadata=metadata, + config=config, + chunk_grid=ChunkGrid.from_metadata(metadata), + codec_pipeline=create_codec_pipeline(metadata=metadata, store=store_path.store), ) - @classmethod - async def _create( - cls, - store: StoreLike, - *, - # v2 and v3 - shape: ShapeLike, - dtype: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], - zarr_format: ZarrFormat = 3, - fill_value: Any | None = DEFAULT_FILL_VALUE, - attributes: dict[str, JSON] | None = None, - # v3 only - chunk_shape: ShapeLike | None = None, - chunk_key_encoding: ( - ChunkKeyEncodingLike - | tuple[Literal["default"], Literal[".", "/"]] - | tuple[Literal["v2"], Literal[".", "/"]] - | None - ) = None, - codecs: Iterable[Codec | dict[str, JSON]] | None = None, - dimension_names: DimensionNamesLike = None, - # v2 only - chunks: ShapeLike | None = None, - dimension_separator: Literal[".", "/"] | None = None, - order: MemoryOrder | None = None, - filters: Iterable[dict[str, JSON] | Numcodec] | None = None, - compressor: CompressorLike = "auto", - # runtime - overwrite: bool = False, - data: npt.ArrayLike | None = None, - config: ArrayConfigLike | None = None, - ) -> AnyAsyncArray: - """Method to create a new asynchronous array instance. - Deprecated in favor of [`zarr.api.asynchronous.create_array`][]. + def with_metadata(self, metadata: T_ArrayMetadata) -> Self: """ + Return a copy of this state with new metadata. - dtype_parsed = parse_dtype(dtype, zarr_format=zarr_format) - store_path = await make_store_path(store) + The codec pipeline is kept: the changes made this way (shape, attributes) do not + change the codecs, and rebuilding the pipeline would repeat its advisory warnings. + """ + return replace(self, metadata=metadata, chunk_grid=ChunkGrid.from_metadata(metadata)) - shape = parse_shapelike(shape) - if chunks is not None and chunk_shape is not None: - raise ValueError("Only one of chunk_shape or chunks can be provided.") +class _StateCell[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: + """ + The one mutable slot holding an array's current `ArrayState`. - # Unify the v2 (chunks) and v3 (chunk_shape) parameter names - _raw_chunks = chunks if chunks is not None else chunk_shape + Every object that views the same array holds the same cell, so a change made + through any of them is seen by all of them. + """ - config_parsed = parse_array_config(config) + __slots__ = ("state",) - result: AnyAsyncArray - if zarr_format == 3: - if dimension_separator is not None: - raise ValueError( - "dimension_separator cannot be used for arrays with zarr_format 3. Use chunk_key_encoding instead." - ) - if filters is not None: - raise ValueError( - "filters cannot be used for arrays with zarr_format 3. Use array-to-array codecs instead." - ) - if compressor != "auto": - raise ValueError( - "compressor cannot be used for arrays with zarr_format 3. Use bytes-to-bytes codecs instead." - ) + state: ArrayState[T_ArrayMetadata] - if order is not None: - _warn_order_kwarg() + def __init__(self, state: ArrayState[T_ArrayMetadata]) -> None: + self.state = state - item_size = 1 - if isinstance(dtype_parsed, HasItemSize): - item_size = dtype_parsed.item_size - if _raw_chunks is None: - outer_chunks = guess_chunks(shape, item_size) - else: - outer_chunks = normalize_chunks_nd(_raw_chunks, shape) - chunk_grid = create_chunk_grid_metadata(outer_chunks) - result = await cls._create_v3( - store_path, - shape=shape, - dtype=dtype_parsed, - fill_value=fill_value, - chunk_key_encoding=chunk_key_encoding, - codecs=codecs, - dimension_names=dimension_names, - attributes=attributes, - overwrite=overwrite, - config=config_parsed, - chunk_grid=chunk_grid, - ) - elif zarr_format == 2: - if codecs is not None: - raise ValueError( - "codecs cannot be used for arrays with zarr_format 2. Use filters and compressor instead." - ) - if chunk_key_encoding is not None: - raise ValueError( - "chunk_key_encoding cannot be used for arrays with zarr_format 2. Use dimension_separator instead." - ) - if dimension_names is not None: - raise ValueError("dimension_names cannot be used for arrays with zarr_format 2.") - if _is_rectilinear_chunks(_raw_chunks): - raise ValueError("Zarr format 2 does not support rectilinear chunk grids.") - item_size = 1 - if isinstance(dtype_parsed, HasItemSize): - item_size = dtype_parsed.item_size - _raw = chunks or chunk_shape - if _raw is None: - outer_chunks = guess_chunks(shape, item_size) - else: - outer_chunks = normalize_chunks_nd(_raw, shape) - _chunks = outer_chunks.chunk_shape +class _ArrayBase[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: + """ + Members shared by `Array` and `AsyncArray`: everything derived from the array's + state without doing IO, plus the private coroutines both classes use for IO. + """ - if order is None: - order_parsed = config_parsed.order - else: - order_parsed = order - config_parsed = replace(config_parsed, order=order) + _cell: _StateCell[T_ArrayMetadata] - result = await cls._create_v2( - store_path, - shape=shape, - dtype=dtype_parsed, - chunks=_chunks, - dimension_separator=dimension_separator, - fill_value=fill_value, - order=order_parsed, - config=config_parsed, - filters=filters, - compressor=compressor, - attributes=attributes, - overwrite=overwrite, - ) - else: - raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + @property + def metadata(self) -> T_ArrayMetadata: + return self._cell.state.metadata - if data is not None: - # insert user-provided data - await result.setitem(..., data) + @property + def store_path(self) -> StorePath: + return self._cell.state.store_path - return result + @property + def config(self) -> ArrayConfig: + """ + The runtime configuration for this array. This is a read-only property. To modify the + runtime configuration, use `with_config` to create a new array with the modified + configuration. - @staticmethod - def _create_metadata_v3( - shape: ShapeLike, - dtype: ZDType[TBaseDType, TBaseScalar], - chunk_grid: ChunkGridMetadata, - fill_value: Any | None = DEFAULT_FILL_VALUE, - chunk_key_encoding: ChunkKeyEncodingLike | None = None, - codecs: Iterable[Codec | dict[str, JSON]] | None = None, - dimension_names: DimensionNamesLike = None, - attributes: dict[str, JSON] | None = None, - ) -> ArrayV3Metadata: - """Create an instance of ArrayV3Metadata.""" - filters: tuple[ArrayArrayCodec, ...] - compressors: tuple[BytesBytesCodec, ...] + Returns + ------- + An `ArrayConfig` object that defines the runtime configuration for the array. + """ + return self._cell.state.config - shape = parse_shapelike(shape) - if codecs is None: - # no data types have default filters - filters = () - serializer = default_serializer_v3(dtype) - compressors = default_compressors_v3(dtype) + @property + def codec_pipeline(self) -> CodecPipeline: + return self._cell.state.codec_pipeline - codecs_parsed = (*filters, serializer, *compressors) - else: - codecs_parsed = tuple(codecs) + @property + def _chunk_grid(self) -> ChunkGrid: + """The chunk grid for this array, bound to the array's shape.""" + return self._cell.state.chunk_grid - chunk_key_encoding_parsed: ChunkKeyEncodingLike - if chunk_key_encoding is None: - chunk_key_encoding_parsed = {"name": "default", "separator": "/"} - else: - chunk_key_encoding_parsed = chunk_key_encoding + @property + def fill_value(self) -> Any: + return self.metadata.fill_value - if fill_value is DEFAULT_FILL_VALUE or fill_value is None: - # Use dtype's default scalar for the DEFAULT_FILL_VALUE sentinel - # For v3, None is converted to DEFAULT_FILL_VALUE behavior - fill_value_parsed = dtype.default_scalar() - else: - fill_value_parsed = fill_value + def __eq__(self, other: object) -> bool: + if not isinstance(other, _ArrayBase) or type(other) is not type(self): + return NotImplemented + return self._cell.state == other._cell.state - return ArrayV3Metadata( - shape=shape, - data_type=dtype, - chunk_grid=chunk_grid, - chunk_key_encoding=chunk_key_encoding_parsed, - fill_value=fill_value_parsed, - codecs=codecs_parsed, # type: ignore[arg-type] - dimension_names=tuple(dimension_names) if dimension_names else None, - attributes=attributes or {}, - ) - - @classmethod - async def _create_v3( - cls, - store_path: StorePath, - *, - shape: ShapeLike, - dtype: ZDType[TBaseDType, TBaseScalar], - chunk_grid: ChunkGridMetadata, - config: ArrayConfig, - fill_value: Any | None = DEFAULT_FILL_VALUE, - chunk_key_encoding: ( - ChunkKeyEncodingLike - | tuple[Literal["default"], Literal[".", "/"]] - | tuple[Literal["v2"], Literal[".", "/"]] - | None - ) = None, - codecs: Iterable[Codec | dict[str, JSON]] | None = None, - dimension_names: DimensionNamesLike = None, - attributes: dict[str, JSON] | None = None, - overwrite: bool = False, - ) -> AsyncArrayV3: - await _prepare_overwrite(store_path, zarr_format=3, overwrite=overwrite) - - if isinstance(chunk_key_encoding, tuple): - chunk_key_encoding = ( - V2ChunkKeyEncoding(separator=chunk_key_encoding[1]) - if chunk_key_encoding[0] == "v2" - else DefaultChunkKeyEncoding(separator=chunk_key_encoding[1]) - ) - - metadata = cls._create_metadata_v3( - shape=shape, - dtype=dtype, - chunk_grid=chunk_grid, - fill_value=fill_value, - chunk_key_encoding=chunk_key_encoding, - codecs=codecs, - dimension_names=dimension_names, - attributes=attributes, - ) - - array = cls(metadata=metadata, store_path=store_path, config=config) - await array._save_metadata(metadata, ensure_parents=True) - return array - - @staticmethod - def _create_metadata_v2( - shape: tuple[int, ...], - dtype: ZDType[TBaseDType, TBaseScalar], - chunks: tuple[int, ...], - order: MemoryOrder, - dimension_separator: Literal[".", "/"] | None = None, - fill_value: Any | None = DEFAULT_FILL_VALUE, - filters: Iterable[dict[str, JSON] | Numcodec] | None = None, - compressor: CompressorLikev2 = None, - attributes: dict[str, JSON] | None = None, - ) -> ArrayV2Metadata: - if dimension_separator is None: - dimension_separator = "." - - # Handle the DEFAULT_FILL_VALUE sentinel - if fill_value is DEFAULT_FILL_VALUE: - fill_value_parsed: Any = dtype.default_scalar() - else: - # For v2, preserve None as-is (backward compatibility) - fill_value_parsed = fill_value - - return ArrayV2Metadata( - shape=shape, - dtype=dtype, - chunks=chunks, - order=order, - dimension_separator=dimension_separator, - fill_value=fill_value_parsed, - compressor=compressor, - filters=filters, - attributes=attributes, - ) - - @classmethod - async def _create_v2( - cls, - store_path: StorePath, - *, - shape: tuple[int, ...], - dtype: ZDType[TBaseDType, TBaseScalar], - chunks: tuple[int, ...], - order: MemoryOrder, - config: ArrayConfig, - dimension_separator: Literal[".", "/"] | None = None, - fill_value: Any | None = DEFAULT_FILL_VALUE, - filters: Iterable[dict[str, JSON] | Numcodec] | None = None, - compressor: CompressorLike = "auto", - attributes: dict[str, JSON] | None = None, - overwrite: bool = False, - ) -> AsyncArrayV2: - await _prepare_overwrite(store_path, zarr_format=2, overwrite=overwrite) - - compressor_parsed: CompressorLikev2 - if compressor == "auto": - compressor_parsed = default_compressor_v2(dtype) - elif isinstance(compressor, BytesBytesCodec): - raise ValueError( - "Cannot use a BytesBytesCodec as a compressor for zarr v2 arrays. " - "Use a numcodecs codec directly instead." - ) - else: - compressor_parsed = compressor - - if filters is None: - filters = default_filters_v2(dtype) - - metadata = cls._create_metadata_v2( - shape=shape, - dtype=dtype, - chunks=chunks, - order=order, - dimension_separator=dimension_separator, - fill_value=fill_value, - filters=filters, - compressor=compressor_parsed, - attributes=attributes, - ) - - array = cls(metadata=metadata, store_path=store_path, config=config) - await array._save_metadata(metadata, ensure_parents=True) - return array - - @classmethod - def from_dict( - cls, - store_path: StorePath, - data: dict[str, JSON], - ) -> AnyAsyncArray: - """ - Create a Zarr array from a dictionary, with support for both Zarr format 2 and 3 metadata. - - Parameters - ---------- - store_path : StorePath - The path within the store where the array should be created. - - data : dict - A dictionary representing the array data. This dictionary should include necessary metadata - for the array, such as shape, dtype, and other attributes. The format of the metadata - will determine whether a Zarr format 2 or 3 array is created. - - Returns - ------- - AsyncArrayV3 or AsyncArrayV2 - The created Zarr array, either using Zarr format 2 or 3 metadata based on the provided data. - - Raises - ------ - ValueError - If the dictionary data is invalid or incompatible with either Zarr format 2 or 3 array creation. - """ - metadata = parse_array_metadata(data) - return cls(metadata=metadata, store_path=store_path) - - @classmethod - async def open( - cls, - store: StoreLike, - zarr_format: ZarrFormat | None = 3, - ) -> AnyAsyncArray: - """ - Async method to open an existing Zarr array from a given store. - - Parameters - ---------- - store : StoreLike - The store containing the Zarr array. See the - [storage documentation in the user guide][user-guide-store-like] - for a description of all valid StoreLike values. - zarr_format : ZarrFormat | None, optional - The Zarr format version (default is 3). - - Returns - ------- - AsyncArray - The opened Zarr array. - - Examples - -------- - ```python - import asyncio - import zarr - from zarr.core.array import AsyncArray - - async def example(): - store = zarr.storage.MemoryStore() - # First create an array to open - await zarr.api.asynchronous.create_array( - store=store, shape=(100, 100), dtype="int32" - ) - # Now open it - async_arr = await AsyncArray.open(store) - return async_arr - - async_arr = asyncio.run(example()) - # - ``` - """ - store_path = await make_store_path(store) - metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format) - # TODO: remove this cast when we have better type hints - _metadata_dict = cast("ArrayMetadataJSON_V3", metadata_dict) - return cls(store_path=store_path, metadata=_metadata_dict) + __hash__: ClassVar[None] = None # type: ignore[assignment] @property def store(self) -> Store: return self.store_path.store - @property - @deprecated("Use AsyncArray.config instead.", category=ZarrDeprecationWarning) - def _config(self) -> ArrayConfig: - return self.config - @property def ndim(self) -> int: """Returns the number of dimensions in the Array. @@ -996,7 +627,7 @@ def serializer(self) -> ArrayBytesCodec | None: ) @property - @deprecated("Use AsyncArray.compressors instead.", category=ZarrDeprecationWarning) + @deprecated("Use compressors instead.", category=ZarrDeprecationWarning) def compressor(self) -> Numcodec | None: """ Compressor that is applied to each chunk of the array. @@ -1059,141 +690,831 @@ def order(self) -> MemoryOrder: return self.config.order @property - def attrs(self) -> dict[str, JSON]: - """Returns the attributes of the array. + def read_only(self) -> bool: + """Returns True if the array is read-only. + + Returns + ------- + bool + True if the array is read-only + """ + # Backwards compatibility for 2.x + return self.store_path.read_only + + @property + def path(self) -> str: + """Storage path. + + Returns + ------- + str + The path to the array in the Zarr store. + """ + return self.store_path.path + + @property + def name(self) -> str: + """Array name following h5py convention. + + Returns + ------- + str + The name of the array. + """ + # follow h5py convention: add leading slash + name = self.path + if not name.startswith("/"): + name = "/" + name + return name + + @property + def basename(self) -> str: + """Final component of name. + + Returns + ------- + str + The basename or final component of the array name. + """ + return self.name.split("/")[-1] + + @property + def cdata_shape(self) -> tuple[int, ...]: + """ + The number of chunks along each dimension. + + When sharding is used, this counts inner chunks (not shards) per dimension. + + Returns + ------- + tuple[int, ...] + The number of chunks along each dimension. + """ + return self._chunk_grid_shape + + @property + def _chunk_grid_shape(self) -> tuple[int, ...]: + """ + The number of chunks along each dimension. + + When sharding is used, this counts inner chunks (not shards) per dimension. + + Returns + ------- + tuple[int, ...] + The number of chunks along each dimension. + """ + if (sharding_codec := _sharding_codec(self.metadata)) is not None: + # When sharding, count inner chunks across the whole array + chunk_shape = sharding_codec.chunk_shape + return tuple(starmap(ceildiv, zip(self.shape, chunk_shape, strict=True))) + return self._chunk_grid.grid_shape + + @property + def _shard_grid_shape(self) -> tuple[int, ...]: + """ + The shape of the shard grid for this array. + + When no shards are present this will automatically fall back to the chunk grid. + + Returns + ------- + tuple[int, ...] + The shape of the shard grid for this array. + """ + # The stored chunk grid is the shard grid when sharding is used, the + # chunk grid otherwise. Works for regular and rectilinear grids alike. + return self._chunk_grid.grid_shape + + @property + def nchunks(self) -> int: + """ + The number of chunks in this array. + + Note that if a sharding codec is used, then the number of chunks may exceed the number of + stored objects supporting this array. + + Returns + ------- + int + The total number of chunks in the array. + """ + return product(self._chunk_grid_shape) + + @property + def _nshards(self) -> int: + """ + The number of shards in this array. + + If no shards are present this will fall back to giving the number of chunks + + Returns + ------- + int + The total number of shards or if absent, chunks in the array. + """ + return product(self._shard_grid_shape) + + def _iter_chunk_coords( + self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None + ) -> Iterator[tuple[int, ...]]: + """ + Create an iterator over the coordinates of chunks in chunk grid space. + + If the `origin` keyword is used, iteration will start at the chunk index specified by `origin`. + The default behavior is to start at the origin of the grid coordinate space. + If the `selection_shape` keyword is used, iteration will be bounded over a contiguous region + ranging from `[origin, origin selection_shape]`, where the upper bound is exclusive as + per python indexing conventions. + + Parameters + ---------- + origin : Sequence[int] | None, default=None + The origin of the selection relative to the array's chunk grid. + selection_shape : Sequence[int] | None, default=None + The shape of the selection in chunk grid coordinates. + + Yields + ------ + chunk_coords: tuple[int, ...] + The coordinates of each chunk in the selection. + """ + return _iter_chunk_coords( + array=self, + origin=origin, + selection_shape=selection_shape, + ) + + def _iter_shard_coords( + self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None + ) -> Iterator[tuple[int, ...]]: + """ + Create an iterator over the coordinates of shards in shard grid space. + + This will fall back to chunk grid space in case no shards are present. + + Note that + + If the `origin` keyword is used, iteration will start at the shard index specified by `origin`. + The default behavior is to start at the origin of the grid coordinate space. + If the `selection_shape` keyword is used, iteration will be bounded over a contiguous region + ranging from `[origin, origin selection_shape]`, where the upper bound is exclusive as + per python indexing conventions. + + Parameters + ---------- + origin : Sequence[int] | None, default=None + The origin of the selection relative to the array's shard grid. + selection_shape : Sequence[int] | None, default=None + The shape of the selection in shard grid coordinates. + + Yields + ------ + chunk_coords: tuple[int, ...] + The coordinates of each shard in the selection or chunk in case of no shard being present. + """ + return _iter_shard_coords( + array=self, + origin=origin, + selection_shape=selection_shape, + ) + + def _iter_shard_keys( + self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None + ) -> Iterator[str]: + """ + Iterate over the keys of the stored objects supporting this array. + + Although only stored objects, e.g. shards should have keys, in case no + shards are present this automatically falls back to chunks. + + Parameters + ---------- + origin : Sequence[int] | None, default=None + The origin of the selection relative to the array's chunk grid. + selection_shape : Sequence[int] | None, default=None + The shape of the selection in shard grid coordinates. + + Yields + ------ + key: str + The storage key of each shard in the selection or in case of no shard + present of each chunk although the latter case is technically incorrect. + """ + # Iterate over the coordinates of chunks in chunk grid space. + return _iter_shard_keys( + array=self, + origin=origin, + selection_shape=selection_shape, + ) + + def _iter_chunk_regions( + self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None + ) -> Iterator[tuple[slice, ...]]: + """ + Iterate over the regions spanned by each chunk. + + Parameters + ---------- + origin : Sequence[int] | None, default=None + The origin of the selection relative to the array's chunk grid. + selection_shape : Sequence[int] | None, default=None + The shape of the selection in chunk grid coordinates. + + Yields + ------ + region: tuple[slice, ...] + A tuple of slice objects representing the region spanned by each chunk in the selection. + """ + return _iter_chunk_regions( + array=self, + origin=origin, + selection_shape=selection_shape, + ) + + def _iter_shard_regions( + self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None + ) -> Iterator[tuple[slice, ...]]: + """ + Iterate over the regions spanned by each shard. + + This will automatically fall back to chunks if no shards are present. + + Parameters + ---------- + origin : Sequence[int] | None, default=None + The origin of the selection relative to the array's shard grid. + selection_shape : Sequence[int] | None, default=None + The shape of the selection in shard grid coordinates. + + Yields + ------ + region: tuple[slice, ...] + A tuple of slice objects representing the region spanned by each shard in the selection or chunk in the + absence of shards. + """ + return _iter_shard_regions(array=self, origin=origin, selection_shape=selection_shape) + + @property + def nbytes(self) -> int: + """ + The total number of bytes that can be stored in the chunks of this array. + + Notes + ----- + This value is calculated by multiplying the number of elements in the array and the size + of each element, the latter of which is determined by the dtype of the array. + For this reason, `nbytes` will likely be inaccurate for arrays with variable-length + dtypes. It is not possible to determine the size of an array with variable-length elements + from the shape and dtype alone. + """ + return self.size * self.dtype.itemsize + + def _info( + self, count_chunks_initialized: int | None = None, count_bytes_stored: int | None = None + ) -> Any: + rectilinear_grid = _stored_rectilinear_grid_or_none(self.metadata) + sharded = _sharding_codec(self.metadata) is not None + # `.chunks` (the inner chunk shape when sharded) is undefined only for a + # non-sharded rectilinear grid, which ArrayInfo renders as ""; + # `.shards` is undefined for a rectilinear shard grid, where the + # "" sentinel keeps the array rendered as sharded. + chunk_shape = self.chunks if (rectilinear_grid is None or sharded) else None + shard_shape: tuple[int, ...] | Literal[""] | None + if rectilinear_grid is None: + shard_shape = self.shards + elif sharded: + shard_shape = "" + else: + shard_shape = None + return ArrayInfo( + _zarr_format=self.metadata.zarr_format, + _data_type=self._zdtype, + _fill_value=self.metadata.fill_value, + _shape=self.shape, + _order=self.order, + _shard_shape=shard_shape, + _chunk_shape=chunk_shape, + _read_only=self.read_only, + _compressors=self.compressors, + _filters=self.filters, + _serializer=self.serializer, + _store_type=type(self.store_path.store).__name__, + _count_bytes=self.nbytes, + _count_bytes_stored=count_bytes_stored, + _count_chunks_initialized=count_chunks_initialized, + ) + + async def _get_selection( + self, + indexer: Indexer, + *, + prototype: BufferPrototype, + out: NDBuffer | None = None, + fields: Fields | None = None, + ) -> NDArrayLikeOrScalar: + state = self._cell.state + return await _get_selection( + state.store_path, + state.metadata, + state.codec_pipeline, + state.config, + state.chunk_grid, + indexer, + prototype=prototype, + out=out, + fields=fields, + ) + + async def _set_selection( + self, + indexer: Indexer, + value: npt.ArrayLike, + *, + prototype: BufferPrototype, + fields: Fields | None = None, + ) -> None: + state = self._cell.state + return await _set_selection( + state.store_path, + state.metadata, + state.codec_pipeline, + state.config, + state.chunk_grid, + indexer, + value, + prototype=prototype, + fields=fields, + ) + + +class AsyncArray[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)](_ArrayBase[T_ArrayMetadata]): + """ + An asynchronous array class representing a chunked array stored in a Zarr store. + + Parameters + ---------- + metadata : ArrayMetadata + The metadata of the array. + store_path : StorePath + The path to the Zarr store. + config : ArrayConfigLike, optional + The runtime configuration of the array, by default None. + + Attributes + ---------- + metadata : ArrayMetadata + The metadata of the array. + store_path : StorePath + The path to the Zarr store. + codec_pipeline : CodecPipeline + The codec pipeline used for encoding and decoding chunks. + config : ArrayConfig + The runtime configuration of the array. + """ + + @overload + def __init__( + self: AsyncArrayV2, + metadata: ArrayV2Metadata | ArrayV2MetadataDict, + store_path: StorePath, + config: ArrayConfigLike | None = None, + ) -> None: ... + + @overload + def __init__( + self: AsyncArrayV3, + metadata: ArrayV3Metadata | ArrayMetadataJSON_V3, + store_path: StorePath, + config: ArrayConfigLike | None = None, + ) -> None: ... + + def __init__( + self, + metadata: ArrayMetadata | ArrayMetadataDict, + store_path: StorePath, + config: ArrayConfigLike | None = None, + ) -> None: + metadata_parsed = parse_array_metadata(metadata) + config_parsed = parse_array_config(config) + self._cell = _StateCell(ArrayState.from_parts(store_path, metadata_parsed, config_parsed)) # type: ignore[arg-type] + + @classmethod + async def _create( + cls, + store: StoreLike, + *, + # v2 and v3 + shape: ShapeLike, + dtype: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], + zarr_format: ZarrFormat = 3, + fill_value: Any | None = DEFAULT_FILL_VALUE, + attributes: dict[str, JSON] | None = None, + # v3 only + chunk_shape: ShapeLike | None = None, + chunk_key_encoding: ( + ChunkKeyEncodingLike + | tuple[Literal["default"], Literal[".", "/"]] + | tuple[Literal["v2"], Literal[".", "/"]] + | None + ) = None, + codecs: Iterable[Codec | dict[str, JSON]] | None = None, + dimension_names: DimensionNamesLike = None, + # v2 only + chunks: ShapeLike | None = None, + dimension_separator: Literal[".", "/"] | None = None, + order: MemoryOrder | None = None, + filters: Iterable[dict[str, JSON] | Numcodec] | None = None, + compressor: CompressorLike = "auto", + # runtime + overwrite: bool = False, + data: npt.ArrayLike | None = None, + config: ArrayConfigLike | None = None, + ) -> AnyAsyncArray: + """Method to create a new asynchronous array instance. + Deprecated in favor of [`zarr.api.asynchronous.create_array`][]. + """ + + dtype_parsed = parse_dtype(dtype, zarr_format=zarr_format) + store_path = await make_store_path(store) + + shape = parse_shapelike(shape) + + if chunks is not None and chunk_shape is not None: + raise ValueError("Only one of chunk_shape or chunks can be provided.") + + # Unify the v2 (chunks) and v3 (chunk_shape) parameter names + _raw_chunks = chunks if chunks is not None else chunk_shape + + config_parsed = parse_array_config(config) + + result: AnyAsyncArray + if zarr_format == 3: + if dimension_separator is not None: + raise ValueError( + "dimension_separator cannot be used for arrays with zarr_format 3. Use chunk_key_encoding instead." + ) + if filters is not None: + raise ValueError( + "filters cannot be used for arrays with zarr_format 3. Use array-to-array codecs instead." + ) + if compressor != "auto": + raise ValueError( + "compressor cannot be used for arrays with zarr_format 3. Use bytes-to-bytes codecs instead." + ) + + if order is not None: + _warn_order_kwarg() + + item_size = 1 + if isinstance(dtype_parsed, HasItemSize): + item_size = dtype_parsed.item_size + if _raw_chunks is None: + outer_chunks = guess_chunks(shape, item_size) + else: + outer_chunks = normalize_chunks_nd(_raw_chunks, shape) + chunk_grid = create_chunk_grid_metadata(outer_chunks) + result = await cls._create_v3( + store_path, + shape=shape, + dtype=dtype_parsed, + fill_value=fill_value, + chunk_key_encoding=chunk_key_encoding, + codecs=codecs, + dimension_names=dimension_names, + attributes=attributes, + overwrite=overwrite, + config=config_parsed, + chunk_grid=chunk_grid, + ) + elif zarr_format == 2: + if codecs is not None: + raise ValueError( + "codecs cannot be used for arrays with zarr_format 2. Use filters and compressor instead." + ) + if chunk_key_encoding is not None: + raise ValueError( + "chunk_key_encoding cannot be used for arrays with zarr_format 2. Use dimension_separator instead." + ) + if dimension_names is not None: + raise ValueError("dimension_names cannot be used for arrays with zarr_format 2.") + if _is_rectilinear_chunks(_raw_chunks): + raise ValueError("Zarr format 2 does not support rectilinear chunk grids.") + + item_size = 1 + if isinstance(dtype_parsed, HasItemSize): + item_size = dtype_parsed.item_size + _raw = chunks or chunk_shape + if _raw is None: + outer_chunks = guess_chunks(shape, item_size) + else: + outer_chunks = normalize_chunks_nd(_raw, shape) + _chunks = outer_chunks.chunk_shape + + if order is None: + order_parsed = config_parsed.order + else: + order_parsed = order + config_parsed = replace(config_parsed, order=order) + + result = await cls._create_v2( + store_path, + shape=shape, + dtype=dtype_parsed, + chunks=_chunks, + dimension_separator=dimension_separator, + fill_value=fill_value, + order=order_parsed, + config=config_parsed, + filters=filters, + compressor=compressor, + attributes=attributes, + overwrite=overwrite, + ) + else: + raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover + + if data is not None: + # insert user-provided data + await result.setitem(..., data) + + return result + + @staticmethod + def _create_metadata_v3( + shape: ShapeLike, + dtype: ZDType[TBaseDType, TBaseScalar], + chunk_grid: ChunkGridMetadata, + fill_value: Any | None = DEFAULT_FILL_VALUE, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + codecs: Iterable[Codec | dict[str, JSON]] | None = None, + dimension_names: DimensionNamesLike = None, + attributes: dict[str, JSON] | None = None, + ) -> ArrayV3Metadata: + """Create an instance of ArrayV3Metadata.""" + filters: tuple[ArrayArrayCodec, ...] + compressors: tuple[BytesBytesCodec, ...] + + shape = parse_shapelike(shape) + if codecs is None: + # no data types have default filters + filters = () + serializer = default_serializer_v3(dtype) + compressors = default_compressors_v3(dtype) + + codecs_parsed = (*filters, serializer, *compressors) + else: + codecs_parsed = tuple(codecs) + + chunk_key_encoding_parsed: ChunkKeyEncodingLike + if chunk_key_encoding is None: + chunk_key_encoding_parsed = {"name": "default", "separator": "/"} + else: + chunk_key_encoding_parsed = chunk_key_encoding + + if fill_value is DEFAULT_FILL_VALUE or fill_value is None: + # Use dtype's default scalar for the DEFAULT_FILL_VALUE sentinel + # For v3, None is converted to DEFAULT_FILL_VALUE behavior + fill_value_parsed = dtype.default_scalar() + else: + fill_value_parsed = fill_value + + return ArrayV3Metadata( + shape=shape, + data_type=dtype, + chunk_grid=chunk_grid, + chunk_key_encoding=chunk_key_encoding_parsed, + fill_value=fill_value_parsed, + codecs=codecs_parsed, # type: ignore[arg-type] + dimension_names=tuple(dimension_names) if dimension_names else None, + attributes=attributes or {}, + ) + + @classmethod + async def _create_v3( + cls, + store_path: StorePath, + *, + shape: ShapeLike, + dtype: ZDType[TBaseDType, TBaseScalar], + chunk_grid: ChunkGridMetadata, + config: ArrayConfig, + fill_value: Any | None = DEFAULT_FILL_VALUE, + chunk_key_encoding: ( + ChunkKeyEncodingLike + | tuple[Literal["default"], Literal[".", "/"]] + | tuple[Literal["v2"], Literal[".", "/"]] + | None + ) = None, + codecs: Iterable[Codec | dict[str, JSON]] | None = None, + dimension_names: DimensionNamesLike = None, + attributes: dict[str, JSON] | None = None, + overwrite: bool = False, + ) -> AsyncArrayV3: + await _prepare_overwrite(store_path, zarr_format=3, overwrite=overwrite) - Returns - ------- - dict - Attributes of the array - """ - return self.metadata.attributes + if isinstance(chunk_key_encoding, tuple): + chunk_key_encoding = ( + V2ChunkKeyEncoding(separator=chunk_key_encoding[1]) + if chunk_key_encoding[0] == "v2" + else DefaultChunkKeyEncoding(separator=chunk_key_encoding[1]) + ) - @property - def read_only(self) -> bool: - """Returns True if the array is read-only. + metadata = cls._create_metadata_v3( + shape=shape, + dtype=dtype, + chunk_grid=chunk_grid, + fill_value=fill_value, + chunk_key_encoding=chunk_key_encoding, + codecs=codecs, + dimension_names=dimension_names, + attributes=attributes, + ) - Returns - ------- - bool - True if the array is read-only - """ - # Backwards compatibility for 2.x - return self.store_path.read_only + array = cls(metadata=metadata, store_path=store_path, config=config) + await array._save_metadata(metadata, ensure_parents=True) + return array - @property - def path(self) -> str: - """Storage path. + @staticmethod + def _create_metadata_v2( + shape: tuple[int, ...], + dtype: ZDType[TBaseDType, TBaseScalar], + chunks: tuple[int, ...], + order: MemoryOrder, + dimension_separator: Literal[".", "/"] | None = None, + fill_value: Any | None = DEFAULT_FILL_VALUE, + filters: Iterable[dict[str, JSON] | Numcodec] | None = None, + compressor: CompressorLikev2 = None, + attributes: dict[str, JSON] | None = None, + ) -> ArrayV2Metadata: + if dimension_separator is None: + dimension_separator = "." - Returns - ------- - str - The path to the array in the Zarr store. - """ - return self.store_path.path + # Handle the DEFAULT_FILL_VALUE sentinel + if fill_value is DEFAULT_FILL_VALUE: + fill_value_parsed: Any = dtype.default_scalar() + else: + # For v2, preserve None as-is (backward compatibility) + fill_value_parsed = fill_value - @property - def name(self) -> str: - """Array name following h5py convention. + return ArrayV2Metadata( + shape=shape, + dtype=dtype, + chunks=chunks, + order=order, + dimension_separator=dimension_separator, + fill_value=fill_value_parsed, + compressor=compressor, + filters=filters, + attributes=attributes, + ) - Returns - ------- - str - The name of the array. - """ - # follow h5py convention: add leading slash - name = self.path - if not name.startswith("/"): - name = "/" + name - return name + @classmethod + async def _create_v2( + cls, + store_path: StorePath, + *, + shape: tuple[int, ...], + dtype: ZDType[TBaseDType, TBaseScalar], + chunks: tuple[int, ...], + order: MemoryOrder, + config: ArrayConfig, + dimension_separator: Literal[".", "/"] | None = None, + fill_value: Any | None = DEFAULT_FILL_VALUE, + filters: Iterable[dict[str, JSON] | Numcodec] | None = None, + compressor: CompressorLike = "auto", + attributes: dict[str, JSON] | None = None, + overwrite: bool = False, + ) -> AsyncArrayV2: + await _prepare_overwrite(store_path, zarr_format=2, overwrite=overwrite) - @property - def basename(self) -> str: - """Final component of name. + compressor_parsed: CompressorLikev2 + if compressor == "auto": + compressor_parsed = default_compressor_v2(dtype) + elif isinstance(compressor, BytesBytesCodec): + raise ValueError( + "Cannot use a BytesBytesCodec as a compressor for zarr v2 arrays. " + "Use a numcodecs codec directly instead." + ) + else: + compressor_parsed = compressor - Returns - ------- - str - The basename or final component of the array name. - """ - return self.name.split("/")[-1] + if filters is None: + filters = default_filters_v2(dtype) - @property - def cdata_shape(self) -> tuple[int, ...]: - """ - The number of chunks along each dimension. + metadata = cls._create_metadata_v2( + shape=shape, + dtype=dtype, + chunks=chunks, + order=order, + dimension_separator=dimension_separator, + fill_value=fill_value, + filters=filters, + compressor=compressor_parsed, + attributes=attributes, + ) - When sharding is used, this counts inner chunks (not shards) per dimension. + array = cls(metadata=metadata, store_path=store_path, config=config) + await array._save_metadata(metadata, ensure_parents=True) + return array - Returns - ------- - tuple[int, ...] - The number of chunks along each dimension. + @classmethod + def from_dict( + cls, + store_path: StorePath, + data: dict[str, JSON], + ) -> AnyAsyncArray: """ - return self._chunk_grid_shape + Create a Zarr array from a dictionary, with support for both Zarr format 2 and 3 metadata. - @property - def _chunk_grid_shape(self) -> tuple[int, ...]: - """ - The number of chunks along each dimension. + Parameters + ---------- + store_path : StorePath + The path within the store where the array should be created. - When sharding is used, this counts inner chunks (not shards) per dimension. + data : dict + A dictionary representing the array data. This dictionary should include necessary metadata + for the array, such as shape, dtype, and other attributes. The format of the metadata + will determine whether a Zarr format 2 or 3 array is created. Returns ------- - tuple[int, ...] - The number of chunks along each dimension. + AsyncArrayV3 or AsyncArrayV2 + The created Zarr array, either using Zarr format 2 or 3 metadata based on the provided data. + + Raises + ------ + ValueError + If the dictionary data is invalid or incompatible with either Zarr format 2 or 3 array creation. """ - if (sharding_codec := _sharding_codec(self.metadata)) is not None: - # When sharding, count inner chunks across the whole array - chunk_shape = sharding_codec.chunk_shape - return tuple(starmap(ceildiv, zip(self.shape, chunk_shape, strict=True))) - return self._chunk_grid.grid_shape + metadata = parse_array_metadata(data) + return cls(metadata=metadata, store_path=store_path) - @property - def _shard_grid_shape(self) -> tuple[int, ...]: + @classmethod + async def open( + cls, + store: StoreLike, + zarr_format: ZarrFormat | None = 3, + ) -> AnyAsyncArray: """ - The shape of the shard grid for this array. + Async method to open an existing Zarr array from a given store. - When no shards are present this will automatically fall back to the chunk grid. + Parameters + ---------- + store : StoreLike + The store containing the Zarr array. See the + [storage documentation in the user guide][user-guide-store-like] + for a description of all valid StoreLike values. + zarr_format : ZarrFormat | None, optional + The Zarr format version (default is 3). Returns ------- - tuple[int, ...] - The shape of the shard grid for this array. - """ - # The stored chunk grid is the shard grid when sharding is used, the - # chunk grid otherwise. Works for regular and rectilinear grids alike. - return self._chunk_grid.grid_shape + AsyncArray + The opened Zarr array. - @property - def nchunks(self) -> int: - """ - The number of chunks in this array. + Examples + -------- + ```python + import asyncio + import zarr + from zarr.core.array import AsyncArray - Note that if a sharding codec is used, then the number of chunks may exceed the number of - stored objects supporting this array. + async def example(): + store = zarr.storage.MemoryStore() + # First create an array to open + await zarr.api.asynchronous.create_array( + store=store, shape=(100, 100), dtype="int32" + ) + # Now open it + async_arr = await AsyncArray.open(store) + return async_arr - Returns - ------- - int - The total number of chunks in the array. + async_arr = asyncio.run(example()) + # + ``` """ - return product(self._chunk_grid_shape) + store_path = await make_store_path(store) + metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format) + # TODO: remove this cast when we have better type hints + _metadata_dict = cast("ArrayMetadataJSON_V3", metadata_dict) + return cls(store_path=store_path, metadata=_metadata_dict) @property - def _nshards(self) -> int: - """ - The number of shards in this array. + @deprecated("Use AsyncArray.config instead.", category=ZarrDeprecationWarning) + def _config(self) -> ArrayConfig: + return self.config - If no shards are present this will fall back to giving the number of chunks + @property + def attrs(self) -> dict[str, JSON]: + """Returns the attributes of the array. Returns ------- - int - The total number of shards or if absent, chunks in the array. + dict + Attributes of the array """ - return product(self._shard_grid_shape) + return self.metadata.attributes @overload def with_config(self: AsyncArrayV2, config: ArrayConfigLike) -> AsyncArrayV2: ... @@ -1309,181 +1630,6 @@ async def example(): async def nbytes_stored(self) -> int: return await _nbytes_stored(self.store_path) - def _iter_chunk_coords( - self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[tuple[int, ...]]: - """ - Create an iterator over the coordinates of chunks in chunk grid space. - - If the `origin` keyword is used, iteration will start at the chunk index specified by `origin`. - The default behavior is to start at the origin of the grid coordinate space. - If the `selection_shape` keyword is used, iteration will be bounded over a contiguous region - ranging from `[origin, origin selection_shape]`, where the upper bound is exclusive as - per python indexing conventions. - - Parameters - ---------- - origin : Sequence[int] | None, default=None - The origin of the selection relative to the array's chunk grid. - selection_shape : Sequence[int] | None, default=None - The shape of the selection in chunk grid coordinates. - - Yields - ------ - chunk_coords: tuple[int, ...] - The coordinates of each chunk in the selection. - """ - return _iter_chunk_coords( - array=self, - origin=origin, - selection_shape=selection_shape, - ) - - def _iter_shard_coords( - self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[tuple[int, ...]]: - """ - Create an iterator over the coordinates of shards in shard grid space. - - This will fall back to chunk grid space in case no shards are present. - - Note that - - If the `origin` keyword is used, iteration will start at the shard index specified by `origin`. - The default behavior is to start at the origin of the grid coordinate space. - If the `selection_shape` keyword is used, iteration will be bounded over a contiguous region - ranging from `[origin, origin selection_shape]`, where the upper bound is exclusive as - per python indexing conventions. - - Parameters - ---------- - origin : Sequence[int] | None, default=None - The origin of the selection relative to the array's shard grid. - selection_shape : Sequence[int] | None, default=None - The shape of the selection in shard grid coordinates. - - Yields - ------ - chunk_coords: tuple[int, ...] - The coordinates of each shard in the selection or chunk in case of no shard being present. - """ - return _iter_shard_coords( - array=self, - origin=origin, - selection_shape=selection_shape, - ) - - def _iter_shard_keys( - self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[str]: - """ - Iterate over the keys of the stored objects supporting this array. - - Although only stored objects, e.g. shards should have keys, in case no - shards are present this automatically falls back to chunks. - - Parameters - ---------- - origin : Sequence[int] | None, default=None - The origin of the selection relative to the array's chunk grid. - selection_shape : Sequence[int] | None, default=None - The shape of the selection in shard grid coordinates. - - Yields - ------ - key: str - The storage key of each shard in the selection or in case of no shard - present of each chunk although the latter case is technically incorrect. - """ - # Iterate over the coordinates of chunks in chunk grid space. - return _iter_shard_keys( - array=self, - origin=origin, - selection_shape=selection_shape, - ) - - def _iter_chunk_regions( - self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[tuple[slice, ...]]: - """ - Iterate over the regions spanned by each chunk. - - Parameters - ---------- - origin : Sequence[int] | None, default=None - The origin of the selection relative to the array's chunk grid. - selection_shape : Sequence[int] | None, default=None - The shape of the selection in chunk grid coordinates. - - Yields - ------ - region: tuple[slice, ...] - A tuple of slice objects representing the region spanned by each chunk in the selection. - """ - return _iter_chunk_regions( - array=self, - origin=origin, - selection_shape=selection_shape, - ) - - def _iter_shard_regions( - self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[tuple[slice, ...]]: - """ - Iterate over the regions spanned by each shard. - - This will automatically fall back to chunks if no shards are present. - - Parameters - ---------- - origin : Sequence[int] | None, default=None - The origin of the selection relative to the array's shard grid. - selection_shape : Sequence[int] | None, default=None - The shape of the selection in shard grid coordinates. - - Yields - ------ - region: tuple[slice, ...] - A tuple of slice objects representing the region spanned by each shard in the selection or chunk in the - absence of shards. - """ - return _iter_shard_regions(array=self, origin=origin, selection_shape=selection_shape) - - @property - def nbytes(self) -> int: - """ - The total number of bytes that can be stored in the chunks of this array. - - Notes - ----- - This value is calculated by multiplying the number of elements in the array and the size - of each element, the latter of which is determined by the dtype of the array. - For this reason, `nbytes` will likely be inaccurate for arrays with variable-length - dtypes. It is not possible to determine the size of an array with variable-length elements - from the shape and dtype alone. - """ - return self.size * self.dtype.itemsize - - async def _get_selection( - self, - indexer: Indexer, - *, - prototype: BufferPrototype, - out: NDBuffer | None = None, - fields: Fields | None = None, - ) -> NDArrayLikeOrScalar: - return await _get_selection( - self.store_path, - self.metadata, - self.codec_pipeline, - self.config, - self._chunk_grid, - indexer, - prototype=prototype, - out=out, - fields=fields, - ) - async def getitem( self, selection: BasicSelection, @@ -1615,26 +1761,6 @@ async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = F """ await save_metadata(self.store_path, metadata, ensure_parents=ensure_parents) - async def _set_selection( - self, - indexer: Indexer, - value: npt.ArrayLike, - *, - prototype: BufferPrototype, - fields: Fields | None = None, - ) -> None: - return await _set_selection( - self.store_path, - self.metadata, - self.codec_pipeline, - self.config, - self._chunk_grid, - indexer, - value, - prototype=prototype, - fields=fields, - ) - async def setitem( self, selection: BasicSelection, @@ -1832,85 +1958,37 @@ async def info_complete(self) -> Any: ------- ArrayInfo - Related - ------- - [zarr.AsyncArray.info][] - A property giving just the statically known information about an array. - """ - return await _info_complete(self) - - def _info( - self, count_chunks_initialized: int | None = None, count_bytes_stored: int | None = None - ) -> Any: - rectilinear_grid = _stored_rectilinear_grid_or_none(self.metadata) - sharded = _sharding_codec(self.metadata) is not None - # `.chunks` (the inner chunk shape when sharded) is undefined only for a - # non-sharded rectilinear grid, which ArrayInfo renders as ""; - # `.shards` is undefined for a rectilinear shard grid, where the - # "" sentinel keeps the array rendered as sharded. - chunk_shape = self.chunks if (rectilinear_grid is None or sharded) else None - shard_shape: tuple[int, ...] | Literal[""] | None - if rectilinear_grid is None: - shard_shape = self.shards - elif sharded: - shard_shape = "" - else: - shard_shape = None - return ArrayInfo( - _zarr_format=self.metadata.zarr_format, - _data_type=self._zdtype, - _fill_value=self.metadata.fill_value, - _shape=self.shape, - _order=self.order, - _shard_shape=shard_shape, - _chunk_shape=chunk_shape, - _read_only=self.read_only, - _compressors=self.compressors, - _filters=self.filters, - _serializer=self.serializer, - _store_type=type(self.store_path.store).__name__, - _count_bytes=self.nbytes, - _count_bytes_stored=count_bytes_stored, - _count_chunks_initialized=count_chunks_initialized, - ) + Related + ------- + [zarr.AsyncArray.info][] - A property giving just the statically known information about an array. + """ + return await _info_complete(self) -# TODO: Array can be a frozen data class again once property setters (e.g. shape) are removed -@dataclass(frozen=False) -class Array[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: +class Array[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)](_ArrayBase[T_ArrayMetadata]): """ A Zarr array. """ _async_array: AsyncArray[T_ArrayMetadata] + def __init__(self, async_array: AsyncArray[T_ArrayMetadata]) -> None: + # `Array` does its own IO against the shared state cell. The `AsyncArray` is kept + # only as the handle returned by `async_array`; it holds the same cell, so a + # change made through either object is seen by both. + self._cell = async_array._cell + self._async_array = async_array + @property def async_array(self) -> AsyncArray[T_ArrayMetadata]: """An asynchronous version of the current array. Useful for batching requests. Returns ------- - An asynchronous array whose metadata + store matches that of this synchronous array. + An asynchronous array that shares its state with this synchronous array. """ return self._async_array - @property - def config(self) -> ArrayConfig: - """ - The runtime configuration for this array. This is a read-only property. To modify the - runtime configuration, use `Array.with_config` to create a new `Array` with the modified - configuration. - - Returns - ------- - An `ArrayConfig` object that defines the runtime configuration for the array. - """ - return self.async_array.config - - @property - def _chunk_grid(self) -> ChunkGrid: - """The chunk grid for this array, bound to the array's shape.""" - return self.async_array._chunk_grid - @classmethod def _create( cls, @@ -1941,7 +2019,7 @@ def _create( # runtime overwrite: bool = False, config: ArrayConfigLike | None = None, - ) -> Self: + ) -> Array[Any]: """Creates a new Array instance from an initialized store. Deprecated in favor of [`zarr.create_array`][]. """ @@ -1973,7 +2051,7 @@ def from_dict( cls, store_path: StorePath, data: dict[str, JSON], - ) -> Self: + ) -> Array[Any]: """ Create a Zarr array from a dictionary. @@ -1994,322 +2072,63 @@ def from_dict( Raises ------ ValueError - If the dictionary data is invalid or missing required fields for array creation. - """ - async_array = AsyncArray.from_dict(store_path=store_path, data=data) - return cls(async_array) - - @classmethod - def open( - cls, - store: StoreLike, - ) -> Self: - """Opens an existing Array from a store. - - Parameters - ---------- - store : StoreLike - Store containing the Array. See the - [storage documentation in the user guide][user-guide-store-like] - for a description of all valid StoreLike values. - - Returns - ------- - Array - Array opened from the store. - """ - async_array = sync(AsyncArray.open(store)) - return cls(async_array) - - @property - def store(self) -> Store: - return self.async_array.store - - @property - def ndim(self) -> int: - """Returns the number of dimensions in the array. - - Returns - ------- - int - The number of dimensions in the array. - """ - return self.async_array.ndim - - @property - def shape(self) -> tuple[int, ...]: - """Returns the shape of the array. - - Returns - ------- - tuple[int, ...] - The shape of the array. - """ - return self.async_array.shape - - @shape.setter - def shape(self, value: tuple[int, ...]) -> None: - """Sets the shape of the array by calling resize.""" - self.resize(value) - - @property - def chunks(self) -> tuple[int, ...]: - """Returns a tuple of integers describing the length of each dimension of a chunk of the array. - If sharding is used the inner chunk shape is returned, which is defined - for any chunk grid (inner chunks are always regular). - - Otherwise, only defined for arrays using a regular chunk grid: for a - non-sharded array with a rectilinear chunk grid, `NotImplementedError` - is raised. Use `read_chunk_sizes` for general access. - - Returns - ------- - tuple - A tuple of integers representing the length of each dimension of a chunk. - """ - return self.async_array.chunks - - @property - def read_chunk_sizes(self) -> tuple[tuple[int, ...], ...]: - """Per-dimension data sizes of chunks used for reading, clipped to the array extent. - - Boundary chunks that extend past the array shape are clipped, so - the last size along a dimension may be smaller than the declared - chunk size. This matches the dask `Array.chunks` convention. - - When sharding is used, returns the inner chunk sizes. - Otherwise, returns the outer chunk sizes (same as `write_chunk_sizes`). - - Returns - ------- - tuple[tuple[int, ...], ...] - One inner tuple per dimension containing the data size of each - chunk (not the encoded buffer size). - - Examples - -------- - Without sharding, `read_chunk_sizes` and `write_chunk_sizes` are the same: - - >>> import zarr - >>> arr = zarr.create_array({}, dtype="i1", shape=(100, 80), chunks=(30, 40)) - >>> arr.read_chunk_sizes - ((30, 30, 30, 10), (40, 40)) - - For a sharded array the two differ: reads are efficient at inner-chunk - granularity, while writes go to storage one shard at a time: - - >>> sharded = zarr.create_array({}, dtype="i1", shape=(40,), chunks=(10,), shards=(20,)) - >>> sharded.read_chunk_sizes - ((10, 10, 10, 10),) - >>> sharded.write_chunk_sizes - ((20, 20),) - """ - return self.async_array.read_chunk_sizes - - @property - def write_chunk_sizes(self) -> tuple[tuple[int, ...], ...]: - """Per-dimension data sizes of storage chunks, clipped to the array extent. - - Always returns the outer chunk sizes, regardless of sharding. - Boundary chunks that extend past the array shape are clipped, so - the last size along a dimension may be smaller than the declared - chunk size. This matches the dask `Array.chunks` convention. - - Returns - ------- - tuple[tuple[int, ...], ...] - One inner tuple per dimension containing the data size of each - chunk (not the encoded buffer size). - - Examples - -------- - Without sharding, `write_chunk_sizes` and `read_chunk_sizes` are the same: - - >>> import zarr - >>> arr = zarr.create_array({}, dtype="i1", shape=(100, 80), chunks=(30, 40)) - >>> arr.write_chunk_sizes - ((30, 30, 30, 10), (40, 40)) - - For a sharded array the two differ: writes go to storage one shard at a - time, while reads are efficient at inner-chunk granularity: - - >>> sharded = zarr.create_array({}, dtype="i1", shape=(40,), chunks=(10,), shards=(20,)) - >>> sharded.write_chunk_sizes - ((20, 20),) - >>> sharded.read_chunk_sizes - ((10, 10, 10, 10),) - """ - return self.async_array.write_chunk_sizes - - @property - def shards(self) -> tuple[int, ...] | None: - """Returns a tuple of integers describing the length of each dimension of a shard of the array. - Returns None if sharding is not used. - - Only defined when the shard grid is regular: for a sharded array with a - rectilinear chunk grid, `NotImplementedError` is raised. Use - `write_chunk_sizes` for general access. - - Returns - ------- - tuple | None - A tuple of integers representing the length of each dimension of a shard or None if sharding is not used. - """ - return self.async_array.shards - - @property - def size(self) -> int: - """Returns the total number of elements in the array. - - Returns - ------- - int - Total number of elements in the array. - """ - return self.async_array.size - - @property - def dtype(self) -> np.dtype[Any]: - """Returns the NumPy data type. - - Returns - ------- - np.dtype - The NumPy data type. - """ - return self.async_array.dtype - - @property - def attrs(self) -> Attributes: - """Returns a [MutableMapping][collections.abc.MutableMapping] containing user-defined attributes. - - Returns - ------- - attrs - A [MutableMapping][collections.abc.MutableMapping] object containing user-defined attributes. - - Notes - ----- - Note that attribute values must be JSON serializable. - """ - return Attributes(self) - - @property - def path(self) -> str: - """Storage path.""" - return self.async_array.path - - @property - def name(self) -> str: - """Array name following h5py convention.""" - return self.async_array.name - - @property - def basename(self) -> str: - """Final component of name.""" - return self.async_array.basename - - @property - def metadata(self) -> ArrayMetadata: - return self.async_array.metadata - - @property - def store_path(self) -> StorePath: - return self.async_array.store_path - - @property - def order(self) -> MemoryOrder: - return self.async_array.order - - @property - def read_only(self) -> bool: - return self.async_array.read_only - - @property - def fill_value(self) -> Any: - return self.metadata.fill_value - - @property - def filters(self) -> tuple[Numcodec, ...] | tuple[ArrayArrayCodec, ...]: - """ - Filters that are applied to each chunk of the array, in order, before serializing that - chunk to bytes. - """ - return self.async_array.filters - - @property - def serializer(self) -> ArrayBytesCodec | None: - """ - Array-to-bytes codec to use for serializing the chunks into bytes. - """ - return self.async_array.serializer - - @property - @deprecated("Use Array.compressors instead.", category=ZarrDeprecationWarning) - def compressor(self) -> Numcodec | None: - """ - Compressor that is applied to each chunk of the array. - - !!! warning "Deprecated" - `array.compressor` is deprecated since v3.0.0 and will be removed in a future release. - Use [`array.compressors`][zarr.Array.compressors] instead. - """ - return self.async_array.compressor - - @property - def compressors(self) -> tuple[Numcodec, ...] | tuple[BytesBytesCodec, ...]: - """ - Compressors that are applied to each chunk of the array. Compressors are applied in order, and after any - filters are applied (if any are specified) and the data is serialized into bytes. - """ - return self.async_array.compressors - - @property - def cdata_shape(self) -> tuple[int, ...]: + If the dictionary data is invalid or missing required fields for array creation. """ - The number of chunks along each dimension. + async_array = AsyncArray.from_dict(store_path=store_path, data=data) + return cls(async_array) - When sharding is used, this counts inner chunks (not shards) per dimension. - """ - return self._chunk_grid_shape + @classmethod + def open( + cls, + store: StoreLike, + ) -> Array[Any]: + """Opens an existing Array from a store. - @property - def _chunk_grid_shape(self) -> tuple[int, ...]: + Parameters + ---------- + store : StoreLike + Store containing the Array. See the + [storage documentation in the user guide][user-guide-store-like] + for a description of all valid StoreLike values. + + Returns + ------- + Array + Array opened from the store. """ - The number of chunks along each dimension. + async_array = sync(AsyncArray.open(store)) + return cls(async_array) - When sharding is used, this counts inner chunks (not shards) per dimension. + @property + def shape(self) -> tuple[int, ...]: + """Returns the shape of the array. Returns ------- tuple[int, ...] - The number of chunks along each dimension. + The shape of the array. """ - return self.async_array._chunk_grid_shape + return self.metadata.shape - @property - def _shard_grid_shape(self) -> tuple[int, ...]: - """ - The shape of the shard grid for this array. - """ - return self.async_array._shard_grid_shape + @shape.setter + def shape(self, value: tuple[int, ...]) -> None: + """Sets the shape of the array by calling resize.""" + self.resize(value) @property - def nchunks(self) -> int: - """ - The number of chunks in this array. + def attrs(self) -> Attributes: + """Returns a [MutableMapping][collections.abc.MutableMapping] containing user-defined attributes. - Note that if a sharding codec is used, then the number of chunks may exceed the number of - stored objects supporting this array. - """ - return self.async_array.nchunks + Returns + ------- + attrs + A [MutableMapping][collections.abc.MutableMapping] object containing user-defined attributes. - @property - def _nshards(self) -> int: - """ - The number of shards in the stored representation of this array. + Notes + ----- + Note that attribute values must be JSON serializable. """ - return self.async_array._nshards + return Attributes(self) @overload def with_config(self: ArrayV2, config: ArrayConfigLike) -> ArrayV2: ... @@ -2334,21 +2153,6 @@ def with_config(self, config: ArrayConfigLike) -> Self: """ return type(self)(self._async_array.with_config(config)) - @property - def nbytes(self) -> int: - """ - The total number of bytes that can be stored in the chunks of this array. - - Notes - ----- - This value is calculated by multiplying the number of elements in the array and the size - of each element, the latter of which is determined by the dtype of the array. - For this reason, `nbytes` will likely be inaccurate for arrays with variable-length - dtypes. It is not possible to determine the size of an array with variable-length elements - from the shape and dtype alone. - """ - return self.async_array.nbytes - @property def nchunks_initialized(self) -> int: """ @@ -2377,7 +2181,7 @@ def nchunks_initialized(self) -> int: >>> arr.nchunks_initialized 6 """ - return sync(self.async_array.nchunks_initialized()) + return sync(_nchunks_initialized(self)) @property def _nshards_initialized(self) -> int: @@ -2399,7 +2203,7 @@ def _nshards_initialized(self) -> int: >>> arr._nshards_initialized 3 """ - return sync(self.async_array._nshards_initialized()) + return sync(_nshards_initialized(self)) def nbytes_stored(self) -> int: """ @@ -2409,125 +2213,7 @@ def nbytes_stored(self) -> int: ------- size : int """ - return sync(self.async_array.nbytes_stored()) - - def _iter_shard_keys( - self, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[str]: - """ - Iterate over the storage keys of each shard, relative to an optional origin, and optionally - limited to a contiguous region in chunk grid coordinates. - - If no shards are present this falls back to chunks, though in this case these are then actually - not storage keys. - - Parameters - ---------- - origin : Sequence[int] | None, default=None - The origin of the selection relative to the array's shard grid. - selection_shape : Sequence[int] | None, default=None - The shape of the selection in shard grid coordinates. - - Yields - ------ - str - The storage key of each shard in the selection or chunk though chunks technically do not have - storage keys. - """ - return self.async_array._iter_shard_keys(origin=origin, selection_shape=selection_shape) - - def _iter_chunk_coords( - self, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[tuple[int, ...]]: - """ - Create an iterator over the coordinates of chunks in chunk grid space. - - If the `origin` keyword is used, iteration will start at the chunk index specified by `origin`. - The default behavior is to start at the origin of the grid coordinate space. - If the `selection_shape` keyword is used, iteration will be bounded over a contiguous region - ranging from `[origin, origin + selection_shape]`, where the upper bound is exclusive as - per python indexing conventions. - - Parameters - ---------- - origin : Sequence[int] | None, default=None - The origin of the selection relative to the array's chunk grid. - selection_shape : Sequence[int] | None, default=None - The shape of the selection in chunk grid coordinates. - - Yields - ------ - tuple[int, ...] - The coordinates of each chunk in the selection. - """ - return self.async_array._iter_chunk_coords(origin=origin, selection_shape=selection_shape) - - def _iter_shard_coords( - self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[tuple[int, ...]]: - """ - Create an iterator over the coordinates of shards in shard grid space. - - If the `origin` keyword is used, iteration will start at the shard index specified by `origin`. - The default behavior is to start at the origin of the grid coordinate space. - If the `selection_shape` keyword is used, iteration will be bounded over a contiguous region - ranging from `[origin, origin selection_shape]`, where the upper bound is exclusive as - per python indexing conventions. - - Parameters - ---------- - origin : Sequence[int] | None, default=None - The origin of the selection relative to the array's shard grid. - selection_shape : Sequence[int] | None, default=None - The shape of the selection in shard grid coordinates. - - Yields - ------ - tuple[int, ...] - The coordinates of each shard in the selection. - """ - return self.async_array._iter_shard_coords(origin=origin, selection_shape=selection_shape) - - def _iter_chunk_regions( - self, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[tuple[slice, ...]]: - """ - Iterate over the regions spanned by each chunk. - - Parameters - ---------- - origin : Sequence[int] | None, default=None - The origin of the selection relative to the array's chunk grid. - selection_shape : Sequence[int] | None, default=None - The shape of the selection in chunk grid coordinates. - - Yields - ------ - tuple[slice, ...] - A tuple of slice objects representing the region spanned by each chunk in the selection. - """ - return self.async_array._iter_chunk_regions(origin=origin, selection_shape=selection_shape) - - def _iter_shard_regions( - self, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[tuple[slice, ...]]: - """ - Iterate over the regions spanned by each shard or chunk if no shard is present. - - Parameters - ---------- - origin : Sequence[int] | None, default=None - The origin of the selection relative to the array's chunk grid. - selection_shape : Sequence[int] | None, default=None - The shape of the selection in chunk grid coordinates. - - Yields - ------ - tuple[slice, ...] - A tuple of slice objects representing the region spanned by each shard or if no shard is present, - chunk in the selection. - """ - return self.async_array._iter_shard_regions(origin=origin, selection_shape=selection_shape) + return sync(_nbytes_stored(self.store_path)) def __array__( self, dtype: npt.DTypeLike | None = None, copy: bool | None = None @@ -2937,7 +2623,7 @@ def get_basic_selection( if prototype is None: prototype = default_buffer_prototype() return sync( - self.async_array._get_selection( + self._get_selection( BasicIndexer(selection, self.shape, self._chunk_grid), out=out, fields=fields, @@ -3046,7 +2732,7 @@ def set_basic_selection( if prototype is None: prototype = default_buffer_prototype() indexer = BasicIndexer(selection, self.shape, self._chunk_grid) - sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + sync(self._set_selection(indexer, value, fields=fields, prototype=prototype)) def get_orthogonal_selection( self, @@ -3175,9 +2861,7 @@ def get_orthogonal_selection( prototype = default_buffer_prototype() indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) return sync( - self.async_array._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype - ) + self._get_selection(indexer=indexer, out=out, fields=fields, prototype=prototype) ) def set_orthogonal_selection( @@ -3292,9 +2976,7 @@ def set_orthogonal_selection( if prototype is None: prototype = default_buffer_prototype() indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) - return sync( - self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) - ) + return sync(self._set_selection(indexer, value, fields=fields, prototype=prototype)) def get_mask_selection( self, @@ -3381,9 +3063,7 @@ def get_mask_selection( prototype = default_buffer_prototype() indexer = MaskIndexer(mask, self.shape, self._chunk_grid) return sync( - self.async_array._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype - ) + self._get_selection(indexer=indexer, out=out, fields=fields, prototype=prototype) ) def set_mask_selection( @@ -3469,7 +3149,7 @@ def set_mask_selection( if prototype is None: prototype = default_buffer_prototype() indexer = MaskIndexer(mask, self.shape, self._chunk_grid) - sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + sync(self._set_selection(indexer, value, fields=fields, prototype=prototype)) def get_coordinate_selection( self, @@ -3558,9 +3238,7 @@ def get_coordinate_selection( prototype = default_buffer_prototype() indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) out_array = sync( - self.async_array._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype - ) + self._get_selection(indexer=indexer, out=out, fields=fields, prototype=prototype) ) if hasattr(out_array, "shape"): @@ -3670,7 +3348,7 @@ def set_coordinate_selection( f"elements with an array of {value.shape[0]} elements." ) - sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + sync(self._set_selection(indexer, value, fields=fields, prototype=prototype)) def get_block_selection( self, @@ -3771,9 +3449,7 @@ def get_block_selection( prototype = default_buffer_prototype() indexer = BlockIndexer(selection, self.shape, self._chunk_grid) return sync( - self.async_array._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype - ) + self._get_selection(indexer=indexer, out=out, fields=fields, prototype=prototype) ) def set_block_selection( @@ -3870,7 +3546,7 @@ def set_block_selection( if prototype is None: prototype = default_buffer_prototype() indexer = BlockIndexer(selection, self.shape, self._chunk_grid) - sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + sync(self._set_selection(indexer, value, fields=fields, prototype=prototype)) @property def vindex(self) -> VIndex: @@ -3934,7 +3610,7 @@ def resize(self, new_shape: ShapeLike) -> None: #>(50, 50) ``` """ - sync(self.async_array.resize(new_shape)) + sync(_resize(self, new_shape)) def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: """Append `data` to `axis`. @@ -3970,7 +3646,7 @@ def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: >>> z.shape (20000, 2000) """ - return sync(self.async_array.append(data, axis=axis)) + return sync(_append(self, data, axis)) def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: """ @@ -3997,8 +3673,8 @@ def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: - The updated attributes will be merged with existing attributes, and any conflicts will be overwritten by the new values. """ - new_array = sync(self.async_array.update_attributes(new_attributes)) - return type(self)(new_array) + sync(_update_attributes(self, new_attributes)) + return self def __repr__(self) -> str: return f"" @@ -4035,7 +3711,7 @@ def info(self) -> Any: Compressors : (ZstdCodec(level=0, checksum=False),) No. bytes : 40 """ - return self.async_array.info + return self._info() def info_complete(self) -> Any: """ @@ -4055,11 +3731,11 @@ def info_complete(self) -> Any: ------- [zarr.Array.info][] - The statically known subset of metadata about an array. """ - return sync(self.async_array.info_complete()) + return sync(_info_complete(self)) async def _shards_initialized( - array: AnyAsyncArray, + array: _AnyArrayBase, ) -> tuple[str, ...]: """ Return the keys of the shards that have been persisted to the storage backend. @@ -4096,6 +3772,8 @@ async def _shards_initialized( ) +type _AnyArrayBase = _ArrayBase[ArrayV2Metadata] | _ArrayBase[ArrayV3Metadata] + type FiltersLike = ( Iterable[dict[str, JSON] | ArrayArrayCodec | Numcodec] | ArrayArrayCodec @@ -5280,7 +4958,7 @@ def _parse_data_params( def _iter_chunk_coords( - array: AnyArray | AnyAsyncArray, + array: _AnyArrayBase, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None, @@ -5311,7 +4989,7 @@ def _iter_chunk_coords( def _iter_shard_coords( - array: AnyArray | AnyAsyncArray, + array: _AnyArrayBase, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None, @@ -5344,7 +5022,7 @@ def _iter_shard_coords( def _iter_shard_keys( - array: AnyArray | AnyAsyncArray, + array: _AnyArrayBase, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None, @@ -5375,7 +5053,7 @@ def _iter_shard_keys( def _iter_shard_regions( - array: AnyArray | AnyAsyncArray, + array: _AnyArrayBase, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None, @@ -5408,7 +5086,7 @@ def _iter_shard_regions( def _iter_chunk_regions( - array: AnyArray | AnyAsyncArray, + array: _AnyArrayBase, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None, @@ -5437,7 +5115,7 @@ def _iter_chunk_regions( async def _nchunks_initialized( - array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + array: _AnyArrayBase, ) -> int: """ Calculate the number of chunks that have been initialized in storage. @@ -5486,7 +5164,7 @@ async def _nchunks_initialized( async def _nshards_initialized( - array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + array: _AnyArrayBase, ) -> int: """ Calculate the number of shards that have been initialized in storage. @@ -5879,8 +5557,78 @@ async def _setitem( ) +def _plan_resize( + state: ArrayState[Any], new_shape: ShapeLike, *, delete_outside_chunks: bool +) -> tuple[ArrayState[Any], tuple[str, ...]]: + """ + Plan a resize without doing any IO. + + Returns the state the array will have after the resize, and the keys of the chunks + that fall outside the new shape and should be deleted. + """ + new_shape = parse_shapelike(new_shape) + if len(new_shape) != len(state.metadata.shape): + raise ValueError( + f"The new shape must have the same number of dimensions as the array. " + f"Got {len(new_shape)} dimension(s), expected {len(state.metadata.shape)}." + ) + new_state = state.with_metadata(state.metadata.update_shape(new_shape)) + + # only compute deletions if the array is shrinking, as the set difference below is + # unbounded in memory + only_growing = all(new >= old for new, old in zip(new_shape, state.metadata.shape, strict=True)) + if not delete_outside_chunks or only_growing: + return new_state, () + outside = set(state.chunk_grid.all_chunk_coords()).difference( + new_state.chunk_grid.all_chunk_coords() + ) + return new_state, tuple(state.metadata.encode_chunk_key(coords) for coords in outside) + + +def _plan_append( + shape: tuple[int, ...], data_shape: tuple[int, ...], axis: int +) -> tuple[tuple[int, ...], tuple[slice, ...]]: + """ + Plan an append without doing any IO. + + Returns the shape of the array after the append, and the selection that the appended + data will be written to. + """ + self_shape_preserved = tuple(s for i, s in enumerate(shape) if i != axis) + data_shape_preserved = tuple(s for i, s in enumerate(data_shape) if i != axis) + if self_shape_preserved != data_shape_preserved: + raise ValueError( + f"shape of data to append is not compatible with the array. " + f"The shape of the data is ({data_shape_preserved})" + f"and the shape of the array is ({self_shape_preserved})." + "All dimensions must match except for the dimension being " + "appended." + ) + new_shape = tuple( + shape[i] if i != axis else shape[i] + data_shape[i] for i in range(len(shape)) + ) + append_selection = tuple( + slice(None) if i != axis else slice(shape[i], new_shape[i]) for i in range(len(shape)) + ) + return new_shape, append_selection + + +def _plan_update_attributes( + state: ArrayState[Any], new_attributes: dict[str, JSON] +) -> ArrayState[Any]: + """ + Plan an attribute update without doing any IO. + + The new attributes are merged into a copy of the current attributes; the current + state is not modified. + """ + return state.with_metadata( + state.metadata.update_attributes({**state.metadata.attributes, **new_attributes}) + ) + + async def _resize( - array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + array: _AnyArrayBase, new_shape: ShapeLike, delete_outside_chunks: bool = True, ) -> None: @@ -5889,7 +5637,7 @@ async def _resize( Parameters ---------- - array : AsyncArray + array : Array | AsyncArray The array to resize. new_shape : ShapeLike The desired new shape of the array. @@ -5897,46 +5645,23 @@ async def _resize( If True (default), chunks that fall outside the new shape will be deleted. If False, the data in those chunks will be preserved. """ - new_shape = parse_shapelike(new_shape) - if len(new_shape) != len(array.metadata.shape): - raise ValueError( - f"The new shape must have the same number of dimensions as the array. " - f"Got {len(new_shape)} dimension(s), expected {len(array.metadata.shape)}." - ) - - new_metadata = array.metadata.update_shape(new_shape) - new_chunk_grid = ChunkGrid.from_metadata(new_metadata) - - # ensure deletion is only run if array is shrinking as the delete_outside_chunks path is unbounded in memory - only_growing = all(new >= old for new, old in zip(new_shape, array.metadata.shape, strict=True)) - - if delete_outside_chunks and not only_growing: - # Remove all chunks outside of the new shape - old_chunk_coords = set(array._chunk_grid.all_chunk_coords()) - new_chunk_coords = set(new_chunk_grid.all_chunk_coords()) - - async def _delete_key(key: str) -> None: - await (array.store_path / key).delete() - - await concurrent_map( - [ - (array.metadata.encode_chunk_key(chunk_coords),) - for chunk_coords in old_chunk_coords.difference(new_chunk_coords) - ], - _delete_key, - zarr_config.get("async.concurrency"), - ) + state = array._cell.state + new_state, outside_keys = _plan_resize( + state, new_shape, delete_outside_chunks=delete_outside_chunks + ) - # Write new metadata - await save_metadata(array.store_path, new_metadata) + async def _delete_key(key: str) -> None: + await (state.store_path / key).delete() - # Update metadata and chunk_grid (in place) - object.__setattr__(array, "metadata", new_metadata) - object.__setattr__(array, "_chunk_grid", new_chunk_grid) + await concurrent_map( + [(key,) for key in outside_keys], _delete_key, zarr_config.get("async.concurrency") + ) + await save_metadata(new_state.store_path, new_state.metadata) + array._cell.state = new_state async def _append( - array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + array: _AnyArrayBase, data: npt.ArrayLike, axis: int = 0, ) -> tuple[int, ...]: @@ -5945,7 +5670,7 @@ async def _append( Parameters ---------- - array : AsyncArray + array : Array | AsyncArray The array to append to. data : npt.ArrayLike Data to be appended. @@ -5966,39 +5691,16 @@ async def _append( if not hasattr(data, "shape"): data = np.asanyarray(data) - self_shape_preserved = tuple(s for i, s in enumerate(array.shape) if i != axis) - data_shape_preserved = tuple(s for i, s in enumerate(data.shape) if i != axis) - if self_shape_preserved != data_shape_preserved: - raise ValueError( - f"shape of data to append is not compatible with the array. " - f"The shape of the data is ({data_shape_preserved})" - f"and the shape of the array is ({self_shape_preserved})." - "All dimensions must match except for the dimension being " - "appended." - ) - # remember old shape - old_shape = array.shape - - # determine new shape - new_shape = tuple( - array.shape[i] if i != axis else array.shape[i] + data.shape[i] - for i in range(len(array.shape)) - ) - - # resize + new_shape, append_selection = _plan_append(array.shape, data.shape, axis) await _resize(array, new_shape) - # store data - append_selection = tuple( - slice(None) if i != axis else slice(old_shape[i], new_shape[i]) - for i in range(len(array.shape)) - ) + state = array._cell.state await _setitem( - array.store_path, - array.metadata, - array.codec_pipeline, - array.config, - array._chunk_grid, + state.store_path, + state.metadata, + state.codec_pipeline, + state.config, + state.chunk_grid, append_selection, data, ) @@ -6007,34 +5709,26 @@ async def _append( async def _update_attributes( - array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + array: _AnyArrayBase, new_attributes: dict[str, JSON], -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +) -> None: """ Update the array's attributes. Parameters ---------- - array : AsyncArray + array : Array | AsyncArray The array whose attributes to update. new_attributes : dict[str, JSON] A dictionary of new attributes to update or add to the array. - - Returns - ------- - AsyncArray - The array with the updated attributes. """ - array.metadata.attributes.update(new_attributes) - - # Write new metadata - await save_metadata(array.store_path, array.metadata) - - return array + new_state = _plan_update_attributes(array._cell.state, new_attributes) + await save_metadata(new_state.store_path, new_state.metadata) + array._cell.state = new_state async def _info_complete( - array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + array: _AnyArrayBase, ) -> Any: """ Return all the information for an array, including dynamic information like storage size. diff --git a/src/zarr/core/attributes.py b/src/zarr/core/attributes.py index 7f29e44365..41358eaf74 100644 --- a/src/zarr/core/attributes.py +++ b/src/zarr/core/attributes.py @@ -18,7 +18,8 @@ def __init__(self, obj: AnyArray | Group) -> None: self._obj = obj def __getitem__(self, key: str) -> JSON: - return self._obj.metadata.attributes[key] + attributes: dict[str, JSON] = self._obj.metadata.attributes + return attributes[key] def __setitem__(self, key: str, value: JSON) -> None: new_attrs = dict(self._obj.metadata.attributes) diff --git a/src/zarr/metadata/migrate_v3.py b/src/zarr/metadata/migrate_v3.py index ad177e19fb..c65c99a64b 100644 --- a/src/zarr/metadata/migrate_v3.py +++ b/src/zarr/metadata/migrate_v3.py @@ -89,7 +89,7 @@ def migrate_to_v3(zarr_v2: AnyArray | Group, output_path: StorePath, dry_run: bo if not zarr_v2.metadata.zarr_format == 2: raise TypeError("Only arrays / groups with zarr v2 metadata can be converted") - if isinstance(zarr_v2.metadata, GroupMetadata): + if isinstance(zarr_v2, Group): _convert_group(zarr_v2, output_path, dry_run) else: _convert_array(zarr_v2, output_path, dry_run) diff --git a/tests/test_array_state.py b/tests/test_array_state.py new file mode 100644 index 0000000000..399f53fe5a --- /dev/null +++ b/tests/test_array_state.py @@ -0,0 +1,162 @@ +"""Tests for the state cell shared by `Array` and `AsyncArray`.""" + +from __future__ import annotations + +import pickle +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +import pytest + +import zarr +from zarr.core.array import ( + Array, + _plan_append, + _plan_resize, + _plan_update_attributes, +) +from zarr.core.sync import sync + +if TYPE_CHECKING: + from zarr.core.common import ZarrFormat + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +@pytest.mark.parametrize("via", ["sync", "async"]) +@pytest.mark.parametrize("op", ["resize", "append", "update_attributes"]) +def test_views_share_state( + zarr_format: ZarrFormat, + via: Literal["sync", "async"], + op: Literal["resize", "append", "update_attributes"], +) -> None: + """ + A change made through an `Array` or through its `async_array` is seen by both, and by + the store. + """ + arr = zarr.create_array({}, shape=(10,), chunks=(2,), dtype="i4", zarr_format=zarr_format) + arr[:] = np.arange(10) + aa = arr.async_array + expected_attrs: dict[str, Any] + + if op == "resize": + expected_shape, expected_attrs = (6,), {} + expected_data = np.arange(6) + if via == "sync": + arr.resize((6,)) + else: + sync(aa.resize((6,))) + elif op == "append": + expected_shape, expected_attrs = (12,), {} + expected_data = np.concatenate([np.arange(10), [7, 7]]) + if via == "sync": + arr.append(np.array([7, 7])) + else: + sync(aa.append(np.array([7, 7]))) + else: + expected_shape, expected_attrs = (10,), {"a": 1} + expected_data = np.arange(10) + if via == "sync": + arr.update_attributes({"a": 1}) + else: + sync(aa.update_attributes({"a": 1})) + + reopened = zarr.open_array(arr.store_path.store, zarr_format=zarr_format) + for view in (arr, aa, reopened): + assert view.shape == expected_shape + assert dict(view.attrs) == expected_attrs + np.testing.assert_array_equal(arr[:], expected_data) + np.testing.assert_array_equal(reopened[:], expected_data) + assert aa._cell is arr._cell + + +@pytest.mark.parametrize( + ("new_shape", "delete_outside_chunks", "expected_keys"), + [ + ((12,), True, set()), + ((5,), False, set()), + ((5,), True, {"c/3", "c/4"}), + ((0,), True, {"c/0", "c/1", "c/2", "c/3", "c/4"}), + ], +) +def test_plan_resize( + new_shape: tuple[int, ...], delete_outside_chunks: bool, expected_keys: set[str] +) -> None: + """ + Planning a resize computes the new state and the keys to delete, and changes nothing. + """ + arr = zarr.create_array({}, shape=(10,), chunks=(2,), dtype="i4") + state = arr._cell.state + + new_state, keys = _plan_resize(state, new_shape, delete_outside_chunks=delete_outside_chunks) + + assert new_state.metadata.shape == new_shape + assert new_state.chunk_grid.grid_shape == (-(-new_shape[0] // 2),) + assert new_state.codec_pipeline is state.codec_pipeline + assert set(keys) == expected_keys + assert arr._cell.state is state + assert arr.shape == (10,) + + +def test_plan_resize_wrong_ndim() -> None: + """ + Planning a resize to a shape with a different number of dimensions raises. + """ + arr = zarr.create_array({}, shape=(10,), chunks=(2,), dtype="i4") + with pytest.raises(ValueError, match="same number of dimensions"): + _plan_resize(arr._cell.state, (10, 10), delete_outside_chunks=True) + + +@pytest.mark.parametrize( + ("shape", "data_shape", "axis", "expected_shape", "expected_selection"), + [ + ((10,), (3,), 0, (13,), (slice(10, 13),)), + ((4, 5), (2, 5), 0, (6, 5), (slice(4, 6), slice(None))), + ((4, 5), (4, 1), 1, (4, 6), (slice(None), slice(5, 6))), + ], +) +def test_plan_append( + shape: tuple[int, ...], + data_shape: tuple[int, ...], + axis: int, + expected_shape: tuple[int, ...], + expected_selection: tuple[slice, ...], +) -> None: + """ + Planning an append computes the new shape and the region the data is written to. + """ + assert _plan_append(shape, data_shape, axis) == (expected_shape, expected_selection) + + +def test_plan_append_incompatible_shape() -> None: + """ + Planning an append of data whose other dimensions do not match the array raises. + """ + with pytest.raises(ValueError, match="not compatible"): + _plan_append((4, 5), (2, 3), 0) + + +def test_plan_update_attributes() -> None: + """ + Planning an attribute update merges into a copy; the current state is not modified. + """ + arr = zarr.create_array({}, shape=(10,), chunks=(2,), dtype="i4", attributes={"a": 1}) + state = arr._cell.state + + new_state = _plan_update_attributes(state, {"b": 2}) + + assert new_state.metadata.attributes == {"a": 1, "b": 2} + assert state.metadata.attributes == {"a": 1} + assert new_state.codec_pipeline is state.codec_pipeline + + +def test_pickle_keeps_views_coherent() -> None: + """ + After a pickle round trip, an `Array` and its `async_array` still share one state cell. + """ + arr = zarr.create_array({}, shape=(10,), chunks=(2,), dtype="i4") + restored: Array[Any] = pickle.loads(pickle.dumps(arr)) + + restored.resize((4,)) + + assert restored.async_array.shape == (4,) + assert restored.async_array._cell is restored._cell From 2edf7d898f1296fccc73c01df88b53c215cd0ed6 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 23 Sep 2026 23:58:53 +0200 Subject: [PATCH 2/2] docs: add changelog fragment for #360 Assisted-by: ClaudeCode:claude-opus-5-5 --- changes/360.misc.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/360.misc.md diff --git a/changes/360.misc.md b/changes/360.misc.md new file mode 100644 index 0000000000..f257e7ea07 --- /dev/null +++ b/changes/360.misc.md @@ -0,0 +1 @@ +`Array` and `AsyncArray` now share one immutable `ArrayState` through a common state cell, so a change made through an `Array` or its `async_array` is seen by both. `Array` no longer delegates its IO to `AsyncArray`, and `update_attributes` no longer mutates the previous metadata in place.