How do I define a new codec in my own Python package which is usable by other Python consumers?

[danking]

Hello! Thank you for your work on Zarrs!

I have tried to define, in Rust, a new Array-to-Bytes codec. I want to distribute this codec as a Python package.

I've implemented a ArrayToBytesCodecTraits as well as a ArrayPartialDecoderTraits. Using these classes, fully inside Rust, after calling ArrayBuilder::array_to_bytes_codec, I'm able to read and write to a Zarr file using my new codec. Yay!

I now want to expose this class, via PyO3, in a small Python package. This is where I am stuck.

Zarrs seems to use inventory for defining new codecs. AFAIK, I can't use that approach because inventory defines, at compile-time, the available codecs. I do not want to distribute an entire copy of the zarrs-python binary library inside my package. I just want to use enough of it to define my codec.

---

The other route is to just directly integrate with the Zarr Python package. Zarr seems to expect codec-authors to subclass the Python class ArrayBytesCodec and define an entrypoint. It seems, in order to create an array with a particular codec, I specify a compressor in create_array.

Unfortunately, I can't find a Zarrs-python struct which corresponds to ArraySpec. I also can't seem to find a way to convert from a Python np.dtype to a Zarrs DataType.

Am I misusing Zarrs-python? In retrospect, it seems like it would have been simpler to write a thin PyO3 shim which uses rust-numpy to consume/produce ndarrays and Vec from/to Python.

[LDeakin]

Hi @danking.

I'm able to read and write to a Zarr file using my new codec. Yay!

That is awesome! If you are happy to share, what is the codec?

Am I misusing Zarrs-python?

Our main motivation with zarrs-python was to accelerate zarr-python, and I don't think we've really considered using it to expand capability. E.g. there are a few codecs that zarrs supports that we do not enable in zarrs-python.

Zarrs seems to use inventory for defining new codecs. AFAIK, I can't use that approach because inventory defines, at compile-time, the available codecs. I do not want to distribute an entire copy of the zarrs-python binary library inside my package. I just want to use enough of it to define my codec.

It might not be too difficult to additionally support runtime registration of extensions in zarrs and support that through zarrs-python. Does that sound sufficient to you? I think zarr-python has or had a validation step that rejected arrays with codecs it doesn't understand. Maybe strict mode circumvents that; otherwise, an upstream patch may be necessary.

Unfortunately, I can't find a Zarrs-python struct which corresponds to ArraySpec. I also can't seem to find a way to convert from a Python np.dtype to a Zarrs DataType.

zarrs-python deals with ArraySpecs etc. internally, but there is no public API to help you with that. While zarrs-python is dedicated to just creating a codec pipeline right now, creating a codec bridge could be a valuable addition. Contributions are welcome :)

In retrospect, it seems like it would have been simpler to write a thin PyO3 shim which uses rust-numpy to consume/produce ndarrays and Vec from/to Python.

@keewis might have some helpful insight here. I believe he wrote a sparse array codec in Rust for direct use in zarr-python.

[danking]

Thanks for the prompt response!

The (very preliminary) code is here: https://github.com/danking/zarr-vortex . I have an implementation using Zarrs as well as one using PyO3 and the Zarr Python library.

Vortex [1] is a new columnar file format. It includes implementations of several of the nouveau numeric compressors, including FastLanes, ALP, and pcodec. It focuses on formats that support random-access to individual elements. There's also eager and active work on on-GPU decompression.

I have some preliminary numbers showing (unsurprisingly) that if you read slices that are misaligned with and smaller than your chunks, Vortex is substantially faster than block compression (I was looking at zstd specifically but this should apply to blosc, etc.).

---

I'm not entirely sure where I want to go with this. To some extent, Vortex (and Parquet, and Lance, ...) doesn't quite fit the Zarr model because Zarr wants to own the file format. I can imagine an alternative Zarr which, as far as storage is concerned, only (1) splits an array into chunks and (2) applies chunk-to-chunk transformations (e.g. transpose). Writing arrays from memory to storage is instead the responsibility of a pluggable file format whose API is something very simple like write(array, url) and read(url, slice) -> array. This moves Zarr closer to the tensorial generalization of a table format (e.g. Apache Iceberg, DeltaLake).

---

This is a bit of a tangent, but I'm curious if you ran into the same issue implementing Zarrs. Zarr's partial decode API only provides a ByteGetter:

class ByteGetter(Protocol):
    async def get(
        self, prototype: BufferPrototype, byte_range: ByteRequest | None = None
    ) -> Buffer | None: ...

Vortex needs to know the size of a file in order to read it. For my demo repo, I just assume ByteGetter is actually a StorePath and use the store to get the size.

---

[1] disclosure: I'm paid to work on Vortex.

[danking]

Relatedly, @LDeakin , do you all assume that Python won't mutate the ndarrays passed to you? There's some discussion threads in pyo3 about this. I currently copy the array given to me by Python which tanks throughput.

[ilan-gold]

@danking zarr-developers/zarr-python#3613 Might also be of interest to you i.e., what to do on the python side. Some discussion about the nd-dness of the buffer (i.e., what to do about the 1d-dness of columnar formats).