Pre-processing only supports .obs and .X attributes

[felix0097]

Report

Currently, our pre-processing code only supports the .obs and .X attributes. If an input anndata object contains other attributes e.g. .obsm or .raw, the preprocessing code becomes very slow.

https://github.com/laminlabs/arrayloaders/blob/e42d36770582adb99767b3e1f5d56f7cdc6fe4db/src/arrayloaders/io.py#L187-L191

We should update the code above or raise an error in this case.

As far as I remember, the code become significantly slower when shuffling the shard via adata_chunk = adata_chunk[idxs, :]. That's why implemented the code in the current way. Perhaps this problem got solved?

Versions

[felix0097]

I think doing adata_chunk[idxs, :].copy() before saving the data solved the issue as well.

[ilan-gold]

I think doing adata_chunk[idxs, :].copy() before saving the data solved the issue as well.

Yea, I think this is a bug. The obsm for example would be wrong because it wasn't shuffled.

[ilan-gold]

If an input anndata object contains other attributes e.g. .obsm or .raw, the preprocessing code becomes very slow.

We should put the option not to carry these over behind some sort of flag then - the previous state of things of lopping these off felt wrong.

[felix0097]

I've looked into this a bit. The .copy() work for .X. However, .obsm is still horribly slow. It runs for hours for like 200k cells. Something is off here. Any thoughts on this @ilan-gold?

[ilan-gold]

You could try going through the obsm and persisting the keys as well in memory if the underlying array is a dask.array.core.Array i.e., adata.obsm["dask"] = adata.obsm["dask"].persist(). Tough to say without profiling though. Things were similarly slow for X and the solution there was persist, no?