feat: add keep_sparse parameter for aggregate and optimize it

[mumichae]

Adjusted the dask call to use dask.delayed, which has massively improved the memory footprint on my end. I also threw in some little adjustments for making use of sparsity and returning sparse matrices, since even aggregated data can be sparse.

• Closes sc.get.aggregate memory leak for Dask array #4074
• Tests included or not required because:

• Release notes not necessary because:

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.01%. Comparing base (7eeba2f) to head (e92b11c).
⚠️ Report is 13 commits behind head on main.
✅ All tests successful. No failed tests found.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4073      +/-   ##
==========================================
+ Coverage   78.99%   79.01%   +0.01%     
==========================================
  Files         128      128              
  Lines       13402    13412      +10     
==========================================
+ Hits        10587    10597      +10     
  Misses       2815     2815              

Flag	Coverage Δ
hatch-test.low-vers	77.27% <85.71%> (+0.02%)	⬆️
hatch-test.pre	78.85% <100.00%> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/scanpy/get/_aggregated.py	95.42% <100.00%> (+0.16%)	⬆️
src/scanpy/tools/_rank_genes_groups.py	93.59% <100.00%> (ø)

[ilan-gold]

Adjusted the dask call to use dask.delayed, which has massively improved the memory footprint on my end. I also threw in some little adjustments for making use of sparsity and returning sparse matrices, since even aggregated data can be sparse.

Are you sure the dask.delayed is responsible for the reduced memory footprint? It seems like the later would be more responsible for that. This would be quite surprising to me if dask.delayed helped with this. Do you have a sense for why?

[ilan-gold]

Interesting, I would have assumed aggregations of sum would produce something for all features - you have features with literally 0 value in a feature across some categories? Maybe with pseudobulk, I could see that.

Maybe we should try out a kernel that creates a coo-like data structure instead as new category-feature combinations are seen which is then csr-ed instead of doing this sparsification. Or category-by-category the sparse-ification happens instead (i.e., allocate a buffer up front per-category, and then sparse-ify in numba). Or a mix between the two - allocate enough memory for a max-sized coo-matrix, fill it up as far as you need, then drop the unused allocation. At the end, concatenate all the results and call tocsr before returning.

I would think sparsifying to be a relatively expensive operation that doesn't do much for a lot of situations (where you don't have such small categories so that feature is non-zero when summed).

[ilan-gold]

Ok I see why we need delayed now - there is 3D CSR matrix so our old trick of adding a dimension i.e, return res[None, :] if unchunked_axis == 1 else res and then summing over hte first dimension doesn't work. It's possible we can work around that. Let's see, otherwise this PR makes a lot of sense. Thanks @mumichae !

[flying-sheep]

@ilan-gold you mean “there is no 3D CSR matrix”, right?

@mumichae needs a release note mentioning the new parameter and optimization

[ilan-gold]

@ilan-gold you mean “there is no 3D CSR matrix”, right?

Yes @flying-sheep but potentially she could add a sparse extra for 3d sparse matrices: https://github.com/dask/dask/blob/ba5045e001d67decd6a83a0f50f666a5e3845250/dask/array/chunk_types.py#L123-L126C25 which are supported by dask: https://sparse.pydata.org/en/latest/examples/dask_example/#create-arrays