feat: serve sorted integer-array selections through zarrs - #189
Draft
selmanozleyen wants to merge 5 commits into
Draft
feat: serve sorted integer-array selections through zarrs#189selmanozleyen wants to merge 5 commits into
selmanozleyen wants to merge 5 commits into
Conversation
A chunk read is described to zarrs as a rectangular subset, so an integer-array selection is refused as discontiguous today and served by zarr-python's pipeline instead. `split_selection_runs` cuts such a selection at its non-consecutive boundaries and emits one (slices, slices) box per run, which is the form the existing slice path already carries into `retrieve_chunks_and_apply_index` -- so this adds no second read path, it makes a refused selection expressible in the one that exists. A pure-slice selection yields a single box, leaving it untouched. Behind `codec_pipeline.integer_array_indexing`, off by default: box count is data-dependent, and a selection whose runs are shorter than an inner chunk decodes that chunk once per run where zarr-python decodes it once in total. Reads only, one integer-array axis only, and non-decreasing indices only -- repeats end a run and are served, while a permutation is refused, since output order and selection order stop agreeing.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Hi, @ilan-gold @LDeakin,
I'd like to expose zarrs-python for the case when the integer indices are sorted instead of falling back. This is possible thanks to zarr-developers/zarr-python#4172. This is actually a realistic and a possibly very useful case for anndata and annbatch! It can't be used right away but it doesn't require much changes to utilize it, I linked the branches under the details.
The setup: 14 sharded stores (Tahoe-100M, CSR
X, 100.6M rows,146e9 nnz, ~89.4M-element shards, 64-row inner chunks).
The loader: Per batch: 1024 random rows
drawn across all 14 stores, (
chunk_size=1, batch_size=1024, preload_nchunks=8192in annbatch), 100 batches = 102,400 rows.
Results
integer_array_indexing=TrueCurrent status: if many runs hit same chunk it regresses due to deduplication of reads. But it might also regress in any local system atm. I only ran these numbers on lustre
Details
After this I will give a shot at #182, and in my AI draft at least it shows some stuff can be improved.
integer_array_indexingonly (branch A)shard_index_cache_size=512plan_readsplan_reads_fetch_threads=32Refs measured:
feat/integer-array-indexingand @1d284dbfeat/planned-reads-onlyfeat/async-csr-row-readfeat/anndata-async-csr-fetch