Patch AnnData.__sizeof__() for backed datasets

anndata/_core/anndata.py

@@ -25,7 +25,7 @@
 from numpy import ma
 from pandas.api.types import infer_dtype, is_string_dtype
 from scipy import sparse
-from scipy.sparse import csr_matrix, issparse
+from scipy.sparse import issparse
 
 from anndata._warnings import ImplicitModificationWarning
 
@@ -592,16 +592,25 @@ def _init_as_actual(
         # layers
         self._layers = Layers(self, layers)
 
-    def __sizeof__(self, show_stratified=None) -> int:
+    def __sizeof__(self, show_stratified=None, with_disk: bool = False) -> int:
         def get_size(X):
-            if issparse(X):
-                X_csr = csr_matrix(X)
-                return X_csr.data.nbytes + X_csr.indptr.nbytes + X_csr.indices.nbytes
+            def cs_to_bytes(X):
+                return X.data.nbytes + X.indptr.nbytes + X.indices.nbytes
+
+            if not with_disk:
+                return X.__sizeof__()
+
+            if isinstance(X, h5py.Dataset):
+                return np.array(X.shape).prod() * X.dtype.itemsize
+            elif isinstance(X, BaseCompressedSparseDataset):
+                return cs_to_bytes(X._to_backed())
+            elif isinstance(X, (sparse.csr_matrix, sparse.csc_matrix)):
+                return cs_to_bytes(X)
             else:
                 return X.__sizeof__()
 
         size = 0
-        attrs = list(["_X", "_obs", "_var"])
+        attrs = list(["X", "_obs", "_var"])
         attrs_multi = list(["_uns", "_obsm", "_varm", "varp", "_obsp", "_layers"])
         for attr in attrs + attrs_multi:
             if attr in attrs_multi:

anndata/tests/test_backed_sparse.py

@@ -212,3 +212,49 @@ def test_anndata_sparse_compat(tmp_path, diskfmt):
     ad._io.specs.write_elem(f, "/", base)
     adata = ad.AnnData(sparse_dataset(f["/"]))
     assert_equal(adata.X, base)
+
+
+def test_dense_sizeof(ondisk_equivalent_adata, diskfmt):
+    _, _, _, dense_disk = ondisk_equivalent_adata
+
+    size_on_disk = np.array(dense_disk.X.shape).prod() * dense_disk.X.dtype.itemsize
+
+    size_nested_objects = 0
+    for x in ("_obs", "_var"):
+        size_nested_objects += getattr(dense_disk, x).__sizeof__()
+    for x in ("_uns", "_obsm", "_varm", "varp", "_obsp", "_layers"):
+        size_nested_objects += sum(
+            [
+                getattr(dense_disk, x)[k].__sizeof__()
+                for k in getattr(dense_disk, x).keys()
+            ]
+        )
+
+    dense_with_disk = dense_disk.__sizeof__(with_disk=True)
+    dense_without_disk = dense_disk.__sizeof__(with_disk=False)
+
+    assert (
+        dense_with_disk - 128 <= size_on_disk + size_nested_objects <= dense_with_disk
+    )
+    if diskfmt == "h5ad":
+        assert dense_without_disk - 128 <= size_nested_objects <= dense_without_disk
+    else:
+        assert_equal(dense_with_disk, dense_without_disk)
+
+
+def test_backed_sizeof(ondisk_equivalent_adata):
+    csr_mem, csr_disk, csc_disk, _ = ondisk_equivalent_adata
+
+    assert_equal(csr_mem.__sizeof__(), csr_disk.__sizeof__())
+    assert_equal(csr_mem.__sizeof__(), csc_disk.__sizeof__())
+    assert_equal(csr_disk.__sizeof__(), csc_disk.__sizeof__())
+
+    assert_equal(
+        csr_mem.__sizeof__(with_disk=True), csr_disk.__sizeof__(with_disk=True)
+    )
+    assert_equal(
+        csr_mem.__sizeof__(with_disk=True), csc_disk.__sizeof__(with_disk=True)
+    )
+    assert_equal(
+        csr_disk.__sizeof__(with_disk=True), csc_disk.__sizeof__(with_disk=True)
+    )

docs/release-notes/0.10.4.md

@@ -3,6 +3,7 @@
 ```{rubric} Bugfix
 ```
 * Only try to use `Categorical.map(na_action=…)` in actually supported Pandas ≥2.1 {pr}`1226` {user}`flying-sheep`
+* `AnnData.__sizeof__()` support for backed datasets {pr}`1222` {user}`Neah-Ko`
 
 ```{rubric} Documentation
 ```