Patch AnnData.__sizeof__() for backed datasets

anndata/_core/anndata.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import collections.abc as cabc
+import sys
 import warnings
 from collections import OrderedDict
 from collections.abc import Iterable, Mapping, MutableMapping, Sequence
@@ -25,7 +26,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,26 +593,35 @@ def _init_as_actual(
         # layers
         self._layers = Layers(self, layers)
 
-    def __sizeof__(self, show_stratified=None) -> 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 __sizeof__(self, show_stratified=None, with_disk: bool = False) -> int:
+        from tqdm import tqdm
+
+        def get_size(X) -> int:
+            def cs_to_bytes(X) -> int:
+                return int(X.data.nbytes + X.indptr.nbytes + X.indices.nbytes)
+
+            if isinstance(X, h5py.Dataset) and with_disk:
+                return int(np.array(X.shape).prod() * X.dtype.itemsize)
+            elif isinstance(X, BaseCompressedSparseDataset) and with_disk:
+                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__()
+                return sys.getsizeof(X)
 
         size = 0
-        attrs = list(["_X", "_obs", "_var"])
-        attrs_multi = list(["_uns", "_obsm", "_varm", "varp", "_obsp", "_layers"])
+        attrs = ["X", "_obs", "_var"]
+        attrs_multi = ["_uns", "_obsm", "_varm", "varp", "_obsp", "_layers"]
         for attr in attrs + attrs_multi:
             if attr in attrs_multi:
                 keys = getattr(self, attr).keys()
-                s = sum([get_size(getattr(self, attr)[k]) for k in keys])
+                s = sum(get_size(getattr(self, attr)[k]) for k in keys)
             else:
                 s = get_size(getattr(self, attr))
             if s > 0 and show_stratified:
-                str_attr = attr.replace("_", ".") + " " * (7 - len(attr))
-                print(f"Size of {str_attr}: {'%3.2f' % (s / (1024 ** 2))} MB")
+                print(
+                    f"Size of {attr.replace('_', '.'):<7}: {tqdm.format_sizeof(s, 'B')}"
+                )
             size += s
         return size
 

anndata/tests/test_backed_sparse.py

@@ -212,3 +212,39 @@ 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:
+        dense_with_disk == dense_without_disk
+
+
+def test_backed_sizeof(ondisk_equivalent_adata):
+    csr_mem, csr_disk, csc_disk, _ = ondisk_equivalent_adata
+
+    csr_mem.__sizeof__() == csr_disk.__sizeof__(with_disk=True)
+    csr_mem.__sizeof__() == csc_disk.__sizeof__(with_disk=True)
+    csr_disk.__sizeof__(with_disk=True) == csc_disk.__sizeof__(with_disk=True)
+    csr_mem.__sizeof__() > csr_disk.__sizeof__()
+    csr_mem.__sizeof__() > csc_disk.__sizeof__()

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}`1230` {user}`Neah-Ko`
 
 ```{rubric} Documentation
 ```