diff --git a/changes/4172.misc.md b/changes/4172.misc.md new file mode 100644 index 0000000000..0be7226476 --- /dev/null +++ b/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. diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index f6eb495cd9..875c22fbd3 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,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 diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 3d80f6364c..04fbdad8c6 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, @@ -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."""