zarrs bindings [do not merge]

[d-v-b]

vibe-coded zarrs bindings. I am not done with this, but I figured it's in a good state for signposting / discussion.

strategy

I wanted to keep the contact surface between zarr-python and zarrs minimal and low-state, so I defined a functional crud API that expresses the core of zarr IO. (that API is not wired up to the top-level zarr.Array class!). The idea is that we can express the stuff users want to do to their data -- create new arrays / groups, write chunks , read chunks, as f(metadata, storage, *parameters).

The crud API supports multiple registered backends, e.g. a default python backend (based on repurposing our existing routines) and the rust backend, when zarrs is available.

The statelessness of the functional API is also a downside if you call read_chunk(metadata, store, ...) repeatedly, because the rust code will re-construct the same chunk decoding machinery each time. I address this with an LRU cache on the rust side. I think "metadata + store + options" is a good cache key but we need to discuss this design further.

An alternative strategy would be to write a zarrs-based expression for the many methods defined on the Array and AsyncArray classes, while ensuring that we avoid crossing the FFI boundary excessively. I avoided this because I imagined it would require covering a huge code surface area and raise tough questions about whether python or rust was owning the life cycle of the object. If people really want a full rust-backed Array class, we can explore that direction.

caveats:

• I have only implemented simple indexing right now. I'm going to add full numpy indexing semantics down the road, using the data structures defined in https://github.com/zarr-developers/ndsel as an FFI-friendly data model.
• Rust handles local file system storage itself. other stores have to cross FFI for every store operation, and I haven't tested pathologies like deleting the store on the python side while rust is working on it. Maybe I am naive but I would really like pure functions here, which means the best solution is for stores to move across the language boundary as plain data, not live python objects. URL pipeline syntax would be a great addition here.

performance

the zarrs backend is faster! here's a benchmark script you can run yourself. It requires the rust toolchain for building the bindings.

I'm seeing ~15x throughput improvement, looks good.

array:       shape=(2048, 2048) dtype=uint16 chunks=(64, 64) shards=(512, 512) compressor=zstd
size:        8.4 MB logical, 10 iterations, LocalStore
correctness: reference and zarrs both match the source data

backend        best (ms)   median (ms)   median MB/s
reference         147.05        157.85            53
zarrs               6.94          8.97           935

impact

these changes require internal changes in the zarr package, as well as a new subpackage for the rust bindings. it adds the rust toolchain to the developer dependencies of the project. it exposes us to changes in the zarrs package, which is outside the zarr-developers org. We definitely need a design plan to limit complexity if we want to pursue this further.

[d-v-b]

@clbarnes related to your interests

[codecov]

Codecov Report

❌ Patch coverage is 20.68127% with 326 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.09%. Comparing base (f12c1dc) to head (004ce2e).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/zarr/crud/_api.py	22.96%	104 Missing ⚠️
src/zarr/crud/_reference.py	23.66%	100 Missing ⚠️
src/zarr/zarrs/_bridge.py	0.00%	54 Missing ⚠️
src/zarr/zarrs/_backend.py	0.00%	45 Missing ⚠️
src/zarr/crud/_registry.py	41.17%	10 Missing ⚠️
src/zarr/zarrs/__init__.py	0.00%	9 Missing ⚠️
src/zarr/crud/_common.py	55.55%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4064      +/-   ##
==========================================
- Coverage   93.52%   91.09%   -2.43%     
==========================================
  Files          90       99       +9     
  Lines       11926    12337     +411     
==========================================
+ Hits        11154    11239      +85     
- Misses        772     1098     +326     

Files with missing lines	Coverage Δ
src/zarr/core/config.py	100.00% <ø> (ø)
src/zarr/crud/__init__.py	100.00% <100.00%> (ø)
src/zarr/crud/_backend.py	100.00% <100.00%> (ø)
src/zarr/crud/_common.py	55.55% <55.55%> (ø)
src/zarr/zarrs/__init__.py	0.00% <0.00%> (ø)
src/zarr/crud/_registry.py	41.17% <41.17%> (ø)
src/zarr/zarrs/_backend.py	0.00% <0.00%> (ø)
src/zarr/zarrs/_bridge.py	0.00% <0.00%> (ø)
src/zarr/crud/_reference.py	23.66% <23.66%> (ø)
src/zarr/crud/_api.py	22.96% <22.96%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[kylebarron]

Returning an ndarray will force a copy out of rust memory. Longer term, I'd perhaps suggest instead making a wrapper class on the Rust side that will expose an ndarray. At least whenever the data type is supported by the buffer protocol that should remove a memory copy. See for example async_tiff.Array

[d-v-b]

we have very adjacent prior art in https://github.com/zarrs/zarrs-python, I should see what they do!

[clbarnes]

You can instantiate numpy arrays on the rust side, either from a vec directly or with an ndarray: https://docs.rs/numpy/latest/numpy/convert/trait.IntoPyArray.html#tymethod.into_pyarray

I imagine it would get more complicated when custom buffer prototypes are passed, because they need to be instantiated on the python side.

[d-v-b]

I am kind of hoping to put the custom buffer prototype thing in the rear view window. It has not worked well.

[d-v-b]

I am open to all ideas here, but "every routine that allocates memory takes a python callable" is probably not going to work for cross-language interop. Another complication is that "contiguous stream of bytes" is a lower-level abstraction than "N-dimensional array".

My rough feeling is that we should partition the chunk encoding / decoding routines by the memory access pattern needed for the array side of the array -> bytes transformation. For dense arrays, rust can hand back a contiguous buffer, or write into a pre-allocated buffer. For sparse arrays, we will need to settle on an memory layout, and we probably can't support writing into a pre-allocated sparse array.

[kylebarron]

Nit: as you might expect from our zulip discussion I'm hesitant to lump all of these together. I'm interested in spending a little time with this thinking about how to represent these traits through a Python API

[kylebarron]

You know I don't love this 😅. I'd like to specialize Sync and Async stores on the rust side.

[d-v-b]

haha I'm not a fan either! ideal scenario is that a store connection has a declarative form that rust can use to reconstruct it (like a URL). Things tied to python memory will remain hard.

[kylebarron]

Fwiw that's how obstore works when passed into a pyo3 binding. It recreates the store on the rust side so no requests need to go through Python

[kylebarron]

I think it's pretty messy to do this in a global LRU cache. There's no reason to have a public LRU cache when you could just cache the metadata in... a class instance 😄

[d-v-b]

in... a class instance 😄

storing the metadata in a class instance isn't the same as globally caching it. The point here is to avoid performing the same computation that is keyed by exactly the metadata document. Even if we were using OOP, I think it would make sense for the procedure that takes metadata and emits chunk encoding / decoding machinery to use a cache.

[kylebarron]

The point here is to avoid performing the same computation that is keyed by exactly the metadata document

You're keying on filesystem - node path - metadata JSON... how would that key ever be shared across more than one Array? Seems like an Array class is a logical place to put that metadata

[d-v-b]

it's possible that we aren't disagreeing! I want a functional API because I think that's actually a very natural way to interact with zarr data -- for a given operation, you specify the metadata document you want to use, the location, and any other parameters.

It also definitely makes sense to provide users with objects that persistently bind the metadata and the storage backend. That's how zarr-python works today. But if you start with the object oriented API, it can be hard to add the functional API later. I think the functional API is actually really important for doing a lot of cool things with Zarr! so i want to see if we can bake it in foundationally, and add the object-oriented API on top.

[kylebarron]

I need to learn more about codec pipelines to understand how they integrate with zarrs

[kylebarron]

I don't like the functional API because it means that you're reparsing metadata that makes sense to just store in a class instance