`read_elem_as_dask`

[ivirshup]

Please describe your wishes and possible alternatives to achieve the desired result.

We should probably add some read_elem_as_dask methods to make reading stored objects from as dask arrays easier.

This largely needs special casing due to:

• Sparse arrays
• Dask not handling h5py object in the process based mode: Distributed writing for H5ad format due to h5py objects being unserializable  #1105

I think we can skip dataframe support for now. It may be worth thinking about how to read objects into GPU memory as well.

I'm not completley set on the API, so input welcome. My initial thought is that we could just have a smaller registry of methods which only supports reads for dense and sparse arrays.

cc: @ilan-gold

[ivirshup]

Starting to remember why I was initially against this kind of thing.

We probably also want:

• read in as cupy
• read in as cupy wrapped by dask

So this becomes three functions quite quickly. But also, I think there's a ton of value in expositing this, so maybe we can start here and transition to a better api in the future.

[ilan-gold]

@ivirshup I agree, I think that we kind of have to start somewhere. I think the fact that we need three functions won't change internally, so exposing it via experimental and giving us a chance to iron out kinks could be good. I understand the maintenance burden, but since they will be experimental, we can always yank them from the public API.

[vitkl]

Looking forward to these changes!

Would this support arbitrary chunk sizes of sparse matrixes with zarr+dask?

I am looking for a solution that would allow loading [random 50 cells, random block of 400 variables] from 50 layers quickly as part of DataLoader (scvi-tools) without entirely loading all 50 sparse matrixes into memory. The anndata object I have has few obs (2k) but lots of vars (500k+), 50 layers and ~20 of other obs/obsm/var/varm loaded as part of minibatch for the same sample of cells and variables. Everything except layers and adata.X is fine to load into memory in full.

[vitkl]

In particular, it would be nice to be able to do the following:

def read_w_sparse_dask_layers(
    group: h5py.Group, obs_chunk: int = 1000, axis: int = 0
) -> ad.AnnData:
    return ad.AnnData(
        X=sparse_dataset_as_dask(sparse_dataset(group["X"]), obs_chunk, axis=axis),
        **{
            k: {
                key: sparse_dataset_as_dask(sparse_dataset(dat), obs_chunk, axis=axis)
                for key, dat in group[k].items()
            }
            if k in group
            else {}
            for k in ["layers"]
        },
        **{
            k: read_elem(group[k]) if k in group else {}
            for k in ["obs", "var", "uns", "obsp", "varp", "obsm", "varm"]
        }
    )

using the code @ivirshup shared with me during scverse hackathon:

def csr_callable(shape: tuple[int, int], dtype) -> sparse.csr_matrix:
    if len(shape) == 0:
        shape = (0, 0)
    if len(shape) == 1:
        shape = (shape[0], 0)
    elif len(shape) == 2:
        pass
    else:
        raise ValueError(shape)

    return sparse.csr_matrix(shape, dtype=dtype)


class CSRCallable:
    """Dummy class to bypass dask checks"""

    def __new__(cls, shape, dtype):
        return csr_callable(shape, dtype)


def make_dask_chunk(x, start: int, end: int) -> da.Array:  #: "SparseDataset",
    def take_slice(x, idx):
        return x[idx]

    return da.from_delayed(
        delayed(take_slice)(x, slice(start, end)),
        dtype=x.dtype,
        shape=(end - start, x.shape[1]),
        meta=CSRCallable,
    )


def make_dask_var_chunk(x, start: int, end: int) -> da.Array:  #: "SparseDataset",
    def take_slice(x, idx):
        return x[:, idx]

    return da.from_delayed(
        delayed(take_slice)(x, slice(start, end)),
        dtype=x.dtype,
        shape=(x.shape[0], end - start),
        meta=CSRCallable,
    )


def sparse_dataset_as_dask(x, stride: int, axis: int = 0):
    n_chunks, rem = divmod(x.shape[axis], stride)

    if axis == 0:
        make_dask_chunk_func = make_dask_chunk
    elif axis == 1:
        make_dask_chunk_func = make_dask_var_chunk
    else:
        raise ValueError("axis must be 0 or 1")

    chunks = []
    cur_pos = 0
    for i in range(n_chunks):
        chunks.append(make_dask_chunk_func(x, cur_pos, cur_pos + stride))
        cur_pos += stride
    if rem:
        chunks.append(make_dask_chunk_func(x, cur_pos, x.shape[axis]))

    return da.concatenate(chunks, axis=axis)

This code is ~1000 times slower at creating minibatches than reading from RAM (2 min instead of a few ms per batch).

[ilan-gold]

This code is ~1000 times slower at creating minibatches than reading from RAM (2 min instead of a few ms per batch).

This makes some sense because RAM is, well, RAM. Even so, zarr + dask make not be the solution. HDF5 is usually memory mapped, which is the secret sauce to making random access fast. Random access with dask is not amazing.

Would this support arbitrary chunk sizes of sparse matrixes with zarr+dask?

Yes it should, but chunking the zarr store on disk is also important!

[vitkl]

This makes some sense because RAM is, well, RAM.

Still, 1000x slower speed makes this approach completely unusable and I was told to expect only ~3-4x slowdown (even 10x would be still usable).

Even so, zarr + dask make not be the solution. HDF5 is usually memory mapped, which is the secret sauce to making random access fast. Random access with dask is not amazing.

Does this mean that you recommend using normal h5ad instead of zarr+dask?
I see that if I simply use anndata.read_h5ad(..., backed='r') all adata.layers slots are loaded into memory. Maybe I am using something incorrectly?

but chunking the zarr store on disk is also important!

How are the sparse matrixes in layers currently chunked when using the code above (both zarr and h5ad)?
My understanding is that sparse matrixes are not affected by the chunking described in this tutorial https://anndata.readthedocs.io/en/latest/tutorials/notebooks/%7Bread%2Cwrite%7D_dispatched.html and I see that they are read in full when I use the code from the tutorial.

Given that I am using batches of vars (eg 0:200, then 15200:15400, then 1200:1400, ...) for a relatively small number of obs (random 50/1600), would it make sense to save the transposed object and then transpose back after loading?