feat: Added flexible backends with the Zarr(s)-Python and Tensorstore support.

[srivarra]

Previously iohub relied on the Zarr-Python internals directly throughout the NGFF layer, and TensorStore was used for a couple of one off features: Position.to_tensorstore() which opens a FOV as a TensorStore handle manually, and for downsampling.

We'd like to be able to "automatically" use TensorStore, or Zarr Python depending on the use case without having to call .to_tensorstore().

We've replaced that with a pluggable backend protocol.

There's a new abstraction layer over in src/iohub/core which defines a Protocol ZarrImplementation[G, A] which is composed of the GroupBackend[G] (open/navigate groups), the ArrayBackend[G, A] (create/open arrays) and an ArrayIO[A] (read, write, downsample, convert).

There are two implementations: ZarrPythonImplementation (with autodetection of the Rust zarrs codec pipeline) and TensorStoreImplementation (needed some light shims to make it handle the filesystem / Zarr Groups like Zarr-Python, it's a bit messy but I couldn't find a better way).

NGFFNDArray wraps any backend handle. All I/O (reads, writes, orthogonal indexing), dask conversion, downsampling, and append delegates to the bound ZarrImplementation. ImageArray subclasses it instead of zarr.Array.

Here's a diagram to show how it's wired up:

                      open_ome_zarr(path, implementation="tensorstore") # or "zarr"
                              │
                              ▼
                      ┌───────────────┐
                      │   NGFFNode    │  (Plate / Well / Position)
                      │  self._impl ──┼──────────────────────────┐
                      └───────┬───────┘                          │
                              │                                  │
                      ┌───────▼───────┐                          │
                      │  NGFFNDArray  │  (ImageArray)            │
                      │  self._impl ──┼──────────────────────────┤
                      └───────────────┘                          │
                                                                 │
                      ┌──────────────────────────────────────────▼──┐
                      │          ZarrImplementation[G, A]           │
                      │              (Protocol)                     │
                      │                                             │
                      │  open_group  · create_array · read · write  │
                      │  group_keys  · open_array   · oindex        │
                      │  array_keys  · get_shape    · to_dask       │
                      │  close       · get_chunks   · downsample    │
                      └──────────────────┬───────────────┬──────────┘
                                         │               │
                 ┌───────────────────────┘               └─────────────────────┐
                 │                                                             │
      ┌──────────▼──────────────┐                            ┌────────────────▼──────────────┐
      │  ZarrPythonImpl         │                            │  TensorStoreImpl              │
      │  G = zarr.Group         │                            │  G = _TsGroup                 │
      │  A = zarr.Array         │                            │  A = ts.TensorStore           │
      │                         │                            │                               │
      │  auto-detects zarrs     │                            │  _detect_zarr_driver()        │
      │  Rust codec pipeline    │                            │  once per open_group          │
      │  shard-region iteration │                            │  → "zarr2" or "zarr3"         │
      │  NumPy downsampling     │                            │  ts.downsample() virtual view │
      └─────────────────────────┘                            └───────────────────────────────┘
                 │                                                             │
      ┌──────────▼──────────────┐                            ┌────────────────▼──────────────┐
      │      zarr.Array         │                            │        _TsGroup shim          │
      │   (zarr-python store)   │                            │  _TsAttrs (reads .zattrs/     │
      └─────────────────────────┘                            │   zarr.json, writes zarr.json)│
                                                             └───────────────────────────────┘

This also now defaults to writing OME-Zarr v0.5 and removed the ability to write OME-Zarr v0.4. (Reading OME-Zarr v0.4 is still allowed).

You can specify the Zarr implementation within open_ome_zarr like so:

open_ome_zarr(..., implementation=Literal["zarr", "tensorstore"]

[srivarra]

While the performance benefits of TensorStore is pretty nice, it's also adding a lot of complexity to handle both of these Zarr implementations.

I think for now this is useful in the short term, but a future release we should completely drop it in favor of just Zarr-Python with the Zarrs-Python codec.

See:

• feat: add ZarrsArray zarrs/zarrs-python#147
• feat/lazy indexing zarr-developers/zarr-python#3678 (Copying over some of tensorstore's Lazy featureset)

Very exciting stuff going on over there!