feat: add ZarrsArray

[LDeakin]

This is a POC for using zarrs' Array from python. Native zarrs performance in Python 🚀🚀🚀. It just layers over zarr.Array and replaces __getitem__ and __setitem__. The advantage over the codec pipeline is that all chunking logic / reassembly happens on the Rust side.

It'd be amazing if there was some way to make this usable by zarr-python (similarly to the codec pipeline) so dask etc can use it.

Also maybe it should be lazy by default, currently opt-in with .lazy

Round trip benchmark

Disclosure: I haven't reviewed this thoroughly, Claude did most of the implementation.

[d-v-b]

I love this!

Also maybe it should be lazy by default, currently opt-in with .lazy

See zarr-developers/zarr-python#3678. Telling claude to copy tensorstore works great!

[ilan-gold]

I am very much so in favor of this because I really abuse this chunk-iteration API and do notice this performance hit. Do you think the performance hit comes from creating many small rust classes? Or is it in zarr-python before we even see the individual chunk selections?

Just wondering out loud for the current state of things, could we make https://github.com/zarrs/zarrs-python/blob/main/python/zarrs/utils.py simpler? I wonder if the creation of a list of selection-out-shape objects or selections/outs/shapes individually specifically could be rust-ified if we could guarantee from zarr-python that the input to it is always a tuple[slice, ...] and we could describe that as Vec<PySlice> or similar for the inputs (and then similar idea for outputs + shape)? I guess I'm just wondering how "strong" we can make the python so that batch_info in our make_chunk_info_for_rust_with_indices function could be turned into one big object that the rust could understand simply instead of our creating many small ones.

I would be very open to dropping this numpy.ndarray support in service of that goal, especially since it predates our fallback support IIRC

Similar/related to this could be an API that the base CodecPipeline in zarr-python exposes for allowing the codec pipeline to construct its own indexer that gets passed down to CodecPipeline.read i.e., a custom batch_info (or some similar change).

Directly to the question of

It'd be amazing if there was some way to make this usable by zarr-python (similarly to the codec pipeline) so dask etc can use it

could be just monkey-patching the Array import from zarr-python i.e., replacing the import-time object, something like setattr(zarr, "Array", OurArray) or something. Very evil and maybe dangerous.

[LDeakin]

Do you think the performance hit comes from creating many small rust classes? Or is it in zarr-python

I'd guess the latter, but I haven't done any profiling. And in regards to everything else you mentioned, I'd say the more we can do in Rust the better, but I have no idea what is feasible.

[LDeakin]

Closing so that zarrista and perhaps a future zarr-python bridge can flourish.

[selmanozleyen]

Hi, it's me again, I am working on bringing native vindex with GIL released. So I found something that might be useful for you guys. I noticed in my branch, even though I added some optimizations on rust level, BasicIndexer would always beat MultiIndexer. Then I did some profiling and noticed there is a algorithmic problem there. So here is the AI written report.

Following up on your question in #147:

"Do you think the performance hit comes from creating many small rust classes? Or is it in zarr-python before we even see the individual chunk selections?"

At least part of it is the latter, and it's measurable in isolation: CoordinateIndexer.__init__
is a large O(n_elements) cost that runs entirely in zarr-python core, before any chunk selections
reach the pipeline. The case is a sorted 1-D coordinate selection made of many contiguous runs
(gathering many disjoint ranges from a big sharded 1-D array). The measurement below is pure
zarr-python — no zarrs, no rust, no I/O — so it's a clean lower bound on the zarr-side cost,
independent of whatever the pipeline does.

1. Isolated, pure zarr-python (no zarrs, no I/O, no storage)

CoordinateIndexer((coords,), shape, chunk_grid) vs the equivalent slice representation
(one BasicIndexer per contiguous run — what arr[start:stop] builds). Grid is a large sharded
1-D array (inner chunk 91549, shard 89351824); each run is ~1450 elements. Build only — nothing
is read. (zarr 3.2.1)

runs	elements	shards hit	Coord ms	Basic ms	ratio
256	371,200	110	3.07	0.42	7.4×
1024	1,484,800	130	9.68	1.48	6.5×
4096	5,939,200	130	35.46	6.30	5.6×
8192	11,878,400	130	71.32	12.91	5.5×

Two things: the coordinate path is 5–7× more expensive, and it scales with the number of
elements, not the number of runs — because it does per-element work over the whole flat selection,
even though the selection only carries O(#runs) of information.

Same script on an HPC compute node (Xeon, slower single-thread) — ratio and scaling hold,
absolutes ~2.5× higher:

runs	Coord ms	Basic ms	ratio
256	3.58	0.84	4.3×
1024	23.92	3.50	6.8×
4096	92.19	14.56	6.3×
8192	180.49	29.69	6.1×

2. Where the time goes (cProfile, 4096-run build ×30)

tottime  function
0.811    indexing.py:__init__            (ravel_multi_index, broadcast, argsort-check, bincount…)
0.129    chunk_grids.py:indices_to_chunks  (coord // chunk, per element)
0.054    indexing.py:boundscheck_indices   (per element)
0.037    numpy diff                        (the np.any(diff<0) sortedness check, per element)
0.029    indexing.py:wraparound_indices    (per element)

Every one of these is a full O(n_elements) numpy pass. For a sorted, contiguous-run 1-D
selection they're all avoidable: no ravel (1-D), no broadcast (1-D), no argsort (already sorted),
boundscheck/wraparound only on run endpoints, and chunk assignment via searchsorted on chunk/shard
boundaries — O(#runs·log) instead of O(n_elements).

Takeaway

• CoordinateIndexer.__init__ for a many-run 1-D selection is tens-to-hundreds of ms of pure
  zarr-python, before the pipeline — and it scales with the element count, when the selection only
  carries O(#runs) of information.
• The win is a fast CoordinateIndexer path for sorted / contiguous-run 1-D selections: skip the
  per-element passes, assign runs to chunks via searchsorted on chunk/shard boundaries.
• I haven't benchmarked the read/pipeline side here, so I'm making no claim about it — this is
  purely the indexer-build cost.

Repro: coordindexer_proof.py — only needs zarr + numpy (no zarrs, no I/O)

"""Proof: for a sorted 1-D coordinate selection made of many contiguous runs, the cost is
zarr-python's `CoordinateIndexer` construction (pure Python, upstream of any codec pipeline),
not the read.

Needs only `zarr` + `numpy`. `CoordinateIndexer.__init__` does several full O(n_elements) numpy
passes (boundscheck, indices_to_chunks, ravel_multi_index, optional argsort, bincount/cumsum) over
the flat coordinate array. We time that against the equivalent slice representation (one
`BasicIndexer` per contiguous run — what `arr[start:stop]` builds), O(#runs). Grid + selection
shape replicate a large sharded 1-D array.
"""

from __future__ import annotations

import argparse
import statistics
import time

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

CHUNK_LEN = 91_549
SHARD_LEN = 89_351_824  # 976 * CHUNK_LEN
RUN_LEN = 1450  # elements per contiguous run
N_SHARDS = 130


def _make_grid():
    n = N_SHARDS * SHARD_LEN
    z = zarr.create_array(
        store=zarr.storage.MemoryStore(),
        shape=(n,), chunks=(CHUNK_LEN,), shards=(SHARD_LEN,), dtype="float32",
    )
    cg = getattr(z._async_array, "_chunk_grid", None) or z.metadata.chunk_grid
    return z, cg, n


def _selection(n_runs: int, n_slots: int, seed: int = 0):
    """A sorted 1-D selection of `n_runs` disjoint contiguous runs of length RUN_LEN."""
    rng = np.random.default_rng(seed)
    starts = np.sort(rng.choice(n_slots, size=n_runs, replace=False))
    runs = [(int(s) * RUN_LEN, (int(s) + 1) * RUN_LEN) for s in starts]
    coords = np.concatenate([np.arange(a, b) for a, b in runs]).astype(np.int64)
    return coords, runs


def _median_ms(fn, k: int = 11) -> tuple[float, float]:
    fn()  # warmup
    ts = []
    for _ in range(k):
        t0 = time.perf_counter()
        fn()
        ts.append((time.perf_counter() - t0) * 1e3)
    return statistics.median(ts), min(ts)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--runs", type=int, nargs="+", default=[256, 1024, 4096, 8192])
    p.add_argument("--cprofile-runs", type=int, default=4096)
    args = p.parse_args()

    z, cg, n = _make_grid()
    shape = (n,)
    n_slots = n // RUN_LEN
    print(f"zarr={zarr.__version__}  array_len={n:,}  shard_len={SHARD_LEN:,}  "
          f"inner_chunk_len={CHUNK_LEN:,}  run_slots={n_slots:,}")
    print("\nCoordinateIndexer(coords)  vs  [BasicIndexer(run) per run]  — build only, NO I/O\n")
    print(f"{'runs':>6} {'elements':>12} {'shards':>7} "
          f"{'Coord ms':>10} {'Basic ms':>10} {'ratio':>7}")
    for nr in args.runs:
        coords, runs = _selection(nr, n_slots)
        n_shards_hit = int(np.unique(coords // SHARD_LEN).size)
        slices = [slice(a, b) for a, b in runs]
        coord_ms, _ = _median_ms(lambda: CoordinateIndexer((coords,), shape, cg))
        basic_ms, _ = _median_ms(
            lambda: [BasicIndexer((s,), shape, cg) for s in slices]  # noqa: B023
        )
        print(f"{nr:>6} {coords.size:>12,} {n_shards_hit:>7} "
              f"{coord_ms:>10.2f} {basic_ms:>10.2f} {coord_ms / basic_ms:>6.1f}x")

    import cProfile
    import pstats
    coords, _ = _selection(args.cprofile_runs, n_slots)
    print(f"\n=== cProfile: CoordinateIndexer build, {args.cprofile_runs} runs "
          f"({coords.size:,} elements), 30 iters ===")
    pr = cProfile.Profile()
    pr.enable()
    for _ in range(30):
        CoordinateIndexer((coords,), shape, cg)
    pr.disable()
    pstats.Stats(pr).sort_stats("tottime").print_stats(12)


if __name__ == "__main__":
    main()