feat: update to anndata 0.11 and memory efficient reads + writes

cellxgene_schema_cli/cellxgene_schema/utils.py

@@ -6,7 +6,10 @@
 from typing import Dict, List, Union
 
 import anndata as ad
+import h5py
 import numpy as np
+from anndata.compat import DaskArray
+from anndata.experimental import read_dispatched, read_elem_as_dask
 from cellxgene_ontology_guide.ontology_parser import OntologyParser
 from scipy import sparse
 from xxhash import xxh3_64_intdigest
@@ -68,7 +71,7 @@
     return adata
 
 
-def get_matrix_format(adata: ad.AnnData, matrix: Union[np.ndarray, sparse.spmatrix]) -> str:
+def get_matrix_format(matrix: DaskArray) -> str:
     """
     Given a matrix, returns the format as one of: csc, csr, coo, dense
     or unknown.
@@ -84,15 +87,11 @@
     # >>> return getattr(matrix, "format_str", "dense)
     #
     matrix_format = "unknown"
-    if adata.n_obs == 0 or adata.n_vars == 0:
+    matrix_slice = matrix[0:1, 0:1].compute()
+    if isinstance(matrix_slice, sparse.spmatrix):
+        matrix_format = matrix_slice.format
+    elif isinstance(matrix_slice, np.ndarray):
         matrix_format = "dense"
-    else:
-        matrix_slice = matrix[0:1, 0:1]
-        if isinstance(matrix_slice, sparse.spmatrix):
-            matrix_format = matrix_slice.format
-        elif isinstance(matrix_slice, np.ndarray):
-            matrix_format = "dense"
-
     assert matrix_format in ["unknown", "csr", "csc", "coo", "dense"]
     return matrix_format
 
@@ -116,6 +115,36 @@
         return getattr(adata, attr)
 
 
+def read_backed(f: h5py.File, chunk_size: int = 10_000) -> ad.AnnData:
+    """
+    Read an AnnData object from a h5py.File object, reading in matrices (dense or sparse) as dask arrays. Does not
+    read full matrices into memory.
+
+    :param f: h5py.File object
+    :param chunk_size: size of chunks to read matrices in
+    :return: ad.AnnData object
+    """
+
+    def callback(func, elem_name: str, elem, iospec):
+        if "layers" in elem_name or ("X" in elem_name and "X_" not in elem_name):
+            if iospec.encoding_type in (
+                "csr_matrix",
+                "csc_matrix",
+            ):
+                n_vars = elem.attrs.get("shape")[1]
+                return read_elem_as_dask(elem, chunks=(chunk_size, n_vars))
+            elif iospec.encoding_type == "array" and len(elem.shape) > 1:
+                return read_elem_as_dask(elem)
+            else:
+                return func(elem)
+        else:
+            return func(elem)
+
+    adata = read_dispatched(f, callback=callback)
+
+    return adata
+
+
 def read_h5ad(h5ad_path: Union[str, bytes, os.PathLike]) -> ad.AnnData:
     """
     Reads h5ad into adata
@@ -124,13 +153,14 @@
     :rtype None
     """
     try:
-        adata = ad.read_h5ad(h5ad_path, backed="r")
+        f = h5py.File(h5ad_path)
+        adata = read_backed(f)
 
         # This code, and AnnData in general, is optimized for row access.
         # Running backed, with CSC, is prohibitively slow. Read the entire
         # AnnData into memory if it is CSC.
-        if (get_matrix_format(adata, adata.X) == "csc") or (
-            (adata.raw is not None) and (get_matrix_format(adata, adata.raw.X) == "csc")
+        if (get_matrix_format(adata.X) == "csc") or (
+            (adata.raw is not None) and (get_matrix_format(adata.raw.X) == "csc")
         ):
             logger.warning("Matrices are in CSC format; loading entire dataset into memory.")
             adata = adata.to_memory()