From e757772a59acb5a6405f6662f3f04e4f876cc4cf Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Tue, 21 Jul 2026 16:21:53 +0200 Subject: [PATCH 1/5] perf(indexing): sorted 1-D fast path --- .../coordinate-indexer-1d-fastpath.feature.md | 7 +++ src/zarr/core/indexing.py | 45 ++++++++++++++ tests/test_indexing.py | 58 +++++++++++++++++++ 3 files changed, 110 insertions(+) create mode 100644 changes/coordinate-indexer-1d-fastpath.feature.md diff --git a/changes/coordinate-indexer-1d-fastpath.feature.md b/changes/coordinate-indexer-1d-fastpath.feature.md new file mode 100644 index 0000000000..809b1f0948 --- /dev/null +++ b/changes/coordinate-indexer-1d-fastpath.feature.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. diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index f6eb495cd9..c32cc85945 100644 --- a/src/zarr/core/indexing.py +++ b/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 (e.g. CSR row + # gathers). The general path below does several full O(n_elements) numpy passes + # (wraparound, boundscheck, indices_to_chunks, broadcast, reshape, ravel_multi_index, + # bincount) 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 bool((np.diff(coords) >= 0).all()) # sorted ascending -> 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 diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 3d80f6364c..d8bc6bfd2c 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -1049,6 +1049,12 @@ def _test_get_coordinate_selection( Expect(input=[29, 15, 8, 1], output=None, id="reversed"), 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 +1147,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.""" From 0595fcd9c334c8fe7969a660bfdfdc426ce18dfc Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Tue, 21 Jul 2026 22:59:42 +0200 Subject: [PATCH 2/5] clean up comments --- src/zarr/core/indexing.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index c32cc85945..a7c52b142f 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1208,11 +1208,10 @@ def __init__( ) # 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 (e.g. CSR row - # gathers). The general path below does several full O(n_elements) numpy passes - # (wraparound, boundscheck, indices_to_chunks, broadcast, reshape, ravel_multi_index, - # bincount) 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 + # (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 From 9ce5ab1ba312be8c0abb4926ee14bef46df27b71 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Tue, 21 Jul 2026 23:06:38 +0200 Subject: [PATCH 3/5] name changelog file --- ...{coordinate-indexer-1d-fastpath.feature.md => 4170.feature.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename changes/{coordinate-indexer-1d-fastpath.feature.md => 4170.feature.md} (100%) diff --git a/changes/coordinate-indexer-1d-fastpath.feature.md b/changes/4170.feature.md similarity index 100% rename from changes/coordinate-indexer-1d-fastpath.feature.md rename to changes/4170.feature.md From 5c081ccb2d4e8729327f26d01f40a7bc9386c064 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Wed, 22 Jul 2026 23:35:34 +0200 Subject: [PATCH 4/5] rename changes file --- changes/{4170.feature.md => 4170.misc.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename changes/{4170.feature.md => 4170.misc.md} (100%) diff --git a/changes/4170.feature.md b/changes/4170.misc.md similarity index 100% rename from changes/4170.feature.md rename to changes/4170.misc.md From 5b717a8f8343e01c98fd292ae619d57b22aafab8 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Wed, 22 Jul 2026 23:44:54 +0200 Subject: [PATCH 5/5] add guard for uint and add test case for it --- src/zarr/core/indexing.py | 3 ++- tests/test_indexing.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index a7c52b142f..16ce0c763a 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1226,7 +1226,8 @@ def __init__( and coords.size > 0 and coords[0] >= 0 and coords[-1] < shape[0] - and bool((np.diff(coords) >= 0).all()) # sorted ascending -> grouped by chunk + and coords[0] <= coords[-1] + and bool((coords[:-1] <= coords[1:]).all()) # sorted -> grouped by chunk ): size = g0.size first = int(coords[0]) // size diff --git a/tests/test_indexing.py b/tests/test_indexing.py index d8bc6bfd2c..7d9dd640ce 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -1047,6 +1047,7 @@ 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