From 288714279893eb92fa8bdce0d4e2980ce54114ba Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Tue, 2 Jun 2026 18:37:21 -0700 Subject: [PATCH 1/9] populated shards primative --- bench/empty_chunks.py | 157 ++++++++++++++++++++++++++++ changes/3929.feature.md | 14 +++ docs/api/zarr/read.md | 6 ++ mkdocs.yml | 1 + src/zarr/__init__.py | 4 + src/zarr/api/asynchronous.py | 4 + src/zarr/api/synchronous.py | 98 ++++++++++++++++++ src/zarr/core/array.py | 182 ++++++++++++++++++++++++++++++-- tests/test_chunk_access.py | 194 +++++++++++++++++++++++++++++++++++ 9 files changed, 652 insertions(+), 8 deletions(-) create mode 100644 bench/empty_chunks.py create mode 100644 changes/3929.feature.md create mode 100644 docs/api/zarr/read.md create mode 100644 tests/test_chunk_access.py diff --git a/bench/empty_chunks.py b/bench/empty_chunks.py new file mode 100644 index 0000000000..50fc7e0d8a --- /dev/null +++ b/bench/empty_chunks.py @@ -0,0 +1,157 @@ +"""Benchmark for sparse-array reads via the chunk-access primitives. + +Compares the stock ``arr[:]`` read against two primitive-based read paths on +sparse arrays (~3% of chunks populated), sweeping chunk count on ``MemoryStore`` +and ``LocalStore``: + +- ``pack``: :func:`zarr.read_regions` + scatter onto a fill-valued array. This + reproduces ``arr[:]`` semantics (a single contiguous array) but only touches + the populated chunks. +- ``stream``: iterate :func:`zarr.read_regions` without packing into one array. + This is the win for pipelines that operate per chunk and never need the whole + array materialized. + +The stock baseline scales with total chunk count; the primitive-based paths scale +with the populated-chunk count, so the speedup grows with sparsity-at-scale. Each +configuration is skipped if the warmup baseline exceeds ``BASELINE_BUDGET_S`` to +keep total runtime bounded. +""" + +from __future__ import annotations + +import tempfile +import time +import timeit +from pathlib import Path + +import numpy as np + +import zarr +from zarr.storage import LocalStore, MemoryStore + +CHUNK_SIZE = 1024 +DTYPE = "int32" +FILL_VALUE = 0 +BASELINE_BUDGET_S = 25.0 # skip configs whose warmup baseline exceeds this + +# (n_chunks, n_populated) — ~3% populated, mirrors the zagg HEALPix report. +SWEEP: list[tuple[int, int]] = [ + (1_024, 32), + (4_096, 128), + (16_384, 512), + (49_152, 1_536), +] + + +def _build_array(store: object, n_chunks: int, n_populated: int) -> zarr.Array: + arr = zarr.create_array( + store=store, + shape=(n_chunks * CHUNK_SIZE,), + chunks=(CHUNK_SIZE,), + dtype=DTYPE, + fill_value=FILL_VALUE, + ) + rng = np.random.default_rng(seed=0) + chunk_indices = rng.choice(n_chunks, size=n_populated, replace=False) + payload = np.arange(CHUNK_SIZE, dtype=DTYPE) + for ci in chunk_indices: + start = int(ci) * CHUNK_SIZE + arr[start : start + CHUNK_SIZE] = payload + return arr + + +def _read_baseline(arr: zarr.Array) -> None: + arr[:] + + +def _read_pack(arr: zarr.Array) -> np.ndarray: + out = np.full(arr.shape, arr.fill_value, dtype=arr.dtype) + for region, data in zarr.read_regions(arr): + out[region] = np.asarray(data) + return out + + +def _read_stream(arr: zarr.Array) -> int: + # Touch each region without materializing a single contiguous array. + total = 0 + for _region, data in zarr.read_regions(arr): + total += int(np.asarray(data).sum()) + return total + + +def _time(fn: object, repeats: int) -> float: + return min(timeit.repeat(fn, repeat=repeats, number=1)) + + +def _adaptive_repeats(warmup_s: float) -> int: + if warmup_s < 0.1: + return 5 + if warmup_s < 1.0: + return 3 + return 1 + + +def _run_one( + store_name: str, store: object, n_chunks: int, n_populated: int +) -> tuple[str, int, int, float, float, float, str]: + arr = _build_array(store, n_chunks, n_populated) + + t0 = time.perf_counter() + _read_baseline(arr) + warmup = time.perf_counter() - t0 + if warmup > BASELINE_BUDGET_S: + return ( + store_name, + n_chunks, + n_populated, + warmup, + float("nan"), + float("nan"), + f"skipped (>{BASELINE_BUDGET_S:.0f}s budget)", + ) + + # warm both primitive paths once + _read_pack(arr) + _read_stream(arr) + + repeats = _adaptive_repeats(warmup) + t_base = _time(lambda: _read_baseline(arr), repeats) + t_pack = _time(lambda: _read_pack(arr), repeats) + t_stream = _time(lambda: _read_stream(arr), repeats) + return store_name, n_chunks, n_populated, t_base, t_pack, t_stream, f"min of {repeats} runs" + + +def main() -> None: + rows = [] + print("Running sweep — this will take a couple of minutes for the largest configs...\n") + for n_chunks, n_populated in SWEEP: + rows.append(_run_one("MemoryStore", MemoryStore(), n_chunks, n_populated)) + with tempfile.TemporaryDirectory() as tmpdir: + rows.append( + _run_one("LocalStore", LocalStore(str(Path(tmpdir))), n_chunks, n_populated) + ) + + print( + f"\n{'store':<14}{'n_chunks':>10}{'populated':>11}" + f"{'arr[:] (s)':>12}{'pack (s)':>11}{'stream (s)':>12}" + f"{'pack x':>9}{'stream x':>10} notes" + ) + print("-" * 100) + for store_name, n_chunks, n_populated, t_base, t_pack, t_stream, note in rows: + print( + f"{store_name:<14}{n_chunks:>10}{n_populated:>11}" + f"{t_base:>12.4f}{_fmt(t_pack):>11}{_fmt(t_stream):>12}" + f"{_speedup(t_base, t_pack):>9}{_speedup(t_base, t_stream):>10} {note}" + ) + + +def _fmt(t: float) -> str: + return "—" if np.isnan(t) else f"{t:.4f}" + + +def _speedup(t_base: float, t: float) -> str: + return "—" if np.isnan(t) or t <= 0 else f"{t_base / t:.1f}x" + + +if __name__ == "__main__": + main() diff --git a/changes/3929.feature.md b/changes/3929.feature.md new file mode 100644 index 0000000000..f5ac547f72 --- /dev/null +++ b/changes/3929.feature.md @@ -0,0 +1,14 @@ +Add two primitives for efficiently reading sparse arrays, where most chunks are +empty and resolve to the fill value. + +- :func:`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. +- :func:`zarr.read_regions` concurrently reads and decodes array regions — by default + only the populated ones — 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. Asynchronous versions are available in +:mod:`zarr.api.asynchronous`; the async :func:`zarr.api.asynchronous.read_regions` +streams each region as soon as its data is available. diff --git a/docs/api/zarr/read.md b/docs/api/zarr/read.md new file mode 100644 index 0000000000..d84bddb8ac --- /dev/null +++ b/docs/api/zarr/read.md @@ -0,0 +1,6 @@ +--- +title: read +--- + +::: zarr.shards_initialized +::: zarr.read_regions diff --git a/mkdocs.yml b/mkdocs.yml index 7a4bfa35ef..7d9eeffc5a 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -39,6 +39,7 @@ nav: - api/zarr/dtype.md - api/zarr/load.md - api/zarr/open.md + - api/zarr/read.md - api/zarr/save.md - api/zarr/codecs.md - api/zarr/codecs/numcodecs.md diff --git a/src/zarr/__init__.py b/src/zarr/__init__.py index cdf3840c3b..d7686d5832 100644 --- a/src/zarr/__init__.py +++ b/src/zarr/__init__.py @@ -27,9 +27,11 @@ open_consolidated, open_group, open_like, + read_regions, save, save_array, save_group, + shards_initialized, tree, zeros, zeros_like, @@ -173,9 +175,11 @@ def set_format(log_format: str) -> None: "open_group", "open_like", "print_debug_info", + "read_regions", "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 7f185535df..9458f9bf4d 100644 --- a/src/zarr/api/asynchronous.py +++ b/src/zarr/api/asynchronous.py @@ -18,6 +18,8 @@ create_array, from_array, get_array_metadata, + read_regions, + shards_initialized, ) from zarr.core.array_spec import ArrayConfigLike, parse_array_config from zarr.core.buffer import NDArrayLike @@ -87,9 +89,11 @@ "open_consolidated", "open_group", "open_like", + "read_regions", "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 8386427b3f..d9f47713f6 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -66,9 +66,11 @@ "open_consolidated", "open_group", "open_like", + "read_regions", "save", "save_array", "save_group", + "shards_initialized", "tree", "zeros", "zeros_like", @@ -1426,3 +1428,99 @@ 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 | AsyncArray[Any]) -> AsyncArray[Any]: + return array._async_array if isinstance(array, Array) else array + + +def shards_initialized( + array: Array | 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 fetch and decode the populated regions, pass the + result of this function (or its regions) 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. + + - ``"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 + -------- + read_regions : Read and decode the populated regions of an array. + """ + return sync(async_api.shards_initialized(_as_async_array(array), strategy=strategy)) + + +def read_regions( + array: Array | AsyncArray[Any], + regions: Iterable[tuple[slice, ...]] | None = None, + *, + concurrency: int | None = None, +) -> list[tuple[tuple[slice, ...], NDArrayLikeOrScalar]]: + """ + Read and decode array regions, returning a list of ``(region, data)`` pairs. + + This is the spatially-resolved companion to [shards_initialized][zarr.shards_initialized]: + each pair associates a region (a tuple of slices into the array) with the decoded data + for that region, letting callers operate on 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, optional + The regions to read. Each region is a tuple of slices, one per array dimension. + If omitted, defaults to the regions spanned by the populated shards of ``array`` + (i.e. every region that holds data). + 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 + -------- + shards_initialized : Discover which shards of an array are populated. + """ + + 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 366c19bb0c..094aaf72f8 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -2,7 +2,7 @@ import json 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 @@ -146,7 +146,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 @@ -3975,17 +3975,183 @@ 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) - ] - store_contents_relative = [ - _relativize_path(path=key, prefix=array.store_path.path) for key in store_contents - ] + store_contents_relative = { + _relativize_path(path=key, prefix=array.store_path.path) + async for key in array.store_path.store.list_prefix(prefix=array.store_path.path) + } return tuple( chunk_key for chunk_key in array._iter_shard_keys() if chunk_key in store_contents_relative ) +# 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 _initialized_shards( + array: AnyAsyncArray, + *, + strategy: Literal["auto", "list", "probe"] = "auto", +) -> list[tuple[tuple[int, ...], str]]: + """ + Discover the populated shards of an array, returning ``(coords, key)`` pairs in + chunk-grid order. This is the shared core of [shards_initialized][zarr.shards_initialized] + (which projects to keys) and [read_regions][zarr.read_regions] (which projects to regions). + """ + coords = list(_iter_shard_coords(array)) + keys = [array.metadata.encode_chunk_key(c) for c in coords] + + if strategy == "auto": + strategy = "probe" if len(coords) <= _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, addressing the case where the prefix + # holds many 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) + } + return [(c, k) for c, k in zip(coords, keys, strict=True) if k 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 / k,) for k in keys], + lambda store_path: store_path.exists(), + zarr_config.get("async.concurrency"), + ) + return [ + (c, k) + for (c, k), is_present in zip(zip(coords, keys, strict=True), present, strict=True) + if is_present + ] + else: + raise ValueError( + f"Unknown strategy {strategy!r}. Expected one of 'auto', 'list', or 'probe'." + ) + + +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 fetch and decode the populated regions, pass the + result of this function (or its regions) 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. + + - ``"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 + -------- + read_regions : Read and decode the populated regions of an array. + """ + if isinstance(array, Array): + array = array._async_array + return tuple(key for _, key in await _initialized_shards(array, strategy=strategy)) + + +async def _initialized_regions( + array: AnyAsyncArray, + *, + strategy: Literal["auto", "list", "probe"] = "auto", +) -> list[tuple[slice, ...]]: + """Return the array regions spanned by each populated shard, in chunk-grid order.""" + shard_shape = array.shards if array.shards is not None else array.chunks + return [ + tuple( + slice(c * s, min((c + 1) * s, dim)) + for c, s, dim in zip(coords, shard_shape, array.shape, strict=True) + ) + for coords, _ in await _initialized_shards(array, strategy=strategy) + ] + + +async def read_regions( + array: AnyArray | AnyAsyncArray, + regions: Iterable[tuple[slice, ...]] | None = None, + *, + concurrency: int | None = None, +) -> AsyncIterator[tuple[tuple[slice, ...], NDArrayLikeOrScalar]]: + """ + Concurrently read and decode array regions, yielding each ``(region, data)`` pair + as soon as its data is available. + + This is the spatially-resolved companion to [shards_initialized][zarr.shards_initialized]: + each yielded value pairs a region (a tuple of slices into the array) with the decoded + data for that region, so callers can stream over 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, optional + The regions to read. Each region is a tuple of slices, one per array dimension. + If omitted, defaults to the regions spanned by the populated shards of ``array`` + (i.e. every region that holds data). + 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 + -------- + shards_initialized : Discover which shards of an array are populated. + """ + if isinstance(array, Array): + array = array._async_array + if concurrency is None: + concurrency = zarr_config.get("async.concurrency") + + region_list = await _initialized_regions(array) if regions is None else list(regions) + + 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 region_list]): + yield await future + + type FiltersLike = ( Iterable[dict[str, JSON] | ArrayArrayCodec | Numcodec] | ArrayArrayCodec diff --git a/tests/test_chunk_access.py b/tests/test_chunk_access.py new file mode 100644 index 0000000000..8ee6041a70 --- /dev/null +++ b/tests/test_chunk_access.py @@ -0,0 +1,194 @@ +"""Tests for the shard-discovery and region-read primitives. + +These cover :func:`zarr.shards_initialized` (discover which shards/chunks of an +array are populated) and :func:`zarr.read_regions` (concurrently read and decode +the populated regions), along with their asynchronous counterparts in +:mod:`zarr.api.asynchronous`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +import zarr +import zarr.api.asynchronous as async_api + +if TYPE_CHECKING: + from collections.abc import Callable + + from zarr.abc.store import Store + + +def _sparse_1d(store: Store) -> tuple[zarr.Array, np.ndarray]: + 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 _dense_1d(store: Store) -> tuple[zarr.Array, np.ndarray]: + 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 _sparse_2d(store: Store) -> tuple[zarr.Array, np.ndarray]: + 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 _sharded_sparse(store: Store) -> tuple[zarr.Array, np.ndarray]: + # 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 _all_empty(store: Store) -> tuple[zarr.Array, np.ndarray]: + arr = zarr.create_array(store=store, shape=(32,), chunks=(8,), dtype="int32", fill_value=7) + return arr, np.asarray(arr[:]) + + +def _all_populated(store: Store) -> tuple[zarr.Array, np.ndarray]: + 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[:]) + + +SETUPS: dict[str, Callable[[Store], tuple[zarr.Array, np.ndarray]]] = { + "sparse_1d": _sparse_1d, + "dense_1d": _dense_1d, + "sparse_2d": _sparse_2d, + "sharded_sparse": _sharded_sparse, + "all_empty": _all_empty, + "all_populated": _all_populated, +} + +STRATEGIES = ["auto", "list", "probe"] + + +def _pack(arr: zarr.Array, regions: list, baseline: np.ndarray) -> np.ndarray: + """Scatter ``(region, data)`` pairs onto a fill-valued array.""" + out = np.full(baseline.shape, arr.fill_value, dtype=baseline.dtype) + for region, data in regions: + out[region] = np.asarray(data) + return out + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize("setup_name", list(SETUPS)) +@pytest.mark.parametrize("strategy", STRATEGIES) +def test_shards_initialized_strategies_agree(store: Store, setup_name: str, strategy: str) -> None: + """Every strategy reports the same set of populated keys, and reports the + expected count for a hand-known layout.""" + arr, _ = SETUPS[setup_name](store) + keys = set(zarr.shards_initialized(arr, strategy=strategy)) + # all strategies must agree + 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, _ = 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, _ = _sparse_1d(store) + with pytest.raises(ValueError, match="Unknown strategy"): + zarr.shards_initialized(arr, strategy="nonsense") # type: ignore[arg-type] + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +def test_list_strategy_ignores_non_chunk_objects(store: Store) -> None: + """The ``list`` strategy must not mistake unrelated objects sharing the + array's prefix (e.g. metadata) for populated chunks.""" + arr, _ = _sparse_1d(store) + # metadata (zarr.json) already lives under the array prefix; add another + # non-chunk object to be sure it is excluded. + keys = set(zarr.shards_initialized(arr, strategy="list")) + assert keys == {"c/1", "c/5"} + assert all(k.startswith("c/") for k in keys) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize("setup_name", list(SETUPS)) +def test_read_regions_reconstructs_baseline(store: Store, setup_name: str) -> None: + """Packing the populated regions onto a fill-valued array reproduces the + full ``arr[:]`` read exactly.""" + arr, baseline = SETUPS[setup_name](store) + regions = zarr.read_regions(arr) + result = _pack(arr, regions, 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", list(SETUPS)) +def test_read_regions_default_count_matches_discovery(store: Store, setup_name: str) -> None: + """With no explicit regions, ``read_regions`` reads exactly the populated + shards discovered by ``shards_initialized``.""" + arr, _ = SETUPS[setup_name](store) + regions = zarr.read_regions(arr) + assert len(regions) == len(zarr.shards_initialized(arr)) + + +@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 = _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 = _sparse_2d(store) + regions = zarr.read_regions(arr, concurrency=1) + assert np.array_equal(_pack(arr, regions, baseline), baseline) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize("setup_name", list(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, _ = SETUPS[setup_name](store) + async_pairs = { + region: np.asarray(data).tobytes() + async for region, data in async_api.read_regions(arr._async_array) + } + sync_pairs = {region: np.asarray(data).tobytes() for region, data in zarr.read_regions(arr)} + assert async_pairs == sync_pairs + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +async def test_shards_initialized_async(store: Store) -> None: + arr, _ = _sparse_1d(store) + keys = await async_api.shards_initialized(arr._async_array) + assert set(keys) == {"c/1", "c/5"} From 82bdbf85af667f2af9f79a4f690fc5f4e4f15d50 Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Tue, 2 Jun 2026 19:01:33 -0700 Subject: [PATCH 2/9] delegation consistency for shards --- src/zarr/core/array.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 094aaf72f8..74bc3225e4 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -3975,13 +3975,11 @@ async def _shards_initialized( [nchunks_initialized][zarr.Array.nchunks_initialized] """ - store_contents_relative = { - _relativize_path(path=key, prefix=array.store_path.path) - async for key in array.store_path.store.list_prefix(prefix=array.store_path.path) - } - return tuple( - chunk_key for chunk_key in array._iter_shard_keys() if chunk_key in store_contents_relative - ) + # Thin wrapper over the shared discovery core, projected to keys. The "list" + # strategy preserves this function's historical behavior (a single prefix + # listing); see [shards_initialized][zarr.shards_initialized] for the public, + # strategy-aware entry point. + return tuple(key for _, key in await _initialized_shards(array, strategy="list")) # When the array has at most this many possible shards, ``shards_initialized`` From a027be53104c020e9246263e7c5a160b42b1232c Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Wed, 3 Jun 2026 16:47:16 -0700 Subject: [PATCH 3/9] easy fixes (tests, file paths, etc) --- bench/empty_chunks.py | 157 ------------------------ changes/3929.feature.md | 14 --- changes/4028.feature.md | 7 ++ docs/api/zarr/read.md | 1 + src/zarr/__init__.py | 2 + src/zarr/api/asynchronous.py | 2 + src/zarr/api/synchronous.py | 55 +++++++-- src/zarr/core/array.py | 142 +++++++++++----------- tests/benchmarks/test_indexing.py | 47 +++++++- tests/test_array.py | 179 +++++++++++++++++++++++++++ tests/test_chunk_access.py | 194 ------------------------------ 11 files changed, 357 insertions(+), 443 deletions(-) delete mode 100644 bench/empty_chunks.py delete mode 100644 changes/3929.feature.md create mode 100644 changes/4028.feature.md delete mode 100644 tests/test_chunk_access.py diff --git a/bench/empty_chunks.py b/bench/empty_chunks.py deleted file mode 100644 index 50fc7e0d8a..0000000000 --- a/bench/empty_chunks.py +++ /dev/null @@ -1,157 +0,0 @@ -"""Benchmark for sparse-array reads via the chunk-access primitives. - -Compares the stock ``arr[:]`` read against two primitive-based read paths on -sparse arrays (~3% of chunks populated), sweeping chunk count on ``MemoryStore`` -and ``LocalStore``: - -- ``pack``: :func:`zarr.read_regions` + scatter onto a fill-valued array. This - reproduces ``arr[:]`` semantics (a single contiguous array) but only touches - the populated chunks. -- ``stream``: iterate :func:`zarr.read_regions` without packing into one array. - This is the win for pipelines that operate per chunk and never need the whole - array materialized. - -The stock baseline scales with total chunk count; the primitive-based paths scale -with the populated-chunk count, so the speedup grows with sparsity-at-scale. Each -configuration is skipped if the warmup baseline exceeds ``BASELINE_BUDGET_S`` to -keep total runtime bounded. -""" - -from __future__ import annotations - -import tempfile -import time -import timeit -from pathlib import Path - -import numpy as np - -import zarr -from zarr.storage import LocalStore, MemoryStore - -CHUNK_SIZE = 1024 -DTYPE = "int32" -FILL_VALUE = 0 -BASELINE_BUDGET_S = 25.0 # skip configs whose warmup baseline exceeds this - -# (n_chunks, n_populated) — ~3% populated, mirrors the zagg HEALPix report. -SWEEP: list[tuple[int, int]] = [ - (1_024, 32), - (4_096, 128), - (16_384, 512), - (49_152, 1_536), -] - - -def _build_array(store: object, n_chunks: int, n_populated: int) -> zarr.Array: - arr = zarr.create_array( - store=store, - shape=(n_chunks * CHUNK_SIZE,), - chunks=(CHUNK_SIZE,), - dtype=DTYPE, - fill_value=FILL_VALUE, - ) - rng = np.random.default_rng(seed=0) - chunk_indices = rng.choice(n_chunks, size=n_populated, replace=False) - payload = np.arange(CHUNK_SIZE, dtype=DTYPE) - for ci in chunk_indices: - start = int(ci) * CHUNK_SIZE - arr[start : start + CHUNK_SIZE] = payload - return arr - - -def _read_baseline(arr: zarr.Array) -> None: - arr[:] - - -def _read_pack(arr: zarr.Array) -> np.ndarray: - out = np.full(arr.shape, arr.fill_value, dtype=arr.dtype) - for region, data in zarr.read_regions(arr): - out[region] = np.asarray(data) - return out - - -def _read_stream(arr: zarr.Array) -> int: - # Touch each region without materializing a single contiguous array. - total = 0 - for _region, data in zarr.read_regions(arr): - total += int(np.asarray(data).sum()) - return total - - -def _time(fn: object, repeats: int) -> float: - return min(timeit.repeat(fn, repeat=repeats, number=1)) - - -def _adaptive_repeats(warmup_s: float) -> int: - if warmup_s < 0.1: - return 5 - if warmup_s < 1.0: - return 3 - return 1 - - -def _run_one( - store_name: str, store: object, n_chunks: int, n_populated: int -) -> tuple[str, int, int, float, float, float, str]: - arr = _build_array(store, n_chunks, n_populated) - - t0 = time.perf_counter() - _read_baseline(arr) - warmup = time.perf_counter() - t0 - if warmup > BASELINE_BUDGET_S: - return ( - store_name, - n_chunks, - n_populated, - warmup, - float("nan"), - float("nan"), - f"skipped (>{BASELINE_BUDGET_S:.0f}s budget)", - ) - - # warm both primitive paths once - _read_pack(arr) - _read_stream(arr) - - repeats = _adaptive_repeats(warmup) - t_base = _time(lambda: _read_baseline(arr), repeats) - t_pack = _time(lambda: _read_pack(arr), repeats) - t_stream = _time(lambda: _read_stream(arr), repeats) - return store_name, n_chunks, n_populated, t_base, t_pack, t_stream, f"min of {repeats} runs" - - -def main() -> None: - rows = [] - print("Running sweep — this will take a couple of minutes for the largest configs...\n") - for n_chunks, n_populated in SWEEP: - rows.append(_run_one("MemoryStore", MemoryStore(), n_chunks, n_populated)) - with tempfile.TemporaryDirectory() as tmpdir: - rows.append( - _run_one("LocalStore", LocalStore(str(Path(tmpdir))), n_chunks, n_populated) - ) - - print( - f"\n{'store':<14}{'n_chunks':>10}{'populated':>11}" - f"{'arr[:] (s)':>12}{'pack (s)':>11}{'stream (s)':>12}" - f"{'pack x':>9}{'stream x':>10} notes" - ) - print("-" * 100) - for store_name, n_chunks, n_populated, t_base, t_pack, t_stream, note in rows: - print( - f"{store_name:<14}{n_chunks:>10}{n_populated:>11}" - f"{t_base:>12.4f}{_fmt(t_pack):>11}{_fmt(t_stream):>12}" - f"{_speedup(t_base, t_pack):>9}{_speedup(t_base, t_stream):>10} {note}" - ) - - -def _fmt(t: float) -> str: - return "—" if np.isnan(t) else f"{t:.4f}" - - -def _speedup(t_base: float, t: float) -> str: - return "—" if np.isnan(t) or t <= 0 else f"{t_base / t:.1f}x" - - -if __name__ == "__main__": - main() diff --git a/changes/3929.feature.md b/changes/3929.feature.md deleted file mode 100644 index f5ac547f72..0000000000 --- a/changes/3929.feature.md +++ /dev/null @@ -1,14 +0,0 @@ -Add two primitives for efficiently reading sparse arrays, where most chunks are -empty and resolve to the fill value. - -- :func:`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. -- :func:`zarr.read_regions` concurrently reads and decodes array regions — by default - only the populated ones — 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. Asynchronous versions are available in -:mod:`zarr.api.asynchronous`; the async :func:`zarr.api.asynchronous.read_regions` -streams each region as soon as its data is available. diff --git a/changes/4028.feature.md b/changes/4028.feature.md new file mode 100644 index 0000000000..1ff1882cab --- /dev/null +++ b/changes/4028.feature.md @@ -0,0 +1,7 @@ +Add three 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.initialized_regions` expresses that discovery as array regions, suitable to pass straight to `read_regions`. +- `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, initialized_regions(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/read.md b/docs/api/zarr/read.md index d84bddb8ac..87e84a1244 100644 --- a/docs/api/zarr/read.md +++ b/docs/api/zarr/read.md @@ -3,4 +3,5 @@ title: read --- ::: zarr.shards_initialized +::: zarr.initialized_regions ::: zarr.read_regions diff --git a/src/zarr/__init__.py b/src/zarr/__init__.py index d7686d5832..ec33511ddf 100644 --- a/src/zarr/__init__.py +++ b/src/zarr/__init__.py @@ -19,6 +19,7 @@ full, full_like, group, + initialized_regions, load, ones, ones_like, @@ -166,6 +167,7 @@ def set_format(log_format: str) -> None: "full", "full_like", "group", + "initialized_regions", "load", "ones", "ones_like", diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py index 9458f9bf4d..5a069ebd34 100644 --- a/src/zarr/api/asynchronous.py +++ b/src/zarr/api/asynchronous.py @@ -18,6 +18,7 @@ create_array, from_array, get_array_metadata, + initialized_regions, read_regions, shards_initialized, ) @@ -81,6 +82,7 @@ "full", "full_like", "group", + "initialized_regions", "load", "ones", "ones_like", diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py index d9f47713f6..331c5bfe3b 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -58,6 +58,7 @@ "full", "full_like", "group", + "initialized_regions", "load", "ones", "ones_like", @@ -1444,8 +1445,8 @@ def shards_initialized( 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 fetch and decode the populated regions, pass the - result of this function (or its regions) to [read_regions][zarr.read_regions]. + arrays it returns chunk keys. To turn these into array regions, use + [initialized_regions][zarr.initialized_regions]. Parameters ---------- @@ -1471,11 +1472,44 @@ def shards_initialized( See Also -------- - read_regions : Read and decode the populated regions of an array. + initialized_regions : 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 initialized_regions( + array: Array | 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.initialized_regions(_as_async_array(array), strategy=strategy)) + + def read_regions( array: Array | AsyncArray[Any], regions: Iterable[tuple[slice, ...]] | None = None, @@ -1485,12 +1519,12 @@ def read_regions( """ Read and decode array regions, returning a list of ``(region, data)`` pairs. - This is the spatially-resolved companion to [shards_initialized][zarr.shards_initialized]: - each pair associates a region (a tuple of slices into the array) with the decoded data - for that region, letting callers operate on 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. + Each pair associates a region (a tuple of slices into the array) with the decoded data + for that region. If ``regions`` is omitted, it defaults to the populated regions of the + array (see [initialized_regions][zarr.initialized_regions]), letting callers operate on + 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 ---------- @@ -1512,7 +1546,8 @@ def read_regions( See Also -------- - shards_initialized : Discover which shards of an array are populated. + initialized_regions : 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]]: diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 74bc3225e4..2138b728a9 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -3975,11 +3975,9 @@ async def _shards_initialized( [nchunks_initialized][zarr.Array.nchunks_initialized] """ - # Thin wrapper over the shared discovery core, projected to keys. The "list" - # strategy preserves this function's historical behavior (a single prefix - # listing); see [shards_initialized][zarr.shards_initialized] for the public, - # strategy-aware entry point. - return tuple(key for _, key in await _initialized_shards(array, strategy="list")) + # 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`` @@ -3989,108 +3987,116 @@ async def _shards_initialized( _PROBE_THRESHOLD = 64 -async def _initialized_shards( - array: AnyAsyncArray, +async def shards_initialized( + array: AnyArray | AnyAsyncArray, *, strategy: Literal["auto", "list", "probe"] = "auto", -) -> list[tuple[tuple[int, ...], str]]: +) -> tuple[str, ...]: """ - Discover the populated shards of an array, returning ``(coords, key)`` pairs in - chunk-grid order. This is the shared core of [shards_initialized][zarr.shards_initialized] - (which projects to keys) and [read_regions][zarr.read_regions] (which projects to regions). + 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 + [initialized_regions][zarr.initialized_regions]. + + 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 + -------- + initialized_regions : The array regions spanned by the populated shards. + read_regions : Read and decode a collection of array regions. """ - coords = list(_iter_shard_coords(array)) - keys = [array.metadata.encode_chunk_key(c) for c in coords] + if isinstance(array, Array): + array = array._async_array + + keys = list(_iter_shard_keys(array)) if strategy == "auto": - strategy = "probe" if len(coords) <= _PROBE_THRESHOLD else "list" + 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, addressing the case where the prefix - # holds many unrelated objects. + # 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) } - return [(c, k) for c, k in zip(coords, keys, strict=True) if k in contents] + 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 / k,) for k in keys], + [(array.store_path / key,) for key in keys], lambda store_path: store_path.exists(), zarr_config.get("async.concurrency"), ) - return [ - (c, k) - for (c, k), is_present in zip(zip(coords, keys, strict=True), present, strict=True) - if is_present - ] + 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 shards_initialized( +async def initialized_regions( array: AnyArray | AnyAsyncArray, *, strategy: Literal["auto", "list", "probe"] = "auto", -) -> tuple[str, ...]: +) -> list[tuple[slice, ...]]: """ - Return the storage keys of the shards that have been persisted to the store. + Return the array regions spanned by 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 fetch and decode the populated regions, pass the - result of this function (or its regions) to [read_regions][zarr.read_regions]. + 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. - - - ``"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. + How to discover which shards exist. See [shards_initialized][zarr.shards_initialized]. Returns ------- - tuple[str, ...] - The storage keys of the populated shards (or chunks, when unsharded), - in chunk-grid order. + list[tuple[slice, ...]] + The regions spanned by the populated shards, in chunk-grid order. See Also -------- - read_regions : Read and decode the populated regions of an array. + 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 - return tuple(key for _, key in await _initialized_shards(array, strategy=strategy)) - - -async def _initialized_regions( - array: AnyAsyncArray, - *, - strategy: Literal["auto", "list", "probe"] = "auto", -) -> list[tuple[slice, ...]]: - """Return the array regions spanned by each populated shard, in chunk-grid order.""" - shard_shape = array.shards if array.shards is not None else array.chunks + 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 [ - tuple( - slice(c * s, min((c + 1) * s, dim)) - for c, s, dim in zip(coords, shard_shape, array.shape, strict=True) - ) - for coords, _ in await _initialized_shards(array, strategy=strategy) + region + for region, key in zip(_iter_shard_regions(array), _iter_shard_keys(array), strict=True) + if key in initialized ] @@ -4104,10 +4110,11 @@ async def read_regions( Concurrently read and decode array regions, yielding each ``(region, data)`` pair as soon as its data is available. - This is the spatially-resolved companion to [shards_initialized][zarr.shards_initialized]: - each yielded value pairs a region (a tuple of slices into the array) with the decoded - data for that region, so callers can stream over only the populated parts of a sparse - array without materializing the full array. + Each yielded value pairs a region (a tuple of slices into the array) with the decoded + data for that region. If ``regions`` is omitted, it defaults to the populated regions + of the array (see [initialized_regions][zarr.initialized_regions]), so callers can + stream over only the populated parts of a sparse array without materializing the full + array. Parameters ---------- @@ -4129,14 +4136,15 @@ async def read_regions( See Also -------- - shards_initialized : Discover which shards of an array are populated. + initialized_regions : 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") - region_list = await _initialized_regions(array) if regions is None else list(regions) + region_list = await initialized_regions(array) if regions is None else list(regions) semaphore = Semaphore(concurrency) diff --git a/tests/benchmarks/test_indexing.py b/tests/benchmarks/test_indexing.py index 385a85b5b5..0090b4304a 100644 --- a/tests/benchmarks/test_indexing.py +++ b/tests/benchmarks/test_indexing.py @@ -11,7 +11,7 @@ import pytest -from zarr import create_array +from zarr import create_array, read_regions indexers = ( (0,) * 3, @@ -277,3 +277,48 @@ def write_with_cache_clear() -> None: data[indexer] = write_data benchmark(write_with_cache_clear) + + +# Sparse-read benchmark: most chunks empty (resolve to the fill value). +sparse_shards = ( + None, + (64,), +) + + +@pytest.mark.parametrize("store", ["memory", "memory_get_latency"], indirect=["store"]) +@pytest.mark.parametrize("shards", sparse_shards, ids=str) +@pytest.mark.parametrize("reader", ["full", "read_regions"], ids=str) +def test_sparse_read( + store: Store, + shards: tuple[int, ...] | None, + reader: str, + benchmark: BenchmarkFixture, +) -> None: + """Benchmark reading a sparse array (most chunks empty) two ways. + + ``full`` is the stock ``arr[:]`` read, which issues a store request for every chunk + including the empty ones; ``read_regions`` touches only the populated chunks. The gap + is largest on ``memory_get_latency``, where each skipped empty-chunk request avoids a + round-trip. + """ + n_chunks = 256 + chunk = 16 + data = create_array( + store=store, + shape=(n_chunks * chunk,), + dtype="uint8", + chunks=(chunk,), + shards=shards, + compressors=None, + filters=None, + fill_value=0, + ) + # populate ~3% of chunks, spread across the array + for ci in range(0, n_chunks, 32): + data[ci * chunk : ci * chunk + chunk] = 1 + + if reader == "full": + benchmark(getitem, data, slice(None)) + else: + benchmark(read_regions, data) diff --git a/tests/test_array.py b/tests/test_array.py index 0d6d2d5906..971533f394 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2374,3 +2374,182 @@ 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 / initialized_regions / read_regions --------------------------- + + +def _ca_sparse_1d(store: Store) -> tuple[Array, np.ndarray]: + 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, np.ndarray]: + 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, np.ndarray]: + 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, np.ndarray]: + # 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, np.ndarray]: + 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, np.ndarray]: + 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, results: list, baseline: np.ndarray) -> np.ndarray: + """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: str) -> 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] + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +def test_list_strategy_ignores_non_chunk_objects(store: Store) -> None: + """The ``list`` strategy must not mistake unrelated objects sharing the array's + prefix (e.g. metadata) for populated chunks.""" + arr, _ = _ca_sparse_1d(store) + keys = set(zarr.shards_initialized(arr, strategy="list")) + assert keys == {"c/1", "c/5"} + assert all(k.startswith("c/") for k in keys) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize("setup_name", list(_CA_SETUPS)) +def test_initialized_regions_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.initialized_regions(arr)) == len(zarr.shards_initialized(arr)) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +def test_initialized_regions_values(store: Store) -> None: + arr, _ = _ca_sparse_1d(store) + assert set(zarr.initialized_regions(arr)) == {(slice(8, 16, 1),), (slice(40, 48, 1),)} + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize("setup_name", list(_CA_SETUPS)) +def test_read_regions_default_reconstructs_baseline(store: Store, setup_name: str) -> None: + """With no explicit regions, ``read_regions`` reads the populated regions; 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) + result = _ca_pack(arr, results, baseline) + assert np.array_equal(result, baseline) + assert result.dtype == baseline.dtype + # the default region set is exactly the initialized regions + assert len(results) == len(zarr.initialized_regions(arr)) + + +@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.initialized_regions(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.initialized_regions(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"} diff --git a/tests/test_chunk_access.py b/tests/test_chunk_access.py deleted file mode 100644 index 8ee6041a70..0000000000 --- a/tests/test_chunk_access.py +++ /dev/null @@ -1,194 +0,0 @@ -"""Tests for the shard-discovery and region-read primitives. - -These cover :func:`zarr.shards_initialized` (discover which shards/chunks of an -array are populated) and :func:`zarr.read_regions` (concurrently read and decode -the populated regions), along with their asynchronous counterparts in -:mod:`zarr.api.asynchronous`. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import numpy as np -import pytest - -import zarr -import zarr.api.asynchronous as async_api - -if TYPE_CHECKING: - from collections.abc import Callable - - from zarr.abc.store import Store - - -def _sparse_1d(store: Store) -> tuple[zarr.Array, np.ndarray]: - 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 _dense_1d(store: Store) -> tuple[zarr.Array, np.ndarray]: - 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 _sparse_2d(store: Store) -> tuple[zarr.Array, np.ndarray]: - 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 _sharded_sparse(store: Store) -> tuple[zarr.Array, np.ndarray]: - # 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 _all_empty(store: Store) -> tuple[zarr.Array, np.ndarray]: - arr = zarr.create_array(store=store, shape=(32,), chunks=(8,), dtype="int32", fill_value=7) - return arr, np.asarray(arr[:]) - - -def _all_populated(store: Store) -> tuple[zarr.Array, np.ndarray]: - 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[:]) - - -SETUPS: dict[str, Callable[[Store], tuple[zarr.Array, np.ndarray]]] = { - "sparse_1d": _sparse_1d, - "dense_1d": _dense_1d, - "sparse_2d": _sparse_2d, - "sharded_sparse": _sharded_sparse, - "all_empty": _all_empty, - "all_populated": _all_populated, -} - -STRATEGIES = ["auto", "list", "probe"] - - -def _pack(arr: zarr.Array, regions: list, baseline: np.ndarray) -> np.ndarray: - """Scatter ``(region, data)`` pairs onto a fill-valued array.""" - out = np.full(baseline.shape, arr.fill_value, dtype=baseline.dtype) - for region, data in regions: - out[region] = np.asarray(data) - return out - - -@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) -@pytest.mark.parametrize("setup_name", list(SETUPS)) -@pytest.mark.parametrize("strategy", STRATEGIES) -def test_shards_initialized_strategies_agree(store: Store, setup_name: str, strategy: str) -> None: - """Every strategy reports the same set of populated keys, and reports the - expected count for a hand-known layout.""" - arr, _ = SETUPS[setup_name](store) - keys = set(zarr.shards_initialized(arr, strategy=strategy)) - # all strategies must agree - 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, _ = 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, _ = _sparse_1d(store) - with pytest.raises(ValueError, match="Unknown strategy"): - zarr.shards_initialized(arr, strategy="nonsense") # type: ignore[arg-type] - - -@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) -def test_list_strategy_ignores_non_chunk_objects(store: Store) -> None: - """The ``list`` strategy must not mistake unrelated objects sharing the - array's prefix (e.g. metadata) for populated chunks.""" - arr, _ = _sparse_1d(store) - # metadata (zarr.json) already lives under the array prefix; add another - # non-chunk object to be sure it is excluded. - keys = set(zarr.shards_initialized(arr, strategy="list")) - assert keys == {"c/1", "c/5"} - assert all(k.startswith("c/") for k in keys) - - -@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) -@pytest.mark.parametrize("setup_name", list(SETUPS)) -def test_read_regions_reconstructs_baseline(store: Store, setup_name: str) -> None: - """Packing the populated regions onto a fill-valued array reproduces the - full ``arr[:]`` read exactly.""" - arr, baseline = SETUPS[setup_name](store) - regions = zarr.read_regions(arr) - result = _pack(arr, regions, 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", list(SETUPS)) -def test_read_regions_default_count_matches_discovery(store: Store, setup_name: str) -> None: - """With no explicit regions, ``read_regions`` reads exactly the populated - shards discovered by ``shards_initialized``.""" - arr, _ = SETUPS[setup_name](store) - regions = zarr.read_regions(arr) - assert len(regions) == len(zarr.shards_initialized(arr)) - - -@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 = _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 = _sparse_2d(store) - regions = zarr.read_regions(arr, concurrency=1) - assert np.array_equal(_pack(arr, regions, baseline), baseline) - - -@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) -@pytest.mark.parametrize("setup_name", list(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, _ = SETUPS[setup_name](store) - async_pairs = { - region: np.asarray(data).tobytes() - async for region, data in async_api.read_regions(arr._async_array) - } - sync_pairs = {region: np.asarray(data).tobytes() for region, data in zarr.read_regions(arr)} - assert async_pairs == sync_pairs - - -@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) -async def test_shards_initialized_async(store: Store) -> None: - arr, _ = _sparse_1d(store) - keys = await async_api.shards_initialized(arr._async_array) - assert set(keys) == {"c/1", "c/5"} From aec3260a7893809bfda6fae90b66bb9920698f7c Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Wed, 3 Jun 2026 17:10:56 -0700 Subject: [PATCH 4/9] switching to caller to provide regions --- src/zarr/api/synchronous.py | 19 +++++++++---------- src/zarr/core/array.py | 21 ++++++++------------- tests/benchmarks/test_indexing.py | 10 +++++----- tests/test_array.py | 10 ++++------ 4 files changed, 26 insertions(+), 34 deletions(-) diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py index 331c5bfe3b..89457ddcf1 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -1512,28 +1512,27 @@ def initialized_regions( def read_regions( array: Array | AsyncArray[Any], - regions: Iterable[tuple[slice, ...]] | None = None, + regions: Iterable[tuple[slice, ...]], *, concurrency: int | None = None, ) -> list[tuple[tuple[slice, ...], NDArrayLikeOrScalar]]: """ - Read and decode array regions, returning a list of ``(region, data)`` pairs. + 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. If ``regions`` is omitted, it defaults to the populated regions of the - array (see [initialized_regions][zarr.initialized_regions]), letting callers operate on - 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. + for that region. The regions to read are supplied by the caller; pass the result of + [initialized_regions][zarr.initialized_regions] 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, optional + regions : iterable of tuple of slice The regions to read. Each region is a tuple of slices, one per array dimension. - If omitted, defaults to the regions spanned by the populated shards of ``array`` - (i.e. every region that holds data). concurrency : int, optional The maximum number of regions read concurrently. Defaults to the ``async.concurrency`` config value. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 2138b728a9..f2fae9b50f 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -4102,28 +4102,25 @@ async def initialized_regions( async def read_regions( array: AnyArray | AnyAsyncArray, - regions: Iterable[tuple[slice, ...]] | None = None, + regions: Iterable[tuple[slice, ...]], *, concurrency: int | None = None, ) -> AsyncIterator[tuple[tuple[slice, ...], NDArrayLikeOrScalar]]: """ - Concurrently read and decode array regions, yielding each ``(region, data)`` pair - as soon as its data is available. + 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. If ``regions`` is omitted, it defaults to the populated regions - of the array (see [initialized_regions][zarr.initialized_regions]), so callers can - stream over only the populated parts of a sparse array without materializing the full - array. + data for that region. The regions to read are supplied by the caller; pass the result + of [initialized_regions][zarr.initialized_regions] 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, optional + regions : iterable of tuple of slice The regions to read. Each region is a tuple of slices, one per array dimension. - If omitted, defaults to the regions spanned by the populated shards of ``array`` - (i.e. every region that holds data). concurrency : int, optional The maximum number of regions read concurrently. Defaults to the ``async.concurrency`` config value. @@ -4144,8 +4141,6 @@ async def read_regions( if concurrency is None: concurrency = zarr_config.get("async.concurrency") - region_list = await initialized_regions(array) if regions is None else list(regions) - semaphore = Semaphore(concurrency) async def _read( @@ -4154,7 +4149,7 @@ async def _read( async with semaphore: return region, await array.getitem(region) - for future in as_completed([ensure_future(_read(region)) for region in region_list]): + for future in as_completed([ensure_future(_read(region)) for region in regions]): yield await future diff --git a/tests/benchmarks/test_indexing.py b/tests/benchmarks/test_indexing.py index 0090b4304a..021c0b1e04 100644 --- a/tests/benchmarks/test_indexing.py +++ b/tests/benchmarks/test_indexing.py @@ -11,7 +11,7 @@ import pytest -from zarr import create_array, read_regions +from zarr import create_array, initialized_regions, read_regions indexers = ( (0,) * 3, @@ -298,9 +298,9 @@ def test_sparse_read( """Benchmark reading a sparse array (most chunks empty) two ways. ``full`` is the stock ``arr[:]`` read, which issues a store request for every chunk - including the empty ones; ``read_regions`` touches only the populated chunks. The gap - is largest on ``memory_get_latency``, where each skipped empty-chunk request avoids a - round-trip. + including the empty ones; ``read_regions`` discovers the populated regions and reads + only those. The gap is largest on ``memory_get_latency``, where each skipped + empty-chunk request avoids a round-trip. """ n_chunks = 256 chunk = 16 @@ -321,4 +321,4 @@ def test_sparse_read( if reader == "full": benchmark(getitem, data, slice(None)) else: - benchmark(read_regions, data) + benchmark(lambda: read_regions(data, initialized_regions(data))) diff --git a/tests/test_array.py b/tests/test_array.py index 971533f394..938a9a4dbe 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2500,16 +2500,14 @@ def test_initialized_regions_values(store: Store) -> None: @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) @pytest.mark.parametrize("setup_name", list(_CA_SETUPS)) -def test_read_regions_default_reconstructs_baseline(store: Store, setup_name: str) -> None: - """With no explicit regions, ``read_regions`` reads the populated regions; scattering - them onto a fill-valued array reproduces the full ``arr[:]`` read exactly.""" +def test_read_initialized_regions_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) + results = zarr.read_regions(arr, zarr.initialized_regions(arr)) result = _ca_pack(arr, results, baseline) assert np.array_equal(result, baseline) assert result.dtype == baseline.dtype - # the default region set is exactly the initialized regions - assert len(results) == len(zarr.initialized_regions(arr)) @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) From cc94ab8203aa5c19828c4f6bd760c6ba8add149c Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Wed, 3 Jun 2026 17:22:43 -0700 Subject: [PATCH 5/9] linting fixes --- src/zarr/api/synchronous.py | 8 ++++---- tests/test_array.py | 22 ++++++++++++++-------- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py index 89457ddcf1..780129f34c 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -1431,12 +1431,12 @@ def zeros_like(a: ArrayLike, **kwargs: Any) -> AnyArray: return Array(sync(async_api.zeros_like(a, **kwargs))) -def _as_async_array(array: Array | AsyncArray[Any]) -> AsyncArray[Any]: +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 | AsyncArray[Any], + array: Array[Any] | AsyncArray[Any], *, strategy: Literal["auto", "list", "probe"] = "auto", ) -> tuple[str, ...]: @@ -1479,7 +1479,7 @@ def shards_initialized( def initialized_regions( - array: Array | AsyncArray[Any], + array: Array[Any] | AsyncArray[Any], *, strategy: Literal["auto", "list", "probe"] = "auto", ) -> list[tuple[slice, ...]]: @@ -1511,7 +1511,7 @@ def initialized_regions( def read_regions( - array: Array | AsyncArray[Any], + array: Array[Any] | AsyncArray[Any], regions: Iterable[tuple[slice, ...]], *, concurrency: int | None = None, diff --git a/tests/test_array.py b/tests/test_array.py index 938a9a4dbe..ca29195c4b 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2379,7 +2379,7 @@ async def test_create_array_chunks_3d( # --- shards_initialized / initialized_regions / read_regions --------------------------- -def _ca_sparse_1d(store: Store) -> tuple[Array, np.ndarray]: +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") @@ -2387,20 +2387,20 @@ def _ca_sparse_1d(store: Store) -> tuple[Array, np.ndarray]: return arr, np.asarray(arr[:]) -def _ca_dense_1d(store: Store) -> tuple[Array, np.ndarray]: +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, np.ndarray]: +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, np.ndarray]: +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 @@ -2410,12 +2410,12 @@ def _ca_sharded_sparse(store: Store) -> tuple[Array, np.ndarray]: return arr, np.asarray(arr[:]) -def _ca_all_empty(store: Store) -> tuple[Array, np.ndarray]: +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, np.ndarray]: +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[:]) @@ -2432,7 +2432,11 @@ def _ca_all_populated(store: Store) -> tuple[Array, np.ndarray]: _CA_STRATEGIES = ["auto", "list", "probe"] -def _ca_pack(arr: Array, results: list, baseline: np.ndarray) -> np.ndarray: +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: @@ -2443,7 +2447,9 @@ def _ca_pack(arr: Array, results: list, baseline: np.ndarray) -> np.ndarray: @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: str) -> None: +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)) From b8d1d8c093925aa6838753ad4a9a2bf2a05a1d42 Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Wed, 3 Jun 2026 17:33:01 -0700 Subject: [PATCH 6/9] bumping codecov to include final three lines... --- tests/test_array.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_array.py b/tests/test_array.py index ca29195c4b..b3c7046572 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2557,3 +2557,16 @@ 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.initialized_regions(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 From 8e601501cdff99cfcebefe3eeb3aaf4dfbdeaf1c Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Thu, 25 Jun 2026 17:50:53 -0700 Subject: [PATCH 7/9] update docstring, tests --- src/zarr/core/array.py | 15 +++++++++++++++ tests/test_array.py | 40 +++++++++++++++++++++++++++++----------- 2 files changed, 44 insertions(+), 11 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 5cbe597532..09826087e9 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -4019,6 +4019,15 @@ async def shards_initialized( 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 -------- initialized_regions : The array regions spanned by the populated shards. @@ -4080,6 +4089,12 @@ async def initialized_regions( 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. diff --git a/tests/test_array.py b/tests/test_array.py index b3c7046572..509cc0336a 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2480,14 +2480,16 @@ def test_shards_initialized_unknown_strategy(store: Store) -> None: zarr.shards_initialized(arr, strategy="nonsense") # type: ignore[arg-type] -@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) -def test_list_strategy_ignores_non_chunk_objects(store: Store) -> None: - """The ``list`` strategy must not mistake unrelated objects sharing the array's - prefix (e.g. metadata) for populated chunks.""" - arr, _ = _ca_sparse_1d(store) - keys = set(zarr.shards_initialized(arr, strategy="list")) - assert keys == {"c/1", "c/5"} - assert all(k.startswith("c/") for k in keys) +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"]) @@ -2499,9 +2501,25 @@ def test_initialized_regions_count_matches_keys(store: Store, setup_name: str) - @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) -def test_initialized_regions_values(store: Store) -> None: - arr, _ = _ca_sparse_1d(store) - assert set(zarr.initialized_regions(arr)) == {(slice(8, 16, 1),), (slice(40, 48, 1),)} +@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_initialized_regions_are_exactly_written( + store: Store, regions: tuple[tuple[slice, ...], ...] +) -> None: + """Writing a set of chunk-aligned regions makes ``initialized_regions`` 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.initialized_regions(arr)) == set(regions) @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) From dc6740b91207d6e41ffc4c4ad891d44cc0d6084a Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Wed, 22 Jul 2026 14:53:06 -0700 Subject: [PATCH 8/9] parity for initialized/uninitialized chunks within shards; fixing merge conflicts --- changes/4028.feature.md | 3 +- src/zarr/__init__.py | 2 + src/zarr/api/asynchronous.py | 2 + src/zarr/api/synchronous.py | 43 +++++++++++++ src/zarr/core/array.py | 116 +++++++++++++++++++++++++++++++++++ tests/test_array.py | 48 +++++++++++++++ 6 files changed, 213 insertions(+), 1 deletion(-) diff --git a/changes/4028.feature.md b/changes/4028.feature.md index 1ff1882cab..2ad3fa0e01 100644 --- a/changes/4028.feature.md +++ b/changes/4028.feature.md @@ -1,7 +1,8 @@ -Add three composable functions for efficiently reading sparse arrays, where most chunks are empty and resolve to the fill value: +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.initialized_regions` expresses that discovery as array regions, suitable to pass straight to `read_regions`. +- `zarr.initialized_chunk_regions` 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 `initialized_regions`. - `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, initialized_regions(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/src/zarr/__init__.py b/src/zarr/__init__.py index ec33511ddf..f4a0f1389a 100644 --- a/src/zarr/__init__.py +++ b/src/zarr/__init__.py @@ -19,6 +19,7 @@ full, full_like, group, + initialized_chunk_regions, initialized_regions, load, ones, @@ -167,6 +168,7 @@ def set_format(log_format: str) -> None: "full", "full_like", "group", + "initialized_chunk_regions", "initialized_regions", "load", "ones", diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py index 68bf69656f..f1f5f54ad0 100644 --- a/src/zarr/api/asynchronous.py +++ b/src/zarr/api/asynchronous.py @@ -18,6 +18,7 @@ create_array, from_array, get_array_metadata, + initialized_chunk_regions, initialized_regions, read_regions, shards_initialized, @@ -82,6 +83,7 @@ "full", "full_like", "group", + "initialized_chunk_regions", "initialized_regions", "load", "ones", diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py index 9c252363cd..8bf0b139e3 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -58,6 +58,7 @@ "full", "full_like", "group", + "initialized_chunk_regions", "initialized_regions", "load", "ones", @@ -1529,6 +1530,48 @@ def initialized_regions( return sync(async_api.initialized_regions(_as_async_array(array), strategy=strategy)) +def initialized_chunk_regions( + 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 [initialized_regions][zarr.initialized_regions], 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 [initialized_regions][zarr.initialized_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]. + 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 + -------- + initialized_regions : The array regions at stored-object (shard) granularity. + read_regions : Read and decode a collection of array regions. + """ + return sync( + async_api.initialized_chunk_regions( + _as_async_array(array), strategy=strategy, concurrency=concurrency + ) + ) + + def read_regions( array: Array[Any] | AsyncArray[Any], regions: Iterable[tuple[slice, ...]], diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 09121532f5..e85e3f07fe 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -4210,6 +4210,122 @@ async def _read( 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 initialized_chunk_regions( + 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 [initialized_regions][zarr.initialized_regions], 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 [initialized_regions][zarr.initialized_regions]. + + 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 + -------- + initialized_regions : 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 initialized_regions(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 + ] + + # 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] | ArrayArrayCodec diff --git a/tests/test_array.py b/tests/test_array.py index 41f9b566d4..b9e2fb8361 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2546,6 +2546,54 @@ def test_read_initialized_regions_reconstructs_baseline(store: Store, setup_name 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_initialized_chunk_regions_unsharded_matches_initialized_regions( + 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.initialized_chunk_regions(arr) == zarr.initialized_regions(arr) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +def test_initialized_chunk_regions_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 ``initialized_regions``.""" + # 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.initialized_regions(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.initialized_chunk_regions(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_initialized_chunk_regions_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.initialized_chunk_regions(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.""" From 3b2aeefb63526ca2320d50d43b88f3f62449d93d Mon Sep 17 00:00:00 2001 From: Shane Grigsby Date: Wed, 22 Jul 2026 15:12:05 -0700 Subject: [PATCH 9/9] standardize names to {chunk_regions,regions}_initialized to match shards_initialized convention --- changes/4028.feature.md | 6 +-- .../functions/chunk_regions_initialized.md | 5 +++ .../functions/initialized_chunk_regions.md | 5 --- .../api/zarr/functions/initialized_regions.md | 5 --- .../api/zarr/functions/regions_initialized.md | 5 +++ mkdocs.yml | 6 +-- src/zarr/__init__.py | 8 ++-- src/zarr/api/asynchronous.py | 8 ++-- src/zarr/api/synchronous.py | 26 ++++++------- src/zarr/core/array.py | 20 +++++----- tests/test_array.py | 38 +++++++++---------- 11 files changed, 66 insertions(+), 66 deletions(-) create mode 100644 docs/api/zarr/functions/chunk_regions_initialized.md delete mode 100644 docs/api/zarr/functions/initialized_chunk_regions.md delete mode 100644 docs/api/zarr/functions/initialized_regions.md create mode 100644 docs/api/zarr/functions/regions_initialized.md diff --git a/changes/4028.feature.md b/changes/4028.feature.md index 2ad3fa0e01..7f53605db5 100644 --- a/changes/4028.feature.md +++ b/changes/4028.feature.md @@ -1,8 +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.initialized_regions` expresses that discovery as array regions, suitable to pass straight to `read_regions`. -- `zarr.initialized_chunk_regions` 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 `initialized_regions`. +- `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, initialized_regions(arr))`. Asynchronous versions are available in `zarr.api.asynchronous`; the async `read_regions` streams each region as soon as its data is available. +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/initialized_chunk_regions.md b/docs/api/zarr/functions/initialized_chunk_regions.md deleted file mode 100644 index 8e8f00b682..0000000000 --- a/docs/api/zarr/functions/initialized_chunk_regions.md +++ /dev/null @@ -1,5 +0,0 @@ ---- -title: zarr.initialized_chunk_regions ---- - -::: zarr.initialized_chunk_regions diff --git a/docs/api/zarr/functions/initialized_regions.md b/docs/api/zarr/functions/initialized_regions.md deleted file mode 100644 index 36ab6844a0..0000000000 --- a/docs/api/zarr/functions/initialized_regions.md +++ /dev/null @@ -1,5 +0,0 @@ ---- -title: zarr.initialized_regions ---- - -::: zarr.initialized_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/mkdocs.yml b/mkdocs.yml index 6325645cca..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 @@ -68,8 +69,6 @@ nav: - ' zarr.full_like': api/zarr/functions/full_like.md - ' zarr.Group': api/zarr/group.md - ' zarr.group': api/zarr/functions/group.md - - ' zarr.initialized_chunk_regions': api/zarr/functions/initialized_chunk_regions.md - - ' zarr.initialized_regions': api/zarr/functions/initialized_regions.md - ' zarr.load': api/zarr/functions/load.md - ' zarr.metadata': api/zarr/metadata.md - ' zarr.ones': api/zarr/functions/ones.md @@ -80,8 +79,9 @@ 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.registry': api/zarr/registry.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 diff --git a/src/zarr/__init__.py b/src/zarr/__init__.py index f4a0f1389a..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, @@ -19,8 +20,6 @@ full, full_like, group, - initialized_chunk_regions, - initialized_regions, load, ones, ones_like, @@ -30,6 +29,7 @@ open_group, open_like, read_regions, + regions_initialized, save, save_array, save_group, @@ -153,6 +153,7 @@ def set_format(log_format: str) -> None: "Group", "__version__", "array", + "chunk_regions_initialized", "config", "consolidate_metadata", "copy", @@ -168,8 +169,6 @@ def set_format(log_format: str) -> None: "full", "full_like", "group", - "initialized_chunk_regions", - "initialized_regions", "load", "ones", "ones_like", @@ -180,6 +179,7 @@ def set_format(log_format: str) -> None: "open_like", "print_debug_info", "read_regions", + "regions_initialized", "save", "save_array", "save_group", diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py index f1f5f54ad0..547f5b94bf 100644 --- a/src/zarr/api/asynchronous.py +++ b/src/zarr/api/asynchronous.py @@ -15,12 +15,12 @@ Array, AsyncArray, CompressorLike, + chunk_regions_initialized, create_array, from_array, get_array_metadata, - initialized_chunk_regions, - initialized_regions, read_regions, + regions_initialized, shards_initialized, ) from zarr.core.array_spec import ArrayConfigLike, parse_array_config @@ -70,6 +70,7 @@ __all__ = [ "array", + "chunk_regions_initialized", "consolidate_metadata", "copy", "copy_all", @@ -83,8 +84,6 @@ "full", "full_like", "group", - "initialized_chunk_regions", - "initialized_regions", "load", "ones", "ones_like", @@ -94,6 +93,7 @@ "open_group", "open_like", "read_regions", + "regions_initialized", "save", "save_array", "save_group", diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py index 8bf0b139e3..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", @@ -58,8 +59,6 @@ "full", "full_like", "group", - "initialized_chunk_regions", - "initialized_regions", "load", "ones", "ones_like", @@ -69,6 +68,7 @@ "open_group", "open_like", "read_regions", + "regions_initialized", "save", "save_array", "save_group", @@ -1466,7 +1466,7 @@ def shards_initialized( 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 - [initialized_regions][zarr.initialized_regions]. + [regions_initialized][zarr.regions_initialized]. Parameters ---------- @@ -1492,13 +1492,13 @@ def shards_initialized( See Also -------- - initialized_regions : The array regions spanned by the populated shards. + 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 initialized_regions( +def regions_initialized( array: Array[Any] | AsyncArray[Any], *, strategy: Literal["auto", "list", "probe"] = "auto", @@ -1527,10 +1527,10 @@ def initialized_regions( shards_initialized : The storage keys of the populated shards. read_regions : Read and decode a collection of array regions. """ - return sync(async_api.initialized_regions(_as_async_array(array), strategy=strategy)) + return sync(async_api.regions_initialized(_as_async_array(array), strategy=strategy)) -def initialized_chunk_regions( +def chunk_regions_initialized( array: Array[Any] | AsyncArray[Any], *, strategy: Literal["auto", "list", "probe"] = "auto", @@ -1539,11 +1539,11 @@ def initialized_chunk_regions( """ Return the array regions spanned by the inner chunks that have been written. - Unlike [initialized_regions][zarr.initialized_regions], which reports at stored-object + 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 [initialized_regions][zarr.initialized_regions]. + identical to [regions_initialized][zarr.regions_initialized]. Parameters ---------- @@ -1562,11 +1562,11 @@ def initialized_chunk_regions( See Also -------- - initialized_regions : The array regions at stored-object (shard) granularity. + regions_initialized : The array regions at stored-object (shard) granularity. read_regions : Read and decode a collection of array regions. """ return sync( - async_api.initialized_chunk_regions( + async_api.chunk_regions_initialized( _as_async_array(array), strategy=strategy, concurrency=concurrency ) ) @@ -1584,7 +1584,7 @@ def read_regions( 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 - [initialized_regions][zarr.initialized_regions] to read only the populated parts of a + [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. @@ -1607,7 +1607,7 @@ def read_regions( See Also -------- - initialized_regions : The array regions spanned by the populated shards. + regions_initialized : The array regions spanned by the populated shards. shards_initialized : The storage keys of the populated shards. """ diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index e85e3f07fe..5b63f2ec6d 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -4037,7 +4037,7 @@ async def shards_initialized( 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 - [initialized_regions][zarr.initialized_regions]. + [regions_initialized][zarr.regions_initialized]. Parameters ---------- @@ -4072,7 +4072,7 @@ async def shards_initialized( See Also -------- - initialized_regions : The array regions spanned by the populated shards. + regions_initialized : The array regions spanned by the populated shards. read_regions : Read and decode a collection of array regions. """ if isinstance(array, Array): @@ -4110,7 +4110,7 @@ async def shards_initialized( ) -async def initialized_regions( +async def regions_initialized( array: AnyArray | AnyAsyncArray, *, strategy: Literal["auto", "list", "probe"] = "auto", @@ -4169,7 +4169,7 @@ async def read_regions( 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 [initialized_regions][zarr.initialized_regions] to read only the populated parts of + of [regions_initialized][zarr.regions_initialized] to read only the populated parts of a sparse array without materializing the full array. Parameters @@ -4190,7 +4190,7 @@ async def read_regions( See Also -------- - initialized_regions : The array regions spanned by the populated shards. + regions_initialized : The array regions spanned by the populated shards. shards_initialized : The storage keys of the populated shards. """ if isinstance(array, Array): @@ -4221,7 +4221,7 @@ def _sharding_codec(array: AnyAsyncArray) -> Any: return None -async def initialized_chunk_regions( +async def chunk_regions_initialized( array: AnyArray | AnyAsyncArray, *, strategy: Literal["auto", "list", "probe"] = "auto", @@ -4230,12 +4230,12 @@ async def initialized_chunk_regions( """ Return the array regions spanned by the *inner chunks* that have been written. - Unlike [initialized_regions][zarr.initialized_regions], which reports at stored-object + 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 [initialized_regions][zarr.initialized_regions]. + 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 @@ -4258,7 +4258,7 @@ async def initialized_chunk_regions( See Also -------- - initialized_regions : The array regions at stored-object (shard) granularity. + regions_initialized : The array regions at stored-object (shard) granularity. read_regions : Read and decode a collection of array regions. """ if isinstance(array, Array): @@ -4268,7 +4268,7 @@ async def initialized_chunk_regions( 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 initialized_regions(array, strategy=strategy) + return await regions_initialized(array, strategy=strategy) if concurrency is None: concurrency = zarr_config.get("async.concurrency") diff --git a/tests/test_array.py b/tests/test_array.py index b9e2fb8361..0f6beed5e1 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2388,7 +2388,7 @@ async def test_create_array_chunks_3d( assert arr.write_chunk_sizes == expected -# --- shards_initialized / initialized_regions / read_regions --------------------------- +# --- shards_initialized / regions_initialized / read_regions --------------------------- def _ca_sparse_1d(store: Store) -> tuple[Array[Any], npt.NDArray[Any]]: @@ -2506,10 +2506,10 @@ async def test_list_strategy_ignores_non_chunk_objects() -> None: @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) @pytest.mark.parametrize("setup_name", list(_CA_SETUPS)) -def test_initialized_regions_count_matches_keys(store: Store, setup_name: str) -> None: +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.initialized_regions(arr)) == len(zarr.shards_initialized(arr)) + assert len(zarr.regions_initialized(arr)) == len(zarr.shards_initialized(arr)) @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) @@ -2523,24 +2523,24 @@ def test_initialized_regions_count_matches_keys(store: Store, setup_name: str) - ], ids=["none", "single", "adjacent", "non_adjacent"], ) -def test_initialized_regions_are_exactly_written( +def test_regions_initialized_are_exactly_written( store: Store, regions: tuple[tuple[slice, ...], ...] ) -> None: - """Writing a set of chunk-aligned regions makes ``initialized_regions`` report + """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.initialized_regions(arr)) == set(regions) + 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_initialized_regions_reconstructs_baseline(store: Store, setup_name: str) -> None: +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.initialized_regions(arr)) + 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 @@ -2548,19 +2548,19 @@ def test_read_initialized_regions_reconstructs_baseline(store: Store, setup_name @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) @pytest.mark.parametrize("setup_name", ["sparse_1d", "dense_1d", "sparse_2d", "all_empty"]) -def test_initialized_chunk_regions_unsharded_matches_initialized_regions( +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.initialized_chunk_regions(arr) == zarr.initialized_regions(arr) + assert zarr.chunk_regions_initialized(arr) == zarr.regions_initialized(arr) @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) -def test_initialized_chunk_regions_sharded_skips_empty_inner_chunks(store: Store) -> None: +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 ``initialized_regions``.""" + 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 @@ -2569,13 +2569,13 @@ def test_initialized_chunk_regions_sharded_skips_empty_inner_chunks(store: Store 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.initialized_regions(arr) == [ + 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.initialized_chunk_regions(arr) == [ + 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)), @@ -2584,13 +2584,13 @@ def test_initialized_chunk_regions_sharded_skips_empty_inner_chunks(store: Store @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) @pytest.mark.parametrize("setup_name", list(_CA_SETUPS)) -def test_read_initialized_chunk_regions_reconstructs_baseline( +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.initialized_chunk_regions(arr)) + results = zarr.read_regions(arr, zarr.chunk_regions_initialized(arr)) assert np.array_equal(_ca_pack(arr, results, baseline), baseline) @@ -2609,7 +2609,7 @@ def test_read_regions_explicit_regions(store: Store) -> None: 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.initialized_regions(arr), concurrency=1) + results = zarr.read_regions(arr, zarr.regions_initialized(arr), concurrency=1) assert np.array_equal(_ca_pack(arr, results, baseline), baseline) @@ -2619,7 +2619,7 @@ async def test_read_regions_async_matches_sync(store: Store, setup_name: str) -> """The async streaming generator yields the same ``(region, data)`` set as the synchronous wrapper.""" arr, _ = _CA_SETUPS[setup_name](store) - regions = zarr.initialized_regions(arr) + 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) @@ -2643,7 +2643,7 @@ async def test_async_chunk_access_accepts_sync_array(store: Store) -> None: underlying async array.""" arr, _ = _ca_sparse_1d(store) keys = await zarr.api.asynchronous.shards_initialized(arr) - regions = await zarr.api.asynchronous.initialized_regions(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),)}