import_anndata_generic: add index_dtype / data_dtype kwargs to avoid mandatory int64 allocations

[sebastianpineda]

Summary

actionet.import_anndata_generic currently hard-codes scipy.io.mmread + .tocsc() for MatrixMarket ingest, which forces int64 arrays for .row, .col, .data, and downstream .indices / .indptr. For large libraries this dominates peak memory even when the values fit int32 (or the nnz fits int32 for indices/indptr).

Proposal: add opt-in index_dtype and data_dtype kwargs so callers can select int32 / float32 at ingest time.

Motivation / numbers

For a real 10x-style Perturb-seq library with ~100M non-zeros:

Stage	Bytes/nnz	Total
mmread COO (int64/int64/int64)	24	~2.4 GB
.tocsc() allocation (before source freed)	24	~2.4 GB
Peak transient	48	~4.8 GB

With int32 indices + int32/float32 data:

Stage	Bytes/nnz	Total
Triplet arrays (int32/int32/int32)	12	~1.2 GB
.tocsr() allocation (int32 indices, int32 indptr)	12	~1.2 GB
Peak transient	24	~2.4 GB

Exact 2x reduction in the biggest single allocation of the ingest step. The values fit int32 — DRAGEN emits %%MatrixMarket matrix coordinate integer general where each entry is a UMI count that is empirically well under 2^31 - 1. For real-field variants, float32 is more than enough for count-like values and unchanged for typical downstream analysis.

Proposed API

def import_anndata_generic(
    input_path: str,
    *,
    mtx_file: str,
    gene_annotations: str,
    sample_annotations: str,
    gene_headers: list[str] | None = None,
    sample_headers: list[str] | None = None,
    sep: str = "\t",
    prefilter: bool = True,
    index_dtype: np.dtype | None = None,   # NEW: forces int32 when nnz fits
    data_dtype: np.dtype | None = None,    # NEW: forces int32/float32 for .data
) -> AnnData:
    ...

Semantics:

• data_dtype=None (default): keep current behaviour.
• data_dtype=np.int32: assert MM field is integer, stream directly into int32.
• data_dtype=np.float32: cast values to float32 during (or immediately after) read.
• index_dtype=None: keep current behaviour (int64 unless scipy's get_index_dtype says otherwise).
• index_dtype=np.int32: raise if nnz > INT32_MAX, otherwise force int32 for .indices / .indptr.

Implementation sketch

scipy.io.mmread has no dtype hooks, so a first pass could keep the current path and cast after construction (still allocates int64 transiently). A cleaner second pass would stream MatrixMarket triplets directly (e.g. via pandas.read_csv with explicit dtype={"row": np.int32, "col": np.int32, "val": data_dtype} and chunksize) and skip mmread entirely — this is what we've prototyped downstream and it drops peak memory as described above.

Happy to open a PR once the API shape is agreed on.

Workaround

We're carrying a local, streaming int32 reader (streaming_mtx.read_dragen_anndata) in a downstream repository that we swap in for import_anndata_generic on the DRAGEN path (coordinate integer general / coordinate real general) and fall back to import_anndata_generic for other variants. That local shim disappears the day this lands upstream.