perf(indexing): speed up sorted 1-D coordinate selections

changes/4172.misc.md

@@ -0,0 +1,7 @@
+Improved `CoordinateIndexer` construction for large, sorted, in-bounds, one-dimensional integer
+coordinate selections over regular chunk grids (e.g. `arr.get_coordinate_selection(sorted_idx)`,
+`arr.vindex[sorted_idx]`, and the gather behind sparse/CSR row selections). When boundary searching
+is estimated to be cheaper than processing every coordinate, per-chunk projections are now built
+with `searchsorted`, making index construction ~15x faster for large gathers. Sparse sorted
+selections spanning many chunks relative to their coordinate count, as well as unsorted, negative,
+multi-dimensional, and irregular-grid selections, continue to use the existing implementation.

src/zarr/core/indexing.py

@@ -22,6 +22,7 @@
 import numpy as np
 import numpy.typing as npt
 
+from zarr.core.chunk_grids import FixedDimension
 from zarr.core.common import ceildiv, product
 from zarr.core.metadata.v2 import ArrayV2Metadata
 from zarr.core.metadata.v3 import ArrayV3Metadata
@@ -1206,6 +1207,59 @@ def __init__(
                 f"got {selection!r}"
             )
 
+        # Optimization for a single sorted, in-bounds, 1-D integer coordinate array over a
+        # regular (fixed-size) chunk grid. The general path below makes several full passes over
+        # the flat selection. For sufficiently dense selections, locating the internal chunk
+        # boundaries with searchsorted is cheaper.
+        if len(selection_normalized) == 1:
+            (coords,) = selection_normalized
+            g0 = dim_grids[0]
+            # coords is an integer ndarray here: is_coordinate_selection() validated above, and
+            # the normalization turned ints/lists into arrays. Only the sorted-1D-over-regular-grid
+            # shape is special-cased; everything else falls through to the general path below.
+            if (
+                isinstance(g0, FixedDimension)
+                and g0.size > 0  # guard the divide below
+                and coords.ndim == 1
+                and coords.size > 0
+                and coords[0] >= 0
+                and coords[-1] < shape[0]
+                and coords[0] <= coords[-1]
+            ):
+                size = g0.size
+                first = int(coords[0]) // size
+                last = int(coords[-1]) // size
+                chunk_span = last - first + 1
+                # searchsorted does O(log n) work per chunk in the spanned range. Fall through
+                # when directly processing the coordinates is expected to be cheaper.
+                if (
+                    chunk_span * coords.size.bit_length() < coords.size
+                    and bool((coords[:-1] <= coords[1:]).all())  # sorted -> grouped by chunk
+                ):
+                    # Search only internal boundaries. Derive the first and last counts from the
+                    # selection bounds so that the boundary after the last chunk cannot overflow.
+                    if first == last:
+                        counts = np.array([coords.size], dtype=np.intp)
+                    else:
+                        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)
+
+                    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, "chunk_rixs", chunk_rixs)
+                    object.__setattr__(self, "chunk_mixs", (chunk_rixs,))
+                    object.__setattr__(self, "dim_grids", dim_grids)
+                    object.__setattr__(self, "shape", coords.shape)
+                    object.__setattr__(self, "drop_axes", ())
+                    return
+
         # handle wraparound, boundscheck
         for dim_sel, dim_len in zip(selection_normalized, shape, strict=True):
             # handle wraparound

tests/test_indexing.py

@@ -14,8 +14,10 @@
 from tests.conftest import Expect, ExpectFail
 from zarr import Array
 from zarr.core.buffer import default_buffer_prototype
+from zarr.core.chunk_grids import ChunkGrid
 from zarr.core.indexing import (
     BasicSelection,
+    CoordinateIndexer,
     CoordinateSelection,
     OrthogonalSelection,
     Selection,
@@ -1047,8 +1049,15 @@ def _test_get_coordinate_selection(
     Expect(input=[3, 25, 8, 17], output=None, id="out-of-order"),
     Expect(input=[1, 8, 15, 29], output=None, id="sorted"),
     Expect(input=[29, 15, 8, 1], output=None, id="reversed"),
+    Expect(input=np.array([29, 15, 8, 1], dtype=np.uint32), output=None, id="reversed-uint"),
     Expect(input=[2, 2, 8, 8], output=None, id="duplicates"),
     Expect(input=np.array([[2, 4], [6, 8]]), output=None, id="multi-dim"),
+    # sorted-1D fast path (chunk_shape=(7,)): boundaries, contiguous runs, single chunk, full
+    Expect(input=[0, 6, 7, 13, 14, 28, 29], output=None, id="sorted-chunk-boundaries"),
+    Expect(input=[0, 1, 2, 8, 9, 10, 21, 22, 23], output=None, id="sorted-contiguous-runs"),
+    Expect(input=[1, 2, 3, 4, 5, 6], output=None, id="sorted-single-chunk"),
+    Expect(input=list(range(30)), output=None, id="sorted-full"),
+    Expect(input=[0, 0, 7, 7, 7, 29], output=None, id="sorted-duplicates-boundaries"),
 ]
 
 # get_coordinate_selection and vindex word their errors differently for these
@@ -1141,6 +1150,107 @@ def test_get_coordinate_selection_1d(
     _test_get_coordinate_selection(a, z, case.input)
 
 
+@pytest.mark.parametrize(
+    ("chunks", "shards"),
+    [((7,), None), ((7,), (21,))],
+    ids=["chunked", "sharded"],
+)
+def test_get_coordinate_selection_1d_fast_path(
+    store: StorePath, chunks: tuple[int, ...], shards: tuple[int, ...] | None
+) -> None:
+    """The sorted-1D-runs fast path in CoordinateIndexer matches numpy on chunked and sharded arrays.
+
+    Exercises the boundary/run/single-chunk/full-array cases that the fast path optimizes, plus
+    the sharded case where the top-level (shard) grid drives chunk assignment.
+    """
+    a = np.arange(210, dtype=int)
+    z = zarr.create_array(
+        store=store / str(uuid4()),
+        shape=a.shape,
+        dtype=a.dtype,
+        chunks=chunks,
+        shards=shards,
+    )
+    z[:] = a
+    rng = np.random.default_rng(0)
+    selections = [
+        np.sort(rng.choice(210, 60, replace=False)),  # scattered sorted
+        np.array([0, 6, 7, 20, 21, 209]),  # chunk/shard boundaries
+        np.concatenate([np.arange(s, s + 5) for s in (0, 33, 100, 180)]),  # contiguous runs
+        np.array([0, 0, 7, 7, 209]),  # sorted with duplicates
+        np.arange(210),  # whole array
+        np.array([5]),  # single element
+    ]
+    for sel in selections:
+        assert_array_equal(a[sel], z.get_coordinate_selection(sel))
+        assert_array_equal(a[sel], z.vindex[sel])
+
+
+def test_coordinate_indexer_1d_last_chunk_boundary_does_not_overflow() -> None:
+    max_intp = np.iinfo(np.intp).max
+    chunk_size = max_intp // 2 + 1
+    coords = np.arange(max_intp - 4, max_intp, dtype=np.intp)
+    chunk_grid = ChunkGrid.from_sizes((max_intp,), (chunk_size,))
+
+    (projection,) = tuple(CoordinateIndexer((coords,), (max_intp,), chunk_grid))
+
+    assert projection.chunk_coords == (1,)
+    assert_array_equal(projection.chunk_selection[0], coords - chunk_size)
+    assert projection.out_selection == slice(0, 4)
+
+
+@pytest.mark.parametrize("coord_dtype", [np.int8, np.uint8, np.uint32])
+def test_coordinate_selection_1d_narrow_dtype_large_chunk(
+    store: StorePath, coord_dtype: type[np.integer[Any]]
+) -> None:
+    source = np.arange(1_000)
+    coords = np.arange(10, dtype=coord_dtype)
+    z = zarr_array_from_numpy_array(store, source, chunk_shape=(1_000,))
+
+    assert_array_equal(z.get_coordinate_selection(coords), source[coords])
+    assert_array_equal(z.vindex[coords], source[coords])
+    assert_array_equal(z[coords], source[coords])
+
+    expected = source.copy()
+    expected[coords] = -1
+    z.set_coordinate_selection(coords, -1)
+    assert_array_equal(z[:], expected)
+    z[:] = source
+    z.vindex[coords] = -1
+    assert_array_equal(z[:], expected)
+
+
+def test_coordinate_indexer_1d_sparse_selection_uses_general_path(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    coords = np.array([0, 99])
+    chunk_grid = ChunkGrid.from_sizes((100,), (1,))
+
+    def unexpected_searchsorted(*args: Any, **kwargs: Any) -> None:
+        pytest.fail("sparse coordinate selection should not call searchsorted")
+
+    monkeypatch.setattr(np, "searchsorted", unexpected_searchsorted)
+    projections = tuple(CoordinateIndexer((coords,), (100,), chunk_grid))
+
+    assert tuple(projection.chunk_coords for projection in projections) == ((0,), (99,))
+
+
+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."""
+    a = np.arange(30, dtype=int)
+    with zarr.config.set({"array.rectilinear_chunks": True}):
+        z = zarr.create_array(
+            store=store / str(uuid4()),
+            shape=a.shape,
+            dtype=a.dtype,
+            chunks=((3, 3, 4, 5, 5, 5, 5),),
+        )
+    z[:] = a
+    for sel in (np.array([1, 8, 15, 29]), np.array([0, 3, 3, 29]), np.arange(30)):
+        assert_array_equal(a[sel], z.get_coordinate_selection(sel))
+
+
 @pytest.mark.parametrize("case", _COORD_1D_BAD_CASES, ids=lambda c: c.id)
 def test_get_coordinate_selection_1d_raises(store: StorePath, case: ExpectFail[Any]) -> None:
     """get_coordinate_selection and vindex both raise IndexError for invalid 1D selections."""