[do not merge] zarrs bindings - #4064
Conversation
|
@clbarnes related to your interests |
Codecov Report❌ Patch coverage is
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
🚀 New features to boost your workflow:
|
| async def list_children(store, path) -> list[tuple[str, dict]] # (path, metadata) | ||
|
|
||
| # chunk-level I/O | ||
| async def decode_chunk(metadata, store, path, chunk_coords, *, selection=None) -> np.ndarray |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
we have very adjacent prior art in https://github.com/zarrs/zarrs-python, I should see what they do!
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
I am kind of hoping to put the custom buffer prototype thing in the rear view window. It has not worked well.
There was a problem hiding this comment.
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.
| 3. **Anything else** → generic `PyStore`: a Rust struct implementing | ||
| `ReadableStorageTraits` / `WritableStorageTraits` / | ||
| `ListableStorageTraits` over a Python callback object. |
There was a problem hiding this comment.
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
| The callback path: the async API function wraps the user's `Store` in a small | ||
| sync Python shim whose methods submit coroutines to zarr-python's existing | ||
| sync event-loop thread (`zarr.core.sync`, | ||
| `asyncio.run_coroutine_threadsafe(...)` + blocking result). Rust calls the | ||
| shim while holding no locks of its own. This makes any conformant `Store` | ||
| (Memory, Zip, Logging, Wrapper, user-defined) work without Rust knowing its | ||
| type. Deadlock safety relies on the existing invariant that code running on | ||
| the zarr sync loop never blocks on these Rust entry points. |
There was a problem hiding this comment.
You know I don't love this 😅. I'd like to specialize Sync and Async stores on the rust side.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
| µs of actual chunk I/O on a warm filesystem. To amortize it across the common | ||
| "open one array, then do many chunk operations" pattern, the chunk/region | ||
| routines memoize the constructed `Array` in a process-global LRU cache | ||
| (capacity 128) keyed on `(filesystem root, node path, metadata JSON)`. |
There was a problem hiding this comment.
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 😄
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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.
| zarr-python already contains everything a pure-Python backend needs: | ||
| `BatchedCodecPipeline` (`src/zarr/core/codec_pipeline.py`), `BasicIndexer` | ||
| (`src/zarr/core/indexing.py`), `save_metadata` (`src/zarr/core/metadata/io.py`), | ||
| metadata parsing (`ArrayV3Metadata.from_dict` / `ArrayV2Metadata.from_dict`), | ||
| and chunk-key encoding (`src/zarr/core/chunk_key_encodings.py`). |
There was a problem hiding this comment.
I need to learn more about codec pipelines to understand how they integrate with zarrs
| but carries only defaults (fields become meaningful in Phase 3). | ||
|
|
||
| ```python | ||
| # node lifecycle |
There was a problem hiding this comment.
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
|
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. |
| Python-store callback bridge: it operates only on stores it can map to a | ||
| zarrista store (currently `LocalStore`) and raises `UnsupportedStoreError` for |
There was a problem hiding this comment.
Why not ObjectStore too?
There was a problem hiding this comment.
🤖 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.
There was a problem hiding this comment.
🤖 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?
There was a problem hiding this comment.
A MemoryStore won't work because the underlying memory can't be extracted
| def _node_path(path: str) -> str: | ||
| """Convert a zarr path (`""`, `"foo/bar"`) to a zarrista node path | ||
| (`"/"`, `"/foo/bar"`).""" | ||
| return f"/{path.strip('/')}" |
There was a problem hiding this comment.
Do zarrista paths currently have to start with /? Should they have to? Should a/b/c be an allowed input?
There was a problem hiding this comment.
🤖 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.
There was a problem hiding this comment.
@d-v-b in terms of API design, should Zarrista support strings without a leading / or require it?
There was a problem hiding this comment.
this is actually a subtle point in the spec. the spec states that, in the hierarchy representation, node names are absolute "/"-delimited paths. But when you address the location of an array / group inside a storage backend, IMO it's more clear to use a relative path, i.e. no leading "/" character.
| if isinstance(store, LocalStore): | ||
| return zarrista.FilesystemStore(str(store.root)) |
There was a problem hiding this comment.
You can pass the underlying store inside an ObjectStore directly into Zarrista.
There was a problem hiding this comment.
Although I guess that does force you into the async backend, and maybe you want to focus on sync operations for now.
There was a problem hiding this comment.
🤖 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.
There was a problem hiding this comment.
I think separately there's a question of should the CrudBackend methods all be async?
There was a problem hiding this comment.
yeah that's a good q, i'm not sure. i suspect we ultimately want async and sync versions, because the best choice is genuinely contingent.
| zarrista accelerates the chunk-level I/O paths (`read_chunk`, `read_subset`, | ||
| `write_chunk`, `delete_chunk`). Node metadata documents are written, read, | ||
| listed and deleted with zarr-python's own machinery (delegated to the | ||
| `ReferenceBackend`); zarrista has no "write this exact metadata document" | ||
| primitive for arrays, and these operations are not performance-critical. | ||
|
|
||
| All methods first resolve the store with `_resolve_store`, so the backend | ||
| consistently rejects stores it cannot ingest (raising | ||
| `UnsupportedStoreError`) rather than half-working on them. |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
🤖 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?
| - Only `LocalStore` is ingestable today (maps to a zarrista | ||
| `FilesystemStore`). obstore- and icechunk-backed stores are future work. |
There was a problem hiding this comment.
And fwiw it should be one line of code to pass an icechunk session, though that's also async
There was a problem hiding this comment.
🤖 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.)
| raise NodeExistsError(str(err)) from err | ||
|
|
||
|
|
||
| class ZarrsBackend: |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
good idea, I'll go with that
Per review on zarr-developers#4064: explicit subclassing states the intent (this class implements CrudBackend) and lets type checkers verify conformance at the definition site instead of only at registration call sites. ReferenceBackend and ZarrsBackend here; ZarristaBackend follows in the next commit alongside its store changes. Assisted-by: ClaudeCode:claude-opus-4.8
…c API Widen the zarrista backend beyond LocalStore, per review on zarr-developers#4064: a `zarr.storage.ObjectStore` unwraps to its inner obstore store and runs through `zarrista.AsyncArray` — natively async, no thread offload (the CrudBackend contract is already async; the sync-API-plus-to_thread path remains for LocalStore only). Memory-backed obstore stores are rejected at the gate, since zarrista cannot ingest them ("expected an async compatible storage object"); an icechunk Session mapping is future work with the same shape. Also: ZarristaBackend now explicitly subclasses CrudBackend (companion to the previous commit), and chunk writes/deletes enforce the zarr-level `read_only` flag, which zarrista I/O would otherwise bypass. obstore joins the zarrista dependency group for the new tests. Assisted-by: ClaudeCode:claude-opus-4.8
| if isinstance(target, zarrista.FilesystemStore): | ||
|
|
||
| def _read() -> bytes: | ||
| array = zarrista.Array.open(target, _node_path(path)) |
There was a problem hiding this comment.
Here you should use Array.from_metadata because you already have the metadata available. That will avoid reading the metadata again from source.
| return _to_bytes(array.retrieve_array_subset(selection), np_dtype) | ||
|
|
||
| return await asyncio.to_thread(_read) | ||
| array = await _open_async(target, path) |
There was a problem hiding this comment.
The above comment means that this is no longer async to open, whenever you already have the metadata (ditto to all the other methods in the backend as well)
| def _to_bytes(decoded: Any, np_dtype: np.dtype[Any]) -> bytes: | ||
| """Reinterpret a zarrista decoded array as C-contiguous native-dtype bytes.""" | ||
| arr = np.asarray(decoded.to_numpy(), dtype=np_dtype) | ||
| return np.ascontiguousarray(arr).tobytes() |
There was a problem hiding this comment.
It's important to note that because bytes is the contract you have, this necessarily copies.
In zarrista I'm exposing Tensor as a core class to minimize data copies. So numpy-compatible fixed-width dtypes can directly be exposed from rust memory without a copy
There was a problem hiding this comment.
decoded: Any 🙃 we can definitely do better than this; this operation should absolutely be zero-copy. I'll look into wrapping the tensor class!
There was a problem hiding this comment.
decoded: Any🙃
What's your point here?
There was a problem hiding this comment.
just that the type annotation is wrong -- the decoded object isn't Any
|
This branch will be re-directed along the following lines:
|
Route `Array`/`AsyncArray` data I/O through a pluggable engine, selected with `engine=` on the array entry points and discoverable via `zarr.list_engines()`. The selection crossing the boundary is a `SelectionRequest`: the user's selection, unresolved, tagged with its dialect. It deliberately does not carry a resolved `Indexer` — `CoordinateIndexer` stores its coordinates chunk-sorted rather than in the user's order, so an engine reconstructing a selection from one would silently reorder results. Both views are derived from the raw selection instead: `request.indexer` for chunk-gather backends, and the raw selection for backends that read boxes. The built-in engine forwards to the existing module-level `_get_selection` / `_set_selection`, so it is the identity engine by construction and future changes to that path reach it without being mirrored. The zarrista backend resolves selections through zarr-indexing's `LazyArray` plus `UnitStepReader`. zarrista natively accepts only integers, step-1 slices and `Ellipsis`, which is exactly that reader's contract, so the engine gets basic, orthogonal, coordinate, mask and block selections — reads and writes, including strided and negative-step — over a box-only backend. Partitioning the read by the chunk grid keeps a sparse selection from pulling its whole bounding box. Writes take one `store_array_subset` call for a dense forward box, and otherwise scatter pointwise grouped by chunk. zarr-indexing stays an optional dependency: nothing under `zarr/core` or `zarr/abc` imports it, only the zarrista backend. Assisted-by: ClaudeCode:claude-opus-4.6
|
🤖 AI text below 🤖 Refactored against current
|
| This PR (before) | Current zarrista |
|---|---|
self._arr.store |
renamed .storage — .store no longer exists |
type(x).__name__ != "Tensor" → raise |
concrete class is FixedLengthTensor; Tensor is now a union alias, so this rejected every read |
zarr_metadata.ArrayMetadataV3 |
renamed ZarrV3ArrayMetadataJSON |
Both hand-rolled read-modify-write chunk loops (~40 lines each) are gone — store_array_subset does that natively now. Also picked up: to_numpy() on all tensor types, and MemoryStore/ZipStore as valid sync stores.
zarr-indexing replaced the hand-rolled selection code
_normalize.py (213 lines) and its 140-line test are deleted. zarrista natively accepts only ints, step-1 slices and Ellipsis — exactly UnitStepReader's contract — so LazyArray supplies the full NumPy dialect for free.
The engine went from basic indexing only to basic, orthogonal, coordinate, mask and block, reads and writes, including strided and negative-step.
This also removes the accepted regression. Serving every selection via its bounding box forced a test_accessed_chunks relaxation; partitioning by the chunk grid means oindex[[0, 999]] reads two chunks, not a thousand rows. The indexing suite now passes unmodified.
Two design changes
The boundary carries a SelectionRequest (raw selection + dialect), not a Region or an IndexTransform. I had written the transform-from-indexer adapter before finding that CoordinateIndexer.selection is stored chunk-sorted (indexing.py:1300) — deriving a transform from it would have silently reordered vindex results. Deriving both views from the raw selection avoids that, and keeps zarr_indexing out of zarr/core entirely, so it stays an optional dependency.
_default.py delegates instead of duplicating. It forwards to main's module-level _get_selection/_set_selection, so the built-in engine is the identity engine by construction. The old copy had already silently drifted from main (#3885, #4205, the v2/v3 metadata.dtype unification) — invisible to git, since it is a new file. array.py grew +317 lines instead of +1446.
Three bugs found and fixed
Each caught by something executable, not by review:
- The async engine deadlocked — it called
sync()from inside a coroutine already on the loop. zarrista's pyo3 futures bind to the loop at creation, so calls are now deferred onto the owning loop and theLazyArrayresolves on a worker thread. - The built-in engine rebuilt the codec pipeline
AsyncArrayhad already built — double work per array, and it double-fired the sharding advisory warning, breakingtest_pipeline_parity. - Scalar writes broadcast to 0-stride arrays that NumPy reports as C-contiguous but zarrista rejects.
Verification
- 7196 passed, 0 failures (full suite, including 134 engine/zarrista tests)
- ruff, mypy, and all pre-commit hooks clean
- Differential suite extended to 17 read + 13 write cases × both engines, with numpy as oracle
- The
open_with_engineexample runs against both engines
Dropped the ~6,500 lines of docs/superpowers/ planning docs, which described the superseded design.
Not done
The ~15x benchmark is unverified against this path. The write path in particular is structurally different (single store_array_subset for dense boxes, pointwise chunk-grouped scatter otherwise) — worth re-running before trusting the number in the description above, which still describes the original design.
Follow-ups, deliberately out of scope: list_engines() is backed by a Literal rather than an entry-point group, so third-party engines can only be passed as instances; and the hierarchy cache in _resolve.py is keyed on id(store).
3 participants
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.Arrayclass!). The idea is that we can express the stuff users want to do to their data -- create new arrays / groups, write chunks , read chunks, asf(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
ArrayandAsyncArrayclasses, 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-backedArrayclass, we can explore that direction.caveats:
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.
impact
these changes require internal changes in the
zarrpackage, 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 thezarrspackage, which is outside thezarr-developersorg. We definitely need a design plan to limit complexity if we want to pursue this further.