(performance): speedup for zarr-based sparse indexing

[ilan-gold]

import scipy.sparse as sp
import numpy as np
import anndata as ad
import zarr
from pathlib import Path

data_path = Path('data/foo.zarr')
if not data_path.exists():
    z = zarr.open_group('data/foo.zarr', mode='w')
    arr = sp.random(1_000_000, 10_000, format="csr", random_state=np.random.default_rng())
    ad.io.write_elem(z, "X", arr)
z = zarr.open_group('data/foo.zarr', mode='r')

adata = ad.AnnData(X=ad.io.sparse_dataset(z["X"]))
idx = np.random.randint(0, adata.shape[0], 1024)
idx.sort()
%time adata[idx].X

On my computer, this is a 3-5X speedup.

• Closes #
• Tests added
• Release note added (or unnecessary)

[scverse-benchmark]

Benchmark changes

Change	Before [d055963]	After [e45a05f]	Ratio	Benchmark (Parameter)
-	1.8306233062330624	1.193271083444494	0.65	readwrite.H5ADReadSuite.track_read_full_memratio('pbmc3k')
-	511±10ms	16.8±0.7ms	0.03	sparse_dataset.SparseCSRContiguousSlice.time_getitem((10000, 10000), ':9000:-1')
-	2.47±0.01s	78.1±3ms	0.03	sparse_dataset.SparseCSRContiguousSlice.time_getitem((10000, 10000), '::-2')
-	521±3ms	48.5±4ms	0.09	sparse_dataset.SparseCSRContiguousSlice.time_getitem((10000, 10000), 'alternating')
-	507±10ms	17.7±0.9ms	0.03	sparse_dataset.SparseCSRContiguousSlice.time_getitem((10000, 10000), 'arange')

Comparison: https://github.com/scverse/anndata/compare/d055963dc70915ad921965a03d2b7342a098dd6b..e45a05f785aa361068c06b92a3070190fd5ed25f
Last changed: Tue, 10 Dec 2024 14:35:46 +0000

More details: https://github.com/scverse/anndata/pull/1790/checks?check_run_id=34197556781

[ilan-gold]

Very good :)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 84.57%. Comparing base (d055963) to head (e45a05f).
Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1790      +/-   ##
==========================================
- Coverage   87.04%   84.57%   -2.47%     
==========================================
  Files          40       40              
  Lines        6089     6095       +6     
==========================================
- Hits         5300     5155     -145     
- Misses        789      940     +151     

Files with missing lines	Coverage Δ
src/anndata/_core/sparse_dataset.py	93.63% <100.00%> (+0.12%)	⬆️

... and 8 files with indirect coverage changes

[flying-sheep]

great! I think for clarity, wherever possible you should unpack an entry of indptr_list into start, end instead of using s[0] and s[1].

you could either do indptr_limits = [(x.indptr[s.start], x.indptr[s.stop + 1]) for s in slices] or so, or just define indptr_slices as above.