perf: read each chunk once when iterating over an array

changes/4559.feature.md

@@ -0,0 +1 @@
+Iterating over an array (`for row in array`) now reads the array one row of chunks (of shards, for a sharded array) at a time and yields the rows from each block, instead of reading `array[i]` separately for every row. A full pass now reads and decodes each chunk once rather than once per row it holds, which makes iteration orders of magnitude faster. The rows yielded are unchanged. At most one block is held in memory at a time, capped by the new `array.iter_block_max_bytes` config value (default 128 MiB). Iterating over a 0-d array now raises `TypeError`, as it does for NumPy arrays, instead of yielding nothing.

docs/user-guide/config.md

@@ -45,6 +45,7 @@ Configuration options include the following:
 - Selections of implementations of codecs, codec pipelines and buffers
 - Enabling GPU support with `zarr.config.enable_gpu()`. See [GPU support](gpu.md) for more.
 - Control request merging when reading multiple chunks from the same shard with `array.sharding_coalesce_max_gap_bytes` and `array.sharding_coalesce_max_bytes`. Reads of nearby chunks are coalesced into a single request to the store when separated by at most `sharding_coalesce_max_gap_bytes` and the resulting merged read is no larger than `sharding_coalesce_max_bytes`.
+- The most data that iterating over an array (`for row in array`) reads at a time, `array.iter_block_max_bytes` (default 128 MiB). Iteration reads one row of chunks (of shards, for a sharded array) along the first dimension at a time. A row of chunks larger than this is read in smaller blocks, and its chunks are then read once per block.
 
 For selecting custom implementations of codecs, pipelines, buffers and ndbuffers,
 first register the implementations in the registry and then select them in the config.

src/zarr/core/array.py

@@ -2764,6 +2764,69 @@ def __getitem__(self, selection: Selection) -> NDArrayLikeOrScalar:
         else:
             return self.get_basic_selection(cast("BasicSelection", pure_selection), fields=fields)
 
+    def __iter__(self) -> Iterator[NDArrayLikeOrScalar]:
+        """Iterate over the array along its first dimension.
+
+        Step `i` yields the same value as `self[i]`. Rather than reading
+        `self[i]` for every `i`, which reads and decodes each chunk once for
+        every row it holds, the array is read one block of rows at a time and
+        the rows of each block are yielded in turn. A block is one row of
+        chunks along the first dimension (of shards, for a sharded array), so
+        a full pass reads each chunk once.
+
+        At most one block is held in memory at a time. A row of chunks larger
+        than the `array.iter_block_max_bytes` config value is read in smaller
+        blocks instead, and its chunks are then read once per block (for a
+        sharded array, blocks are kept to whole rows of inner chunks where
+        possible, so each inner chunk is still read once).
+
+        Raises
+        ------
+        TypeError
+            If the array is 0-dimensional, as for NumPy arrays.
+
+        Examples
+        --------
+        >>> import zarr
+        >>> import numpy as np
+        >>> z = zarr.create_array({}, shape=(4, 2), chunks=(2, 2), dtype="int32")
+        >>> z[:] = np.arange(8).reshape(4, 2)
+        >>> [row.tolist() for row in z]
+        [[0, 1], [2, 3], [4, 5], [6, 7]]
+        """
+        if self.ndim == 0:
+            raise TypeError("iteration over a 0-d array")
+        return self._iter_rows()
+
+    def _iter_rows(self) -> Iterator[NDArrayLikeOrScalar]:
+        for start, stop in self._iter_row_blocks():
+            block: Any = self[start:stop]
+            for i in range(stop - start):
+                # Copy, so that each row is an independent array like `self[i]`
+                # and does not keep the whole block alive.
+                yield block[i, ...].copy(order="K")
+
+    def _iter_row_blocks(self) -> Iterator[tuple[int, int]]:
+        """The `(start, stop)` bounds along the first dimension of the blocks read by `__iter__`.
+
+        Each block lies within one row of storage chunks (shards, for a sharded
+        array). A row of storage chunks larger than `array.iter_block_max_bytes`
+        is split into smaller blocks, which for a sharded array are kept to
+        whole rows of inner chunks where possible.
+        """
+        row_nbytes = math.prod(self.shape[1:]) * self.dtype.itemsize
+        max_rows = max(1, zarr_config.get("array.iter_block_max_bytes") // max(1, row_nbytes))
+        if (sharding_codec := _sharding_codec(self.metadata)) is not None:
+            inner_rows = sharding_codec.chunk_shape[0]
+            if max_rows > inner_rows:
+                max_rows -= max_rows % inner_rows
+        dim = self._chunk_grid.dimensions[0]
+        for i in range(dim.nchunks):
+            start = dim.chunk_offset(i)
+            stop = start + dim.data_size(i)
+            for block_start in range(start, stop, max_rows):
+                yield block_start, min(block_start + max_rows, stop)
+
     def __setitem__(self, selection: Selection, value: npt.ArrayLike) -> None:
         """Modify data for an item or region of the array.
 

src/zarr/core/config.py

@@ -103,6 +103,7 @@ def enable_gpu(self) -> ConfigSet:
                 "rectilinear_chunks": False,
                 "sharding_coalesce_max_gap_bytes": 1 << 20,  # 1 MiB
                 "sharding_coalesce_max_bytes": 16 << 20,  # 16 MiB
+                "iter_block_max_bytes": 128 << 20,  # 128 MiB
                 # Decoding data with the pickle codec can execute arbitrary code,
                 # so it must be explicitly enabled.
                 "allow_unsafe_pickle_decode": False,

tests/benchmarks/test_indexing.py

@@ -52,6 +52,29 @@ def test_slice_indexing(
     benchmark(getitem, data, indexer)
 
 
+@pytest.mark.parametrize("store", ["memory", "memory_get_latency"], indirect=["store"])
+@pytest.mark.parametrize("shards", shards, ids=str)
+def test_iterate(
+    store: Store,
+    shards: tuple[int, ...] | None,
+    benchmark: BenchmarkFixture,
+) -> None:
+    """Benchmark iterating over the rows of an array with `for row in array`."""
+    data = create_array(
+        store=store,
+        shape=(105,) * 3,
+        dtype="uint8",
+        chunks=(10,) * 3,
+        shards=shards,
+        compressors=None,
+        filters=None,
+        fill_value=0,
+    )
+
+    data[:] = 1
+    benchmark(list, data)
+
+
 # Benchmark for Morton order optimization with power-of-2 shards
 # Morton order is used internally by sharding codec for chunk iteration
 morton_shards = (

tests/test_config.py

@@ -59,6 +59,7 @@ def test_config_defaults_set() -> None:
                     "rectilinear_chunks": False,
                     "sharding_coalesce_max_gap_bytes": 1 << 20,
                     "sharding_coalesce_max_bytes": 16 << 20,
+                    "iter_block_max_bytes": 128 << 20,
                     "allow_unsafe_pickle_decode": False,
                 },
                 "async": {"concurrency": 10, "timeout": None},

tests/test_indexing.py

@@ -2627,3 +2627,231 @@ def test_sparse_indexer_projection_reconstructs_coordinates(
         assert_array_equal(reconstructed, coords)
         assert_array_equal(covered, np.ones_like(coords))
         assert len(indexer.chunk_run_ends) == len(expected_chunks)
+
+
+def _assert_iter_matches_getitem(z: Array) -> None:
+    """Iterating over `z` yields the same arrays as `z[i]` for each `i`."""
+    rows = list(z)
+    assert len(rows) == z.shape[0]
+    for i, row in enumerate(rows):
+        expected = z[i]
+        assert isinstance(row, np.ndarray)
+        assert isinstance(expected, np.ndarray)
+        assert row.dtype == expected.dtype
+        assert row.shape == expected.shape
+        assert row.flags.c_contiguous == expected.flags.c_contiguous
+        assert row.flags.f_contiguous == expected.flags.f_contiguous
+        assert_array_equal(row, expected)
+
+
+@pytest.mark.parametrize("zarr_format", [2, 3])
+@pytest.mark.parametrize("order", ["C", "F"])
+@pytest.mark.parametrize(
+    ("shape", "chunks"),
+    [
+        ((20,), (5,)),
+        ((23,), (5,)),
+        ((23,), (23,)),
+        ((23,), (100,)),
+        ((11, 7), (3, 4)),
+        ((11, 7, 5), (4, 3, 5)),
+        ((4, 0), (2, 1)),
+    ],
+)
+def test_iter(
+    shape: tuple[int, ...], chunks: tuple[int, ...], order: str, zarr_format: ZarrFormat
+) -> None:
+    """Iterating over an array yields `z[i]` for each `i`, matching NumPy iteration."""
+    a = np.arange(np.prod(shape), dtype="int32").reshape(shape)
+    z = zarr.create_array(
+        MemoryStore(),
+        shape=shape,
+        chunks=chunks,
+        dtype=a.dtype,
+        zarr_format=zarr_format,
+        config={"order": order},
+    )
+    z[...] = a
+    _assert_iter_matches_getitem(z)
+    for row, expected in zip(z, a, strict=True):
+        assert_array_equal(row, expected)
+
+
+@pytest.mark.parametrize(
+    ("shape", "chunks", "shards"),
+    [
+        ((40, 6), (5, 3), (20, 6)),
+        ((47, 6), (5, 3), (20, 6)),
+        ((40,), (4,), (8,)),
+    ],
+)
+def test_iter_sharded(
+    shape: tuple[int, ...], chunks: tuple[int, ...], shards: tuple[int, ...]
+) -> None:
+    """Iteration over a sharded array, including a partial last shard, yields `z[i]`."""
+    a = np.arange(np.prod(shape), dtype="float64").reshape(shape)
+    z = zarr.create_array(MemoryStore(), shape=shape, chunks=chunks, shards=shards, dtype=a.dtype)
+    z[...] = a
+    _assert_iter_matches_getitem(z)
+    assert_array_equal(np.stack(list(z)), a)
+
+
+def test_iter_rectilinear() -> None:
+    """Iteration follows the chunk edges of a rectilinear chunk grid, including after a resize."""
+    with zarr.config.set({"array.rectilinear_chunks": True}):
+        z = zarr.create_array(MemoryStore(), shape=(10, 4), chunks=[[3, 5, 2], [4]], dtype="i4")
+        z[...] = np.arange(40).reshape(10, 4)
+        assert list(z._iter_row_blocks()) == [(0, 3), (3, 8), (8, 10)]
+        _assert_iter_matches_getitem(z)
+        z.resize((7, 4))
+        assert list(z._iter_row_blocks()) == [(0, 3), (3, 7)]
+        _assert_iter_matches_getitem(z)
+
+        sharded = zarr.create_array(
+            MemoryStore(), shape=(10,), chunks=(2,), shards=[[4, 6]], dtype="i4"
+        )
+        sharded[...] = np.arange(10)
+        assert list(sharded._iter_row_blocks()) == [(0, 4), (4, 10)]
+        _assert_iter_matches_getitem(sharded)
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    ["bool", "uint8", "complex128", "<U3", [("a", "i4"), ("b", "f8")], "str"],
+    ids=str,
+)
+def test_iter_dtypes(dtype: Any) -> None:
+    """Iteration yields rows of the array's dtype, for fixed-size, structured and string dtypes."""
+    z = zarr.create_array(MemoryStore(), shape=(7, 2), chunks=(3, 2), dtype=dtype)
+    if np.dtype(dtype).fields is None:
+        z[...] = np.arange(14).reshape(7, 2).astype(dtype)
+    _assert_iter_matches_getitem(z)
+
+
+@pytest.mark.parametrize("fill_value", [0, 7])
+def test_iter_missing_chunks(fill_value: int) -> None:
+    """Rows in chunks that were never written are filled with the fill value."""
+    z = zarr.create_array(
+        MemoryStore(), shape=(10, 4), chunks=(3, 2), dtype="i2", fill_value=fill_value
+    )
+    z[3:6, 2:] = 1
+    expected = np.full((10, 4), fill_value, dtype="i2")
+    expected[3:6, 2:] = 1
+    _assert_iter_matches_getitem(z)
+    assert_array_equal(np.stack(list(z)), expected)
+
+
+def test_iter_0d() -> None:
+    """Iterating over a 0-d array raises `TypeError`, as for NumPy arrays."""
+    z = zarr.create_array(MemoryStore(), shape=(), dtype="i4")
+    with pytest.raises(TypeError, match="iteration over a 0-d array"):
+        iter(z)
+    with pytest.raises(TypeError):
+        iter(np.zeros(()))
+
+
+def test_iter_empty() -> None:
+    """Iterating over an array with an empty first dimension yields nothing."""
+    z = zarr.create_array(MemoryStore(), shape=(0, 3), chunks=(2, 2), dtype="i4")
+    assert list(z) == []
+
+
+def test_iter_rows_are_independent() -> None:
+    """Each row is its own array, as `z[i]` is: changing one row changes no other row."""
+    z = zarr.create_array(MemoryStore(), shape=(4, 3), chunks=(4, 3), dtype="i4")
+    rows = list(z)
+    rows[0][...] = 1
+    assert all(r.base is None for r in rows)
+    assert_array_equal(np.stack(rows[1:]), np.zeros((3, 3), dtype="i4"))
+
+
+@pytest.mark.parametrize(
+    ("shape", "chunks", "shards"),
+    [
+        ((95, 30), (10, 10), None),
+        ((80, 30), (10, 10), (40, 30)),
+        ((80, 30), (10, 10), (20, 10)),
+    ],
+)
+async def test_iter_reads_each_chunk_once(
+    shape: tuple[int, ...], chunks: tuple[int, ...], shards: tuple[int, ...] | None
+) -> None:
+    """A full pass of iteration reads each chunk (or shard) from the store exactly once."""
+    store = await CountingDict.open()
+    z = zarr.create_array(
+        StorePath(store) / "arr", shape=shape, chunks=chunks, shards=shards, dtype="f8"
+    )
+    z[...] = np.arange(np.prod(shape)).reshape(shape)
+    counts_before = store.counter.copy()
+    rows = list(z)
+    reads = store.counter - counts_before
+    grid_shape = tuple(ceildiv_int(s, c) for s, c in zip(shape, shards or chunks, strict=True))
+    expected_keys = {
+        ("__getitem__", "c/" + "/".join(map(str, ix))) for ix in np.ndindex(*grid_shape)
+    }
+    assert set(reads) == expected_keys
+    assert all(count == 1 for count in reads.values())
+    assert_array_equal(np.stack(rows), z[...])
+
+
+async def test_iter_reads_lazily() -> None:
+    """Iteration reads one row of chunks at a time, as the rows are consumed."""
+    store = await CountingDict.open()
+    z = zarr.create_array(StorePath(store) / "arr", shape=(30, 4), chunks=(10, 2), dtype="f8")
+    z[...] = 1
+    counts_before = store.counter.copy()
+    it = iter(z)
+    assert store.counter == counts_before
+    for _ in range(10):
+        next(it)
+    reads = store.counter - counts_before
+    assert reads == {("__getitem__", "c/0/0"): 1, ("__getitem__", "c/0/1"): 1}
+
+
+def test_iter_block_max_bytes() -> None:
+    """A row of chunks larger than `array.iter_block_max_bytes` is read in smaller blocks."""
+    a = np.arange(23 * 4, dtype="int64").reshape(23, 4)  # 32 bytes per row
+    z = zarr.create_array(MemoryStore(), shape=a.shape, chunks=(10, 2), dtype=a.dtype)
+    z[...] = a
+    assert list(z._iter_row_blocks()) == [(0, 10), (10, 20), (20, 23)]
+    with zarr.config.set({"array.iter_block_max_bytes": 100}):
+        assert list(z._iter_row_blocks()) == [
+            (0, 3),
+            (3, 6),
+            (6, 9),
+            (9, 10),
+            (10, 13),
+            (13, 16),
+            (16, 19),
+            (19, 20),
+            (20, 23),
+        ]
+        _assert_iter_matches_getitem(z)
+    with zarr.config.set({"array.iter_block_max_bytes": 1}):
+        assert list(z._iter_row_blocks()) == [(i, i + 1) for i in range(23)]
+        _assert_iter_matches_getitem(z)
+
+
+def test_iter_block_max_bytes_sharded() -> None:
+    """Blocks of a sharded array that are smaller than a shard row hold whole inner chunk rows."""
+    a = np.arange(40 * 4, dtype="int64").reshape(40, 4)  # 32 bytes per row
+    z = zarr.create_array(
+        MemoryStore(), shape=a.shape, chunks=(4, 2), shards=(20, 4), dtype=a.dtype
+    )
+    z[...] = a
+    assert list(z._iter_row_blocks()) == [(0, 20), (20, 40)]
+    # 10 rows fit, which is rounded down to 2 inner chunk rows of 4 rows each.
+    with zarr.config.set({"array.iter_block_max_bytes": 10 * 32}):
+        assert list(z._iter_row_blocks()) == [
+            (0, 8),
+            (8, 16),
+            (16, 20),
+            (20, 28),
+            (28, 36),
+            (36, 40),
+        ]
+        _assert_iter_matches_getitem(z)
+    # Fewer rows than an inner chunk row fit.
+    with zarr.config.set({"array.iter_block_max_bytes": 3 * 32}):
+        assert list(z._iter_row_blocks())[:3] == [(0, 3), (3, 6), (6, 9)]
+        _assert_iter_matches_getitem(z)

tests/test_properties.py

@@ -331,6 +331,22 @@ def test_block_indexing(data: st.DataObject) -> None:
     assert_array_equal(nparray, zarray[:])
 
 
+@settings(deadline=None)
+@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning")
+@given(data=st.data())
+def test_iter(data: st.DataObject) -> None:
+    # Iteration reads whole rows of chunks (shards, for sharded arrays) at a time.
+    # Whatever the grid, it must yield the same rows as NumPy iteration over the
+    # array's contents, and the same arrays as integer indexing.
+    zarray, nparray = data.draw(block_test_arrays())
+    rows = list(zarray)
+    assert len(rows) == len(nparray)
+    for i, (row, expected) in enumerate(zip(rows, nparray, strict=True)):
+        assert_array_equal(expected, row)
+        assert row.shape == zarray[i].shape
+        assert row.dtype == zarray[i].dtype
+
+
 @given(store=stores, meta=array_metadata())  # type: ignore[misc]
 @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning")
 async def test_roundtrip_array_metadata_from_store(