fix(indexing): make sparse selections O(npoints) instead of O(nchunks)

changes/4174.bugfix.md

@@ -0,0 +1,7 @@
+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.

src/zarr/core/indexing.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import itertools
-import math
 import numbers
 from collections.abc import Iterator, Sequence
 from dataclasses import dataclass
@@ -21,6 +20,7 @@
 
 import numpy as np
 import numpy.typing as npt
+from typing_extensions import deprecated
 
 from zarr.core.chunk_grids import FixedDimension
 from zarr.core.common import ceildiv, product
@@ -31,6 +31,7 @@
     BoundsCheckError,
     NegativeStepError,
     VindexInvalidSelectionError,
+    ZarrDeprecationWarning,
 )
 
 if TYPE_CHECKING:
@@ -740,6 +741,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."""
@@ -751,9 +769,9 @@ class IntArrayDimIndexer:
     order: Order
     dim_sel: npt.NDArray[np.intp]
     dim_out_sel: npt.NDArray[np.intp]
-    chunk_nitems: int
     dim_chunk_ixs: npt.NDArray[np.intp]
-    chunk_nitems_cumsum: 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,
@@ -793,23 +811,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)
@@ -819,21 +835,44 @@ 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)
+
+    @property
+    @deprecated(
+        "IntArrayDimIndexer.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(
+        "IntArrayDimIndexer.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
 
-        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:
+            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]
             if self.order == Order.INCREASING:
                 dim_out_sel = slice(start, stop)
             else:
@@ -1171,7 +1210,9 @@ class CoordinateIndexer(Indexer):
     sel_shape: tuple[int, ...]
     selection: CoordinateSelectionNormalized
     sel_sort: npt.NDArray[np.intp] | None
-    chunk_nitems_cumsum: npt.NDArray[np.intp]
+    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, ...]
@@ -1188,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))
@@ -1244,15 +1284,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)
@@ -1297,39 +1337,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", ())
 
+    @property
+    @deprecated(
+        "CoordinateIndexer.chunk_nitems_cumsum is deprecated: it materializes a dense array "
+        "with one entry per chunk in the array. Use chunk_rixs and chunk_run_ends, which "
+        "cover only the occupied chunks.",
+        category=ZarrDeprecationWarning,
+    )
+    def chunk_nitems_cumsum(self) -> npt.NDArray[np.intp]:
+        dense = np.zeros(product(self.cdata_shape), dtype=np.intp)
+        dense[self.chunk_rixs] = np.diff(self.chunk_run_ends, prepend=0)
+        return np.cumsum(dense)
+
     def __iter__(self) -> Iterator[ChunkProjection]:
         # 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:
+            if i == 0:
                 start = 0
             else:
-                start = self.chunk_nitems_cumsum[chunk_rix - 1]
-            stop = self.chunk_nitems_cumsum[chunk_rix]
+                start = self.chunk_run_ends[i - 1]
+            stop = self.chunk_run_ends[i]
             out_selection: slice | npt.NDArray[np.intp]
             if self.sel_sort is None:
                 out_selection = slice(start, stop)

tests/test_indexing.py

@@ -19,6 +19,7 @@
     BasicSelection,
     CoordinateIndexer,
     CoordinateSelection,
+    IntArrayDimIndexer,
     OrthogonalSelection,
     Selection,
     _ArrayIndexingOrder,
@@ -1235,6 +1236,110 @@ def unexpected_searchsorted(*args: Any, **kwargs: Any) -> None:
     assert tuple(projection.chunk_coords for projection in projections) == ((0,), (99,))
 
 
+def test_deprecated_dense_indexer_attributes() -> None:
+    """`chunk_nitems` / `chunk_nitems_cumsum` keep their original dense per-chunk semantics
+    for external code that introspected indexers, but warn: they allocate O(nchunks) and the
+    indexers no longer use them internally (see gh-4174)."""
+    chunk_grid = ChunkGrid.from_sizes((20,), (3,))  # 7 chunks
+    coords = np.array([1, 4, 4, 19])
+    expected_nitems = np.bincount(coords // 3, minlength=7)
+
+    indexer = CoordinateIndexer((coords,), (20,), chunk_grid)
+    with pytest.warns(zarr.errors.ZarrDeprecationWarning):
+        assert_array_equal(indexer.chunk_nitems_cumsum, np.cumsum(expected_nitems))
+
+    (dim_grid,) = chunk_grid._dimensions
+    dim_indexer = IntArrayDimIndexer(coords, 20, dim_grid)
+    with pytest.warns(zarr.errors.ZarrDeprecationWarning):
+        assert_array_equal(dim_indexer.chunk_nitems, expected_nitems)
+    with pytest.warns(zarr.errors.ZarrDeprecationWarning):
+        assert_array_equal(dim_indexer.chunk_nitems_cumsum, np.cumsum(expected_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,
+    not the total number of chunks in the array. This array has 2**44 chunks, so any dense
+    per-chunk allocation fails before the assertions are reached. See gh-4174."""
+    z = zarr.create_array(
+        store=store / str(uuid4()),
+        shape=(2**22, 2**22),
+        chunks=(1, 1),
+        dtype="int32",
+        fill_value=-1,
+    )
+    # unsorted coords with duplicates force the argsort branch of the general path
+    rows = np.array([7, 0, 3, 0])
+    cols = np.array([1, 5, 2**22 - 1, 5])
+    z.set_coordinate_selection((rows, cols), np.array([10, 20, 30, 20], dtype="int32"))
+    assert_array_equal(
+        z.get_coordinate_selection((rows, cols)), np.array([10, 20, 30, 20], dtype="int32")
+    )
+    # points in untouched chunks come back as fill_value
+    assert_array_equal(
+        z.get_coordinate_selection((np.array([0, 42]), np.array([0, 42]))),
+        np.array([-1, -1], dtype="int32"),
+    )
+    # orthogonal integer-array selections: increasing, decreasing, and unsorted order
+    for coords in ([3, 5, 6], [6, 5, 3], [5, 6, 3]):
+        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)
+
+
+def test_coordinate_indexer_many_chunks() -> None:
+    """Both CoordinateIndexer paths (sorted-1D fast path and general path) produce correct
+    projections on a grid whose chunk count (2**62) is far too large for any dense per-chunk
+    array. See gh-4174."""
+    dim_len = 2**62
+    chunk_grid = ChunkGrid.from_sizes((dim_len,), (1,))
+
+    # sorted coords, sparse relative to their span: the general path without a sort
+    coords = np.array([3, 4, 4, dim_len - 1])
+    projections = tuple(CoordinateIndexer((coords,), (dim_len,), chunk_grid))
+    assert tuple(p.chunk_coords for p in projections) == ((3,), (4,), (dim_len - 1,))
+    assert [p.out_selection for p in projections] == [slice(0, 1), slice(1, 3), slice(3, 4)]
+    for p in projections:
+        assert_array_equal(p.chunk_selection[0], np.zeros(len(p.chunk_selection[0]), dtype=int))
+
+    # sorted coords, dense relative to their span: the searchsorted fast path
+    coords = np.concatenate([np.full(100, 3), np.full(100, 7)])
+    projections = tuple(CoordinateIndexer((coords,), (dim_len,), chunk_grid))
+    assert tuple(p.chunk_coords for p in projections) == ((3,), (7,))
+    assert [p.out_selection for p in projections] == [slice(0, 100), slice(100, 200)]
+
+    # unsorted coords: the general (argsort) path
+    coords = np.array([dim_len - 1, 5])
+    projections = tuple(CoordinateIndexer((coords,), (dim_len,), chunk_grid))
+    assert tuple(p.chunk_coords for p in projections) == ((5,), (dim_len - 1,))
+    assert [list(p.out_selection) for p in projections] == [[1], [0]]
+
+
+@pytest.mark.parametrize(
+    ("coords", "expected_out_sels"),
+    [
+        pytest.param([3, 5, 5], [slice(0, 1), slice(1, 3)], id="increasing"),
+        pytest.param([5, 3, 1], [[2], [1], [0]], id="decreasing"),
+        pytest.param([5, 1, 3], [[1], [2], [0]], id="unordered"),
+    ],
+)
+def test_int_array_dim_indexer_many_chunks(coords: list[int], expected_out_sels: list[Any]) -> None:
+    """IntArrayDimIndexer produces correct projections for every ordering of the selection
+    on a dimension with 2**62 chunks, where any dense per-chunk array would fail. See gh-4174."""
+    dim_len = 2**62
+    (dim_grid,) = ChunkGrid.from_sizes((dim_len,), (1,))._dimensions
+    indexer = IntArrayDimIndexer(np.array(coords), dim_len, dim_grid)
+
+    projections = tuple(indexer)
+    assert [p.dim_chunk_ix for p in projections] == sorted(set(coords))
+    for p, expected_out_sel in zip(projections, expected_out_sels, strict=True):
+        # chunk size is 1, so every selected point maps to offset 0 within its chunk
+        assert_array_equal(p.dim_chunk_sel, np.zeros(len(p.dim_chunk_sel), dtype=int))
+        if isinstance(expected_out_sel, slice):
+            assert p.dim_out_sel == expected_out_sel
+        else:
+            assert list(p.dim_out_sel) == expected_out_sel
+
+
 def test_get_coordinate_selection_1d_irregular_grid(store: StorePath) -> None:
     """Coordinate selections on an irregular (rectilinear) chunk grid bypass the sorted-1D fast
     path (which requires a regular grid) and still match numpy via the general path."""