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 iszarr-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) numpypasses (boundscheck, indices_to_chunks, ravel_multi_index, optional argsort, bincount/cumsum) overthe flat coordinate array. We time that against the equivalent slice representation (one`BasicIndexer` per contiguous run — what `arr[start:stop]` builds), O(#runs). Grid + selectionshape replicate a large sharded 1-D array."""from __future__ importannotationsimportargparseimportstatisticsimporttimeimportnumpyasnpimportzarrfromzarr.core.indexingimportBasicIndexer, CoordinateIndexerCHUNK_LEN=91_549SHARD_LEN=89_351_824# 976 * CHUNK_LENRUN_LEN=1450# elements per contiguous runN_SHARDS=130def_make_grid():
n=N_SHARDS*SHARD_LENz=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) orz.metadata.chunk_gridreturnz, cg, ndef_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) forsinstarts]
coords=np.concatenate([np.arange(a, b) fora, binruns]).astype(np.int64)
returncoords, runsdef_median_ms(fn, k: int=11) ->tuple[float, float]:
fn() # warmupts= []
for_inrange(k):
t0=time.perf_counter()
fn()
ts.append((time.perf_counter() -t0) *1e3)
returnstatistics.median(ts), min(ts)
defmain() ->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_LENprint(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}")
fornrinargs.runs:
coords, runs=_selection(nr, n_slots)
n_shards_hit=int(np.unique(coords//SHARD_LEN).size)
slices= [slice(a, b) fora, binruns]
coord_ms, _=_median_ms(lambda: CoordinateIndexer((coords,), shape, cg))
basic_ms, _=_median_ms(
lambda: [BasicIndexer((s,), shape, cg) forsinslices] # 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")
importcProfileimportpstatscoords, _=_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_inrange(30):
CoordinateIndexer((coords,), shape, cg)
pr.disable()
pstats.Stats(pr).sort_stats("tottime").print_stats(12)
if__name__=="__main__":
main()