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

[kylebarron]

In case it's useful for evaluation, I started writing a direct, standalone binding of zarrs to Python: https://github.com/developmentseed/zarrista

[d-v-b]

In case it's useful for evaluation, I started writing a direct, standalone binding of zarrs to Python: https://github.com/developmentseed/zarrista

this is awesome kyle! Extremely useful to have a comparison point for the work here.

[kylebarron]

Why not ObjectStore too?

[d-v-b]

🤖 AI text below 🤖

No reason beyond increment size — LocalStore was the smallest store that let the differential suite run against the other backends. zarr.storage.ObjectStore exposes the underlying obstore store as .store, and the CrudBackend contract is already async, so the natural next step is mapping it onto zarrista's async API (AsyncArray.open_async(store.store, ...)). Planned as a follow-up, and an icechunk Session mapping has the same shape.

[d-v-b]

🤖 AI text below 🤖

Done in 6c76610: zarr.storage.ObjectStore now unwraps to its inner obstore store and runs through zarrista.AsyncArray — natively async, no thread offload. Verified round-trips (chunk read/write + cross-chunk subset) against an obstore LocalStore, cross-checked with zarr-python and the reference backend.

One empirical finding: zarrista 0.1.0b5 rejects an obstore MemoryStore at open time with TypeError: expected an async compatible storage object, while other obstore stores pass through fine — so the backend rejects memory-backed ObjectStores at its gate for consistency. Is that expected on the zarrista side, or worth an issue?

[kylebarron]

A MemoryStore won't work because the underlying memory can't be extracted

[kylebarron]

Do zarrista paths currently have to start with /? Should they have to? Should a/b/c be an allowed input?

[d-v-b]

🤖 AI text below 🤖

Checked against 0.1.0b5: yes, the leading slash is required today — Array.open(store, "a") raises ValueError: invalid node path a, and "/a/" (trailing slash) fails too; only "/a" opens. That matches zarrs' NodePath, but zarr-python's convention is store-key style ("" for the root, "a/b/c" for children), hence this shim. If zarrista accepted bare a/b/c and normalized internally, this adapter would disappear — weak vote for that from the interop side, but the shim is cheap either way.

[kylebarron]

@d-v-b in terms of API design, should Zarrista support strings without a leading / or require it?

[kylebarron]

You can pass the underlying store inside an ObjectStore directly into Zarrista.

[kylebarron]

Although I guess that does force you into the async backend, and maybe you want to focus on sync operations for now.

[d-v-b]

🤖 AI text below 🤖

Yes — and the async constraint is actually a non-issue here: the CrudBackend methods are all async def already. The sync zarrista API + asyncio.to_thread in this file was just the smallest first increment for LocalStore; the obstore path would use zarrista.AsyncArray natively with no thread offload, which is a better fit for the contract, not a worse one. .store on zarr.storage.ObjectStore is the unwrap.

[kylebarron]

I think separately there's a question of should the CrudBackend methods all be async?

[kylebarron]

With a CRUD backend this means that zarrista is only doing chunk-by-chunk operations? I'm thinking perhaps zarrista should own more of the indexing strategy than just chunk-by-chunk reading

[d-v-b]

🤖 AI text below 🤖

Reads are not chunk-by-chunk: read_subset hands the whole region to zarrista's retrieve_array_subset, so zarrista owns the multi-chunk read decomposition. Writes are chunk-level today because the CRUD write contract is write_chunk — the shared facade does the region decomposition and read-modify-write for partial boundary chunks in Python.

Agreed the backend should be able to own more of this. The natural evolution is an optional backend-level write_region, with the facade falling back to chunk-wise RMW for backends that don't provide it. zarrista would need a region-write primitive to exploit that (the stubs expose store_chunk/store_encoded_chunk but no subset write) — is exposing zarrs' store_array_subset something you'd entertain on the zarrista side?

[kylebarron]

And fwiw it should be one line of code to pass an icechunk session, though that's also async

[d-v-b]

🤖 AI text below 🤖

Good to know. Once the obstore/async mapping lands, a Session mapping is the same shape — keeping it in scope for that follow-up. (One caveat from the zarrista docs worth flagging: sessions backed by in_memory_storage() won't work since the Rust side reconstructs the session as a separate instance, so tests will need file- or S3-backed storage.)

[kylebarron]

I'd recommend subclassing from CrudBackend. That gives the reader more information over the intent here (for ZarrsBackend to implement CrudBackend) as well as telling the type checker that it should give you type errors if you don't accurately implement CrudBackend

[d-v-b]

good idea, I'll go with that

[d-v-b]

🤖 AI text below 🤖

Done in be59fb4 (and 6c76610 for ZarristaBackend): all three backends now subclass CrudBackend explicitly, and the protocol docstring notes that conformance is checked at the definition site.