Migrating to zarr v3

[JulienBrn]

Currently spike interface uses zarr v2 which does not have a compatible API with zarr v3 which has different features (mainly sharding).
My motivation behind sharding is that I like being able to load a single channel pretty fast, so I usually put channel_chunk_size fairly low, but by doing so, I get a huge number of files (384x2x3600 for a 2h recording...).

I've tried to migrate the code of zarrextractors.py and here is my conclusion:

• migrating to v3 is mostly easy by changing group.create_dataset(key=name,..., data=data) by group.create_array(key=name, ..., shape=data.shape, dtype=data.dtype) followed by group[name][:] = data
• we have a problem for structured numpy arrays which were handled in v2 and not in v3. These seem in the process of being handled by zarr, see this pr. Either we wish to rely on this future implementation or we can have a custom convention for handling them. A solution could be to create a group with a name like "_structuredarraygrp[name]" and put each individual array in it.
• However, for reading the structured arrays, the code is less obvious: either we remove the lazyness of zarr and we just create the numpy array from the group each time, or we need to find a way to provide a lazyarray that handles structured dtypes. Perhaps by using dask ?
• adding sharding is extremely easy as its just a parameter of create_array. However, one needs to modify the functions that process arguments.

Is there any interest in such a migration (enough for me to submit a pr) ?

[alejoe91]

Hi @JulienBrn

Thank you for this! Definitely there is interest in this. For the structured numpy array, I guess it's only used by the Sorting object, so we could handle it internally as we do now. Generally, the sorting objects can be easily loaded in memory so it shouldn't be a problem. Ok to add sharding args to the write functions :)

Please go ahead and make I PR! We'd be happy to review it and collaborate on it

[JulienBrn]

I've successfully made zarr v3 work for one of my recording with sharding.

Core code

The core code is the following:

def write_nparray_to_zarr(group: zarr.Group, key: str, data: np.array, compressors=None, filters=None):    
    if data.dtype.fields:
        subgrp = group.create_group(key)
        subgrp.attrs["is_structured_array"] = True
        for nd, d in data.dtype.fields.items():
            subgrp.create_array(name=nd, shape=data.shape, dtype=d[0], compressors=compressors, filters=filters)
            subgrp[nd][:]=data[nd]
    else:
        group.create_array(name=key, shape=data.shape, dtype=data.dtype, compressors=compressors, filters=filters)
        group[key][:]=data

def load_array(grp: zarr.Group | zarr.Array):
    if "is_structured_array" in grp.attrs and grp.attrs["is_structured_array"]:
        l=[]
        shape = None
        for n, a in  grp.arrays():
            if not shape:
                shape = a.shape
            elif shape != a.shape:
                raise Exception("All arrays of structured array should have same shape")
            if a.dtype ==  np.dtypes.StringDType():
                l.append((n, object))
            else:
                l.append((n, a.dtype))
        res = np.empty(shape, dtype=[(t[0], t[1]) for t in l])
        for n, a in  grp.arrays():
            if a.dtype ==  np.dtypes.StringDType():
                res[n] = np.array([str(i) for i in np.array(a).reshape(-1)], dtype=str).reshape(shape)
            else:
                res[n] = a
        return res
    elif isinstance(grp, zarr.Array):
        if grp.dtype == np.dtypes.StringDType(): #Seems like the rest of spikeinterface does not like this dtype
            return np.array(grp).astype(object).astype(str)
        return grp
    else:
        raise Exception("Problem expected a zarr group representing a numpy structured array.")

Warnings to handle

I have several warnings due to zarr 3 specification issues that I do not exactly know how to address. Currently, I simply added the following:

warnings.filterwarnings(action="ignore", message=".*the Zarr format 3 specification.*", append=True)
warnings.filterwarnings(action="ignore", message=".*the Zarr version 3 specification.*", append=True)

The actual warnings are:

UserWarning: The codec `vlen-utf8` is currently not part in the Zarr format 3 specification. It may not be supported by other zarr implementations and may change in the future.

UserWarning: Numcodecs codecs are not in the Zarr version 3 specification and may not be supported by other zarr implementations.

My understanding of these warnings:

• the first warning is because zarr 3.0.8 handles np.StringDType() but maybe not in a zarr3 compliant way. This should be handled in zarr 3.1
• the second is because I use from numcodecs.zarr3 import Delta and this is maybe python specific.

As one of the reasons for using the zarr format is long term (hopefully 15years+) preservation of data with compatibility with other languages and possibly tools in the future (the format is rather recent). Basically, we would like to use it as a modern version of hdf5 with additional features, so these warnings are concerning. Or maybe we are wrong to use the zarr format for this purpose ?

SpikeInterface Compatibility with old zarr files

How do we wish to handle zarr files in spikeinterface that were generated with zarr v2 ?

Conclusion

Before a submit a PR, I believe these questions should be addressed.

[samuelgarcia]

Hi @JulienBrn
thanks you for working on this.
About the structure dtype, the solution you propose is consistent in the zarr itself because all fileds are separated.
But it has a cost when writing and reading : double the memory and big reorganizing.
Another simple way way would be to store the structure dtype somewhere in the zarr attr and then save
arr.view('u1') and when we read back we can do arr_u1.view(original_dtype).

@alejoe91
A propose a call to discuss this.

[d-v-b]

FYI zarr python will will support structured dtypes for zarr v3 in version 3.1.0, which should be released this week

[h-mayorquin]

This will crash with nwb dependency which has not move forward because of the structured dtypes situation and other small problems.

[JulienBrn]

Hi @JulienBrn thanks you for working on this. About the structure dtype, the solution you propose is consistent in the zarr itself because all fileds are separated. But it has a cost when writing and reading : double the memory and big reorganizing. Another simple way way would be to store the structure dtype somewhere in the zarr attr and then save arr.view('u1') and when we read back we can do arr_u1.view(original_dtype).

My understanding is that the structured dtypes are only for small arrays (for example contact vectors). So we are indeed losing a bit of performance, but I would not expect this to be a problem. But encoding as u1 could be another way to handle this.

FYI zarr python will will support structured dtypes for zarr v3 in version 3.1.0, which should be released this week

If 3.1.0 is indeed for this week, this would be extremely nice. Though is the way structured dtypes are going to be implemented in zarr python specific to the python backend or are structured dtypes supported in all zarr backends ? If the latter, doesn't this mean possible issues with reading those files from other programming languages ? If so, are we at least expecting structured dtypes to be in the zarr specification in the future ? If not, isn't using a spikeinterface solution more reliable in the long run (the zarr file would match the zarr specification and spikeinterface is the one interpreting it differently).
I'm also wondering (seeing the warnings) if this is also an issue for the np.StringDType and the Delta codec from numcodecs.

This will crash with nwb dependency which has not move forward because of the structured dtypes situation and other small problems.

Maybe we could implement a zarr2 backend and a zarr3 backend solution for compatibility with other tools such as nwb, leaving for now "zarr" export to "zarr2" backend ? xarray seems to have that features by specifying zarr_format =2 .

But perhaps we are going in the wrong direction and should just leverage that all this has probably already been solved (or will be) in xarray, and, instead of duplicating solutions, we should just write the to_xarray and from_xarray functions (leveraging dask arrays for laziness). By doing so, users would be free to go through xarray for other backends such as zarr. Then spikeinterface would only need to support its binary format and some neuroscience specific formats (perhaps relying on neo). This would perhaps add other benefits than just storing as xarray has many other features.
And maybe I'm wrong about the difficulty of doing this. Note: this relates to issue #3929

[JulienBrn]

zarr version 3.1.0 has been released. I still have to check if it handles the AoS arrays of spikeinterface, but I believe it should. If so, the code to support zarr 3 should be fairly straightforward.

However, I do not know how we would handle backward compatibility (allowing zarr 2 files to be loaded) as I believe this would require both zarr3 and zarr2 to be installed in the same environment.

Furthermore, I still do not fully grasp what in python zarr 3.1 is specific to the python backend and what is in the zarr specifications.

[d-v-b]

zarr python 3.x supports zarr v2 and zarr v3 formats