From e757772a59acb5a6405f6662f3f04e4f876cc4cf Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Tue, 21 Jul 2026 16:21:53 +0200 Subject: [PATCH 01/11] 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 02/11] 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 03/11] 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 04/11] 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 05/11] 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 From 3d95ea59a0ec29bcf9c1ca7881ead0b8b4d10d34 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Wed, 29 Jul 2026 13:26:26 +0200 Subject: [PATCH 06/11] apply suggestions --- changes/{4170.misc.md => 4172.misc.md} | 0 src/zarr/core/indexing.py | 4 ++-- 2 files changed, 2 insertions(+), 2 deletions(-) rename changes/{4170.misc.md => 4172.misc.md} (100%) diff --git a/changes/4170.misc.md b/changes/4172.misc.md similarity index 100% rename from changes/4170.misc.md rename to changes/4172.misc.md diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index 16ce0c763a..e98986be06 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1240,14 +1240,14 @@ def __init__( 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, "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 or (1,)) + object.__setattr__(self, "shape", coords.shape) object.__setattr__(self, "drop_axes", ()) return From 220224a3f08c898cf5f393b7fa568141f125e948 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Wed, 29 Jul 2026 14:14:11 +0200 Subject: [PATCH 07/11] add test_coordinate_indexer_1d_last_chunk_boundary_does_not_overflow and it's fix --- src/zarr/core/indexing.py | 8 +++++--- tests/test_indexing.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index e98986be06..1ca2fa9403 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1232,9 +1232,11 @@ def __init__( 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)) + # Search only internal boundaries. Derive the first and last counts from the + # selection bounds so that the boundary after the last chunk cannot overflow. + 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 diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 7d9dd640ce..40f9c50bcd 100644 --- a/tests/test_indexing.py +++ b/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, @@ -1184,6 +1186,19 @@ def test_get_coordinate_selection_1d_fast_path( 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.array([max_intp - 1], 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, 1) + + 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.""" From 11f77faba8d04b4d19387411c2b35b5aef3e7fd7 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Wed, 29 Jul 2026 14:24:16 +0200 Subject: [PATCH 08/11] add path whenever the requests are sparse --- src/zarr/core/indexing.py | 44 +++++++++++++++++++++------------------ tests/test_indexing.py | 19 +++++++++++++++-- 2 files changed, 41 insertions(+), 22 deletions(-) diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index 1ca2fa9403..babc002df3 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1232,26 +1232,30 @@ def __init__( size = g0.size first = int(coords[0]) // size last = int(coords[-1]) // size - # Search only internal boundaries. Derive the first and last counts from the - # selection bounds so that the boundary after the last chunk cannot overflow. - 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 + chunk_span = last - first + 1 + # searchsorted does O(log n) work per boundary. Fall through when directly + # processing the coordinates is expected to be cheaper. + if chunk_span * coords.size.bit_length() < coords.size: + # Search only internal boundaries. Derive the first and last counts from the + # selection bounds so that the boundary after the last chunk cannot overflow. + 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): diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 40f9c50bcd..bc7f3270a0 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -1189,14 +1189,29 @@ def test_get_coordinate_selection_1d_fast_path( 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.array([max_intp - 1], dtype=np.intp) + 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, 1) + assert projection.out_selection == slice(0, 4) + + +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: From eaabba19c6d954b60e8410d49bbc997fd2cc0fa2 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Wed, 29 Jul 2026 14:38:54 +0200 Subject: [PATCH 09/11] first == last edge case and do the cost check inside the active path not before (duh) --- src/zarr/core/indexing.py | 15 ++++++++++----- tests/test_indexing.py | 21 +++++++++++++++++++++ 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index babc002df3..ebd8bc2464 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1227,7 +1227,6 @@ def __init__( 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 @@ -1235,12 +1234,18 @@ def __init__( chunk_span = last - first + 1 # searchsorted does O(log n) work per boundary. Fall through when directly # processing the coordinates is expected to be cheaper. - if chunk_span * coords.size.bit_length() < coords.size: + 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. - 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) + 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 diff --git a/tests/test_indexing.py b/tests/test_indexing.py index bc7f3270a0..04fbdad8c6 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -1199,6 +1199,27 @@ def test_coordinate_indexer_1d_last_chunk_boundary_does_not_overflow() -> None: 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: From b89605aa7a0a5ca0145d09691736b6e3edb277a7 Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Wed, 29 Jul 2026 14:42:46 +0200 Subject: [PATCH 10/11] update the changelog file --- changes/4172.misc.md | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/changes/4172.misc.md b/changes/4172.misc.md index 809b1f0948..56dc4a1014 100644 --- a/changes/4172.misc.md +++ b/changes/4172.misc.md @@ -1,7 +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. +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, +unsorted, negative, multi-dimensional, and irregular-grid selections continue to use the existing +implementation. From 975a5f10debe2bc33270879bc03ff6b5c03913ce Mon Sep 17 00:00:00 2001 From: selmanozleyen Date: Wed, 29 Jul 2026 14:43:51 +0200 Subject: [PATCH 11/11] rewording --- changes/4172.misc.md | 6 +++--- src/zarr/core/indexing.py | 14 ++++++-------- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/changes/4172.misc.md b/changes/4172.misc.md index 56dc4a1014..0be7226476 100644 --- a/changes/4172.misc.md +++ b/changes/4172.misc.md @@ -2,6 +2,6 @@ Improved `CoordinateIndexer` construction for large, sorted, in-bounds, one-dime 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, -unsorted, negative, multi-dimensional, and irregular-grid selections continue to use the existing -implementation. +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. diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index ebd8bc2464..875c22fbd3 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1207,12 +1207,10 @@ 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. + # 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] @@ -1232,8 +1230,8 @@ def __init__( first = int(coords[0]) // size last = int(coords[-1]) // size chunk_span = last - first + 1 - # searchsorted does O(log n) work per boundary. Fall through when directly - # processing the coordinates is expected to be cheaper. + # 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