Merge branch 'main' into feature/concat_hdf5_virtual_datasets

Author: selmanozleyen

docs/release-notes/2084.bugfix.md

@@ -0,0 +1 @@
+Allow writing of views of {class}`dask.array.Array` {user}`ilan-gold`

src/anndata/_io/specs/methods.py

@@ -492,6 +492,7 @@ def write_chunked_dense_array_to_group(
 )
 
 
+@_REGISTRY.register_write(ZarrGroup, views.DaskArrayView, IOSpec("array", "0.2.0"))
 @_REGISTRY.register_write(ZarrGroup, DaskArray, IOSpec("array", "0.2.0"))
 def write_basic_dask_zarr(
     f: ZarrGroup,
@@ -514,6 +515,7 @@ def write_basic_dask_zarr(
 
 # Adding this separately because h5py isn't serializable
 # https://github.com/pydata/xarray/issues/4242
+@_REGISTRY.register_write(H5Group, views.DaskArrayView, IOSpec("array", "0.2.0"))
 @_REGISTRY.register_write(H5Group, DaskArray, IOSpec("array", "0.2.0"))
 def write_basic_dask_h5(
     f: H5Group,
@@ -803,21 +805,7 @@ def write_sparse_dataset(
     f[k].attrs["encoding-version"] = "0.1.0"
 
 
-@_REGISTRY.register_write(H5Group, (DaskArray, CupyArray), IOSpec("array", "0.2.0"))
-@_REGISTRY.register_write(ZarrGroup, (DaskArray, CupyArray), IOSpec("array", "0.2.0"))
-@_REGISTRY.register_write(
-    H5Group, (DaskArray, CupyCSRMatrix), IOSpec("csr_matrix", "0.1.0")
-)
-@_REGISTRY.register_write(
-    H5Group, (DaskArray, CupyCSCMatrix), IOSpec("csc_matrix", "0.1.0")
-)
-@_REGISTRY.register_write(
-    ZarrGroup, (DaskArray, CupyCSRMatrix), IOSpec("csr_matrix", "0.1.0")
-)
-@_REGISTRY.register_write(
-    ZarrGroup, (DaskArray, CupyCSCMatrix), IOSpec("csc_matrix", "0.1.0")
-)
-def write_cupy_dask_sparse(f, k, elem, _writer, dataset_kwargs=MappingProxyType({})):
+def write_cupy_dask(f, k, elem, _writer, dataset_kwargs=MappingProxyType({})):
     _writer.write_elem(
         f,
         k,
@@ -826,18 +814,6 @@ def write_cupy_dask_sparse(f, k, elem, _writer, dataset_kwargs=MappingProxyType(
     )
 
 
-@_REGISTRY.register_write(
-    H5Group, (DaskArray, sparse.csr_matrix), IOSpec("csr_matrix", "0.1.0")
-)
-@_REGISTRY.register_write(
-    H5Group, (DaskArray, sparse.csc_matrix), IOSpec("csc_matrix", "0.1.0")
-)
-@_REGISTRY.register_write(
-    ZarrGroup, (DaskArray, sparse.csr_matrix), IOSpec("csr_matrix", "0.1.0")
-)
-@_REGISTRY.register_write(
-    ZarrGroup, (DaskArray, sparse.csc_matrix), IOSpec("csc_matrix", "0.1.0")
-)
 def write_dask_sparse(
     f: GroupStorageType,
     k: str,
@@ -886,6 +862,26 @@ def chunk_slice(start: int, stop: int) -> tuple[slice | None, slice | None]:
         disk_mtx.append(elem[chunk_slice(chunk_start, chunk_stop)].compute())
 
 
+for array_type, group_type in product(
+    [DaskArray, views.DaskArrayView], [H5Group, ZarrGroup]
+):
+    for cupy_array_type, spec in [
+        (CupyArray, IOSpec("array", "0.2.0")),
+        (CupyCSCMatrix, IOSpec("csc_matrix", "0.1.0")),
+        (CupyCSRMatrix, IOSpec("csr_matrix", "0.1.0")),
+    ]:
+        _REGISTRY.register_write(group_type, (array_type, cupy_array_type), spec)(
+            write_cupy_dask
+        )
+    for scipy_sparse_type, spec in [
+        (sparse.csr_matrix, IOSpec("csr_matrix", "0.1.0")),
+        (sparse.csc_matrix, IOSpec("csc_matrix", "0.1.0")),
+    ]:
+        _REGISTRY.register_write(group_type, (array_type, scipy_sparse_type), spec)(
+            write_dask_sparse
+        )
+
+
 @_REGISTRY.register_read(H5Group, IOSpec("csc_matrix", "0.1.0"))
 @_REGISTRY.register_read(H5Group, IOSpec("csr_matrix", "0.1.0"))
 @_REGISTRY.register_read(ZarrGroup, IOSpec("csc_matrix", "0.1.0"))

tests/test_dask.py

@@ -17,6 +17,7 @@
 from anndata.experimental.merge import as_group
 from anndata.tests.helpers import (
     GEN_ADATA_DASK_ARGS,
+    as_cupy_sparse_dask_array,
     as_dense_cupy_dask_array,
     as_dense_dask_array,
     as_sparse_dask_array,
@@ -310,6 +311,36 @@ def test_dask_to_memory_unbacked(array_func, mem_type):
     assert isinstance(orig.uns["da"]["da"], DaskArray)
 
 
+@pytest.mark.parametrize(
+    "array_func",
+    [
+        pytest.param(as_dense_dask_array, id="dense_dask_array"),
+        pytest.param(as_sparse_dask_array, id="sparse_dask_array"),
+        pytest.param(
+            as_dense_cupy_dask_array,
+            id="cupy_dense_dask_array",
+            marks=pytest.mark.gpu,
+        ),
+        pytest.param(
+            as_cupy_sparse_dask_array,
+            id="cupy_sparse_dask_array",
+            marks=pytest.mark.gpu,
+        ),
+    ],
+)
+def test_dask_to_disk_view(array_func, diskfmt, tmp_path):
+    random_state = np.random.default_rng()
+    orig = ad.AnnData(
+        # need to change type for cupy
+        array_func(random_state.binomial(100, 0.005, (20, 15)).astype("float32"))
+    )
+    orig = orig[orig.shape[0] // 2]
+    path = tmp_path / f"test.{diskfmt}"
+    getattr(orig, f"write_{diskfmt}")(path)
+    roundtrip = getattr(ad.io, f"read_{diskfmt}")(path)
+    assert_equal(roundtrip, orig)
+
+
 # Test if dask arrays turn into numpy arrays after to_memory is called
 def test_dask_to_memory_copy_unbacked():
     import numpy as np