perf: use name in map_blocks to bypass tokenization (#2121)

Author: ilan-gold

benchmarks/benchmarks/sparse_dataset.py

@@ -7,7 +7,7 @@
 from dask.array.core import Array as DaskArray
 from scipy import sparse
 
-from anndata import AnnData
+from anndata import AnnData, concat
 from anndata._core.sparse_dataset import sparse_dataset
 from anndata._io.specs import write_elem
 from anndata.experimental import read_elem_lazy
@@ -77,3 +77,34 @@ def peakmem_getitem_adata(self, *_):
         res = self.adata[self.index]
         if isinstance(res, DaskArray):
             res.compute()
+
+
+class SparseCSRDask:
+    filepath = "data.zarr"
+
+    def setup_cache(self):
+        X = sparse.random(
+            10_000,
+            10_000,
+            density=0.01,
+            format="csr",
+            random_state=np.random.default_rng(42),
+        )
+        g = zarr.group(self.filepath)
+        write_elem(g, "X", X)
+
+    def setup(self):
+        self.group = zarr.group(self.filepath)
+        self.adata = AnnData(X=read_elem_lazy(self.group["X"]))
+
+    def time_concat(self):
+        concat([self.adata for i in range(100)])
+
+    def peakmem_concat(self):
+        concat([self.adata for i in range(100)])
+
+    def time_read(self):
+        AnnData(X=read_elem_lazy(self.group["X"]))
+
+    def peakmem_read(self):
+        AnnData(X=read_elem_lazy(self.group["X"]))

docs/release-notes/2121.perf.md

@@ -0,0 +1 @@
+Use `name` on {func}`dask.array.map_blocks` internally when concatenating {class}`anndata.experimental.backed.Dataset2D` objects whose categoricals/nullable types must be converted to dask arrays {user}`ilan-gold`

src/anndata/_core/merge.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import uuid
 from collections import OrderedDict
 from collections.abc import Callable, Mapping, MutableSet
 from functools import partial, reduce, singledispatch
@@ -1250,6 +1251,7 @@ def get_chunk(block_info=None):
             chunks=chunk_size,
             meta=np.array([], dtype=dtype),
             dtype=dtype,
+            name=f"{uuid.uuid4()}/{base_path_or_zarr_group}/{elem_name}-{dtype}",
         )
 
     return da.from_array(col.values, chunks=-1)  # in-memory