feat: serve sorted integer-array selections through zarrs - #189
Draft
selmanozleyen wants to merge 4 commits into
Draft
feat: serve sorted integer-array selections through zarrs#189selmanozleyen wants to merge 4 commits into
selmanozleyen wants to merge 4 commits into
Conversation
selmanozleyen
force-pushed
the
feat/integer-array-indexing
branch
9 times, most recently
from
August 28, 2026 11:44
1e054a4 to
bf4e570
Compare
A dependency bump and the API migration it forces, and nothing else. Kept on its own because a dependency migration reviewed together with a behaviour change means neither can be reverted alone. All of it is one upstream idea -- the #411 refactor -- which is that a codec chain is BOUND to a data type and fill value once, where it is built, instead of being handed both at every call. `CodecChain` is now unbound; `decode`, `encode`, `decode_into`, `partial_decoder` and `recommended_concurrency` live on `CodecChainBound`, reached by `with_context`. So the chain is bound in the constructor and five call sites stop passing a pair they no longer need. `zarrs::array::StoragePartialDecoder` is also gone, and nothing named replaced it: the (storage, key) TUPLE is the store-backed BytesPartialDecoderTraits implementation now. No behaviour change, with one exception worth naming rather than hiding: a codec chain that cannot BIND to its data type now fails in the constructor as a TypeError, where before it surfaced at the first read as a RuntimeError. Binding has to happen somewhere, and the constructor is the only place it can. The codec metadata is still PARSED where it was, so an array with both bad codecs and a bad fill value still reports the codecs. The dependency is a git rev rather than a version, because 0.24 is not released yet, and it carries a [patch.crates-io] block for zarrs_storage. That patch is load-bearing, not cosmetic: zarrs_opendal and zarrs_object_store track RELEASED zarrs, so they pull zarrs_storage from crates.io while zarrs comes from git -- two copies of one crate, two distinct AsyncReadableStorageTraits, and trait bounds that cannot be satisfied. Patching that single crate collapses the graph; every other zarrs crate already resolves through the git checkout, verified from a clean lockfile. Both the pin and the patch are commented with exactly what to do on release day, and nothing in src/ changes then.
selmanozleyen
force-pushed
the
feat/integer-array-indexing
branch
from
August 28, 2026 13:25
bf4e570 to
a52c2aa
Compare
zarr-python enforces read-only in the store itself -- `Store._check_writable`, reached from the concrete store's `_set`. This pipeline never gets there: it is handed a `StoreConfig` and builds its own Rust store, writable whatever mode the array was opened in, so a write to a `mode="r"` array succeeds here and raises through the default pipeline. Silent divergence, and silent data loss for anyone relying on the mode. The flag has to be carried from Python because nothing on the Rust side can infer it: the mode is zarr-python's, not the store's. One guard, at the only entry point Python can write through. `store_chunk_bytes` and `store_chunk_subset_bytes` are reached only from it. The tests open STRICT. zarr's own refusal message is byte-identical to the guard's, so with a fallback available they would pass whether the guard fired or zarr-python served the write.
…t all A sorted integer-array selection is mostly runs of consecutive indices, and a run is a slice. Split into runs, each one is a contiguous box the existing read path already knows how to serve. Without this a fancy index that is not contiguous is declined outright and the whole read falls back to zarr-python's pipeline -- so the shape a row-sampling data loader produces never reached zarrs, however well zarrs could have served it. Reads and writes stop sharing one description function, because they now want different ones. A read is split into runs; a WRITE must not be. Splitting a write can put two items on one chunk key, and the write path is read-modify-write, so two items on one key race. Indices are normalised to int64 first, which two dtypes needed for different reasons. A uint8 selection like [255, 0] differenced to 1, read as consecutive, and returned an empty slice -- wrong data, with no error. And a BOOLEAN MASK is not an index array at all: cast blindly it becomes [1, 1, ...], which is non-decreasing and exactly as long as the output slice, so it passed every test here and read element 1 once per True. A mask's POSITIONS are what it means, so that is what it is turned into -- which also lets a mask take this path rather than falling back.
selmanozleyen
force-pushed
the
feat/integer-array-indexing
branch
from
August 28, 2026 15:06
a52c2aa to
ee7d47e
Compare
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,
This PR is currently blocked by #190 and #191
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