arr_1d[np.sort(idx)] can be much faster (CoordinateIndexer fast path)

[selmanozleyen]

Describe the new feature you'd like

Hi,

Fancy-indexing a 1-D Zarr array with a sorted integer array (arr[sorted_idx]) is dominated by CoordinateIndexer.__init__. It runs several checks over the flat selection. For a sorted selection this is avoidable: the per-chunk grouping can be computed with a single searchsorted over the touched chunk boundaries in O(#chunks_touched · log n).

arr[idx], vindex[idx], and get_coordinate_selection all construct a CoordinateIndexer, so this covers the natural indexing syntax. And we have batched reads of CSR data stored in zarr for ML training, where each batch is a large sorted gather. And it is definetely a hot path.

Evidence

Index construction only (no I/O), regular sharded 1-D array, ~1450-element runs (table written in md with AI):

runs	elements	CoordinateIndexer build
1024	1.48M	8.7 ms → 0.5 ms
4096	5.94M	35 ms → 2.1 ms
8192	11.9M	71 ms → 4.5 ms

Minimal repro (only needs zarr + numpy):

import numpy as np, zarr, time
from zarr.core.indexing import CoordinateIndexer

n = 130 * 89_351_824
z = zarr.create_array(
    store=zarr.storage.MemoryStore(),
    shape=(n,),
    chunks=(91_549,),
    shards=(89_351_824,),
    dtype="f4"
)
cg = z._async_array._chunk_grid
starts = np.sort(np.random.default_rng(0).choice(n // 1450, 4096, replace=False))
coords = np.concatenate([np.arange(s * 1450, (s + 1) * 1450) for s in starts])

t = time.perf_counter()
CoordinateIndexer((coords,), (n,), cg)
print((time.perf_counter() - t) * 1e3, "ms")

Suggestion

If we added this codeblock around here in the constructor:

zarr-python/src/zarr/core/indexing.py

Line 1208 in ecc2d77

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

kinda related convo: zarrs/zarrs-python#147 (comment)

[d-v-b]

thanks for this work! if you have the time, could you explore whether the array indexing logic defined in #3906 benefits from the same performance optimization?

[selmanozleyen]

I don't know if you mean "if the changes in CoordinateIndexer benefit us" or "if we used the sorted-coordinate grouping trick it would benefit us". But the answer is yes to both!

• For ordinary arrays, feat: add IndexTransform library for composable, lazy coordinate mappings #3906 already calls CoordinateIndexer, so it receives the improvement directly from the new CoordinateIndexer.

• For lazy arrays, the proposed change is to partition a sorted 1-D ArrayMap once using searchsorted, rather than scanning the complete coordinate array in transform.intersect(...) for every touched shard.

I compared the current #3906 head with the new CoordinateIndexer and the #3906 after my proposal.

All timings measure index planning and shard projection/resolution only; no I/O.

(disclaimer table and dropdown prettification is processed through AI)

Indexing path	Runs	Elements	Current #3906	Proposed #3906 + new CoordinateIndexer
Ordinary 1-D coordinate indexing	1,024	1.48M	11.2 ms	1.73 ms
Ordinary 1-D coordinate indexing	4,096	5.94M	47.7 ms	9.85 ms
Ordinary 1-D coordinate indexing	8,192	11.9M	95.8 ms	14.7 ms
Lazy 1-D coordinate indexing	1,024	1.48M	190 ms	2.04 ms
Lazy 1-D coordinate indexing	4,096	5.94M	1,810 ms	15.6 ms
Lazy 1-D coordinate indexing	8,192	11.9M	2,957 ms	30.2 ms

Suggested code for #3906:

def iter_sorted_1d(transform, chunk_grid):
    m = transform.output[0]
    coords = m.offset + m.stride * m.index_array.reshape(-1)
    dg = chunk_grid._dimensions[0]

    first = dg.index_to_chunk(int(coords[0]))
    last = dg.index_to_chunk(int(coords[-1]))

    edges = np.fromiter(
        (dg.chunk_offset(c) for c in range(first, last + 2)),
        dtype=np.intp,
        count=last - first + 2,
    )
    offsets = np.searchsorted(coords, edges)

    for relative_shard in np.flatnonzero(np.diff(offsets)):
        shard = first + int(relative_shard)
        start = int(offsets[relative_shard])
        stop = int(offsets[relative_shard + 1])

        values = m.index_array.reshape(-1)[start:stop]

        restricted = IndexTransform(
            domain=IndexDomain((0,), (stop - start,)),
            output=(
                ArrayMap(
                    values,
                    offset=m.offset,
                    stride=m.stride,
                ),
            ),
        )

        local = restricted.translate((-dg.chunk_offset(shard),))

        yield (
            (shard,),
            local,
            np.arange(start, stop, dtype=np.intp),
        )

Results for the lazy transform resolver alone:

runs,elements,transform_build_ms,current_resolve_ms,proposed_resolve_ms
1024,1484800,0.669,189.594,1.373
4096,5939200,6.182,1804.259,9.374
8192,11878400,12.295,2944.312,17.916

Reproducer

Run this from a checkout of #3906:

from __future__ import annotations

import gc
import time
from dataclasses import dataclass

import numpy as np

from zarr_transforms.chunk_resolution import iter_chunk_transforms
from zarr_transforms.domain import IndexDomain
from zarr_transforms.output_map import ArrayMap
from zarr_transforms.transform import IndexTransform, selection_to_transform


@dataclass(frozen=True)
class FixedDimension:
    size: int

    def index_to_chunk(self, idx: int) -> int:
        return int(idx) // self.size

    def chunk_offset(self, chunk_ix: int) -> int:
        return chunk_ix * self.size

    def chunk_size(self, chunk_ix: int) -> int:
        return self.size

    def indices_to_chunks(self, indices: np.ndarray) -> np.ndarray:
        return indices // self.size


@dataclass(frozen=True)
class Grid:
    _dimensions: tuple[FixedDimension, ...]


def iter_sorted_1d(transform, chunk_grid):
    m = transform.output[0]
    assert isinstance(m, ArrayMap)
    assert m.input_dimension is None
    assert m.stride > 0

    coords = m.offset + m.stride * m.index_array.reshape(-1)
    assert coords.size > 0
    assert bool((np.diff(coords) >= 0).all())

    dg = chunk_grid._dimensions[0]
    first = dg.index_to_chunk(int(coords[0]))
    last = dg.index_to_chunk(int(coords[-1]))

    edges = np.fromiter(
        (dg.chunk_offset(c) for c in range(first, last + 2)),
        dtype=np.intp,
        count=last - first + 2,
    )
    offsets = np.searchsorted(coords, edges)

    for relative_shard in np.flatnonzero(np.diff(offsets)):
        shard = first + int(relative_shard)
        start = int(offsets[relative_shard])
        stop = int(offsets[relative_shard + 1])

        values = m.index_array.reshape(-1)[start:stop]

        restricted = IndexTransform(
            domain=IndexDomain((0,), (stop - start,)),
            output=(
                ArrayMap(
                    values,
                    offset=m.offset,
                    stride=m.stride,
                ),
            ),
        )

        local = restricted.translate((-dg.chunk_offset(shard),))

        yield (
            (shard,),
            local,
            np.arange(start, stop, dtype=np.intp),
        )


def timed(fn, repeats=3):
    results = []

    for _ in range(repeats):
        gc.collect()
        start = time.perf_counter()
        fn()
        results.append(time.perf_counter() - start)

    return min(results)


n = 130 * 89_351_824

# This is the outer shard grid from the original example.
grid = Grid((FixedDimension(89_351_824),))

print(
    "runs,elements,transform_build_ms,"
    "current_resolve_ms,proposed_resolve_ms"
)

for nruns in (1024, 4096, 8192):
    rng = np.random.default_rng(0)
    starts = np.sort(rng.choice(n // 1450, nruns, replace=False))
    coords = np.concatenate(
        [
            np.arange(start * 1450, (start + 1) * 1450)
            for start in starts
        ]
    )

    def construct():
        selection_to_transform(
            (coords,),
            IndexTransform.from_shape((n,)),
            "vectorized",
        )

    transform = selection_to_transform(
        (coords,),
        IndexTransform.from_shape((n,)),
        "vectorized",
    )

    def current_resolve():
        for _ in iter_chunk_transforms(transform, grid):
            pass

    def proposed_resolve():
        for _ in iter_sorted_1d(transform, grid):
            pass

    construct_ms = timed(construct) * 1e3
    current_ms = timed(current_resolve, repeats=1) * 1e3
    proposed_ms = timed(proposed_resolve) * 1e3

    assert sum(1 for _ in iter_chunk_transforms(transform, grid)) == sum(
        1 for _ in iter_sorted_1d(transform, grid)
    )

    print(
        f"{nruns},{coords.size},{construct_ms:.3f},"
        f"{current_ms:.3f},{proposed_ms:.3f}"
    )

uv run python benchmark_sorted_1d_transform.py

The ordinary-array measurements were produced by running this benchmark once on #3906 head and once on the new CoordinateIndexer branch:

import gc
import time

import numpy as np
import zarr
from zarr.core.indexing import CoordinateIndexer


def timed(fn, repeats=7):
    results = []

    for _ in range(repeats):
        gc.collect()
        start = time.perf_counter()
        fn()
        results.append(time.perf_counter() - start)

    return min(results)


n = 130 * 89_351_824

z = zarr.create_array(
    store=zarr.storage.MemoryStore(),
    shape=(n,),
    chunks=(91_549,),
    shards=(89_351_824,),
    dtype="f4",
)

chunk_grid = z._async_array._chunk_grid

print("runs,elements,build_and_iterate_ms")

for nruns in (1024, 4096, 8192):
    rng = np.random.default_rng(0)
    starts = np.sort(rng.choice(n // 1450, nruns, replace=False))
    coords = np.concatenate(
        [
            np.arange(start * 1450, (start + 1) * 1450)
            for start in starts
        ]
    )

    def build_and_iterate():
        indexer = CoordinateIndexer((coords,), (n,), chunk_grid)
        for _ in indexer:
            pass

    elapsed = timed(build_and_iterate)

    print(f"{nruns},{coords.size},{elapsed * 1e3:.3f}")

[d-v-b]

that's great, would you mind opening a PR against my fork, targeting the lazy indexing branch (refactor/simplify-indexing)? No worries if not, but it would be cool to credit you in the commit history

[selmanozleyen]

that's great, would you mind opening a PR against my fork, targeting the lazy indexing branch (refactor/simplify-indexing)? No worries if not, but it would be cool to credit you in the commit history

Yeah ofc I can!

[selmanozleyen]

But should I also open a separate PR for the changes CoordinateIndexer?

[d-v-b]

yes that would be great, thank you!