[bhaskargurram-ai]

Summary

Closes #2529.

zarr.Array has no __iter__, so for row in z falls back to Python's sequence protocol. That protocol calls z[0], z[1], and so on until IndexError. Each call reads and decodes every chunk that row touches. A chunk holding n rows is therefore read and decoded n times per pass, and on a remote store each row becomes a fresh set of requests. zarr v2 avoided this with a chunk-aware Array.__iter__/islice.

This PR adds Array.__iter__, following the direction discussed in the issue: chunk-aware iteration with a size check.

• One row of chunks at a time. Iteration walks the chunk grid along the first dimension and reads one row of chunks at a time with self[start:stop]. For a sharded array, that row is a row of shards. It then yields the rows of that block. A full pass reads each chunk (or shard) once.
• Chunk edges come from the array's chunk grid (_chunk_grid.dimensions[0]), not from chunks[0]. So rectilinear grids, a partial last chunk, and arrays resized below their last edge all work. Bounds are generated lazily, so nothing is built per chunk up front.
• Size check: array.iter_block_max_bytes, default 128 MiB. A row of chunks larger than this is read in smaller blocks, and its chunks are decoded once per block instead of once per row. For sharded arrays the smaller blocks are rounded to whole rows of inner chunks, so each inner chunk is still decoded once. This is the guard for the "chunks span the whole first axis" case raised in the thread.
• Same rows as before. Each step yields what z[i] returns: the same dtype, the same shape, and the same memory order (C or F, following config.order). For a 1-D array that is a 0-d ndarray, as today. Each row is copied out of its block, so it is an independent array like z[i] and does not keep the block alive.
• One behaviour change: iterating over a 0-d array now raises TypeError("iteration over a 0-d array"), as NumPy and zarr v2 do. Today it silently yields nothing, because z[0] raises IndexError and the sequence protocol stops.
• Memory: at most one block is held at a time. For a 30.5 MiB array with 0.76 MiB row blocks, the peak traced allocation during a full pass is 2.5 MiB. With the limit set to 256 KiB it is 1.5 MiB.

Benchmark

Each cell times for x in z over the whole array. The last column is the "read everything once, then iterate in NumPy" reference.

Store	Case (float64, zstd)	main	this PR	Speedup	z[:] then iterate
memory	1-D (10,000), chunks 1,000	4.19 s	8.1 ms	514x	2.2 ms
memory	1-D (10,000), one chunk	5.46 s	3.5 ms	1,577x	0.9 ms
memory	1-D (10,500), chunks 1,000 (partial last chunk)	7.56 s	7.3 ms	1,033x	2.4 ms
memory	2-D (2000, 100), chunks (500, 100)	2.35 s	5.4 ms	438x	3.7 ms
memory	2-D (2000, 1000), chunks (100, 100)	7.11 s	80.0 ms	89x	73.5 ms
memory	1-D (10,000), chunks 100, shards 1,000	11.53 s	26.2 ms	441x	21.9 ms
local	1-D (10,000), chunks 1,000	8.80 s	17.3 ms	510x	4.2 ms
local	1-D (10,000), one chunk	9.48 s	4.5 ms	2,115x	1.2 ms
local	1-D (10,500), chunks 1,000 (partial last chunk)	11.16 s	9.0 ms	1,235x	3.0 ms
local	2-D (2000, 100), chunks (500, 100)	4.24 s	7.2 ms	586x	4.4 ms
local	2-D (2000, 1000), chunks (100, 100)	10.91 s	107.0 ms	102x	90.8 ms
local	1-D (10,000), chunks 100, shards 1,000	14.81 s	26.6 ms	558x	17.1 ms

How these were measured:

• Hardware and versions: Intel Xeon @ 2.10 GHz, 4 cores, Linux; Python 3.12.3, NumPy 2.5.4; zarr main at 957a5f6 vs this branch.
• LocalStore was on tmpfs.
• "main" is the faster of two single runs. The machine was shared, and the two runs differed by up to 2x, so the speedups are conservative.
• "this PR" is the best of three runs.

The gap that remains against z[:] comes from copying each row out of its block, roughly 0.5 µs per row.

Benchmark script

"""Benchmark `for row in arr` on zarr arrays (issue #2529)."""
import argparse, os, platform, tempfile, time

import numpy as np
import zarr
from zarr.storage import LocalStore, MemoryStore

CASES = [
    ("1-D, chunks=1,000", (10_000,), (1_000,), None),
    ("1-D, one chunk", (10_000,), (10_000,), None),
    ("1-D, chunks=1,000 (last chunk partial)", (10_500,), (1_000,), None),
    ("2-D (2000, 100), chunks (500, 100)", (2_000, 100), (500, 100), None),
    ("2-D (2000, 1000), chunks (100, 100)", (2_000, 1_000), (100, 100), None),
    ("1-D sharded, chunks=100, shards=1,000", (10_000,), (100,), (1_000,)),
]

def make(store, shape, chunks, shards):
    z = zarr.create_array(store, shape=shape, chunks=chunks, shards=shards, dtype="f8")
    z[...] = np.random.default_rng(0).standard_normal(shape)
    return z

def best(fn, repeat):
    times = []
    for _ in range(repeat):
        t0 = time.perf_counter(); fn(); times.append(time.perf_counter() - t0)
    return min(times)

def iterate(z):
    assert sum(1 for _ in z) == z.shape[0]

def read_then_iterate(z):
    assert sum(1 for _ in z[:]) == z.shape[0]

p = argparse.ArgumentParser()
p.add_argument("--repeat", type=int, default=3)
args = p.parse_args()
print(f"zarr {zarr.__version__}, numpy {np.__version__}, python {platform.python_version()}")
base = os.environ.get("BENCH_DIR") or tempfile.mkdtemp()
for store_name in ("memory", "local"):
    for i, (name, shape, chunks, shards) in enumerate(CASES):
        store = MemoryStore() if store_name == "memory" else LocalStore(os.path.join(base, f"case{i}"))
        z = make(store, shape, chunks, shards)
        print(store_name, name, best(lambda: iterate(z), args.repeat),
              best(lambda: read_then_iterate(z), args.repeat))

Tests

Correctness is checked against z[i], which is the old code path, and against NumPy iteration over the same data. Store reads are counted with the existing CountingDict store in tests/test_indexing.py.

• tests/test_indexing.py
  • test_iter: zarr formats 2 and 3, C and F order; chunks dividing the shape, a partial last chunk, one chunk, a chunk larger than the shape, 2-D and 3-D, and a zero-length trailing axis. For every row it checks the value, dtype, shape and C/F contiguity against z[i], and the values against NumPy iteration.
  • test_iter_sharded: sharded arrays, including a partial last shard.
  • test_iter_rectilinear: rectilinear chunk edges, the same array after a resize, and a rectilinear shard grid.
  • test_iter_dtypes: bool, uint8, complex128, <U3, a structured dtype, and variable-length strings.
  • test_iter_missing_chunks: chunks that were never written yield the fill value.
  • test_iter_0d: raises TypeError, as NumPy does.
  • test_iter_empty: an empty first dimension yields nothing.
  • test_iter_rows_are_independent: each row is its own array, not a view into its block.
  • test_iter_reads_each_chunk_once: unsharded, sharded with full-width shards, and sharded with a 2-D shard grid. Each chunk or shard key is read from the store exactly once per pass, and no other key is read.
  • test_iter_reads_lazily: creating the iterator reads nothing. Consuming the first chunk row reads only that row's chunks, once each.
  • test_iter_block_max_bytes and test_iter_block_max_bytes_sharded: the block bounds for small limits, including the round-down to whole inner-chunk rows for sharded arrays. Values are still correct when blocks are split.
• tests/test_properties.py::test_iter: a Hypothesis test over block_test_arrays(), covering regular, sharded and rectilinear grids, random dtypes and the available stores. Iteration must match NumPy iteration, and each row must match z[i] in shape and dtype.
• tests/benchmarks/test_indexing.py::test_iterate: a CodSpeed benchmark on memory and latency stores, with and without sharding.
• tests/test_config.py: the defaults now include the new key.

On main, the read-count, laziness, 0-d and block-bound tests fail. The value-parity tests pass on both, as they should.

For reviewers

• The size check. Is a config value the right shape for it, and is 128 MiB a sensible default? v2 had no limit. Without one, a layout whose chunks span the whole first axis loads the whole array at the first step. With the limit, those chunks are decoded about row_of_chunks_bytes / 128 MiB times per pass, not once per row.
• Sharded arrays. Blocks are rows of shards when they fit under the limit. That means one GET per shard; with a full shard selected, the sharding codec does a single read. The alternative is rows of inner chunks, which bounds memory more tightly but re-reads each shard's index once per inner-chunk row.
• The 0-d change. Iterating a 0-d array now raises TypeError instead of yielding nothing.
• Copying rows. Each row is copied out of its block (.copy(order="K")), so rows stay independent arrays like z[i], keep the same memory order, and don't pin their block in memory. Yielding views would be a little faster (about 0.5 µs per row) and is how NumPy and zarr v2 behaved.
• Out of scope: AsyncArray stays non-iterable (an __aiter__ could reuse _iter_row_blocks later), the v2 islice(start, end) method is not restored, and __len__ is not added. Happy to add any of these here or in a follow-up.
• What I'm confident in: the rows yielded are the same as before for every 1-D or higher array I could construct. The Hypothesis test and the z[i] parity tests cover this.

Author attestation

• I am a human, these are my changes, and I have reviewed and understood every change and can explain why each is correct.

TODO

• Add unit tests and/or doctests in docstrings
• Add docstrings and API docs for any new/modified user-facing classes and functions
• New/modified features documented in docs/user-guide/*.md
• Changes documented as a new file in changes/
• GitHub Actions have all passed
• Test coverage is 100% (Codecov passes)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.74%. Comparing base (e28e7dc) to head (5fd48b7).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #4559   +/-   ##
=======================================
  Coverage   94.73%   94.74%           
=======================================
  Files          94       94           
  Lines       13706    13728   +22     
=======================================
+ Hits        12985    13007   +22     
  Misses        721      721           

Files with missing lines	Coverage Δ
src/zarr/core/array.py	98.46% <100.00%> (+0.02%)	⬆️
src/zarr/core/config.py	100.00% <ø> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.