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

changes/4170.misc.md

@@ -0,0 +1,7 @@
+Added a fast path to `CoordinateIndexer` for 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). These now build their
+per-chunk projections with a single `searchsorted` over the touched chunk boundaries instead of
+several passes over every selected element, making index construction ~15x faster for large
+gathers. Unsorted, negative, multi-dimensional, and irregular-grid selections are unaffected and
+continue to use the existing path.

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,50 @@ def __init__(
                 f"got {selection!r}"
             )
 
+        # Fast path: a single sorted, in-bounds, 1-D integer coordinate array over a regular
+        # (fixed-size) chunk grid -- the common "gather many disjoint ranges" case.
+        # The path below does several full O(n_elements) numpy passes over the flat selection;
+        # when the selection carries only O(#runs) of information that is wasteful.
+        # Here we locate chunk boundaries with a single searchsorted
+        # over the touched chunk edges -> O(#chunks_touched * log n_elements), skipping the passes.
+        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]
+                and bool((coords[:-1] <= coords[1:]).all())  # sorted -> grouped by chunk
+            ):
+                size = g0.size
+                first = int(coords[0]) // size
+                last = int(coords[-1]) // size
+                # count selected points per chunk in [first, last] via boundary searchsorted
+                edges = np.arange(first, last + 2, dtype=np.intp) * size
+                counts = np.diff(np.searchsorted(coords, edges))
+                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 or (1,))
+                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 or (1,))
+                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

@@ -1047,8 +1047,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 +1148,58 @@ 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_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."""