Add array storage helpers

pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
     'asciitree',
     'numpy>=1.24',
     'fasteners',
-    'numcodecs>=0.10.0',
+    'numcodecs>=0.13.0',
     'fsspec>2024',
     'crc32c',
     'typing_extensions',
@@ -273,6 +273,7 @@ filterwarnings = [
     "error:::zarr.*",
     "ignore:PY_SSIZE_T_CLEAN will be required.*:DeprecationWarning",
     "ignore:The loop argument is deprecated since Python 3.8.*:DeprecationWarning",
+    "ignore:.*may be removed in an early zarr-python v3 release.:DeprecationWarning",
     "ignore:Creating a zarr.buffer.gpu.*:UserWarning",
     "ignore:Duplicate name:UserWarning",  # from ZipFile
 ]

src/zarr/abc/store.py

@@ -1,15 +1,22 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from asyncio import gather
-from collections.abc import AsyncGenerator, Iterable
-from typing import Any, NamedTuple, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Any, NamedTuple, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Iterable
+    from typing import Any, TypeAlias
 
-from typing_extensions import Self
+    from typing_extensions import Self
 
-from zarr.core.buffer import Buffer, BufferPrototype
-from zarr.core.common import AccessModeLiteral, BytesLike
+    from zarr.core.buffer import Buffer, BufferPrototype
+    from zarr.core.common import AccessModeLiteral, BytesLike
 
 __all__ = ["Store", "AccessMode", "ByteGetter", "ByteSetter", "set_or_delete"]
 
+ByteRangeRequest: TypeAlias = tuple[int | None, int | None]
+
 
 class AccessMode(NamedTuple):
     str: AccessModeLiteral
@@ -94,7 +101,7 @@ async def get(
         self,
         key: str,
         prototype: BufferPrototype,
-        byte_range: tuple[int | None, int | None] | None = None,
+        byte_range: ByteRangeRequest | None = None,
     ) -> Buffer | None:
         """Retrieve the value associated with a given key.
 
@@ -113,13 +120,13 @@ async def get(
     async def get_partial_values(
         self,
         prototype: BufferPrototype,
-        key_ranges: list[tuple[str, tuple[int | None, int | None]]],
+        key_ranges: Iterable[tuple[str, ByteRangeRequest]],
     ) -> list[Buffer | None]:
         """Retrieve possibly partial values from given key_ranges.
 
         Parameters
         ----------
-        key_ranges : list[tuple[str, tuple[int, int]]]
+        key_ranges : list[tuple[str, tuple[int | None, int | None]]]
             Ordered set of key, range pairs, a key may occur multiple times with different ranges
 
         Returns
@@ -189,7 +196,9 @@ def supports_partial_writes(self) -> bool:
         ...
 
     @abstractmethod
-    async def set_partial_values(self, key_start_values: list[tuple[str, int, BytesLike]]) -> None:
+    async def set_partial_values(
+        self, key_start_values: Iterable[tuple[str, int, BytesLike]]
+    ) -> None:
         """Store values at a given key, starting at byte range_start.
 
         Parameters
@@ -253,21 +262,32 @@ def close(self) -> None:
         """Close the store."""
         self._is_open = False
 
+    async def _get_many(
+        self, requests: Iterable[tuple[str, BufferPrototype, ByteRangeRequest | None]]
+    ) -> AsyncGenerator[tuple[str, Buffer | None], None]:
+        """
+        Retrieve a collection of objects from storage. In general this method does not guarantee
+        that objects will be retrieved in the order in which they were requested, so this method
+        yields tuple[str, Buffer | None] instead of just Buffer | None
+        """
+        for req in requests:
+            yield (req[0], await self.get(*req))
+
 
 @runtime_checkable
 class ByteGetter(Protocol):
     async def get(
-        self, prototype: BufferPrototype, byte_range: tuple[int, int | None] | None = None
+        self, prototype: BufferPrototype, byte_range: ByteRangeRequest | None = None
     ) -> Buffer | None: ...
 
 
 @runtime_checkable
 class ByteSetter(Protocol):
     async def get(
-        self, prototype: BufferPrototype, byte_range: tuple[int, int | None] | None = None
+        self, prototype: BufferPrototype, byte_range: ByteRangeRequest | None = None
     ) -> Buffer | None: ...
 
-    async def set(self, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: ...
+    async def set(self, value: Buffer, byte_range: ByteRangeRequest | None = None) -> None: ...
 
     async def delete(self) -> None: ...
 

src/zarr/codecs/sharding.py

@@ -17,7 +17,7 @@
     Codec,
     CodecPipeline,
 )
-from zarr.abc.store import ByteGetter, ByteSetter
+from zarr.abc.store import ByteGetter, ByteRangeRequest, ByteSetter
 from zarr.codecs.bytes import BytesCodec
 from zarr.codecs.crc32c_ import Crc32cCodec
 from zarr.core.array_spec import ArraySpec
@@ -78,7 +78,7 @@ class _ShardingByteGetter(ByteGetter):
     chunk_coords: ChunkCoords
 
     async def get(
-        self, prototype: BufferPrototype, byte_range: tuple[int, int | None] | None = None
+        self, prototype: BufferPrototype, byte_range: ByteRangeRequest | None = None
     ) -> Buffer | None:
         assert byte_range is None, "byte_range is not supported within shards"
         assert (
@@ -91,7 +91,7 @@ async def get(
 class _ShardingByteSetter(_ShardingByteGetter, ByteSetter):
     shard_dict: ShardMutableMapping
 
-    async def set(self, value: Buffer, byte_range: tuple[int, int] | None = None) -> None:
+    async def set(self, value: Buffer, byte_range: ByteRangeRequest | None = None) -> None:
         assert byte_range is None, "byte_range is not supported within shards"
         self.shard_dict[self.chunk_coords] = value
 

src/zarr/core/array.py

@@ -2,11 +2,13 @@
 
 import json
 from asyncio import gather
+from collections.abc import Iterator, Sequence
 from dataclasses import dataclass, field, replace
 from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 import numpy.typing as npt
+from typing_extensions import deprecated
 
 from zarr._compat import _deprecate_positional_args
 from zarr.abc.codec import Codec, CodecPipeline
@@ -50,6 +52,8 @@
     OrthogonalSelection,
     Selection,
     VIndex,
+    _iter_grid,
+    ceildiv,
     check_fields,
     check_no_multi_fields,
     is_pure_fancy_indexing,
@@ -59,7 +63,7 @@
 )
 from zarr.core.metadata.v2 import ArrayV2Metadata
 from zarr.core.metadata.v3 import ArrayV3Metadata
-from zarr.core.sync import sync
+from zarr.core.sync import collect_aiterator, sync
 from zarr.registry import get_pipeline_class
 from zarr.store import StoreLike, StorePath, make_store_path
 from zarr.store.common import (
@@ -399,10 +403,12 @@ def shape(self) -> ChunkCoords:
     def chunks(self) -> ChunkCoords:
         if isinstance(self.metadata.chunk_grid, RegularChunkGrid):
             return self.metadata.chunk_grid.chunk_shape
-        else:
-            raise TypeError(
-                f"chunk attribute is only available for RegularChunkGrid, this array has a {self.metadata.chunk_grid}"
-            )
+
+        msg = (
+            f"The `chunks` attribute is only defined for arrays using `RegularChunkGrid`."
+            f"This array has a {self.metadata.chunk_grid} instead."
+        )
+        raise NotImplementedError(msg)
 
     @property
     def size(self) -> int:
@@ -443,6 +449,55 @@ def basename(self) -> str | None:
             return self.name.split("/")[-1]
         return None
 
+    @property
+    @deprecated("AsyncArray.cdata_shape may be removed in an early zarr-python v3 release.")
+    def cdata_shape(self) -> ChunkCoords:
+        """
+        The shape of the chunk grid for this array.
+        """
+        return tuple(ceildiv(s, c) for s, c in zip(self.shape, self.chunks, strict=False))
+
+    @property
+    @deprecated("AsyncArray.nchunks may be removed in an early zarr-python v3 release.")
+    def nchunks(self) -> int:
+        """
+        The number of chunks in the stored representation of this array.
+        """
+        return product(self.cdata_shape)
+
+    def _iter_chunk_coords(self, origin: Sequence[int] | None = None) -> Iterator[ChunkCoords]:
+        """
+        Produce an iterator over the coordinates of each chunk, in chunk grid space, relative to
+        an optional origin.
+        """
+        return _iter_grid(self.cdata_shape, origin=origin)
+
+    def _iter_chunk_keys(self, origin: Sequence[int] | None = None) -> Iterator[str]:
+        """
+        Return an iterator over the storage keys of each chunk, relative to an optional origin.
+        """
+        for k in self._iter_chunk_coords(origin=origin):
+            yield self.metadata.encode_chunk_key(k)
+
+    def _iter_chunk_regions(self) -> Iterator[tuple[slice, ...]]:
+        """
+        Iterate over the regions spanned by each chunk.
+        """
+        for cgrid_position in self._iter_chunk_coords():
+            out: tuple[slice, ...] = ()
+            for c_pos, c_shape in zip(cgrid_position, self.chunks, strict=False):
+                start = c_pos * c_shape
+                stop = start + c_shape
+                out += (slice(start, stop, 1),)
+            yield out
+
+    @property
+    def nbytes(self) -> int:
+        """
+        The number of bytes that can be stored in this array.
+        """
+        return self.nchunks * self.dtype.itemsize
+
     async def _get_selection(
         self,
         indexer: Indexer,
@@ -751,6 +806,50 @@ def read_only(self) -> bool:
     def fill_value(self) -> Any:
         return self.metadata.fill_value
 
+    @property
+    @deprecated(
+        "cdata_shape is transitional and will be removed in an early zarr-python v3 release."
+    )
+    def cdata_shape(self) -> ChunkCoords:
+        """
+        The shape of the chunk grid for this array.
+        """
+        return tuple(ceildiv(s, c) for s, c in zip(self.shape, self.chunks, strict=False))
+
+    @property
+    @deprecated("nchunks is transitional and will be removed in an early zarr-python v3 release.")
+    def nchunks(self) -> int:
+        """
+        The number of chunks in the stored representation of this array.
+        """
+        return self._async_array.nchunks
+
+    def _iter_chunks(self, origin: Sequence[int] | None = None) -> Iterator[ChunkCoords]:
+        """
+        Produce an iterator over the coordinates of each chunk, in chunk grid space, relative
+        to an optional origin.
+        """
+        yield from self._async_array._iter_chunk_coords(origin=origin)
+
+    @property
+    def nbytes(self) -> int:
+        """
+        The number of bytes that can be stored in this array.
+        """
+        return self._async_array.nbytes
+
+    def _iter_chunk_keys(self, origin: Sequence[int] | None = None) -> Iterator[str]:
+        """
+        Return an iterator over the keys of each chunk, relative to an optional origin
+        """
+        yield from self._async_array._iter_chunk_keys(origin=origin)
+
+    def _iter_chunk_regions(self) -> Iterator[tuple[slice, ...]]:
+        """
+        Iterate over the regions spanned by each chunk.
+        """
+        yield from self._async_array._iter_chunk_regions()
+
     def __array__(
         self, dtype: npt.DTypeLike | None = None, copy: bool | None = None
     ) -> NDArrayLike:
@@ -2082,3 +2181,27 @@ def info(self) -> None:
         return sync(
             self._async_array.info(),
         )
+
+
+@deprecated(
+    "nchunks_initialized is transitional and will be removed in an early zarr-python v3 release."
+)
+def nchunks_initialized(array: Array) -> int:
+    return len(chunks_initialized(array))
+
+
+def chunks_initialized(array: Array) -> tuple[str, ...]:
+    """
+    Return the keys of all the chunks that exist in storage.
+    """
+    # todo: make this compose with the underlying async iterator
+    store_contents = list(
+        collect_aiterator(array.store_path.store.list_prefix(prefix=array.store_path.path))
+    )
+    out: list[str] = []
+
+    for chunk_key in array._iter_chunk_keys():
+        if chunk_key in store_contents:
+            out.append(chunk_key)
+
+    return tuple(out)

src/zarr/core/common.py

@@ -45,7 +45,7 @@ def product(tup: ChunkCoords) -> int:
 
 
 async def concurrent_map(
-    items: list[T], func: Callable[..., Awaitable[V]], limit: int | None = None
+    items: Iterable[T], func: Callable[..., Awaitable[V]], limit: int | None = None
 ) -> list[V]:
     if limit is None:
         return await asyncio.gather(*[func(*item) for item in items])