Anndata and zarr

[Amitbergman]

Question

Hey all,

I have a few anndata datasets with sparse csr X matrices (each is with ~10M cells and 40K genes, with parity of about 5%).

I want to be able to quickly load whole rows from these datasets (say given a query, load all rows based on a condition on the obs table).
Currently I am taking the anndata object and converting it to tileDB, but I recently encountered the zarr file format, and specifically the support of zarr v3 in anndata.

I have a few questions regarding zarr:

1. Is Zarr v3 would be a good fit for our use case? Should I expect improvement over tileDB?
2. Are there some guidelines on what codec to use? Chunk sizes?
3. Are there some guidelines as to how to benefit from concurrency? I see dask being used in many places together with zarr.

Thanks!

[ilan-gold]

Is Zarr v3 would be a good fit for our use case? Should I expect improvement over tileDB?

Yes, it's a great fit, especially with a bit of work (I would love to review PRs and get help). I would definitely have a look at our zarr v3 guide and our "lazy loading" notebook, to learn more on what is going on when you save as zarr and how to access the data, at least as a first step. Certainly trying out zarr should be extremely easy.

Whether or not you see an improvement over tileDB is a different question, and not one I'm really equipped to answer now, nor one I want to answer quite yet. Instead I'll try to highlight the current limitations so if you wish to benchmark or compare, you'll be aware what possible improvements there are. These limitations are exactly what I want to work out before seeing what happens in comparison, although you're free (and welcome) to try zarr vs tileDB out for yourself right now. I'll split up things into the biggest performance factors I see right now, should you try things out:

Asynchronous access

At the moment our data access pattern in anndata.io.sparse_dataset is synchronous and blocking (up to the fact that zarr is concurrent/parallel under-the-hood i.e., in ever zarr_array[my_selection] call). However, when wrapped in dask via read_elem_lazy, it's possible (even likely) you get performance improvements, although it is tough for me to 100% say.

There is definitely a route out of this state of things beginning in #1927. In short

1. Remove the dependency on scipy.sparse in feat: allow gpu io in sparse_dataset by removing scipy inheritance #1927 which frees up our API and generally removes a bad practice of inheritance/private methods
2. Make asynchronous versions of the data fetching functions in feat: allow gpu io in sparse_dataset by removing scipy inheritance #1927, wrapping them in zarr.core.syncto get back the data in synchronous callers
3. Make custom-indexing via Custom Indexing API for Pipeline-level Parallelism zarr-developers/zarr-python#3175 a real possibility so we can maximize throughput and have a clean API with sparse arrays
4. (feat): add async versions of {write,read}_{elem,dispatched} #1902 would allow us to begin thinking about asynchronous access for the whole AnnData object as well as exposing asynchronous APIs we may not wish to expose in the first step

I don't assess #1927 + asynchronous methods to be particularly difficult TBH although point 3 will require discussion with the zarr devs. Point 4 is a bit challenging but certainly doable. I expect these things, even one at a time to greatly improve throughput. If you're wondering why I don't do this right now (a) I just finished up proving that this sequence of steps is actually worth it (more to come on that soon) and (b) I don't work with remote data i.e., like tileDB, so this sort of thing isn't a priority for me at the moment, although I find it very exciting.

On Disk Format

Are there some guidelines on what codec to use? Chunk sizes?

I generally recommend blosc because it used to be the default in zarr v2 and a paper that looked into this for a different use-case found it to be the best. Chunk sizes for sparse data only have meaning if you are referring to indices and data in which case I would not simply write the ~number of elements in a row to those elements as the chunk size for two reasons:

1. the number of elements in a sparse row varies and the chunk sizes are fixed, and
2. even if there were variable chunk sizes, you would be limited by the page size of your computer to see any reasonable performance justification for such a small size. Consider writing 5000 data points to data of 4 byte data type (float 64, int 64 etc.). That's 20k bytes at a 5x compression ratio gives you about a page of data (~4KB). If you write just over that amount (say 4KB + 1 byte of data), you won't benefit from the small chunk size because your computer has to fetch in page sizes of 4KB and so you'll fetch almost twice the data. If you write less data to a chunk, you may as well have written more because your computer will fetch 4KB of data regardless. Also, writing would be pretty slow at such a small size. So I would recommend somewhere in the ball park 30K data points (this is backed up by experience, more to come soon hopefully).

I would also highly recommend sharded data - see the guide for more info there and a quickstart function.

Are there some guidelines as to how to benefit from concurrency? I see dask being used in many places together with zarr.

Dask is very useful for computation i.e., downstream analysis in scanpy. But if all you are doing is fetching data over a remote connection, I am not 100% sure. I imagine in the absence of some of the above 4 points, dask might be in general faster, but I imagine that to drastically change when the above 4 points are completed. Definitely try out read_lazy and try swapping out X with anndata.io.sparse_dataset if you wish to compare.

Final thoughts

Even if we are not "there" yet in terms of remote data or local throughput, the pieces are laid out and will be done within the next 6 months for high performance, I imagine. Having help would be immense, especially with #1927 and #1902. We already offer APIs for working with "remote" data as shown in this notebook, which might be a good starting point if you're looking to compare. The first step was making it possible via the read_lazy API + anndata.io.sparse_dataset to read data remotely, but now we can really make things sing. I really want someone with a lot of data to take the plunge because I feel very strongly we have a very strong thing going here, and it's only going to get stronger.

[flying-sheep]

Some good writing there @ilan-gold, you should probably move that somewhere where more people will see it!