add tests and make this an optional argument

Author: selmanozleyen

src/anndata/experimental/merge.py

@@ -208,11 +208,13 @@ def write_concat_dense(  # noqa: PLR0917
     })
 
 
-def write_concat_sparse(  # noqa: PLR0917
+def write_concat_sparse(  # noqa: PLR0913, PLR0917
     datasets: Sequence[BaseCompressedSparseDataset],
     output_group: ZarrGroup | H5Group,
     output_path: ZarrGroup | H5Group,
+    *,
     max_loaded_elems: int,
+    virtual_concat: bool = False,
     axis: Literal[0, 1] = 0,
     reindexers: Reindexer | None = None,
     fill_value: Any = None,
@@ -235,6 +237,9 @@ def write_concat_sparse(  # noqa: PLR0917
         A reindexer object that defines the reindexing operation to be applied.
     fill_value
         The fill value to use for missing elements. Defaults to None.
+    virtual_concat
+        Whether to use virtual concatenation for sparse arrays.
+
     """
     elems = None
     use_reindexing = True
@@ -245,7 +250,7 @@ def write_concat_sparse(  # noqa: PLR0917
         elems = _gen_slice_to_append(
             datasets, reindexers, max_loaded_elems, axis, fill_value
         )
-    if datasets[0].backend == "hdf5" and not use_reindexing:
+    if datasets[0].backend == "hdf5" and not use_reindexing and virtual_concat:
         BaseCompressedSparseDataset.virtual_concat_hdf5(
             datasets, output_group, output_path
         )
@@ -278,6 +283,7 @@ def _write_concat_mappings(  # noqa: PLR0913, PLR0917
     index: pd.Index = None,
     reindexers: list[Reindexer] | None = None,
     fill_value: Any = None,
+    virtual_concat: bool = False,
 ):
     """
     Write a list of mappings to a zarr/h5 group.
@@ -298,6 +304,7 @@ def _write_concat_mappings(  # noqa: PLR0913, PLR0917
             reindexers=reindexers,
             fill_value=fill_value,
             max_loaded_elems=max_loaded_elems,
+            virtual_concat=virtual_concat,
         )
 
 
@@ -306,6 +313,7 @@ def _write_concat_arrays(  # noqa: PLR0913, PLR0917
     output_group: ZarrGroup | H5Group,
     output_path: str | Path,
     max_loaded_elems: int,
+    virtual_concat: bool = False,
     axis: Literal[0, 1] = 0,
     reindexers: list[Reindexer] | None = None,
     fill_value: Any = None,
@@ -331,17 +339,24 @@ def _write_concat_arrays(  # noqa: PLR0913, PLR0917
                 arrays,
                 output_group,
                 output_path,
-                max_loaded_elems,
-                axis,
-                reindexers,
-                fill_value,
+                max_loaded_elems=max_loaded_elems,
+                virtual_concat=virtual_concat,
+                axis=axis,
+                reindexers=reindexers,
+                fill_value=fill_value,
             )
         else:
             msg = f"Concat of following not supported: {[a.format for a in arrays]}"
             raise NotImplementedError(msg)
     else:
         write_concat_dense(
-            arrays, output_group, output_path, axis, reindexers, fill_value
+            arrays,
+            output_group,
+            output_path,
+            virtual_concat=virtual_concat,
+            axis=axis,
+            reindexers=reindexers,
+            fill_value=fill_value,
         )
 
 
@@ -354,6 +369,7 @@ def _write_concat_sequence(  # noqa: PLR0913, PLR0917
     index: pd.Index | None = None,
     reindexers: list[Reindexer] | None = None,
     fill_value: Any = None,
+    virtual_concat: bool = False,
     join: Join_T = "inner",
 ):
     """
@@ -388,11 +404,12 @@ def _write_concat_sequence(  # noqa: PLR0913, PLR0917
             arrays,
             output_group,
             output_path,
-            max_loaded_elems,
-            axis,
-            reindexers,
-            fill_value,
-            join,
+            max_loaded_elems=max_loaded_elems,
+            virtual_concat=virtual_concat,
+            axis=axis,
+            reindexers=reindexers,
+            fill_value=fill_value,
+            join=join,
         )
     else:
         msg = f"Concatenation of these types is not yet implemented: {[type(a) for a in arrays]} with axis={axis}."
@@ -470,6 +487,7 @@ def concat_on_disk(  # noqa: PLR0913
     out_file: PathLike[str] | str | H5Group | ZarrGroup,
     *,
     max_loaded_elems: int = 100_000_000,
+    virtual_concat: bool = False,
     axis: Literal["obs", 0, "var", 1] = 0,
     join: Join_T = "inner",
     merge: StrategiesLiteral | Callable[[Collection[Mapping]], Mapping] | None = None,
@@ -499,10 +517,9 @@ def concat_on_disk(  # noqa: PLR0913
     see the Dask documentation, as the Dask concatenation function is used
     to concatenate dense arrays in this function.
 
-    For sparse arrays, if the backend is hdf5 and there is no reindexing,
+    For sparse arrays, if the backend is hdf5 and there is no reindexing and
+    `virtual_concat` is True,
     the virtual concatenation is used using the `h5py` virtual dataset support.
-    This will create soft links to the source files instead of copying the whole content.
-    Be aware that this will make the output file dependent on the source files.
 
     Params
     ------
@@ -517,6 +534,11 @@ def concat_on_disk(  # noqa: PLR0913
         sparse arrays. Note that this number also includes the empty entries.
         Set to 100m by default meaning roughly 400mb will be loaded
         to memory simultaneously.
+    virtual_concat
+        Whether to use virtual concatenation for sparse arrays.
+        This will create soft links to the source files instead of copying the whole content.
+        Be aware that this will make the output file dependent on the source files.
+        This is False by default.
     axis
         Which axis to concatenate along.
     join
@@ -664,6 +686,7 @@ def concat_on_disk(  # noqa: PLR0913
             label=label,
             index_unique=index_unique,
             fill_value=fill_value,
+            virtual_concat=virtual_concat,
             merge=merge,
         )
 
@@ -681,6 +704,7 @@ def _concat_on_disk_inner(  # noqa: PLR0913
     label: str | None,
     index_unique: str | None,
     fill_value: Any | None,
+    virtual_concat: bool = False,
     merge: Callable[[Collection[Mapping]], Mapping],
 ) -> None:
     """Internal helper to minimize the amount of indented code within the context manager"""
@@ -754,6 +778,7 @@ def _concat_on_disk_inner(  # noqa: PLR0913
         reindexers=reindexers,
         fill_value=fill_value,
         max_loaded_elems=max_loaded_elems,
+        virtual_concat=virtual_concat,
     )
 
     # Write Layers and {axis_name}m
@@ -774,6 +799,7 @@ def _concat_on_disk_inner(  # noqa: PLR0913
             intersect_keys(maps),
             m,
             max_loaded_elems=max_loaded_elems,
+            virtual_concat=virtual_concat,
             axis=m_axis,
             index=m_index,
             reindexers=m_reindexers,

tests/test_concatenate_disk.py

@@ -98,7 +98,9 @@ def assert_eq_concat_on_disk(
     adatas,
     tmp_path: Path,
     file_format: Literal["zarr", "h5ad"],
+    *,
     max_loaded_elems: int | None = None,
+    virtual_concat: bool = False,
     *args,
     merge_strategy: merge.StrategiesLiteral | None = None,
     **kwargs,
@@ -110,7 +112,7 @@ def assert_eq_concat_on_disk(
     out_name = tmp_path / f"out.{file_format}"
     if max_loaded_elems is not None:
         kwargs["max_loaded_elems"] = max_loaded_elems
-    concat_on_disk(paths, out_name, *args, merge=merge_strategy, **kwargs)
+    concat_on_disk(paths, out_name, *args, virtual_concat=virtual_concat, merge=merge_strategy, **kwargs)
     with as_group(out_name, mode="r") as rg:
         res2 = read_elem(rg)
     assert_equal(res1, res2, exact=False)
@@ -181,6 +183,48 @@ def test_anndatas(
     )
 
 
+def test_anndatas_virtual_concat(
+    *,
+    tmp_path: Path,
+):
+    axis = 0
+    max_loaded_elems = 1_000_000
+    file_format = "h5ad"
+    array_type = "sparse"
+    join_type = "inner"
+    _, off_axis_name = _resolve_axis(1 - axis)
+    random_axes = {0, 1}
+    sparse_fmt = "csr" if axis == 0 else "csc"
+    kw = GEN_ADATA_OOC_CONCAT_ARGS
+
+    adatas = []
+    for i in range(3):
+        M, N = (np.random.randint(5, 10) if a in random_axes else 50 for a in (0, 1))
+        a = gen_adata(
+            (M, N),
+            X_type=get_array_type(array_type, axis),
+            sparse_fmt=sparse_fmt,
+            obs_dtypes=[pd.CategoricalDtype(ordered=False)],
+            var_dtypes=[pd.CategoricalDtype(ordered=False)],
+            **kw,
+        )
+        # ensure some names overlap, others do not, for the off-axis so that inner/outer is properly tested
+        off_names = getattr(a, f"{off_axis_name}_names").array
+        off_names[1::2] = f"{i}-" + off_names[1::2]
+        setattr(a, f"{off_axis_name}_names", off_names)
+        adatas.append(a)
+
+    assert_eq_concat_on_disk(
+        adatas,
+        tmp_path,
+        file_format,
+        max_loaded_elems=max_loaded_elems,
+        virtual_concat=True,
+        axis=axis,
+        join=join_type,
+    )
+
+
 def test_concat_ordered_categoricals_retained(tmp_path, file_format):
     a = AnnData(
         X=np.ones((5, 1)),