tests for nchunks_initialized, chunks_initialized; add selection_shape kwarg to grid iteration; make chunk grid iterators consistent for array and async array

Author: d-v-b

src/zarr/core/array.py

@@ -461,25 +461,83 @@ def nchunks(self) -> int:
         """
         return product(self.cdata_shape)
 
-    def _iter_chunk_coords(self, origin: Sequence[int] | None = None) -> Iterator[ChunkCoords]:
+    @property
+    def nchunks_initialized(self) -> int:
+        """
+        The number of chunks that have been persisted in storage.
         """
-        Produce an iterator over the coordinates of each chunk, in chunk grid space, relative to
-        an optional origin.
+        return nchunks_initialized(self)
+
+    def _iter_chunk_coords(
+        self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None
+    ) -> Iterator[ChunkCoords]:
         """
-        return _iter_grid(self.cdata_shape, origin=origin)
+        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.
 
-    def _iter_chunk_keys(self, origin: Sequence[int] | None = None) -> Iterator[str]:
+        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: ChunkCoords
+            The coordinates of each chunk in the selection.
         """
-        Return an iterator over the storage keys of each chunk, relative to an optional origin.
+        return _iter_grid(self.cdata_shape, origin=origin, selection_shape=selection_shape)
+
+    def _iter_chunk_keys(
+        self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None
+    ) -> Iterator[str]:
+        """
+        Iterate over the storage keys of each chunk, relative to an optional origin, and optionally
+        limited to a contiguous region in chunk grid coordinates.
+
+        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
+        ------
+        key: str
+            The storage key of each chunk in the selection.
         """
-        for k in self._iter_chunk_coords(origin=origin):
+        # Iterate over the coordinates of chunks in chunk grid space.
+        for k in self._iter_chunk_coords(origin=origin, selection_shape=selection_shape):
+            # Encode the chunk key from the chunk coordinates.
             yield self.metadata.encode_chunk_key(k)
 
-    def _iter_chunk_regions(self) -> Iterator[tuple[slice, ...]]:
+    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.
         """
-        for cgrid_position in self._iter_chunk_coords():
+        for cgrid_position in self._iter_chunk_coords(
+            origin=origin, selection_shape=selection_shape
+        ):
             out: tuple[slice, ...] = ()
             for c_pos, c_shape in zip(cgrid_position, self.chunks, strict=False):
                 start = c_pos * c_shape
@@ -816,12 +874,32 @@ def nchunks(self) -> int:
         """
         return self._async_array.nchunks
 
-    def _iter_chunks(self, origin: Sequence[int] | None = None) -> Iterator[ChunkCoords]:
+    def _iter_chunk_coords(
+        self, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None
+    ) -> Iterator[ChunkCoords]:
         """
-        Produce an iterator over the coordinates of each chunk, in chunk grid space, relative
-        to an optional origin.
+        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: ChunkCoords
+            The coordinates of each chunk in the selection.
         """
-        yield from self._async_array._iter_chunk_coords(origin=origin)
+        yield from self._async_array._iter_chunk_coords(
+            origin=origin, selection_shape=selection_shape
+        )
 
     @property
     def nbytes(self) -> int:
@@ -830,17 +908,57 @@ def nbytes(self) -> int:
         """
         return self._async_array.nbytes
 
-    def _iter_chunk_keys(self, origin: Sequence[int] | None = None) -> Iterator[str]:
+    @property
+    def nchunks_initialized(self) -> int:
+        """
+        The number of chunks that have been initialized in the stored representation of this array.
+        """
+        return self._async_array.nchunks_initialized
+
+    def _iter_chunk_keys(
+        self, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None
+    ) -> Iterator[str]:
         """
-        Return an iterator over the keys of each chunk, relative to an optional origin
+        Iterate over the storage keys of each chunk, relative to an optional origin, and optionally
+        limited to a contiguous region in chunk grid coordinates.
+
+        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
+        ------
+        key: str
+            The storage key of each chunk in the selection.
         """
-        yield from self._async_array._iter_chunk_keys(origin=origin)
+        yield from self._async_array._iter_chunk_keys(
+            origin=origin, selection_shape=selection_shape
+        )
 
-    def _iter_chunk_regions(self) -> Iterator[tuple[slice, ...]]:
+    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.
         """
-        yield from self._async_array._iter_chunk_regions()
+        yield from self._async_array._iter_chunk_regions(
+            origin=origin, selection_shape=selection_shape
+        )
 
     def __array__(
         self, dtype: npt.DTypeLike | None = None, copy: bool | None = None
@@ -2175,17 +2293,46 @@ def info(self) -> None:
         )
 
 
-def nchunks_initialized(array: Array) -> int:
+def nchunks_initialized(array: AsyncArray | Array) -> int:
     """
     Calculate the number of chunks that have been initialized, i.e. the number of chunks that have
     been persisted to the storage backend.
+
+    Parameters
+    ----------
+    array : Array
+        The array to inspect.
+
+    Returns
+    -------
+    nchunks_initialized : int
+        The number of chunks that have been initialized.
+
+    See Also
+    --------
+    chunks_initialized
     """
     return len(chunks_initialized(array))
 
 
-def chunks_initialized(array: Array) -> tuple[str, ...]:
+def chunks_initialized(array: Array | AsyncArray) -> tuple[str, ...]:
     """
     Return the keys of the chunks that have been persisted to the storage backend.
+
+    Parameters
+    ----------
+    array : Array
+        The array to inspect.
+
+    Returns
+    -------
+    chunks_initialized : tuple[str, ...]
+        The keys of the chunks that have been initialized.
+
+    See Also
+    --------
+    nchunks_initialized
+
     """
     # TODO: make this compose with the underlying async iterator
     store_contents = list(

src/zarr/core/indexing.py

@@ -93,23 +93,24 @@ def ceildiv(a: float, b: float) -> int:
 
 
 def _iter_grid(
-    shape: Sequence[int],
+    grid_shape: Sequence[int],
     *,
     origin: Sequence[int] | None = None,
+    selection_shape: Sequence[int] | None = None,
     order: _ArrayIndexingOrder = "lexicographic",
 ) -> Iterator[ChunkCoords]:
     """
-    Iterate over the elements of grid of integers.
-
-    Takes a grid shape expressed as a sequence of integers and an optional origin and
-    yields tuples bounded by [origin, origin + grid_shape].
+    Iterate over the elements of grid of integers, with the option to restrict the domain of
+    iteration to those from a contiguous subregion of that grid.
 
     Parameters
     ---------
-    shape: Sequence[int]
+    grid_shape: Sequence[int]
         The size of the domain to iterate over.
     origin: Sequence[int] | None, default=None
-        The first coordinate of the domain.
+        The first coordinate of the domain to return.
+    selection_shape: Sequence[int] | None, default=None
+        The shape of the selection.
     order: Literal["lexicographic"], default="lexicographic"
         The linear indexing order to use.
 
@@ -129,22 +130,37 @@ def _iter_grid(
 
     >>> tuple(iter_grid((2,3)), origin=(1,1))
     ((1, 1), (1, 2), (1, 3), (2, 1), (2, 2), (2, 3))
+
+    >>> tuple(iter_grid((2,3)), origin=(1,1), selection_shape=(2,2))
+    ((1, 1), (1, 2), (1, 3), (2, 1))
     """
     if origin is None:
-        origin_parsed = (0,) * len(shape)
+        origin_parsed = (0,) * len(grid_shape)
     else:
-        if len(origin) != len(shape):
+        if len(origin) != len(grid_shape):
             msg = (
                 "Shape and origin parameters must have the same length."
-                f"Got {len(shape)} elements in shape, but {len(origin)} elements in origin."
+                f"Got {len(grid_shape)} elements in shape, but {len(origin)} elements in origin."
             )
             raise ValueError(msg)
         origin_parsed = tuple(origin)
-
-    if order == "lexicographic":
-        yield from itertools.product(
-            *(range(o, o + s) for o, s in zip(origin_parsed, shape, strict=True))
+    if selection_shape is None:
+        selection_shape_parsed = tuple(
+            g - o for o, g in zip(origin_parsed, grid_shape, strict=True)
         )
+    else:
+        selection_shape_parsed = tuple(selection_shape)
+    if order == "lexicographic":
+        dimensions: tuple[range, ...] = ()
+        for idx, (o, gs, ss) in enumerate(
+            zip(origin_parsed, grid_shape, selection_shape_parsed, strict=True)
+        ):
+            if o + ss > gs:
+                raise IndexError(
+                    f"Invalid selection shape ({selection_shape}) for origin ({origin}) and grid shape ({grid_shape}) at axis {idx}."
+                )
+            dimensions += (range(o, o + ss),)
+        yield from itertools.product(*(dimensions))
 
     else:
         msg = f"Indexing order {order} is not supported at this time."  # type: ignore[unreachable]

tests/v3/test_array.py

@@ -1,12 +1,15 @@
 import pickle
+from itertools import accumulate
 from typing import Literal
 
 import numpy as np
 import pytest
 
+from src.zarr.core.array import chunks_initialized
 from zarr import Array, AsyncArray, Group
 from zarr.core.buffer.cpu import NDBuffer
 from zarr.core.common import ZarrFormat
+from zarr.core.sync import sync
 from zarr.errors import ContainsArrayError, ContainsGroupError
 from zarr.store import LocalStore, MemoryStore
 from zarr.store.common import StorePath
@@ -232,3 +235,55 @@ def test_serializable_sync_array(store: LocalStore, zarr_format: ZarrFormat) ->
 
     assert actual == expected
     np.testing.assert_array_equal(actual[:], expected[:])
+
+
+@pytest.mark.parametrize("test_cls", [Array, AsyncArray])
+def test_nchunks_initialized(test_cls: type[Array] | type[AsyncArray]) -> None:
+    """
+    Test that nchunks_initialized accurately returns the number of stored chunks.
+    """
+    store = MemoryStore({}, mode="w")
+    arr = Array.create(store, shape=(100,), chunks=(10,), dtype="i4")
+
+    # write chunks one at a time
+    for idx, region in enumerate(arr._iter_chunk_regions()):
+        arr[region] = 1
+        expected = idx + 1
+        if test_cls == Array:
+            observed = arr.nchunks_initialized
+        else:
+            observed = arr._async_array.nchunks_initialized
+        assert observed == expected
+
+    # delete chunks
+    for idx, key in enumerate(arr._iter_chunk_keys()):
+        sync(arr.store_path.store.delete(key))
+        if test_cls == Array:
+            observed = arr.nchunks_initialized
+        else:
+            observed = arr._async_array.nchunks_initialized
+        expected = arr.nchunks - idx - 1
+        assert observed == expected
+
+
+@pytest.mark.parametrize("test_cls", [Array, AsyncArray])
+def test_chunks_initialized(test_cls: type[Array] | type[AsyncArray]) -> None:
+    """
+    Test that chunks_initialized accurately returns the keys of stored chunks.
+    """
+    store = MemoryStore({}, mode="w")
+    arr = Array.create(store, shape=(100,), chunks=(10,), dtype="i4")
+
+    chunks_accumulated = tuple(
+        accumulate(tuple(map(lambda v: tuple(v.split(" ")), arr._iter_chunk_keys())))
+    )
+    for keys, region in zip(chunks_accumulated, arr._iter_chunk_regions(), strict=False):
+        arr[region] = 1
+
+        if test_cls == Array:
+            observed = sorted(chunks_initialized(arr))
+        else:
+            observed = sorted(chunks_initialized(arr._async_array))
+
+        expected = sorted(keys)
+        assert observed == expected