fix(indexing): exact ceildiv everywhere and sparse boolean-axis selections

Review follow-ups for the sparse-selection change:

- Fix ceildiv itself rather than bypassing it at one call site. It went
  through float division, so ceildiv(2**62 - 1, 1) returned 2**62; four
  other callers (slice arithmetic, the regular-grid check, dask-style
  chunk sizes) had the same latent error. Integers now divide exactly;
  floats keep the ceil-of-quotient path. FixedDimension goes back to
  calling it.

- BoolArrayDimIndexer still allocated a dense per-chunk count array and
  ran a Python loop over every chunk, so an orthogonal boolean selection
  on a finely chunked axis paid O(nchunks) in time and memory on top of
  the mask. It now derives the occupied chunks and run ends from the
  selected positions with sorted_run_ends, like the other two indexers,
  and exposes the same deprecated dense properties. A boolean axis with
  2**22 chunks goes from a multi-second loop to ~2 ms.

- Changelog: state precisely which selection kinds are covered, that
  reading the deprecated properties rebuilds the dense array, and the
  ceildiv correction.

Assisted-by: ClaudeCode:claude-fable-5-1

Author: d-v-b

changes/4174.bugfix.md

@@ -9,13 +9,19 @@ The intermediate `ChunksTuple` representation was removed in the process, and
 size) with explicit edge-length sequences, matching what the normalizer and
 the rectilinear grid spec already accepted.
 
-Coordinate, mask, and orthogonal integer-array selections no longer allocate memory
-proportional to the total number of chunks in the array. Previously a selection of a few
-points on an array with billions of chunks could raise `MemoryError`; indexing cost now
-scales with the number of selected points instead. The internal indexers now store
-per-occupied-chunk run offsets (`chunk_run_ends`); the dense `chunk_nitems` /
-`chunk_nitems_cumsum` attributes of `CoordinateIndexer` and `IntArrayDimIndexer` are
-deprecated but keep their original semantics.
+Coordinate selections (including `vindex` masks) and orthogonal selections with
+integer or boolean arrays no longer allocate memory proportional to the total number of
+chunks in the array, nor loop over every chunk in Python. Previously a selection of a few
+points on an array with billions of chunks could raise `MemoryError`, and a boolean array
+along a finely chunked axis paid a Python-level pass per chunk; indexing cost now scales
+with the number of selected points (plus, for a boolean array, the length of the array
+itself). The internal indexers now store per-occupied-chunk run offsets
+(`chunk_run_ends`); the dense `chunk_nitems` / `chunk_nitems_cumsum` attributes of
+`CoordinateIndexer`, `IntArrayDimIndexer` and `BoolArrayDimIndexer` are deprecated but
+keep their original semantics — note that reading them rebuilds the dense per-chunk array
+they used to hold, so on a huge array they still cost what they always did.
 
-Uniform grid chunk counts now use exact integer arithmetic, so sparse selections
-near the end of dimensions larger than floating-point precision remain in bounds.
+`zarr.core.common.ceildiv` now divides integers exactly instead of through a float, so
+chunk counts of dimensions longer than `2**53` are correct; previously
+`ceildiv(2**62 - 1, 1)` rounded to `2**62`, which put sparse selections near the end of
+such a dimension out of bounds.

src/zarr/core/chunk_grids.py

@@ -62,7 +62,7 @@ def __post_init__(self) -> None:
         if self.size == 0:
             n = 0
         else:
-            n = -(-self.extent // self.size)
+            n = ceildiv(self.extent, self.size)
         object.__setattr__(self, "nchunks", n)
         object.__setattr__(self, "ngridcells", n)
 

src/zarr/core/common.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import math
+import numbers
 import warnings
 from collections.abc import Iterable, Mapping, Sequence
 from enum import Enum
@@ -92,8 +93,21 @@ def product(tup: tuple[int, ...]) -> int:
 
 
 def ceildiv(a: float, b: float) -> int:
+    """Ceiling of ``a / b``.
+
+    Integer inputs are divided exactly: ``math.ceil(a / b)`` goes through a float and
+    loses precision above ``2**53``, so e.g. ``ceildiv(2**62 - 1, 1)`` came out as
+    ``2**62``. Chunk counts of large arrays depend on this being exact.
+    """
     if a == 0:
         return 0
+    # Widen to ``object`` for the check: mypy models ``int`` as a ``float`` and does not
+    # know that it satisfies ``Integral``, so it would consider this branch unreachable.
+    a_obj: object = a
+    b_obj: object = b
+    if isinstance(a_obj, numbers.Integral) and isinstance(b_obj, numbers.Integral):
+        a_int, b_int = int(a_obj), int(b_obj)
+        return -(-a_int // b_int)
     return math.ceil(a / b)
 
 

src/zarr/core/indexing.py

@@ -629,10 +629,10 @@ class BoolArrayDimIndexer:
     dim_grid: DimensionGrid
     nchunks: int
 
-    chunk_nitems: npt.NDArray[Any]
-    chunk_nitems_cumsum: npt.NDArray[Any]
     nitems: int
     dim_chunk_ixs: npt.NDArray[np.intp]
+    # 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,
@@ -653,33 +653,51 @@ def __init__(
         g = dim_grid
         nchunks = g.nchunks
 
-        # precompute number of selected items for each chunk
-        chunk_nitems = np.zeros(nchunks, dtype="i8")
-        for dim_chunk_ix in range(nchunks):
-            dim_offset = g.chunk_offset(dim_chunk_ix)
-            chunk_len = g.data_size(dim_chunk_ix)
-            chunk_nitems[dim_chunk_ix] = np.count_nonzero(
-                dim_sel[dim_offset : dim_offset + chunk_len]
-            )
-        chunk_nitems_cumsum = np.cumsum(chunk_nitems)
-        nitems = chunk_nitems_cumsum[-1]
-        dim_chunk_ixs = np.nonzero(chunk_nitems)[0]
+        # The selected positions are already in increasing order, so their chunk ids are
+        # non-decreasing and the occupied chunks with their run ends follow in O(nitems).
+        # The previous per-chunk count loop was O(nchunks) in both time and memory (gh-4174).
+        sel_positions = np.nonzero(dim_sel)[0]
+        dim_chunk_ixs, chunk_run_ends = sorted_run_ends(g.indices_to_chunks(sel_positions))
+        nitems = int(sel_positions.size)
 
         # store attributes
         object.__setattr__(self, "dim_sel", dim_sel)
         object.__setattr__(self, "dim_len", dim_len)
         object.__setattr__(self, "dim_grid", dim_grid)
         object.__setattr__(self, "nchunks", nchunks)
-        object.__setattr__(self, "chunk_nitems", chunk_nitems)
-        object.__setattr__(self, "chunk_nitems_cumsum", chunk_nitems_cumsum)
         object.__setattr__(self, "nitems", nitems)
         object.__setattr__(self, "dim_chunk_ixs", dim_chunk_ixs)
+        object.__setattr__(self, "chunk_run_ends", chunk_run_ends)
+
+    @property
+    @deprecated(
+        "BoolArrayDimIndexer.chunk_nitems is deprecated: it materializes a dense array with "
+        "one entry per chunk along the dimension. Use dim_chunk_ixs and chunk_run_ends, "
+        "which cover only the occupied chunks.",
+        category=ZarrDeprecationWarning,
+    )
+    def chunk_nitems(self) -> npt.NDArray[np.intp]:
+        dense = np.zeros(self.nchunks, dtype=np.intp)
+        dense[self.dim_chunk_ixs] = np.diff(self.chunk_run_ends, prepend=0)
+        return dense
+
+    @property
+    @deprecated(
+        "BoolArrayDimIndexer.chunk_nitems_cumsum is deprecated: it materializes a dense array "
+        "with one entry per chunk along the dimension. Use dim_chunk_ixs and chunk_run_ends, "
+        "which cover only the occupied chunks.",
+        category=ZarrDeprecationWarning,
+    )
+    def chunk_nitems_cumsum(self) -> npt.NDArray[np.intp]:
+        dense = np.zeros(self.nchunks, dtype=np.intp)
+        dense[self.dim_chunk_ixs] = np.diff(self.chunk_run_ends, prepend=0)
+        return np.cumsum(dense)
 
     def __iter__(self) -> Iterator[ChunkDimProjection]:
         g = self.dim_grid
 
         # iterate over chunks with at least one item
-        for dim_chunk_ix in self.dim_chunk_ixs:
+        for i, dim_chunk_ix in enumerate(self.dim_chunk_ixs):
             # find region in chunk
             dim_offset = g.chunk_offset(dim_chunk_ix)
             chunk_len = g.data_size(dim_chunk_ix)
@@ -693,11 +711,11 @@ def __iter__(self) -> Iterator[ChunkDimProjection]:
                 dim_chunk_sel = tmp
 
             # find region in output
-            if dim_chunk_ix == 0:
+            if i == 0:
                 start = 0
             else:
-                start = self.chunk_nitems_cumsum[dim_chunk_ix - 1]
-            stop = self.chunk_nitems_cumsum[dim_chunk_ix]
+                start = self.chunk_run_ends[i - 1]
+            stop = self.chunk_run_ends[i]
             dim_out_sel = slice(start, stop)
             is_complete_chunk = False  # TODO
 

tests/test_common.py

@@ -10,6 +10,7 @@
 from zarr.core.common import (
     ANY_ACCESS_MODE,
     AccessModeLiteral,
+    ceildiv,
     concurrent_iter,
     parse_bool,
     parse_int,
@@ -24,6 +25,27 @@
     from typing import Any, Literal
 
 
+@pytest.mark.parametrize(
+    ("a", "b", "expected"),
+    [
+        (0, 3, 0),
+        (7, 3, 3),
+        (9, 3, 3),
+        (2**62 - 1, 1, 2**62 - 1),
+        (2**62 + 1, 2, 2**61 + 1),
+        (2**60 + 3, 1, 2**60 + 3),
+        (np.int64(2**62 - 1), np.int64(1), 2**62 - 1),
+        (7.5, 2, 4),
+    ],
+)
+def test_ceildiv(a: float, b: float, expected: int) -> None:
+    """Integer inputs divide exactly, beyond the 2**53 range where float division rounds;
+    float inputs keep the ceil-of-quotient semantics."""
+    result = ceildiv(a, b)
+    assert result == expected
+    assert isinstance(result, int)
+
+
 @pytest.mark.parametrize("data", [(0, 0, 0, 0), (1, 3, 4, 5, 6), (2, 4)])
 def test_product(data: tuple[int, ...]) -> None:
     assert product(data) == np.prod(data)

tests/test_indexing.py

@@ -19,6 +19,7 @@
 from zarr.core.chunk_grids import ChunkGrid
 from zarr.core.indexing import (
     BasicSelection,
+    BoolArrayDimIndexer,
     CoordinateIndexer,
     CoordinateSelection,
     IntArrayDimIndexer,
@@ -1358,6 +1359,15 @@ def test_deprecated_dense_indexer_attributes() -> None:
     with pytest.warns(zarr.errors.ZarrDeprecationWarning):
         assert_array_equal(dim_indexer.chunk_nitems_cumsum, np.cumsum(expected_nitems))
 
+    mask = np.zeros(20, dtype=bool)
+    mask[np.unique(coords)] = True
+    expected_mask_nitems = np.bincount(np.unique(coords) // 3, minlength=7)
+    bool_indexer = BoolArrayDimIndexer(mask, 20, dim_grid)
+    with pytest.warns(zarr.errors.ZarrDeprecationWarning):
+        assert_array_equal(bool_indexer.chunk_nitems, expected_mask_nitems)
+    with pytest.warns(zarr.errors.ZarrDeprecationWarning):
+        assert_array_equal(bool_indexer.chunk_nitems_cumsum, np.cumsum(expected_mask_nitems))
+
 
 def test_sparse_selections_on_arrays_with_many_chunks(store: StorePath) -> None:
     """Coordinate and orthogonal selections must scale with the number of selected points,
@@ -1387,6 +1397,13 @@ def test_sparse_selections_on_arrays_with_many_chunks(store: StorePath) -> None:
         vals = np.arange(1, len(coords) + 1, dtype="int32")
         z.set_orthogonal_selection((np.array(coords), 0), vals)
         assert_array_equal(z.get_orthogonal_selection((np.array(coords), 0)), vals)
+    # an orthogonal boolean selection along an axis with 2**22 chunks: the mask itself is
+    # O(dim_len), but the indexer must not add a per-chunk pass on top of it
+    mask = np.zeros(2**22, dtype=bool)
+    mask[[3, 5, 6, 2**22 - 1]] = True
+    vals = np.array([1, 2, 3, 4], dtype="int32")
+    z.set_orthogonal_selection((mask, 0), vals)
+    assert_array_equal(z.get_orthogonal_selection((mask, 0)), vals)
 
 
 def test_coordinate_indexer_many_chunks() -> None: