diff --git a/changes/4028.feature.md b/changes/4028.feature.md new file mode 100644 index 0000000000..7f53605db5 --- /dev/null +++ b/changes/4028.feature.md @@ -0,0 +1,8 @@ +Add four composable functions for efficiently reading sparse arrays, where most chunks are empty and resolve to the fill value: + +- `zarr.shards_initialized` discovers which shards (or chunks, for unsharded arrays) of an array have actually been written to the store, via either a single prefix listing or concurrent per-key existence probes (selected by `strategy=`). +- `zarr.regions_initialized` expresses that discovery as array regions, suitable to pass straight to `read_regions`. +- `zarr.chunk_regions_initialized` refines that discovery to inner-chunk granularity for sharded arrays: it reads only each populated shard's index (never the chunk data) and reports the regions of the inner chunks that were actually written, skipping empty inner chunks. For unsharded arrays it is identical to `regions_initialized`. +- `zarr.read_regions` concurrently reads and decodes a caller-supplied collection of array regions, yielding each `(region, data)` pair spatially resolved to its location in the array. + +Together these let callers skip the per-empty-chunk store round-trips that dominate `arr[:]` on sparse arrays — e.g. `read_regions(arr, regions_initialized(arr))`. Asynchronous versions are available in `zarr.api.asynchronous`; the async `read_regions` streams each region as soon as its data is available. diff --git a/docs/api/zarr/functions/chunk_regions_initialized.md b/docs/api/zarr/functions/chunk_regions_initialized.md new file mode 100644 index 0000000000..2451fb661a --- /dev/null +++ b/docs/api/zarr/functions/chunk_regions_initialized.md @@ -0,0 +1,5 @@ +--- +title: zarr.chunk_regions_initialized +--- + +::: zarr.chunk_regions_initialized diff --git a/docs/api/zarr/functions/read_regions.md b/docs/api/zarr/functions/read_regions.md new file mode 100644 index 0000000000..4662af7bf5 --- /dev/null +++ b/docs/api/zarr/functions/read_regions.md @@ -0,0 +1,5 @@ +--- +title: zarr.read_regions +--- + +::: zarr.read_regions diff --git a/docs/api/zarr/functions/regions_initialized.md b/docs/api/zarr/functions/regions_initialized.md new file mode 100644 index 0000000000..e3dc361aa4 --- /dev/null +++ b/docs/api/zarr/functions/regions_initialized.md @@ -0,0 +1,5 @@ +--- +title: zarr.regions_initialized +--- + +::: zarr.regions_initialized diff --git a/docs/api/zarr/functions/shards_initialized.md b/docs/api/zarr/functions/shards_initialized.md new file mode 100644 index 0000000000..466378ae78 --- /dev/null +++ b/docs/api/zarr/functions/shards_initialized.md @@ -0,0 +1,5 @@ +--- +title: zarr.shards_initialized +--- + +::: zarr.shards_initialized diff --git a/mkdocs.yml b/mkdocs.yml index 46bfc1764c..41cbcde059 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -53,6 +53,7 @@ nav: - ' zarr.codecs': api/zarr/codecs.md - ' zarr.codecs.numcodecs': api/zarr/codecs/numcodecs.md - ' zarr.config': api/zarr/config.md + - ' zarr.chunk_regions_initialized': api/zarr/functions/chunk_regions_initialized.md - ' zarr.consolidate_metadata': api/zarr/functions/consolidate_metadata.md - ' zarr.create': api/zarr/functions/create.md - ' zarr.create_array': api/zarr/functions/create_array.md @@ -78,10 +79,13 @@ nav: - ' zarr.open_group': api/zarr/functions/open_group.md - ' zarr.open_like': api/zarr/functions/open_like.md - ' zarr.print_debug_info': api/zarr/functions/print_debug_info.md + - ' zarr.read_regions': api/zarr/functions/read_regions.md + - ' zarr.regions_initialized': api/zarr/functions/regions_initialized.md - ' zarr.registry': api/zarr/registry.md - ' zarr.save': api/zarr/functions/save.md - ' zarr.save_array': api/zarr/functions/save_array.md - ' zarr.save_group': api/zarr/functions/save_group.md + - ' zarr.shards_initialized': api/zarr/functions/shards_initialized.md - ' zarr.storage': api/zarr/storage.md - ' zarr.testing': - api/zarr/testing/index.md diff --git a/src/zarr/__init__.py b/src/zarr/__init__.py index cdf3840c3b..85f07b92dc 100644 --- a/src/zarr/__init__.py +++ b/src/zarr/__init__.py @@ -5,6 +5,7 @@ from zarr._version import version as __version__ from zarr.api.synchronous import ( array, + chunk_regions_initialized, consolidate_metadata, copy, copy_all, @@ -27,9 +28,12 @@ open_consolidated, open_group, open_like, + read_regions, + regions_initialized, save, save_array, save_group, + shards_initialized, tree, zeros, zeros_like, @@ -149,6 +153,7 @@ def set_format(log_format: str) -> None: "Group", "__version__", "array", + "chunk_regions_initialized", "config", "consolidate_metadata", "copy", @@ -173,9 +178,12 @@ def set_format(log_format: str) -> None: "open_group", "open_like", "print_debug_info", + "read_regions", + "regions_initialized", "save", "save_array", "save_group", + "shards_initialized", "tree", "zeros", "zeros_like", diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py index c751d6a31c..547f5b94bf 100644 --- a/src/zarr/api/asynchronous.py +++ b/src/zarr/api/asynchronous.py @@ -15,9 +15,13 @@ Array, AsyncArray, CompressorLike, + chunk_regions_initialized, create_array, from_array, get_array_metadata, + read_regions, + regions_initialized, + shards_initialized, ) from zarr.core.array_spec import ArrayConfigLike, parse_array_config from zarr.core.buffer import NDArrayLike @@ -66,6 +70,7 @@ __all__ = [ "array", + "chunk_regions_initialized", "consolidate_metadata", "copy", "copy_all", @@ -87,9 +92,12 @@ "open_consolidated", "open_group", "open_like", + "read_regions", + "regions_initialized", "save", "save_array", "save_group", + "shards_initialized", "tree", "zeros", "zeros_like", diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py index 3231837a04..17cbfd2e70 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -45,6 +45,7 @@ __all__ = [ "array", + "chunk_regions_initialized", "consolidate_metadata", "copy", "copy_all", @@ -66,9 +67,12 @@ "open_consolidated", "open_group", "open_like", + "read_regions", + "regions_initialized", "save", "save_array", "save_group", + "shards_initialized", "tree", "zeros", "zeros_like", @@ -1445,3 +1449,174 @@ def zeros_like(a: ArrayLike, **kwargs: Any) -> AnyArray: The new array. """ return Array(sync(async_api.zeros_like(a, **kwargs))) + + +def _as_async_array(array: Array[Any] | AsyncArray[Any]) -> AsyncArray[Any]: + return array._async_array if isinstance(array, Array) else array + + +def shards_initialized( + array: Array[Any] | AsyncArray[Any], + *, + strategy: Literal["auto", "list", "probe"] = "auto", +) -> tuple[str, ...]: + """ + Return the storage keys of the shards that have been persisted to the store. + + This reports storage at the granularity of stored objects: for sharded arrays it + returns shard keys (the objects that actually exist in the store), and for unsharded + arrays it returns chunk keys. To turn these into array regions, use + [regions_initialized][zarr.regions_initialized]. + + Parameters + ---------- + array : Array or AsyncArray + The array to inspect. + strategy : {"auto", "list", "probe"}, default "auto" + How to discover which shards exist. + + - ``"list"`` issues a single ``store.list_prefix`` call and keeps the keys that + belong to this array's shard grid (ignoring metadata and any other objects + under the same prefix). + - ``"probe"`` checks the existence of each possible shard key individually and + concurrently. This avoids listing a prefix that may hold many unrelated + objects, and is faster when the array has few possible shards. + - ``"auto"`` uses ``"probe"`` when the array has at most a small number of + possible shards and ``"list"`` otherwise. + + Returns + ------- + tuple[str, ...] + The storage keys of the populated shards (or chunks, when unsharded), + in chunk-grid order. + + See Also + -------- + regions_initialized : The array regions spanned by the populated shards. + read_regions : Read and decode a collection of array regions. + """ + return sync(async_api.shards_initialized(_as_async_array(array), strategy=strategy)) + + +def regions_initialized( + array: Array[Any] | AsyncArray[Any], + *, + strategy: Literal["auto", "list", "probe"] = "auto", +) -> list[tuple[slice, ...]]: + """ + Return the array regions spanned by the shards that have been persisted to the store. + + This is [shards_initialized][zarr.shards_initialized] expressed as regions: it filters + the array's shard regions down to those whose shard is populated. The result is + suitable to pass directly to [read_regions][zarr.read_regions]. + + Parameters + ---------- + array : Array or AsyncArray + The array to inspect. + strategy : {"auto", "list", "probe"}, default "auto" + How to discover which shards exist. See [shards_initialized][zarr.shards_initialized]. + + Returns + ------- + list[tuple[slice, ...]] + The regions spanned by the populated shards, in chunk-grid order. + + See Also + -------- + shards_initialized : The storage keys of the populated shards. + read_regions : Read and decode a collection of array regions. + """ + return sync(async_api.regions_initialized(_as_async_array(array), strategy=strategy)) + + +def chunk_regions_initialized( + array: Array[Any] | AsyncArray[Any], + *, + strategy: Literal["auto", "list", "probe"] = "auto", + concurrency: int | None = None, +) -> list[tuple[slice, ...]]: + """ + Return the array regions spanned by the inner chunks that have been written. + + Unlike [regions_initialized][zarr.regions_initialized], which reports at stored-object + (shard) granularity, this reports at chunk granularity: for a sharded array it reads each + populated shard's index (not its data) and reports the regions of the inner chunks that + were actually written, skipping empty inner chunks. For an unsharded array this is + identical to [regions_initialized][zarr.regions_initialized]. + + Parameters + ---------- + array : Array or AsyncArray + The array to inspect. + strategy : {"auto", "list", "probe"}, default "auto" + How to discover which shards exist. See [shards_initialized][zarr.shards_initialized]. + concurrency : int, optional + The maximum number of shard indexes read concurrently. Defaults to the + ``async.concurrency`` config value. + + Returns + ------- + list[tuple[slice, ...]] + The regions spanned by the populated inner chunks, in chunk-grid order. + + See Also + -------- + regions_initialized : The array regions at stored-object (shard) granularity. + read_regions : Read and decode a collection of array regions. + """ + return sync( + async_api.chunk_regions_initialized( + _as_async_array(array), strategy=strategy, concurrency=concurrency + ) + ) + + +def read_regions( + array: Array[Any] | AsyncArray[Any], + regions: Iterable[tuple[slice, ...]], + *, + concurrency: int | None = None, +) -> list[tuple[tuple[slice, ...], NDArrayLikeOrScalar]]: + """ + Read and decode a collection of array regions, returning a list of ``(region, data)`` + pairs. + + Each pair associates a region (a tuple of slices into the array) with the decoded data + for that region. The regions to read are supplied by the caller; pass the result of + [regions_initialized][zarr.regions_initialized] to read only the populated parts of a + sparse array without materializing the full array. For lazy, streaming consumption use + the asynchronous [zarr.api.asynchronous.read_regions][] instead, which yields each pair + as soon as its data is available. + + Parameters + ---------- + array : Array or AsyncArray + The array to read from. + regions : iterable of tuple of slice + The regions to read. Each region is a tuple of slices, one per array dimension. + concurrency : int, optional + The maximum number of regions read concurrently. Defaults to the + ``async.concurrency`` config value. + + Returns + ------- + list[tuple[tuple[slice, ...], NDArrayLikeOrScalar]] + Each region paired with its decoded data, in completion order (not necessarily + the order of ``regions``). + + See Also + -------- + regions_initialized : The array regions spanned by the populated shards. + shards_initialized : The storage keys of the populated shards. + """ + + async def _collect() -> list[tuple[tuple[slice, ...], NDArrayLikeOrScalar]]: + return [ + item + async for item in async_api.read_regions( + _as_async_array(array), regions, concurrency=concurrency + ) + ] + + return sync(_collect()) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index f75ef72415..5b63f2ec6d 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -2,7 +2,7 @@ import math import warnings -from asyncio import gather +from asyncio import Semaphore, as_completed, ensure_future, gather from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field, replace from itertools import starmap @@ -147,7 +147,7 @@ from zarr.storage._utils import _relativize_path if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import AsyncIterator, Iterator from typing import Self import numpy.typing as npt @@ -4014,20 +4014,317 @@ async def _shards_initialized( [nchunks_initialized][zarr.Array.nchunks_initialized] """ - store_contents = [ - x async for x in array.store_path.store.list_prefix(prefix=array.store_path.path) + # Thin wrapper over the public, strategy-aware entry point. The "list" strategy + # preserves this function's historical behavior (a single prefix listing). + return await shards_initialized(array, strategy="list") + + +# When the array has at most this many possible shards, ``shards_initialized`` +# probes each key individually rather than listing the prefix. Probing avoids +# paying for a prefix listing that may contain many unrelated objects, and is +# cheap when there are few keys to check. +_PROBE_THRESHOLD = 64 + + +async def shards_initialized( + array: AnyArray | AnyAsyncArray, + *, + strategy: Literal["auto", "list", "probe"] = "auto", +) -> tuple[str, ...]: + """ + Return the storage keys of the shards that have been persisted to the store. + + This reports storage at the granularity of stored objects: for sharded arrays it + returns shard keys (the objects that actually exist in the store), and for unsharded + arrays it returns chunk keys. To turn these into array regions, use + [regions_initialized][zarr.regions_initialized]. + + Parameters + ---------- + array : Array or AsyncArray + The array to inspect. + strategy : {"auto", "list", "probe"}, default "auto" + How to discover which shards exist. + + - ``"list"`` issues a single ``store.list_prefix`` call and keeps the keys that + belong to this array's shard grid (ignoring metadata and any other objects + under the same prefix). + - ``"probe"`` checks the existence of each possible shard key individually and + concurrently. This avoids listing a prefix that may hold many unrelated + objects, and is faster when the array has few possible shards. + - ``"auto"`` uses ``"probe"`` when the array has at most a small number of + possible shards and ``"list"`` otherwise. + + Returns + ------- + tuple[str, ...] + The storage keys of the populated shards (or chunks, when unsharded), + in chunk-grid order. + + Notes + ----- + This reports membership at the granularity of *stored objects* and does not + introspect the contents of a shard. For a sharded array, a shard is reported as + initialized if its object exists in the store, even if some (or all) of its inner + chunks were never written — those positions still read back as the fill value. In + other words, an "initialized" region may contain empty inner chunks. For unsharded + arrays there is one object per chunk, so no such ambiguity arises. + + See Also + -------- + regions_initialized : The array regions spanned by the populated shards. + read_regions : Read and decode a collection of array regions. + """ + if isinstance(array, Array): + array = array._async_array + + keys = list(_iter_shard_keys(array)) + + if strategy == "auto": + strategy = "probe" if len(keys) <= _PROBE_THRESHOLD else "list" + + if strategy == "list": + # A single prefix listing, filtered to keys that belong to this array's shard + # grid. Non-chunk objects under the same prefix (metadata, etc.) are excluded by + # the intersection, handling prefixes that also hold unrelated objects. + contents = { + _relativize_path(path=key, prefix=array.store_path.path) + async for key in array.store_path.store.list_prefix(prefix=array.store_path.path) + # obstore can include a directory marker whose key matches the listed prefix; + # it is not an initialized shard and must be excluded before relativizing. + if array.store_path.path == "" or key != array.store_path.path + } + return tuple(key for key in keys if key in contents) + elif strategy == "probe": + # Per-key existence checks, concurrently. Preferable when the prefix may + # contain many unrelated objects, or when there are few keys to check. + present = await concurrent_map( + [(array.store_path / key,) for key in keys], + lambda store_path: store_path.exists(), + zarr_config.get("async.concurrency"), + ) + return tuple(key for key, is_present in zip(keys, present, strict=True) if is_present) + else: + raise ValueError( + f"Unknown strategy {strategy!r}. Expected one of 'auto', 'list', or 'probe'." + ) + + +async def regions_initialized( + array: AnyArray | AnyAsyncArray, + *, + strategy: Literal["auto", "list", "probe"] = "auto", +) -> list[tuple[slice, ...]]: + """ + Return the array regions spanned by the shards that have been persisted to the store. + + This is [shards_initialized][zarr.shards_initialized] expressed as regions: it filters + the array's shard regions down to those whose shard is populated. The result is + suitable to pass directly to [read_regions][zarr.read_regions]. + + Parameters + ---------- + array : Array or AsyncArray + The array to inspect. + strategy : {"auto", "list", "probe"}, default "auto" + How to discover which shards exist. See [shards_initialized][zarr.shards_initialized]. + + Returns + ------- + list[tuple[slice, ...]] + The regions spanned by the populated shards, in chunk-grid order. + + Notes + ----- + Regions are reported at stored-object granularity: a sharded region may contain + empty inner chunks that were never written (they read back as the fill value). See + [shards_initialized][zarr.shards_initialized] for details. + + See Also + -------- + shards_initialized : The storage keys of the populated shards. + read_regions : Read and decode a collection of array regions. + """ + if isinstance(array, Array): + array = array._async_array + initialized = frozenset(await shards_initialized(array, strategy=strategy)) + # "regions that are initialized" is a filter over the array's shard regions, keyed on + # whether the corresponding shard key is populated. + return [ + region + for region, key in zip(_iter_shard_regions(array), _iter_shard_keys(array), strict=True) + if key in initialized ] - store_contents_relative = [ - _relativize_path(path=key, prefix=array.store_path.path) - for key in store_contents - # obstore can include a directory marker whose key matches the listed prefix; - # it is not an initialized shard and must be excluded before relativizing. - if array.store_path.path == "" or key != array.store_path.path + + +async def read_regions( + array: AnyArray | AnyAsyncArray, + regions: Iterable[tuple[slice, ...]], + *, + concurrency: int | None = None, +) -> AsyncIterator[tuple[tuple[slice, ...], NDArrayLikeOrScalar]]: + """ + Concurrently read and decode a collection of array regions, yielding each + ``(region, data)`` pair as soon as its data is available. + + Each yielded value pairs a region (a tuple of slices into the array) with the decoded + data for that region. The regions to read are supplied by the caller; pass the result + of [regions_initialized][zarr.regions_initialized] to read only the populated parts of + a sparse array without materializing the full array. + + Parameters + ---------- + array : Array or AsyncArray + The array to read from. + regions : iterable of tuple of slice + The regions to read. Each region is a tuple of slices, one per array dimension. + concurrency : int, optional + The maximum number of regions read concurrently. Defaults to the + ``async.concurrency`` config value. + + Yields + ------ + tuple[tuple[slice, ...], NDArrayLikeOrScalar] + A region and its decoded data, in completion order (not necessarily the order + of ``regions``). + + See Also + -------- + regions_initialized : The array regions spanned by the populated shards. + shards_initialized : The storage keys of the populated shards. + """ + if isinstance(array, Array): + array = array._async_array + if concurrency is None: + concurrency = zarr_config.get("async.concurrency") + + semaphore = Semaphore(concurrency) + + async def _read( + region: tuple[slice, ...], + ) -> tuple[tuple[slice, ...], NDArrayLikeOrScalar]: + async with semaphore: + return region, await array.getitem(region) + + for future in as_completed([ensure_future(_read(region)) for region in regions]): + yield await future + + +def _sharding_codec(array: AnyAsyncArray) -> Any: + """Return the array's top-level ``ShardingCodec`` instance, or ``None`` if the array + is not sharded.""" + from zarr.codecs.sharding import ShardingCodec + + codecs: tuple[Codec, ...] = getattr(array.metadata, "codecs", ()) + if len(codecs) == 1 and isinstance(codecs[0], ShardingCodec): + return codecs[0] + return None + + +async def chunk_regions_initialized( + array: AnyArray | AnyAsyncArray, + *, + strategy: Literal["auto", "list", "probe"] = "auto", + concurrency: int | None = None, +) -> list[tuple[slice, ...]]: + """ + Return the array regions spanned by the *inner chunks* that have been written. + + Unlike [regions_initialized][zarr.regions_initialized], which reports at stored-object + (shard) granularity, this reports at chunk granularity: for a sharded array it looks + *inside* each populated shard — reading only the shard index, not the chunk data — and + reports the regions of the inner chunks that were actually written, skipping empty inner + chunks. For an unsharded array, stored objects already are chunks, so this is identical + to [regions_initialized][zarr.regions_initialized]. + + The cost scales with the number of *populated shards*: discovering them is a single + listing (see [shards_initialized][zarr.shards_initialized]), and each populated shard + then costs one small range read of its index (the chunk data is never fetched). + + Parameters + ---------- + array : Array or AsyncArray + The array to inspect. + strategy : {"auto", "list", "probe"}, default "auto" + How to discover which shards exist. See [shards_initialized][zarr.shards_initialized]. + concurrency : int, optional + The maximum number of shard indexes read concurrently. Defaults to the + ``async.concurrency`` config value. + + Returns + ------- + list[tuple[slice, ...]] + The regions spanned by the populated inner chunks, in chunk-grid order. + + See Also + -------- + regions_initialized : The array regions at stored-object (shard) granularity. + read_regions : Read and decode a collection of array regions. + """ + if isinstance(array, Array): + array = array._async_array + + sharding_codec = _sharding_codec(array) + if sharding_codec is None: + # No sharding: each stored object is a chunk, so chunk granularity == shard + # granularity. Reuse the cheap listing-based path verbatim. + return await regions_initialized(array, strategy=strategy) + + if concurrency is None: + concurrency = zarr_config.get("async.concurrency") + + # array.chunks is the inner-chunk shape when sharding; array.shards is the shard shape. + assert array.shards is not None # guaranteed by sharding_codec is not None + chunks_per_shard = tuple(s // c for s, c in zip(array.shards, array.chunks, strict=True)) + chunk_grid_shape = array._chunk_grid_shape + + # 1. Discover populated shards (a single listing) and recover their grid coordinates. + populated_keys = frozenset(await shards_initialized(array, strategy=strategy)) + populated_shard_coords = [ + coord + for coord, key in zip(_iter_shard_coords(array), _iter_shard_keys(array), strict=True) + if key in populated_keys ] - return tuple( - chunk_key for chunk_key in array._iter_shard_keys() if chunk_key in store_contents_relative + + # 2. Read each populated shard's index (only) and expand its populated inner chunks to + # global chunk coordinates. The shard data itself is never fetched. + async def _populated_chunk_coords(shard_coord: tuple[int, ...]) -> list[tuple[int, ...]]: + byte_getter = array.store_path / array.metadata.encode_chunk_key(shard_coord) + index = await sharding_codec._load_shard_index_maybe(byte_getter, chunks_per_shard) + if index is None: + return [] + origin = tuple(s * cps for s, cps in zip(shard_coord, chunks_per_shard, strict=True)) + coords: list[tuple[int, ...]] = [] + for local in zip( + *(d.tolist() for d in np.nonzero(index.get_full_chunk_map())), strict=True + ): + coord = tuple(o + lo for o, lo in zip(origin, local, strict=True)) + # Guard against an edge shard whose index references chunks past the array bounds. + if all(ci < gi for ci, gi in zip(coord, chunk_grid_shape, strict=True)): + coords.append(coord) + return coords + + per_shard = await concurrent_map( + [(coord,) for coord in populated_shard_coords], + _populated_chunk_coords, + concurrency, ) + # 3. Turn the populated chunk coordinates into array regions, in chunk-grid order. + chunk_shape = array.chunks + ndim = len(array.shape) + unit = (1,) * ndim + return [ + next( + iter( + _iter_regions( + array.shape, chunk_shape, origin=coord, selection_shape=unit, trim_excess=True + ) + ) + ) + for coord in sorted({coord for coords in per_shard for coord in coords}) + ] + type FiltersLike = ( Iterable[dict[str, JSON] | ArrayArrayCodec | Numcodec] diff --git a/tests/test_array.py b/tests/test_array.py index b1a7a3c0f2..0f6beed5e1 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2386,3 +2386,265 @@ async def test_create_array_chunks_3d( shape = (10, 12, 15) arr = await create_array(store={}, shape=shape, chunks=chunk_input, dtype="float64") assert arr.write_chunk_sizes == expected + + +# --- shards_initialized / regions_initialized / read_regions --------------------------- + + +def _ca_sparse_1d(store: Store) -> tuple[Array[Any], npt.NDArray[Any]]: + arr = zarr.create_array(store=store, shape=(64,), chunks=(8,), dtype="int32", fill_value=42) + # populate two non-adjacent chunks (chunks 1 and 5) + arr[8:16] = np.arange(8, dtype="int32") + arr[40:48] = np.arange(100, 108, dtype="int32") + return arr, np.asarray(arr[:]) + + +def _ca_dense_1d(store: Store) -> tuple[Array[Any], npt.NDArray[Any]]: + arr = zarr.create_array(store=store, shape=(32,), chunks=(8,), dtype="int32", fill_value=0) + arr[:] = np.arange(32, dtype="int32") + return arr, np.asarray(arr[:]) + + +def _ca_sparse_2d(store: Store) -> tuple[Array[Any], npt.NDArray[Any]]: + arr = zarr.create_array(store=store, shape=(8, 8), chunks=(2, 2), dtype="int32", fill_value=-1) + arr[0:2, 0:2] = np.ones((2, 2), dtype="int32") + arr[4:6, 4:6] = np.full((2, 2), 7, dtype="int32") + return arr, np.asarray(arr[:]) + + +def _ca_sharded_sparse(store: Store) -> tuple[Array[Any], npt.NDArray[Any]]: + # chunks (2, 2) within shards (4, 4): the shard grid is 2x2 over the 8x8 array. + arr = zarr.create_array( + store=store, shape=(8, 8), chunks=(2, 2), shards=(4, 4), dtype="int32", fill_value=42 + ) + arr[0:2, 0:2] = np.ones((2, 2), dtype="int32") # shard (0, 0) + arr[4:6, 4:6] = np.full((2, 2), 7, dtype="int32") # shard (1, 1) + return arr, np.asarray(arr[:]) + + +def _ca_all_empty(store: Store) -> tuple[Array[Any], npt.NDArray[Any]]: + arr = zarr.create_array(store=store, shape=(32,), chunks=(8,), dtype="int32", fill_value=7) + return arr, np.asarray(arr[:]) + + +def _ca_all_populated(store: Store) -> tuple[Array[Any], npt.NDArray[Any]]: + arr = zarr.create_array(store=store, shape=(32,), chunks=(8,), dtype="int32", fill_value=0) + arr[:] = np.arange(32, dtype="int32") + return arr, np.asarray(arr[:]) + + +_CA_SETUPS = { + "sparse_1d": _ca_sparse_1d, + "dense_1d": _ca_dense_1d, + "sparse_2d": _ca_sparse_2d, + "sharded_sparse": _ca_sharded_sparse, + "all_empty": _ca_all_empty, + "all_populated": _ca_all_populated, +} +_CA_STRATEGIES = ["auto", "list", "probe"] + + +def _ca_pack( + arr: Array[Any], + results: list[tuple[tuple[slice, ...], Any]], + baseline: npt.NDArray[Any], +) -> npt.NDArray[Any]: + """Scatter ``(region, data)`` pairs onto a fill-valued array.""" + out = np.full(baseline.shape, arr.fill_value, dtype=baseline.dtype) + for region, data in results: + out[region] = np.asarray(data) + return out + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize("setup_name", list(_CA_SETUPS)) +@pytest.mark.parametrize("strategy", _CA_STRATEGIES) +def test_shards_initialized_strategies_agree( + store: Store, setup_name: str, strategy: Literal["auto", "list", "probe"] +) -> None: + """Every strategy reports the same set of populated keys.""" + arr, _ = _CA_SETUPS[setup_name](store) + keys = set(zarr.shards_initialized(arr, strategy=strategy)) + assert keys == set(zarr.shards_initialized(arr, strategy="auto")) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize( + ("setup_name", "expected_count"), + [ + ("sparse_1d", 2), + ("dense_1d", 4), + ("sparse_2d", 2), + ("sharded_sparse", 2), + ("all_empty", 0), + ("all_populated", 4), + ], +) +def test_shards_initialized_counts(store: Store, setup_name: str, expected_count: int) -> None: + arr, _ = _CA_SETUPS[setup_name](store) + assert len(zarr.shards_initialized(arr)) == expected_count + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +def test_shards_initialized_unknown_strategy(store: Store) -> None: + arr, _ = _ca_sparse_1d(store) + with pytest.raises(ValueError, match="Unknown strategy"): + zarr.shards_initialized(arr, strategy="nonsense") # type: ignore[arg-type] + + +async def test_list_strategy_ignores_non_chunk_objects() -> None: + """The ``list`` strategy must ignore objects that share the array's prefix but are + not chunks (metadata, stray writes). With no chunks written, an unrelated object + under the prefix must not be reported as an initialized shard.""" + store = MemoryStore() + arr = zarr.create_array(store=store, shape=(64,), chunks=(8,), dtype="int32", fill_value=0) + # No chunks written; drop a non-chunk object under the array's prefix. + await store.set("foo", default_buffer_prototype().buffer.from_bytes(b"blablabla")) + keys = await zarr.api.asynchronous.shards_initialized(arr._async_array, strategy="list") + assert keys == () + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize("setup_name", list(_CA_SETUPS)) +def test_regions_initialized_count_matches_keys(store: Store, setup_name: str) -> None: + """There is one initialized region per populated shard key.""" + arr, _ = _CA_SETUPS[setup_name](store) + assert len(zarr.regions_initialized(arr)) == len(zarr.shards_initialized(arr)) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize( + "regions", + [ + (), + ((slice(8, 16, 1),),), + ((slice(0, 8, 1),), (slice(8, 16, 1),)), + ((slice(8, 16, 1),), (slice(40, 48, 1),)), + ], + ids=["none", "single", "adjacent", "non_adjacent"], +) +def test_regions_initialized_are_exactly_written( + store: Store, regions: tuple[tuple[slice, ...], ...] +) -> None: + """Writing a set of chunk-aligned regions makes ``regions_initialized`` report + exactly those regions, and nothing else.""" + arr = zarr.create_array(store=store, shape=(64,), chunks=(8,), dtype="int32", fill_value=0) + for region in regions: + arr[region] = 1 # non-fill value so the chunk is persisted + assert set(zarr.regions_initialized(arr)) == set(regions) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize("setup_name", list(_CA_SETUPS)) +def test_read_regions_initialized_reconstructs_baseline(store: Store, setup_name: str) -> None: + """Reading the initialized regions and scattering them onto a fill-valued array + reproduces the full ``arr[:]`` read exactly.""" + arr, baseline = _CA_SETUPS[setup_name](store) + results = zarr.read_regions(arr, zarr.regions_initialized(arr)) + result = _ca_pack(arr, results, baseline) + assert np.array_equal(result, baseline) + assert result.dtype == baseline.dtype + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize("setup_name", ["sparse_1d", "dense_1d", "sparse_2d", "all_empty"]) +def test_chunk_regions_initialized_unsharded_matches_regions_initialized( + store: Store, setup_name: str +) -> None: + """For unsharded arrays each stored object is a chunk, so chunk granularity equals + stored-object granularity.""" + arr, _ = _CA_SETUPS[setup_name](store) + assert zarr.chunk_regions_initialized(arr) == zarr.regions_initialized(arr) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +def test_chunk_regions_initialized_sharded_skips_empty_inner_chunks(store: Store) -> None: + """Within a populated shard, only the inner chunks that were actually written are + reported — finer than the shard-level ``regions_initialized``.""" + # chunks (2, 2) within shards (4, 4): each shard is a 2x2 grid of inner chunks. + arr = zarr.create_array( + store=store, shape=(8, 8), chunks=(2, 2), shards=(4, 4), dtype="int32", fill_value=0 + ) + arr[0:2, 0:2] = 1 # one inner chunk of shard (0, 0) + arr[4:6, 4:6] = 2 # two inner chunks of shard (1, 1) + arr[6:8, 6:8] = 3 + # Shard granularity: two coarse 4x4 shard regions. + assert zarr.regions_initialized(arr) == [ + (slice(0, 4, 1), slice(0, 4, 1)), + (slice(4, 8, 1), slice(4, 8, 1)), + ] + # Chunk granularity: exactly the three written 2x2 inner chunks (empty inner chunks + # within the populated shards are skipped). + assert zarr.chunk_regions_initialized(arr) == [ + (slice(0, 2, 1), slice(0, 2, 1)), + (slice(4, 6, 1), slice(4, 6, 1)), + (slice(6, 8, 1), slice(6, 8, 1)), + ] + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize("setup_name", list(_CA_SETUPS)) +def test_read_chunk_regions_initialized_reconstructs_baseline( + store: Store, setup_name: str +) -> None: + """Reading the populated inner-chunk regions and scattering them reproduces the full + ``arr[:]`` read exactly, for both sharded and unsharded arrays.""" + arr, baseline = _CA_SETUPS[setup_name](store) + results = zarr.read_regions(arr, zarr.chunk_regions_initialized(arr)) + assert np.array_equal(_ca_pack(arr, results, baseline), baseline) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +def test_read_regions_explicit_regions(store: Store) -> None: + """Explicit regions are read and returned with their decoded data.""" + arr, baseline = _ca_sparse_1d(store) + explicit = [(slice(8, 16),), (slice(40, 48),)] + regions = dict(zarr.read_regions(arr, explicit)) + assert set(regions) == set(explicit) + assert np.array_equal(np.asarray(regions[(slice(8, 16),)]), baseline[8:16]) + assert np.array_equal(np.asarray(regions[(slice(40, 48),)]), baseline[40:48]) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +def test_read_regions_concurrency_one(store: Store) -> None: + """A concurrency limit of 1 produces the same result as the default.""" + arr, baseline = _ca_sparse_2d(store) + results = zarr.read_regions(arr, zarr.regions_initialized(arr), concurrency=1) + assert np.array_equal(_ca_pack(arr, results, baseline), baseline) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize("setup_name", list(_CA_SETUPS)) +async def test_read_regions_async_matches_sync(store: Store, setup_name: str) -> None: + """The async streaming generator yields the same ``(region, data)`` set as the + synchronous wrapper.""" + arr, _ = _CA_SETUPS[setup_name](store) + regions = zarr.regions_initialized(arr) + async_pairs = { + region: np.asarray(data).tobytes() + async for region, data in zarr.api.asynchronous.read_regions(arr._async_array, regions) + } + sync_pairs = { + region: np.asarray(data).tobytes() for region, data in zarr.read_regions(arr, regions) + } + assert async_pairs == sync_pairs + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +async def test_shards_initialized_async(store: Store) -> None: + arr, _ = _ca_sparse_1d(store) + keys = await zarr.api.asynchronous.shards_initialized(arr._async_array) + assert set(keys) == {"c/1", "c/5"} + + +@pytest.mark.parametrize("store", ["memory"], indirect=["store"]) +async def test_async_chunk_access_accepts_sync_array(store: Store) -> None: + """The async API also accepts a synchronous ``Array``, unwrapping it to the + underlying async array.""" + arr, _ = _ca_sparse_1d(store) + keys = await zarr.api.asynchronous.shards_initialized(arr) + regions = await zarr.api.asynchronous.regions_initialized(arr) + results = [pair async for pair in zarr.api.asynchronous.read_regions(arr, regions)] + assert set(keys) == {"c/1", "c/5"} + assert set(regions) == {(slice(8, 16, 1),), (slice(40, 48, 1),)} + assert len(results) == 2