Skip to content

Subset cells and genes while unpacking IVCSR/IVCSC, without a format change #67

Open

Description

@aarmey

Follow-up to #24. Rather than a new on-disk format with a byte-offset index, keep the current IVCSR format and subset cells at the point of unpacking. This needs no format change and works on every existing file.

Background (measured on lupus.ivcsr.h5ad, 1.26M cells, 1.06 GB packed_indices, warm cache)

Without a byte index, a reader has to read every byte of packed_indices. But it only has to decode, and allocate output for, the rows it keeps.

A prototype that finds the start of each selected row with a parallel pass over the varints (reusing _ivcsc._realign_forward / _count_all), then decodes only those rows, matched a full decode followed by a subset exactly. Times:

full X decode 0.1% of cells 1% 10%
Blosc2 file (old default) 1.97 s 1.47 s 1.49 s 1.54 s
Same format, packed_indices stored unfiltered 0.84 s 0.20 s 0.22 s 0.28 s

Each part of the subset path:

Step Time
Read the whole stream (Blosc2, serial h5py filter) ~1.23 s
Find the start of each selected row ~0.10 s
Read major_ptr/values/value_ptr ~0.06 s
Decode the selected rows ≤ 0.04 s

So reading and decompressing the stream is the bottleneck, not decoding. #68 (zstd-1 + parallel chunk decode) addresses that: it reads the full lupus stream in ~0.08 s on 16 threads, against 1.50 s for Blosc2 through h5py, and the stream is 31% smaller (0.73 GB). With it, a 1% read should take about 0.25 s, which still needs measuring.

Plan

  1. Selective decode kernel (_ivcsc.py):
    • Given sorted target rows, find each row's starting byte. Each parallel chunk does one pass, counting varint terminator bytes (& 0x80 == 0) against the varint rank value_ptr[major_ptr[r]]. This generalizes _group_chunk_boundaries.
    • Decode only those rows, in parallel, into a compact VCSRArray.
    • Rebuild major_ptr/value_ptr for the subset from major_ptr/value_ptr (no decode needed).
  2. Gene subsetting in the same pass: drop and renumber minor-axis indices through a lookup table during decode (as _rapid_load._fill_kept does), so selecting both cells and genes still decodes once.
  3. Loader API: load_packed(path, obs_idx=None, var_idx=None).
    • Read obs/var first, so a selection can be a boolean mask, labels, integers, or a callable on obs, as with load_and_normalize's obs_filter.
    • Slice the metadata to match.
    • obs_idx=None keeps today's full decode.
  4. Bounded memory: decompress the stream in slabs of zstd chunks with the parallel reader (_compression.read_array), rather than holding all of it. Carry the row/rank state across slab boundaries, so peak memory is one slab plus the selected rows' output. This matters at the scale of Parse 10M (22 GB stream; a full decode allocates ~44 GB of indices).
  5. IVCSC (gene-major) files: here a cell selection is on the minor axis, so it works as a filter during a full decode (step 2), and gene selection is the fast path. All current files under /opt/IVCSC are IVCSR.
  6. Hook for Design a composable lazy Selection object for VCSCAnnData #25: this is where a lazy Selection materializes, in one pass over the file. Since cell totals can be computed from values/value_ptr alone, metadata-only filters (Create a statistics sidecar as an on-disk format feature #26) never touch packed_indices.

Not in scope

  • A stored per-row byte index (byte_ptr). It would remove the locate scan and make random reads proportional to the selection, and it can be added later as an extra dataset that current readers ignore. A prototype read 1% of cells in 0.08 s, touching 1.00% of the bytes. Worth revisiting if the scan becomes the bottleneck at 10M+ cells.
  • Changing the IVCSC/IVCSR encoding.

Acceptance criteria

  • load_packed(path, obs_idx=...) returns exactly load_packed(path)[obs_idx, :] (and likewise with var_idx) for IVCSR and IVCSC, including empty rows, multi-byte varints, and unsorted or duplicate indices (or a documented error for those).
  • Peak memory for a 1% selection is bounded by the slab size plus the output, not by the full decoded matrix.
  • A benchmark in benchmarks/ covers 0.1/1/10% random selections and a sample-clustered selection on a zstd-compressed file.
  • load_and_normalize's obs_filter path (or its replacement after Remove load_and_normalize; load, subset, then normalize #66) uses the same kernel instead of _decode_selected_rows.

References

  • src/vsparse/_ivcsc.py (_realign_forward, _count_all, _group_chunk_boundaries)
  • src/vsparse/_rapid_load.py (_decode_selected_rows, _fill_kept)
  • src/vsparse/_compression.py (read_array, parallel zstd chunk decode)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions