fix(indexing): avoid full-grid allocations for sparse selections

changes/4356.bugfix.md

@@ -0,0 +1,16 @@
+Coordinate selections (including `vindex` masks) and orthogonal integer-array
+selections no longer allocate counts proportional to the total number of chunks.
+Previously, selecting a few points on an array with billions of chunks could raise
+`MemoryError`. The indexers now store occupied chunk IDs and run-end offsets
+(`chunk_run_ends`); unsorted selections can still require sorting. The dense
+`chunk_nitems` / `chunk_nitems_cumsum` attributes of `CoordinateIndexer` and
+`IntArrayDimIndexer` remain available without deprecation warnings. Each array is
+materialized on first access and cached, preserving its identity and in-place
+mutations on subsequent reads. These attributes remain dataclass fields, so
+`dataclasses.asdict()` materializes them; `repr` omits them to avoid dense allocation
+during inspection. Accessing a dense attribute still allocates memory proportional
+to the full chunk grid.
+
+The new `zarr.core.common.ceildiv_int` helper uses exact integer arithmetic for
+chunk counts and slice calculations. For example, `ceildiv_int(2**62 - 1, 1)`
+returns `2**62 - 1`. The existing `ceildiv` helper retains its floating-point behavior.

src/zarr/core/array.py

@@ -69,7 +69,7 @@
     ZarrFormat,
     _default_zarr_format,
     _warn_order_kwarg,
-    ceildiv,
+    ceildiv_int,
     concurrent_map,
     parse_shapelike,
     product,
@@ -195,7 +195,7 @@ def _chunk_sizes_from_shape(
     """Compute dask-style chunk sizes from an array shape and uniform chunk shape."""
     result: list[tuple[int, ...]] = []
     for s, c in zip(array_shape, chunk_shape, strict=True):
-        nchunks = ceildiv(s, c)
+        nchunks = ceildiv_int(s, c)
         sizes = tuple(min(c, s - i * c) for i in range(nchunks))
         result.append(sizes)
     return tuple(result)
@@ -1147,7 +1147,7 @@ def _chunk_grid_shape(self) -> tuple[int, ...]:
         if (sharding_codec := _sharding_codec(self.metadata)) is not None:
             # When sharding, count inner chunks across the whole array
             chunk_shape = sharding_codec.chunk_shape
-            return tuple(starmap(ceildiv, zip(self.shape, chunk_shape, strict=True)))
+            return tuple(starmap(ceildiv_int, zip(self.shape, chunk_shape, strict=True)))
         return self._chunk_grid.grid_shape
 
     @property

src/zarr/core/chunk_grids.py

@@ -25,7 +25,7 @@
 import zarr
 from zarr.core.common import (
     ShapeLike,
-    ceildiv,
+    ceildiv_int,
     parse_shapelike,
 )
 from zarr.errors import ZarrUserWarning
@@ -62,7 +62,7 @@ def __post_init__(self) -> None:
         if self.size == 0:
             n = 0
         else:
-            n = ceildiv(self.extent, self.size)
+            n = ceildiv_int(self.extent, self.size)
         object.__setattr__(self, "nchunks", n)
         object.__setattr__(self, "ngridcells", n)
 
@@ -465,7 +465,9 @@ def from_sizes(
                 if (
                     edges_list[0] > 0
                     and all(e == edges_list[0] for e in edges_list)
-                    and (extent == edge_sum or len(edges_list) == ceildiv(extent, edges_list[0]))
+                    and (
+                        extent == edge_sum or len(edges_list) == ceildiv_int(extent, edges_list[0])
+                    )
                 ):
                     dims.append(FixedDimension(size=edges_list[0], extent=extent))
                 else:

src/zarr/core/common.py

@@ -92,11 +92,17 @@ def product(tup: tuple[int, ...]) -> int:
 
 
 def ceildiv(a: float, b: float) -> int:
+    """Ceiling of ``a / b`` using floating-point division; zero when ``a`` is zero."""
     if a == 0:
         return 0
     return math.ceil(a / b)
 
 
+def ceildiv_int(a: int, b: int) -> int:
+    """Ceiling of integer division using exact Python integer arithmetic."""
+    return -(-int(a) // int(b))
+
+
 def concurrent_iter[T: tuple[Any, ...], V](
     items: Iterable[T],
     func: Callable[..., Awaitable[V]],

src/zarr/core/indexing.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
 import itertools
-import math
 import numbers
 from collections.abc import Iterator, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
 from functools import lru_cache
 from types import EllipsisType
@@ -23,7 +22,7 @@
 import numpy.typing as npt
 
 from zarr.core.chunk_grids import FixedDimension
-from zarr.core.common import ceildiv, product
+from zarr.core.common import ceildiv_int, product
 from zarr.core.metadata.v2 import ArrayV2Metadata
 from zarr.core.metadata.v3 import ArrayV3Metadata
 from zarr.errors import (
@@ -207,7 +206,7 @@ def _iter_regions(
     # ((slice(0, 1, 1), slice(0, 2, 1)), (slice(1, 2, 1), slice(0, 2, 1)))
     ```
     """
-    grid_shape = tuple(itertools.starmap(ceildiv, zip(domain_shape, region_shape, strict=True)))
+    grid_shape = tuple(itertools.starmap(ceildiv_int, zip(domain_shape, region_shape, strict=True)))
     for grid_position in _iter_grid(
         grid_shape=grid_shape, origin=origin, selection_shape=selection_shape, order=order
     ):
@@ -415,7 +414,7 @@ def __init__(
 
         object.__setattr__(self, "dim_len", dim_len)
         object.__setattr__(self, "dim_grid", dim_grid)
-        object.__setattr__(self, "nitems", max(0, ceildiv((stop - start), step)))
+        object.__setattr__(self, "nitems", max(0, ceildiv_int((stop - start), step)))
         object.__setattr__(self, "nchunks", dim_grid.nchunks)
 
     def __iter__(self) -> Iterator[ChunkDimProjection]:
@@ -441,7 +440,7 @@ def __iter__(self) -> Iterator[ChunkDimProjection]:
                 if remainder:
                     dim_chunk_sel_start += self.step - remainder
                 # compute number of previous items, provides offset into output array
-                dim_out_offset = ceildiv((dim_offset - self.start), self.step)
+                dim_out_offset = ceildiv_int((dim_offset - self.start), self.step)
             else:
                 # selection starts within current chunk
                 dim_chunk_sel_start = self.start - dim_offset
@@ -455,7 +454,7 @@ def __iter__(self) -> Iterator[ChunkDimProjection]:
                 dim_chunk_sel_stop = self.stop - dim_offset
 
             dim_chunk_sel = slice(dim_chunk_sel_start, dim_chunk_sel_stop, self.step)
-            dim_chunk_nitems = ceildiv((dim_chunk_sel_stop - dim_chunk_sel_start), self.step)
+            dim_chunk_nitems = ceildiv_int((dim_chunk_sel_stop - dim_chunk_sel_start), self.step)
 
             # If there are no elements on the selection within this chunk, then skip
             if dim_chunk_nitems == 0:
@@ -744,6 +743,23 @@ def boundscheck_indices(x: npt.NDArray[Any], dim_len: int) -> None:
         raise BoundsCheckError(msg)
 
 
+def sorted_run_ends(
+    a: npt.NDArray[Any],
+) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]:
+    """Group a sorted 1-D integer array into runs of equal values.
+
+    Returns `(values, run_ends)` where `values` holds the distinct values in order and
+    `run_ends[i]` is the exclusive end offset of run `i` in `a`. Cost is O(len(a)),
+    independent of the range of values — unlike a dense `np.bincount` histogram, which
+    allocates O(max value) memory (see gh-4174).
+    """
+    if a.size == 0:
+        return np.empty(0, dtype=np.intp), np.empty(0, dtype=np.intp)
+    run_starts = np.concatenate(([0], np.nonzero(np.diff(a))[0] + 1))
+    run_ends = np.append(run_starts[1:], a.size).astype(np.intp, copy=False)
+    return a[run_starts].astype(np.intp, copy=False), run_ends
+
+
 @dataclass(frozen=True)
 class IntArrayDimIndexer:
     """Integer array selection against a single dimension."""
@@ -755,9 +771,12 @@ class IntArrayDimIndexer:
     order: Order
     dim_sel: npt.NDArray[np.intp]
     dim_out_sel: npt.NDArray[np.intp]
-    chunk_nitems: int
+    # Dense compatibility arrays are populated on first access, not during indexing.
+    chunk_nitems: npt.NDArray[np.intp] = field(repr=False)
     dim_chunk_ixs: npt.NDArray[np.intp]
-    chunk_nitems_cumsum: npt.NDArray[np.intp]
+    chunk_nitems_cumsum: npt.NDArray[np.intp] = field(repr=False)
+    # end offset of each occupied chunk's run of selected items, aligned with dim_chunk_ixs
+    chunk_run_ends: npt.NDArray[np.intp]
 
     def __init__(
         self,
@@ -802,23 +821,21 @@ def __init__(
 
         if order == Order.INCREASING:
             dim_out_sel = None
+            dim_sel_chunk_sorted = dim_sel_chunk
         elif order == Order.DECREASING:
             dim_sel = dim_sel[::-1]
             # TODO should be possible to do this without creating an arange
             dim_out_sel = np.arange(nitems - 1, -1, -1)
+            dim_sel_chunk_sorted = dim_sel_chunk[::-1]
         else:
             # sort indices to group by chunk
             dim_out_sel = np.argsort(dim_sel_chunk)
             dim_sel = np.take(dim_sel, dim_out_sel)
+            dim_sel_chunk_sorted = dim_sel_chunk[dim_out_sel]
 
-        # precompute number of selected items for each chunk
-        chunk_nitems = np.bincount(dim_sel_chunk, minlength=nchunks)
-
-        # find chunks that we need to visit
-        dim_chunk_ixs = np.nonzero(chunk_nitems)[0]
-
-        # compute offsets into the output array
-        chunk_nitems_cumsum = np.cumsum(chunk_nitems)
+        # the chunks to visit and, per occupied chunk, the end offset of its run of
+        # selected items — O(nitems), never O(nchunks)
+        dim_chunk_ixs, chunk_run_ends = sorted_run_ends(dim_sel_chunk_sorted)
 
         # store attributes
         object.__setattr__(self, "dim_len", dim_len)
@@ -828,21 +845,32 @@ def __init__(
         object.__setattr__(self, "order", order)
         object.__setattr__(self, "dim_sel", dim_sel)
         object.__setattr__(self, "dim_out_sel", dim_out_sel)
-        object.__setattr__(self, "chunk_nitems", chunk_nitems)
         object.__setattr__(self, "dim_chunk_ixs", dim_chunk_ixs)
-        object.__setattr__(self, "chunk_nitems_cumsum", chunk_nitems_cumsum)
+        object.__setattr__(self, "chunk_run_ends", chunk_run_ends)
+
+    def __getattr__(self, name: str) -> npt.NDArray[np.intp]:
+        if name not in ("chunk_nitems", "chunk_nitems_cumsum"):
+            raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}")
+        dense = np.zeros(self.nchunks, dtype=np.intp)
+        dense[self.dim_chunk_ixs] = np.diff(self.chunk_run_ends, prepend=0)
+        if name == "chunk_nitems_cumsum":
+            np.cumsum(dense, out=dense)
+        object.__setattr__(self, name, dense)
+        return dense
 
     def __iter__(self) -> Iterator[ChunkDimProjection]:
         g = self.dim_grid
+        dense_cumsum = self.__dict__.get("chunk_nitems_cumsum")
 
-        for dim_chunk_ix in self.dim_chunk_ixs:
+        for i, dim_chunk_ix in enumerate(self.dim_chunk_ixs):
             dim_out_sel: slice | npt.NDArray[np.intp]
             # find region in output
-            if dim_chunk_ix == 0:
-                start = 0
+            if dense_cumsum is None:
+                start = 0 if i == 0 else self.chunk_run_ends[i - 1]
+                stop = self.chunk_run_ends[i]
             else:
-                start = self.chunk_nitems_cumsum[dim_chunk_ix - 1]
-            stop = self.chunk_nitems_cumsum[dim_chunk_ix]
+                start = 0 if dim_chunk_ix == 0 else dense_cumsum[dim_chunk_ix - 1]
+                stop = dense_cumsum[dim_chunk_ix]
             if self.order == Order.INCREASING:
                 dim_out_sel = slice(start, stop)
             else:
@@ -1180,7 +1208,11 @@ class CoordinateIndexer(Indexer):
     sel_shape: tuple[int, ...]
     selection: CoordinateSelectionNormalized
     sel_sort: npt.NDArray[np.intp] | None
-    chunk_nitems_cumsum: npt.NDArray[np.intp]
+    # Exclude the lazy dense field from repr to keep inspection sparse.
+    chunk_nitems_cumsum: npt.NDArray[np.intp] = field(repr=False)
+    cdata_shape: tuple[int, ...]
+    # end offset of each occupied chunk's run of selected points, aligned with chunk_rixs
+    chunk_run_ends: npt.NDArray[np.intp]
     chunk_rixs: npt.NDArray[np.intp]
     chunk_mixs: tuple[npt.NDArray[np.intp], ...]
     shape: tuple[int, ...]
@@ -1197,7 +1229,6 @@ def __init__(
             cdata_shape = (1,)
         else:
             cdata_shape = tuple(g.nchunks for g in dim_grids)
-        nchunks = math.prod(cdata_shape)
 
         # some initial normalization
         selection_normalized = cast("CoordinateSelectionNormalized", ensure_tuple(selection))
@@ -1262,15 +1293,15 @@ def __init__(
                         edges = np.arange(first + 1, last + 1, dtype=coords.dtype) * size
                         cuts = np.searchsorted(coords, edges)
                         counts = np.diff(cuts, prepend=0, append=coords.size)
-                    chunk_rixs = (first + np.nonzero(counts)[0]).astype(np.intp)
-                    chunk_nitems = np.zeros(nchunks, dtype=np.intp)
-                    chunk_nitems[first : last + 1] = counts
-                    chunk_nitems_cumsum = np.cumsum(chunk_nitems)
+                    occupied = np.nonzero(counts)[0]
+                    chunk_rixs = (first + occupied).astype(np.intp)
+                    chunk_run_ends = np.cumsum(counts[occupied])
 
                     object.__setattr__(self, "sel_shape", coords.shape)
                     object.__setattr__(self, "selection", (coords,))
                     object.__setattr__(self, "sel_sort", None)
-                    object.__setattr__(self, "chunk_nitems_cumsum", chunk_nitems_cumsum)
+                    object.__setattr__(self, "cdata_shape", cdata_shape)
+                    object.__setattr__(self, "chunk_run_ends", chunk_run_ends)
                     object.__setattr__(self, "chunk_rixs", chunk_rixs)
                     object.__setattr__(self, "chunk_mixs", (chunk_rixs,))
                     object.__setattr__(self, "dim_grids", dim_grids)
@@ -1315,39 +1346,51 @@ def __init__(
             # optimisation, only sort if needed
             sel_sort = np.argsort(chunks_raveled_indices)
             selection_broadcast = tuple(dim_sel[sel_sort] for dim_sel in selection_broadcast)
+            chunks_raveled_indices = chunks_raveled_indices[sel_sort]
         else:
             sel_sort = None
 
         shape = selection_broadcast[0].shape or (1,)
 
-        # precompute number of selected items for each chunk
-        chunk_nitems = np.bincount(chunks_raveled_indices, minlength=nchunks)
-        chunk_nitems_cumsum = np.cumsum(chunk_nitems)
-        # locate the chunks we need to process
-        chunk_rixs = np.nonzero(chunk_nitems)[0]
+        # the chunks to visit and, per occupied chunk, the end offset of its run of
+        # selected points — O(npoints), never O(nchunks)
+        chunk_rixs, chunk_run_ends = sorted_run_ends(chunks_raveled_indices)
 
         # unravel chunk indices
         chunk_mixs = np.unravel_index(chunk_rixs, cdata_shape)
 
         object.__setattr__(self, "sel_shape", sel_shape)
         object.__setattr__(self, "selection", selection_broadcast)
         object.__setattr__(self, "sel_sort", sel_sort)
-        object.__setattr__(self, "chunk_nitems_cumsum", chunk_nitems_cumsum)
+        object.__setattr__(self, "cdata_shape", cdata_shape)
+        object.__setattr__(self, "chunk_run_ends", chunk_run_ends)
         object.__setattr__(self, "chunk_rixs", chunk_rixs)
         object.__setattr__(self, "chunk_mixs", chunk_mixs)
         object.__setattr__(self, "dim_grids", dim_grids)
         object.__setattr__(self, "shape", shape)
         object.__setattr__(self, "drop_axes", ())
 
+    def __getattr__(self, name: str) -> npt.NDArray[np.intp]:
+        if name != "chunk_nitems_cumsum":
+            raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}")
+        dense = np.zeros(product(self.cdata_shape), dtype=np.intp)
+        dense[self.chunk_rixs] = np.diff(self.chunk_run_ends, prepend=0)
+        np.cumsum(dense, out=dense)
+        object.__setattr__(self, name, dense)
+        return dense
+
     def __iter__(self) -> Iterator[ChunkProjection]:
+        dense_cumsum = self.__dict__.get("chunk_nitems_cumsum")
         # iterate over chunks
-        for i, chunk_rix in enumerate(self.chunk_rixs):
+        for i in range(len(self.chunk_rixs)):
             chunk_coords = tuple(m[i] for m in self.chunk_mixs)
-            if chunk_rix == 0:
-                start = 0
+            if dense_cumsum is None:
+                start = 0 if i == 0 else self.chunk_run_ends[i - 1]
+                stop = self.chunk_run_ends[i]
             else:
-                start = self.chunk_nitems_cumsum[chunk_rix - 1]
-            stop = self.chunk_nitems_cumsum[chunk_rix]
+                chunk_rix = self.chunk_rixs[i]
+                start = 0 if chunk_rix == 0 else dense_cumsum[chunk_rix - 1]
+                stop = dense_cumsum[chunk_rix]
             out_selection: slice | npt.NDArray[np.intp]
             if self.sel_sort is None:
                 out_selection = slice(start, stop)