From 79a6ad788630055320a8e09f92699c9b253a9749 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 26 Jan 2026 17:39:20 +0100 Subject: [PATCH 1/2] lift array methods to separate functions --- src/zarr/core/array.py | 951 ++++++++++++++++++++++++++++++++--------- 1 file changed, 748 insertions(+), 203 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 00536a1ec0..0e68288c68 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -1337,13 +1337,7 @@ async def example(): result = asyncio.run(example()) ``` """ - if self.shards is None: - chunks_per_shard = 1 - else: - chunks_per_shard = product( - tuple(a // b for a, b in zip(self.shards, self.chunks, strict=True)) - ) - return (await self._nshards_initialized()) * chunks_per_shard + return await _nchunks_initialized(self) async def _nshards_initialized(self) -> int: """ @@ -1381,10 +1375,10 @@ async def example(): result = asyncio.run(example()) ``` """ - return len(await _shards_initialized(self)) + return await _nshards_initialized(self) async def nbytes_stored(self) -> int: - return await self.store_path.store.getsize_prefix(self.store_path.path) + return await _nbytes_stored(self.store_path) def _iter_chunk_coords( self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None @@ -1549,49 +1543,16 @@ async def _get_selection( out: NDBuffer | None = None, fields: Fields | None = None, ) -> NDArrayLikeOrScalar: - # check fields are sensible - out_dtype = check_fields(fields, self.dtype) - - # setup output buffer - if out is not None: - if isinstance(out, NDBuffer): - out_buffer = out - else: - raise TypeError(f"out argument needs to be an NDBuffer. Got {type(out)!r}") - if out_buffer.shape != indexer.shape: - raise ValueError( - f"shape of out argument doesn't match. Expected {indexer.shape}, got {out.shape}" - ) - else: - out_buffer = prototype.nd_buffer.empty( - shape=indexer.shape, - dtype=out_dtype, - order=self.order, - ) - if product(indexer.shape) > 0: - # need to use the order from the metadata for v2 - _config = self._config - if self.metadata.zarr_format == 2: - _config = replace(_config, order=self.order) - - # reading chunks and decoding them - await self.codec_pipeline.read( - [ - ( - self.store_path / self.metadata.encode_chunk_key(chunk_coords), - self.metadata.get_chunk_spec(chunk_coords, _config, prototype=prototype), - chunk_selection, - out_selection, - is_complete_chunk, - ) - for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexer - ], - out_buffer, - drop_axes=indexer.drop_axes, - ) - if isinstance(indexer, BasicIndexer) and indexer.shape == (): - return out_buffer.as_scalar() - return out_buffer.as_ndarray_like() + return await _get_selection( + self.store_path, + self.metadata, + self.codec_pipeline, + self._config, + indexer, + prototype=prototype, + out=out, + fields=fields, + ) async def getitem( self, @@ -1636,14 +1597,14 @@ async def example(): value = asyncio.run(example()) ``` """ - if prototype is None: - prototype = default_buffer_prototype() - indexer = BasicIndexer( + return await _getitem( + self.store_path, + self.metadata, + self.codec_pipeline, + self._config, selection, - shape=self.metadata.shape, - chunk_grid=self.metadata.chunk_grid, + prototype=prototype, ) - return await self._get_selection(indexer, prototype=prototype) async def get_orthogonal_selection( self, @@ -1653,11 +1614,15 @@ async def get_orthogonal_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: - if prototype is None: - prototype = default_buffer_prototype() - indexer = OrthogonalIndexer(selection, self.shape, self.metadata.chunk_grid) - return await self._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype + return await _get_orthogonal_selection( + self.store_path, + self.metadata, + self.codec_pipeline, + self._config, + selection, + out=out, + fields=fields, + prototype=prototype, ) async def get_mask_selection( @@ -1668,11 +1633,15 @@ async def get_mask_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: - if prototype is None: - prototype = default_buffer_prototype() - indexer = MaskIndexer(mask, self.shape, self.metadata.chunk_grid) - return await self._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype + return await _get_mask_selection( + self.store_path, + self.metadata, + self.codec_pipeline, + self._config, + mask, + out=out, + fields=fields, + prototype=prototype, ) async def get_coordinate_selection( @@ -1683,18 +1652,17 @@ async def get_coordinate_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: - if prototype is None: - prototype = default_buffer_prototype() - indexer = CoordinateIndexer(selection, self.shape, self.metadata.chunk_grid) - out_array = await self._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype + return await _get_coordinate_selection( + self.store_path, + self.metadata, + self.codec_pipeline, + self._config, + selection, + out=out, + fields=fields, + prototype=prototype, ) - if hasattr(out_array, "shape"): - # restore shape - out_array = np.array(out_array).reshape(indexer.sel_shape) - return out_array - async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = False) -> None: """ Asynchronously save the array metadata. @@ -1709,56 +1677,15 @@ async def _set_selection( prototype: BufferPrototype, fields: Fields | None = None, ) -> None: - # check fields are sensible - check_fields(fields, self.dtype) - fields = check_no_multi_fields(fields) - - # check value shape - if np.isscalar(value): - array_like = prototype.buffer.create_zero_length().as_array_like() - if isinstance(array_like, np._typing._SupportsArrayFunc): - # TODO: need to handle array types that don't support __array_function__ - # like PyTorch and JAX - array_like_ = cast("np._typing._SupportsArrayFunc", array_like) - value = np.asanyarray(value, dtype=self.dtype, like=array_like_) - else: - if not hasattr(value, "shape"): - value = np.asarray(value, self.dtype) - # assert ( - # value.shape == indexer.shape - # ), f"shape of value doesn't match indexer shape. Expected {indexer.shape}, got {value.shape}" - if not hasattr(value, "dtype") or value.dtype.name != self.dtype.name: - if hasattr(value, "astype"): - # Handle things that are already NDArrayLike more efficiently - value = value.astype(dtype=self.dtype, order="A") - else: - value = np.array(value, dtype=self.dtype, order="A") - value = cast("NDArrayLike", value) - - # We accept any ndarray like object from the user and convert it - # to an NDBuffer (or subclass). From this point onwards, we only pass - # Buffer and NDBuffer between components. - value_buffer = prototype.nd_buffer.from_ndarray_like(value) - - # need to use the order from the metadata for v2 - _config = self._config - if self.metadata.zarr_format == 2: - _config = replace(_config, order=self.metadata.order) - - # merging with existing data and encoding chunks - await self.codec_pipeline.write( - [ - ( - self.store_path / self.metadata.encode_chunk_key(chunk_coords), - self.metadata.get_chunk_spec(chunk_coords, _config, prototype), - chunk_selection, - out_selection, - is_complete_chunk, - ) - for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexer - ], - value_buffer, - drop_axes=indexer.drop_axes, + return await _set_selection( + self.store_path, + self.metadata, + self.codec_pipeline, + self._config, + indexer, + value, + prototype=prototype, + fields=fields, ) async def setitem( @@ -1800,14 +1727,15 @@ async def setitem( - This method is asynchronous and should be awaited. - Supports basic indexing, where the selection is contiguous and does not involve advanced indexing. """ - if prototype is None: - prototype = default_buffer_prototype() - indexer = BasicIndexer( + return await _setitem( + self.store_path, + self.metadata, + self.codec_pipeline, + self._config, selection, - shape=self.metadata.shape, - chunk_grid=self.metadata.chunk_grid, + value, + prototype=prototype, ) - return await self._set_selection(indexer, value, prototype=prototype) @property def oindex(self) -> AsyncOIndex[T_ArrayMetadata]: @@ -1849,32 +1777,7 @@ async def resize(self, new_shape: ShapeLike, delete_outside_chunks: bool = True) ----- - This method is asynchronous and should be awaited. """ - new_shape = parse_shapelike(new_shape) - assert len(new_shape) == len(self.metadata.shape) - new_metadata = self.metadata.update_shape(new_shape) - - if delete_outside_chunks: - # Remove all chunks outside of the new shape - old_chunk_coords = set(self.metadata.chunk_grid.all_chunk_coords(self.metadata.shape)) - new_chunk_coords = set(self.metadata.chunk_grid.all_chunk_coords(new_shape)) - - async def _delete_key(key: str) -> None: - await (self.store_path / key).delete() - - await concurrent_map( - [ - (self.metadata.encode_chunk_key(chunk_coords),) - for chunk_coords in old_chunk_coords.difference(new_chunk_coords) - ], - _delete_key, - zarr_config.get("async.concurrency"), - ) - - # Write new metadata - await self._save_metadata(new_metadata) - - # Update metadata (in place) - object.__setattr__(self, "metadata", new_metadata) + return await _resize(self, new_shape, delete_outside_chunks) async def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: """Append `data` to `axis`. @@ -1895,40 +1798,7 @@ async def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: The size of all dimensions other than `axis` must match between this array and `data`. """ - # ensure data is array-like - if not hasattr(data, "shape"): - data = np.asanyarray(data) - - self_shape_preserved = tuple(s for i, s in enumerate(self.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 = self.shape - - # determine new shape - new_shape = tuple( - self.shape[i] if i != axis else self.shape[i] + data.shape[i] - for i in range(len(self.shape)) - ) - - # resize - await self.resize(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(self.shape)) - ) - await self.setitem(append_selection, data) - - return new_shape + return await _append(self, data, axis) async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: """ @@ -1956,11 +1826,7 @@ async 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. """ - self.metadata.attributes.update(new_attributes) - - # Write new metadata - await self._save_metadata(self.metadata) - + await _update_attributes(self, new_attributes) return self def __repr__(self) -> str: @@ -2017,10 +1883,7 @@ async def info_complete(self) -> Any: ------- [zarr.AsyncArray.info][] - A property giving just the statically known information about an array. """ - return self._info( - await self._nshards_initialized(), - await self.store_path.store.getsize_prefix(self.store_path.path), - ) + return await _info_complete(self) def _info( self, count_chunks_initialized: int | None = None, count_bytes_stored: int | None = None @@ -5518,3 +5381,685 @@ def _iter_chunk_regions( return _iter_regions( array.shape, array.chunks, origin=origin, selection_shape=selection_shape, trim_excess=True ) + + +async def _nchunks_initialized( + array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], +) -> int: + """ + Calculate the number of chunks that have been initialized in storage. + + This value is calculated as the product of the number of initialized shards and the number + of chunks per shard. For arrays that do not use sharding, the number of chunks per shard is + effectively 1, and in that case the number of chunks initialized is the same as the number + of stored objects associated with an array. + + Parameters + ---------- + array : AsyncArray + The array to inspect. + + Returns + ------- + nchunks_initialized : int + The number of chunks that have been initialized. + """ + if array.shards is None: + chunks_per_shard = 1 + else: + chunks_per_shard = product( + tuple(a // b for a, b in zip(array.shards, array.chunks, strict=True)) + ) + return (await _nshards_initialized(array)) * chunks_per_shard + + +async def _nshards_initialized( + array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], +) -> int: + """ + Calculate the number of shards that have been initialized in storage. + + This is the number of shards that have been persisted to the storage backend. + + Parameters + ---------- + array : AsyncArray + The array to inspect. + + Returns + ------- + nshards_initialized : int + The number of shards that have been initialized. + """ + return len(await _shards_initialized(array)) + + +async def _nbytes_stored( + store_path: StorePath, +) -> int: + """ + Calculate the number of bytes stored for an array. + + Parameters + ---------- + store_path : StorePath + The store path of the array. + + Returns + ------- + nbytes_stored : int + The number of bytes stored. + """ + return await store_path.store.getsize_prefix(store_path.path) + + +async def _get_selection( + store_path: StorePath, + metadata: ArrayMetadata, + codec_pipeline: CodecPipeline, + config: ArrayConfig, + indexer: Indexer, + *, + prototype: BufferPrototype, + out: NDBuffer | None = None, + fields: Fields | None = None, +) -> NDArrayLikeOrScalar: + """ + Get a selection from an array. + + Parameters + ---------- + store_path : StorePath + The store path of the array. + metadata : ArrayMetadata + The array metadata. + codec_pipeline : CodecPipeline + The codec pipeline for encoding/decoding. + config : ArrayConfig + The array configuration. + indexer : Indexer + The indexer specifying the selection. + prototype : BufferPrototype + A buffer prototype to use for the retrieved data. + out : NDBuffer | None, optional + An output buffer to write the data to. + fields : Fields | None, optional + Fields to select from structured arrays. + + Returns + ------- + NDArrayLikeOrScalar + The selected data. + """ + # Get dtype from metadata + if metadata.zarr_format == 2: + zdtype = metadata.dtype + else: + zdtype = metadata.data_type + dtype = zdtype.to_native_dtype() + + # Determine memory order + if metadata.zarr_format == 2: + order = metadata.order + else: + order = config.order + + # check fields are sensible + out_dtype = check_fields(fields, dtype) + + # setup output buffer + if out is not None: + if isinstance(out, NDBuffer): + out_buffer = out + else: + raise TypeError(f"out argument needs to be an NDBuffer. Got {type(out)!r}") + if out_buffer.shape != indexer.shape: + raise ValueError( + f"shape of out argument doesn't match. Expected {indexer.shape}, got {out.shape}" + ) + else: + out_buffer = prototype.nd_buffer.empty( + shape=indexer.shape, + dtype=out_dtype, + order=order, + ) + if product(indexer.shape) > 0: + # need to use the order from the metadata for v2 + _config = config + if metadata.zarr_format == 2: + _config = replace(_config, order=order) + + # reading chunks and decoding them + await codec_pipeline.read( + [ + ( + store_path / metadata.encode_chunk_key(chunk_coords), + metadata.get_chunk_spec(chunk_coords, _config, prototype=prototype), + chunk_selection, + out_selection, + is_complete_chunk, + ) + for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexer + ], + out_buffer, + drop_axes=indexer.drop_axes, + ) + if isinstance(indexer, BasicIndexer) and indexer.shape == (): + return out_buffer.as_scalar() + return out_buffer.as_ndarray_like() + + +async def _getitem( + store_path: StorePath, + metadata: ArrayMetadata, + codec_pipeline: CodecPipeline, + config: ArrayConfig, + selection: BasicSelection, + *, + prototype: BufferPrototype | None = None, +) -> NDArrayLikeOrScalar: + """ + Retrieve a subset of the array's data based on the provided selection. + + Parameters + ---------- + store_path : StorePath + The store path of the array. + metadata : ArrayMetadata + The array metadata. + codec_pipeline : CodecPipeline + The codec pipeline for encoding/decoding. + config : ArrayConfig + The array configuration. + selection : BasicSelection + A selection object specifying the subset of data to retrieve. + prototype : BufferPrototype, optional + A buffer prototype to use for the retrieved data (default is None). + + Returns + ------- + NDArrayLikeOrScalar + The retrieved subset of the array's data. + """ + if prototype is None: + prototype = default_buffer_prototype() + indexer = BasicIndexer( + selection, + shape=metadata.shape, + chunk_grid=metadata.chunk_grid, + ) + return await _get_selection( + store_path, metadata, codec_pipeline, config, indexer, prototype=prototype + ) + + +async def _get_orthogonal_selection( + store_path: StorePath, + metadata: ArrayMetadata, + codec_pipeline: CodecPipeline, + config: ArrayConfig, + selection: OrthogonalSelection, + *, + out: NDBuffer | None = None, + fields: Fields | None = None, + prototype: BufferPrototype | None = None, +) -> NDArrayLikeOrScalar: + """ + Get an orthogonal selection from the array. + + Parameters + ---------- + store_path : StorePath + The store path of the array. + metadata : ArrayMetadata + The array metadata. + codec_pipeline : CodecPipeline + The codec pipeline for encoding/decoding. + config : ArrayConfig + The array configuration. + selection : OrthogonalSelection + The orthogonal selection specification. + out : NDBuffer | None, optional + An output buffer to write the data to. + fields : Fields | None, optional + Fields to select from structured arrays. + prototype : BufferPrototype | None, optional + A buffer prototype to use for the retrieved data. + + Returns + ------- + NDArrayLikeOrScalar + The selected data. + """ + if prototype is None: + prototype = default_buffer_prototype() + indexer = OrthogonalIndexer(selection, metadata.shape, metadata.chunk_grid) + return await _get_selection( + store_path, + metadata, + codec_pipeline, + config, + indexer=indexer, + out=out, + fields=fields, + prototype=prototype, + ) + + +async def _get_mask_selection( + store_path: StorePath, + metadata: ArrayMetadata, + codec_pipeline: CodecPipeline, + config: ArrayConfig, + mask: MaskSelection, + *, + out: NDBuffer | None = None, + fields: Fields | None = None, + prototype: BufferPrototype | None = None, +) -> NDArrayLikeOrScalar: + """ + Get a mask selection from the array. + + Parameters + ---------- + store_path : StorePath + The store path of the array. + metadata : ArrayMetadata + The array metadata. + codec_pipeline : CodecPipeline + The codec pipeline for encoding/decoding. + config : ArrayConfig + The array configuration. + mask : MaskSelection + The boolean mask specifying the selection. + out : NDBuffer | None, optional + An output buffer to write the data to. + fields : Fields | None, optional + Fields to select from structured arrays. + prototype : BufferPrototype | None, optional + A buffer prototype to use for the retrieved data. + + Returns + ------- + NDArrayLikeOrScalar + The selected data. + """ + if prototype is None: + prototype = default_buffer_prototype() + indexer = MaskIndexer(mask, metadata.shape, metadata.chunk_grid) + return await _get_selection( + store_path, + metadata, + codec_pipeline, + config, + indexer=indexer, + out=out, + fields=fields, + prototype=prototype, + ) + + +async def _get_coordinate_selection( + store_path: StorePath, + metadata: ArrayMetadata, + codec_pipeline: CodecPipeline, + config: ArrayConfig, + selection: CoordinateSelection, + *, + out: NDBuffer | None = None, + fields: Fields | None = None, + prototype: BufferPrototype | None = None, +) -> NDArrayLikeOrScalar: + """ + Get a coordinate selection from the array. + + Parameters + ---------- + store_path : StorePath + The store path of the array. + metadata : ArrayMetadata + The array metadata. + codec_pipeline : CodecPipeline + The codec pipeline for encoding/decoding. + config : ArrayConfig + The array configuration. + selection : CoordinateSelection + The coordinate selection specification. + out : NDBuffer | None, optional + An output buffer to write the data to. + fields : Fields | None, optional + Fields to select from structured arrays. + prototype : BufferPrototype | None, optional + A buffer prototype to use for the retrieved data. + + Returns + ------- + NDArrayLikeOrScalar + The selected data. + """ + if prototype is None: + prototype = default_buffer_prototype() + indexer = CoordinateIndexer(selection, metadata.shape, metadata.chunk_grid) + out_array = await _get_selection( + store_path, + metadata, + codec_pipeline, + config, + indexer=indexer, + out=out, + fields=fields, + prototype=prototype, + ) + + if hasattr(out_array, "shape"): + # restore shape + out_array = np.array(out_array).reshape(indexer.sel_shape) + return out_array + + +async def _set_selection( + store_path: StorePath, + metadata: ArrayMetadata, + codec_pipeline: CodecPipeline, + config: ArrayConfig, + indexer: Indexer, + value: npt.ArrayLike, + *, + prototype: BufferPrototype, + fields: Fields | None = None, +) -> None: + """ + Set a selection in an array. + + Parameters + ---------- + store_path : StorePath + The store path of the array. + metadata : ArrayMetadata + The array metadata. + codec_pipeline : CodecPipeline + The codec pipeline for encoding/decoding. + config : ArrayConfig + The array configuration. + indexer : Indexer + The indexer specifying the selection. + value : npt.ArrayLike + The values to write. + prototype : BufferPrototype + A buffer prototype to use. + fields : Fields | None, optional + Fields to select from structured arrays. + """ + # Get dtype from metadata + if metadata.zarr_format == 2: + zdtype = metadata.dtype + else: + zdtype = metadata.data_type + dtype = zdtype.to_native_dtype() + + # check fields are sensible + check_fields(fields, dtype) + fields = check_no_multi_fields(fields) + + # check value shape + if np.isscalar(value): + array_like = prototype.buffer.create_zero_length().as_array_like() + if isinstance(array_like, np._typing._SupportsArrayFunc): + # TODO: need to handle array types that don't support __array_function__ + # like PyTorch and JAX + array_like_ = cast("np._typing._SupportsArrayFunc", array_like) + value = np.asanyarray(value, dtype=dtype, like=array_like_) + else: + if not hasattr(value, "shape"): + value = np.asarray(value, dtype) + # assert ( + # value.shape == indexer.shape + # ), f"shape of value doesn't match indexer shape. Expected {indexer.shape}, got {value.shape}" + if not hasattr(value, "dtype") or value.dtype.name != dtype.name: + if hasattr(value, "astype"): + # Handle things that are already NDArrayLike more efficiently + value = value.astype(dtype=dtype, order="A") + else: + value = np.array(value, dtype=dtype, order="A") + value = cast("NDArrayLike", value) + + # We accept any ndarray like object from the user and convert it + # to an NDBuffer (or subclass). From this point onwards, we only pass + # Buffer and NDBuffer between components. + value_buffer = prototype.nd_buffer.from_ndarray_like(value) + + # Determine memory order + if metadata.zarr_format == 2: + order = metadata.order + else: + order = config.order + + # need to use the order from the metadata for v2 + _config = config + if metadata.zarr_format == 2: + _config = replace(_config, order=order) + + # merging with existing data and encoding chunks + await codec_pipeline.write( + [ + ( + store_path / metadata.encode_chunk_key(chunk_coords), + metadata.get_chunk_spec(chunk_coords, _config, prototype), + chunk_selection, + out_selection, + is_complete_chunk, + ) + for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexer + ], + value_buffer, + drop_axes=indexer.drop_axes, + ) + + +async def _setitem( + store_path: StorePath, + metadata: ArrayMetadata, + codec_pipeline: CodecPipeline, + config: ArrayConfig, + selection: BasicSelection, + value: npt.ArrayLike, + prototype: BufferPrototype | None = None, +) -> None: + """ + Set values in the array using basic indexing. + + Parameters + ---------- + store_path : StorePath + The store path of the array. + metadata : ArrayMetadata + The array metadata. + codec_pipeline : CodecPipeline + The codec pipeline for encoding/decoding. + config : ArrayConfig + The array configuration. + selection : BasicSelection + The selection defining the region of the array to set. + value : npt.ArrayLike + The values to be written into the selected region of the array. + prototype : BufferPrototype or None, optional + A prototype buffer that defines the structure and properties of the array chunks being modified. + If None, the default buffer prototype is used. + """ + if prototype is None: + prototype = default_buffer_prototype() + indexer = BasicIndexer( + selection, + shape=metadata.shape, + chunk_grid=metadata.chunk_grid, + ) + return await _set_selection( + store_path, metadata, codec_pipeline, config, indexer, value, prototype=prototype + ) + + +async def _resize( + array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + new_shape: ShapeLike, + delete_outside_chunks: bool = True, +) -> None: + """ + Resize an array to a new shape. + + Parameters + ---------- + array : AsyncArray + The array to resize. + new_shape : ShapeLike + The desired new shape of the array. + delete_outside_chunks : bool, optional + 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) + assert len(new_shape) == len(array.metadata.shape) + new_metadata = array.metadata.update_shape(new_shape) + + if delete_outside_chunks: + # Remove all chunks outside of the new shape + old_chunk_coords = set(array.metadata.chunk_grid.all_chunk_coords(array.metadata.shape)) + new_chunk_coords = set(array.metadata.chunk_grid.all_chunk_coords(new_shape)) + + 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"), + ) + + # Write new metadata + await save_metadata(array.store_path, new_metadata) + + # Update metadata (in place) + object.__setattr__(array, "metadata", new_metadata) + + +async def _append( + array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + data: npt.ArrayLike, + axis: int = 0, +) -> tuple[int, ...]: + """ + Append data to an array along the specified axis. + + Parameters + ---------- + array : AsyncArray + The array to append to. + data : npt.ArrayLike + Data to be appended. + axis : int + Axis along which to append. + + Returns + ------- + new_shape : tuple[int, ...] + The new shape of the array after appending. + + Notes + ----- + The size of all dimensions other than `axis` must match between the + array and `data`. + """ + # ensure data is array-like + 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 + 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)) + ) + await _setitem( + array.store_path, + array.metadata, + array.codec_pipeline, + array._config, + append_selection, + data, + ) + + return new_shape + + +async def _update_attributes( + array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + new_attributes: dict[str, JSON], +) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + """ + Update the array's attributes. + + Parameters + ---------- + 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 + + +async def _info_complete( + array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], +) -> Any: + """ + Return all the information for an array, including dynamic information like storage size. + + Parameters + ---------- + array : AsyncArray + The array to get info for. + + Returns + ------- + ArrayInfo + Complete information about the array including: + - The count of chunks initialized + - The sum of the bytes written + """ + return array._info( + await _nshards_initialized(array), + await array.store_path.store.getsize_prefix(array.store_path.path), + ) From 9399b7fb502c17678e594f17febcf78ce48a0c1c Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 28 Jan 2026 20:21:04 +0100 Subject: [PATCH 2/2] create experimental array class --- src/zarr/experimental/array.py | 955 +++++++++++++++++++++++++++++++++ 1 file changed, 955 insertions(+) create mode 100644 src/zarr/experimental/array.py diff --git a/src/zarr/experimental/array.py b/src/zarr/experimental/array.py new file mode 100644 index 0000000000..4f5bda7302 --- /dev/null +++ b/src/zarr/experimental/array.py @@ -0,0 +1,955 @@ +from __future__ import annotations + +from itertools import starmap +from typing import TYPE_CHECKING, Any + +import numpy as np + +from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec +from zarr.abc.numcodec import Numcodec +from zarr.core._info import ArrayInfo +from zarr.core.array import ( + _append, + _get_coordinate_selection, + _get_mask_selection, + _get_orthogonal_selection, + _getitem, + _info_complete, + _iter_chunk_coords, + _iter_chunk_regions, + _iter_shard_coords, + _iter_shard_keys, + _iter_shard_regions, + _nbytes_stored, + _nchunks_initialized, + _nshards_initialized, + _resize, + _setitem, + _update_attributes, + create_codec_pipeline, + get_array_metadata, + parse_array_metadata, +) +from zarr.core.array_spec import ArrayConfig, ArrayConfigLike, parse_array_config +from zarr.core.buffer import ( + BufferPrototype, + NDArrayLikeOrScalar, + NDBuffer, +) +from zarr.core.common import ( + JSON, + MemoryOrder, + ShapeLike, + ZarrFormat, + ceildiv, + product, +) +from zarr.core.indexing import ( + BasicSelection, + CoordinateSelection, + Fields, + MaskSelection, + OrthogonalSelection, +) +from zarr.core.metadata import ( + ArrayMetadata, + ArrayMetadataDict, + ArrayV2Metadata, + ArrayV3Metadata, +) +from zarr.core.sync import sync +from zarr.storage._common import StorePath, make_store_path + +if TYPE_CHECKING: + from collections.abc import Iterator, Sequence + from typing import Self + + import numpy.typing as npt + + from zarr.abc.codec import CodecPipeline + from zarr.abc.store import Store + from zarr.storage import StoreLike + + +class Array: + """ + A unified Zarr array class with both synchronous and asynchronous methods. + + This class combines the functionality of AsyncArray and Array into a single class. + For each operation, there is both a synchronous method (e.g., `getitem`) and an + asynchronous method (e.g., `getitem_async`). + + Parameters + ---------- + metadata : ArrayV2Metadata | ArrayV3Metadata + 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 : ArrayV2Metadata | ArrayV3Metadata + 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. + """ + + metadata: ArrayV2Metadata | ArrayV3Metadata + store_path: StorePath + codec_pipeline: CodecPipeline + config: ArrayConfig + + def __init__( + self, + store_path: StorePath, + metadata: ArrayMetadata | ArrayMetadataDict, + *, + codec_pipeline: CodecPipeline | None = None, + config: ArrayConfigLike | None = None, + ) -> None: + metadata_parsed = parse_array_metadata(metadata) + config_parsed = parse_array_config(config) + + if codec_pipeline is None: + codec_pipeline = create_codec_pipeline(metadata=metadata_parsed, store=store_path.store) + + self.metadata = metadata_parsed + self.store_path = store_path + self.config = config_parsed + self.codec_pipeline = codec_pipeline + + # ------------------------------------------------------------------------- + # Class methods: open + # ------------------------------------------------------------------------- + + @classmethod + async def open_async( + cls, + store: StoreLike, + *, + config: ArrayConfigLike | None = None, + codec_pipeline: CodecPipeline | None = None, + zarr_format: ZarrFormat | None = 3, + ) -> Array: + """ + Async method to open an existing Zarr array from a given store. + + Parameters + ---------- + store : StoreLike + The store containing the Zarr array. + zarr_format : ZarrFormat | None, optional + The Zarr format version (default is 3). + + Returns + ------- + Array + The opened Zarr array. + """ + store_path = await make_store_path(store) + metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format) + return cls( + store_path=store_path, + metadata=metadata_dict, + codec_pipeline=codec_pipeline, + config=config, + ) + + @classmethod + def open( + cls, + store: StoreLike, + *, + config: ArrayConfigLike | None = None, + codec_pipeline: CodecPipeline | None = None, + zarr_format: ZarrFormat | None = 3, + ) -> Array: + """ + Open an existing Zarr array from a given store. + + Parameters + ---------- + store : StoreLike + The store containing the Zarr array. + zarr_format : ZarrFormat | None, optional + The Zarr format version (default is 3). + + Returns + ------- + Array + The opened Zarr array. + """ + return sync(cls.open_async(store, zarr_format=zarr_format)) + + # ------------------------------------------------------------------------- + # Properties (all synchronous, derived from metadata/store_path) + # ------------------------------------------------------------------------- + + @property + def store(self) -> Store: + """The store containing the array data.""" + return self.store_path.store + + @property + def ndim(self) -> int: + """Returns the number of dimensions in the Array.""" + return len(self.metadata.shape) + + @property + def shape(self) -> tuple[int, ...]: + """Returns the shape of the Array.""" + return self.metadata.shape + + @property + def chunks(self) -> tuple[int, ...]: + """Returns the chunk shape of the Array.""" + return self.metadata.chunks + + @property + def shards(self) -> tuple[int, ...] | None: + """Returns the shard shape of the Array, or None if sharding is not used.""" + return self.metadata.shards + + @property + def size(self) -> int: + """Returns the total number of elements in the array.""" + return np.prod(self.metadata.shape).item() + + @property + def filters(self) -> tuple[Numcodec, ...] | tuple[ArrayArrayCodec, ...]: + """Filters applied to each chunk before serialization.""" + if self.metadata.zarr_format == 2: + filters = self.metadata.filters + if filters is None: + return () + return filters + return tuple( + codec for codec in self.metadata.inner_codecs if isinstance(codec, ArrayArrayCodec) + ) + + @property + def serializer(self) -> ArrayBytesCodec | None: + """Array-to-bytes codec for serializing chunks.""" + if self.metadata.zarr_format == 2: + return None + return next( + codec for codec in self.metadata.inner_codecs if isinstance(codec, ArrayBytesCodec) + ) + + @property + def compressors(self) -> tuple[Numcodec, ...] | tuple[BytesBytesCodec, ...]: + """Compressors applied to each chunk after serialization.""" + if self.metadata.zarr_format == 2: + if self.metadata.compressor is not None: + return (self.metadata.compressor,) + return () + return tuple( + codec for codec in self.metadata.inner_codecs if isinstance(codec, BytesBytesCodec) + ) + + @property + def _zdtype(self) -> Any: + """The zarr-specific representation of the array data type.""" + if self.metadata.zarr_format == 2: + return self.metadata.dtype + else: + return self.metadata.data_type + + @property + def dtype(self) -> np.dtype[Any]: + """Returns the data type of the array.""" + return self._zdtype.to_native_dtype() + + @property + def order(self) -> MemoryOrder: + """Returns the memory order of the array.""" + if self.metadata.zarr_format == 2: + return self.metadata.order + else: + return self.config.order + + @property + def attrs(self) -> dict[str, JSON]: + """Returns the attributes of the array.""" + return self.metadata.attributes + + @property + def read_only(self) -> bool: + """Returns True if the array is read-only.""" + return self.store_path.read_only + + @property + def path(self) -> str: + """Storage path.""" + return self.store_path.path + + @property + def name(self) -> str: + """Array name following h5py convention.""" + name = self.path + if not name.startswith("/"): + name = "/" + name + return name + + @property + def basename(self) -> str: + """Final component of name.""" + return self.name.split("/")[-1] + + @property + def cdata_shape(self) -> tuple[int, ...]: + """The shape of the chunk grid for this array.""" + return self._chunk_grid_shape + + @property + def _chunk_grid_shape(self) -> tuple[int, ...]: + """The shape of the chunk grid for this array.""" + return tuple(starmap(ceildiv, zip(self.shape, self.chunks, strict=True))) + + @property + def _shard_grid_shape(self) -> tuple[int, ...]: + """The shape of the shard grid for this array.""" + if self.shards is None: + shard_shape = self.chunks + else: + shard_shape = self.shards + return tuple(starmap(ceildiv, zip(self.shape, shard_shape, strict=True))) + + @property + def nchunks(self) -> int: + """The number of chunks in this array.""" + return product(self._chunk_grid_shape) + + @property + def _nshards(self) -> int: + """The number of shards in this array.""" + return product(self._shard_grid_shape) + + @property + def nbytes(self) -> int: + """The total number of bytes that would be stored if all chunks were initialized.""" + return self.size * self.dtype.itemsize + + @property + def info(self) -> ArrayInfo: + """Return the statically known information for an array.""" + return self._info() + + def _info( + self, count_chunks_initialized: int | None = None, count_bytes_stored: int | None = None + ) -> ArrayInfo: + 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=self.shards, + _chunk_shape=self.chunks, + _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, + ) + + # ------------------------------------------------------------------------- + # Iteration methods (synchronous) + # ------------------------------------------------------------------------- + + def _iter_chunk_coords( + self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None + ) -> Iterator[tuple[int, ...]]: + """Iterate over chunk coordinates in chunk grid space.""" + 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, ...]]: + """Iterate over shard coordinates in shard grid space.""" + 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 stored objects supporting this array.""" + 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 chunk regions in array index space.""" + 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 shard regions in array index space.""" + return _iter_shard_regions(array=self, origin=origin, selection_shape=selection_shape) + + # ------------------------------------------------------------------------- + # nchunks_initialized: async and sync + # ------------------------------------------------------------------------- + + async def nchunks_initialized_async(self) -> int: + """ + Asynchronously calculate the number of chunks that have been initialized. + + Returns + ------- + int + The number of chunks that have been initialized. + """ + return await _nchunks_initialized(self) + + def nchunks_initialized(self) -> int: + """ + Calculate the number of chunks that have been initialized. + + Returns + ------- + int + The number of chunks that have been initialized. + """ + return sync(self.nchunks_initialized_async()) + + # ------------------------------------------------------------------------- + # _nshards_initialized: async and sync + # ------------------------------------------------------------------------- + + async def _nshards_initialized_async(self) -> int: + """ + Asynchronously calculate the number of shards that have been initialized. + + Returns + ------- + int + The number of shards that have been initialized. + """ + return await _nshards_initialized(self) + + def _nshards_initialized(self) -> int: + """ + Calculate the number of shards that have been initialized. + + Returns + ------- + int + The number of shards that have been initialized. + """ + return sync(self._nshards_initialized_async()) + + # ------------------------------------------------------------------------- + # nbytes_stored: async and sync + # ------------------------------------------------------------------------- + + async def nbytes_stored_async(self) -> int: + """ + Asynchronously calculate the number of bytes stored for this array. + + Returns + ------- + int + The number of bytes stored. + """ + return await _nbytes_stored(self.store_path) + + def nbytes_stored(self) -> int: + """ + Calculate the number of bytes stored for this array. + + Returns + ------- + int + The number of bytes stored. + """ + return sync(self.nbytes_stored_async()) + + # ------------------------------------------------------------------------- + # getitem: async and sync + # ------------------------------------------------------------------------- + + async def getitem_async( + self, + selection: BasicSelection, + *, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """ + Asynchronously retrieve a subset of the array's data based on the provided selection. + + Parameters + ---------- + selection : BasicSelection + A selection object specifying the subset of data to retrieve. + prototype : BufferPrototype, optional + A buffer prototype to use for the retrieved data. + + Returns + ------- + NDArrayLikeOrScalar + The retrieved subset of the array's data. + """ + return await _getitem( + self.store_path, + self.metadata, + self.codec_pipeline, + self.config, + selection, + prototype=prototype, + ) + + def getitem( + self, + selection: BasicSelection, + *, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """ + Retrieve a subset of the array's data based on the provided selection. + + Parameters + ---------- + selection : BasicSelection + A selection object specifying the subset of data to retrieve. + prototype : BufferPrototype, optional + A buffer prototype to use for the retrieved data. + + Returns + ------- + NDArrayLikeOrScalar + The retrieved subset of the array's data. + """ + return sync(self.getitem_async(selection, prototype=prototype)) + + def __getitem__(self, selection: BasicSelection) -> NDArrayLikeOrScalar: + """Retrieve data using indexing syntax.""" + return self.getitem(selection) + + # ------------------------------------------------------------------------- + # setitem: async and sync + # ------------------------------------------------------------------------- + + async def setitem_async( + self, + selection: BasicSelection, + value: npt.ArrayLike, + prototype: BufferPrototype | None = None, + ) -> None: + """ + Asynchronously set values in the array using basic indexing. + + Parameters + ---------- + selection : BasicSelection + The selection defining the region of the array to set. + value : npt.ArrayLike + The values to be written into the selected region. + prototype : BufferPrototype, optional + A buffer prototype to use. + """ + return await _setitem( + self.store_path, + self.metadata, + self.codec_pipeline, + self.config, + selection, + value, + prototype=prototype, + ) + + def setitem( + self, + selection: BasicSelection, + value: npt.ArrayLike, + prototype: BufferPrototype | None = None, + ) -> None: + """ + Set values in the array using basic indexing. + + Parameters + ---------- + selection : BasicSelection + The selection defining the region of the array to set. + value : npt.ArrayLike + The values to be written into the selected region. + prototype : BufferPrototype, optional + A buffer prototype to use. + """ + sync(self.setitem_async(selection, value, prototype=prototype)) + + def __setitem__(self, selection: BasicSelection, value: npt.ArrayLike) -> None: + """Set data using indexing syntax.""" + self.setitem(selection, value) + + # ------------------------------------------------------------------------- + # get_orthogonal_selection: async and sync + # ------------------------------------------------------------------------- + + async def get_orthogonal_selection_async( + self, + selection: OrthogonalSelection, + *, + out: NDBuffer | None = None, + fields: Fields | None = None, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """ + Asynchronously get an orthogonal selection from the array. + + Parameters + ---------- + selection : OrthogonalSelection + The orthogonal selection specification. + out : NDBuffer | None, optional + An output buffer to write the data to. + fields : Fields | None, optional + Fields to select from structured arrays. + prototype : BufferPrototype | None, optional + A buffer prototype to use for the retrieved data. + + Returns + ------- + NDArrayLikeOrScalar + The selected data. + """ + return await _get_orthogonal_selection( + self.store_path, + self.metadata, + self.codec_pipeline, + self.config, + selection, + out=out, + fields=fields, + prototype=prototype, + ) + + def get_orthogonal_selection( + self, + selection: OrthogonalSelection, + *, + out: NDBuffer | None = None, + fields: Fields | None = None, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """ + Get an orthogonal selection from the array. + + Parameters + ---------- + selection : OrthogonalSelection + The orthogonal selection specification. + out : NDBuffer | None, optional + An output buffer to write the data to. + fields : Fields | None, optional + Fields to select from structured arrays. + prototype : BufferPrototype | None, optional + A buffer prototype to use for the retrieved data. + + Returns + ------- + NDArrayLikeOrScalar + The selected data. + """ + return sync( + self.get_orthogonal_selection_async( + selection, out=out, fields=fields, prototype=prototype + ) + ) + + # ------------------------------------------------------------------------- + # get_mask_selection: async and sync + # ------------------------------------------------------------------------- + + async def get_mask_selection_async( + self, + mask: MaskSelection, + *, + out: NDBuffer | None = None, + fields: Fields | None = None, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """ + Asynchronously get a mask selection from the array. + + Parameters + ---------- + mask : MaskSelection + The boolean mask specifying the selection. + out : NDBuffer | None, optional + An output buffer to write the data to. + fields : Fields | None, optional + Fields to select from structured arrays. + prototype : BufferPrototype | None, optional + A buffer prototype to use for the retrieved data. + + Returns + ------- + NDArrayLikeOrScalar + The selected data. + """ + return await _get_mask_selection( + self.store_path, + self.metadata, + self.codec_pipeline, + self.config, + mask, + out=out, + fields=fields, + prototype=prototype, + ) + + def get_mask_selection( + self, + mask: MaskSelection, + *, + out: NDBuffer | None = None, + fields: Fields | None = None, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """ + Get a mask selection from the array. + + Parameters + ---------- + mask : MaskSelection + The boolean mask specifying the selection. + out : NDBuffer | None, optional + An output buffer to write the data to. + fields : Fields | None, optional + Fields to select from structured arrays. + prototype : BufferPrototype | None, optional + A buffer prototype to use for the retrieved data. + + Returns + ------- + NDArrayLikeOrScalar + The selected data. + """ + return sync( + self.get_mask_selection_async(mask, out=out, fields=fields, prototype=prototype) + ) + + # ------------------------------------------------------------------------- + # get_coordinate_selection: async and sync + # ------------------------------------------------------------------------- + + async def get_coordinate_selection_async( + self, + selection: CoordinateSelection, + *, + out: NDBuffer | None = None, + fields: Fields | None = None, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """ + Asynchronously get a coordinate selection from the array. + + Parameters + ---------- + selection : CoordinateSelection + The coordinate selection specification. + out : NDBuffer | None, optional + An output buffer to write the data to. + fields : Fields | None, optional + Fields to select from structured arrays. + prototype : BufferPrototype | None, optional + A buffer prototype to use for the retrieved data. + + Returns + ------- + NDArrayLikeOrScalar + The selected data. + """ + return await _get_coordinate_selection( + self.store_path, + self.metadata, + self.codec_pipeline, + self.config, + selection, + out=out, + fields=fields, + prototype=prototype, + ) + + def get_coordinate_selection( + self, + selection: CoordinateSelection, + *, + out: NDBuffer | None = None, + fields: Fields | None = None, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """ + Get a coordinate selection from the array. + + Parameters + ---------- + selection : CoordinateSelection + The coordinate selection specification. + out : NDBuffer | None, optional + An output buffer to write the data to. + fields : Fields | None, optional + Fields to select from structured arrays. + prototype : BufferPrototype | None, optional + A buffer prototype to use for the retrieved data. + + Returns + ------- + NDArrayLikeOrScalar + The selected data. + """ + return sync( + self.get_coordinate_selection_async( + selection, out=out, fields=fields, prototype=prototype + ) + ) + + # ------------------------------------------------------------------------- + # resize: async and sync + # ------------------------------------------------------------------------- + + async def resize_async(self, new_shape: ShapeLike, delete_outside_chunks: bool = True) -> None: + """ + Asynchronously resize the array to a new shape. + + Parameters + ---------- + new_shape : ShapeLike + The desired new shape of the array. + delete_outside_chunks : bool, optional + If True (default), chunks that fall outside the new shape will be deleted. + """ + return await _resize(self, new_shape, delete_outside_chunks) + + def resize(self, new_shape: ShapeLike, delete_outside_chunks: bool = True) -> None: + """ + Resize the array to a new shape. + + Parameters + ---------- + new_shape : ShapeLike + The desired new shape of the array. + delete_outside_chunks : bool, optional + If True (default), chunks that fall outside the new shape will be deleted. + """ + sync(self.resize_async(new_shape, delete_outside_chunks)) + + # ------------------------------------------------------------------------- + # append: async and sync + # ------------------------------------------------------------------------- + + async def append_async(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: + """ + Asynchronously append data to the array along the specified axis. + + Parameters + ---------- + data : npt.ArrayLike + Data to be appended. + axis : int + Axis along which to append. + + Returns + ------- + tuple[int, ...] + The new shape of the array after appending. + """ + return await _append(self, data, axis) + + def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: + """ + Append data to the array along the specified axis. + + Parameters + ---------- + data : npt.ArrayLike + Data to be appended. + axis : int + Axis along which to append. + + Returns + ------- + tuple[int, ...] + The new shape of the array after appending. + """ + return sync(self.append_async(data, axis)) + + # ------------------------------------------------------------------------- + # update_attributes: async and sync + # ------------------------------------------------------------------------- + + async def update_attributes_async(self, new_attributes: dict[str, JSON]) -> Self: + """ + Asynchronously update the array's attributes. + + Parameters + ---------- + new_attributes : dict[str, JSON] + A dictionary of new attributes to update or add. + + Returns + ------- + Array + The array with the updated attributes. + """ + await _update_attributes(self, new_attributes) + return self + + def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: + """ + Update the array's attributes. + + Parameters + ---------- + new_attributes : dict[str, JSON] + A dictionary of new attributes to update or add. + + Returns + ------- + Array + The array with the updated attributes. + """ + return sync(self.update_attributes_async(new_attributes)) + + # ------------------------------------------------------------------------- + # info_complete: async and sync + # ------------------------------------------------------------------------- + + async def info_complete_async(self) -> ArrayInfo: + """ + Asynchronously return all the information for an array, including dynamic information. + + Returns + ------- + ArrayInfo + Complete information about the array including chunks initialized and bytes stored. + """ + return await _info_complete(self) + + def info_complete(self) -> ArrayInfo: + """ + Return all the information for an array, including dynamic information. + + Returns + ------- + ArrayInfo + Complete information about the array including chunks initialized and bytes stored. + """ + return sync(self.info_complete_async()) + + # ------------------------------------------------------------------------- + # __repr__ + # ------------------------------------------------------------------------- + + def __repr__(self) -> str: + return f""