Release 0.6.1: sparse outputs at their real size, and split in one pass

CHANGELOG.md

@@ -3,6 +3,42 @@
 Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no
 `v` prefix.
 
+## 0.6.1
+
+Makes the outputs of `subset`, `split` and `concat` the size anndata writes
+(they were about twice that), and makes `split` read each matrix once rather
+than once per group. In 0.6.0 a hand-written anndata loop split the ci-tier
+benchmark 7x faster than `split` did; `split` is now the faster of the two.
+
+### Fixed
+
+- **`subset`, `split` and `concat` wrote sparse matrices about twice as large
+  as they needed to be.** Every sparse output widened `indices` and `indptr`
+  to int64, whatever the source used, and was written without the source's
+  compression, so with an lzf input the `data` and `indices` came out
+  uncompressed. Outputs now keep the source's index dtypes and storage
+  settings. `concat` widens only when the combined matrix could not
+  be addressed otherwise. `convert` already did this; the code next to it
+  did not.
+
+- **`split` read the whole matrix once per group.** It ran a separate subset
+  for each group. When groups were interleaved through the file, as samples
+  usually are, each subset read nearly all of X, so a split into k groups
+  read and decompressed X k times: 56.7 s against 7.8 s for a hand-written
+  anndata loop in the 0.6.0 benchmark. `split` now reads each matrix once
+  and shares every block out among all the outputs. It holds up to a quarter
+  of the file-descriptor limit open at once, and makes one pass per batch
+  beyond that. `subset` shares the same writer, which now gathers rows in
+  one vectorised step instead of a Python loop over rows, and skips blocks
+  that hold none of the selected rows. A new guard in
+  `tests/test_performance.py` checks that split's reads of X do not grow
+  with the number of groups.
+
+  Dataframe columns are no longer read through an h5py fancy index, which
+  cost 0.38 s per column per group on 50,000 cells. They are now read as
+  contiguous blocks and selected in memory. With both changes, the ci-tier
+  benchmark on a laptop splits in 1.6 s against anndata's 2.3 s.
+
 ## 0.6.0
 
 Adds `adata convert`, and fixes four quadratic paths that made `concat` and

pyproject.toml

@@ -3,7 +3,7 @@
 # so the distribution is published as `pyadata-cli`. The import package
 # and the command are both still `adata`.
 name = "pyadata-cli"
-version = "0.6.0"
+version = "0.6.1"
 description = "Streaming CLI for exploring and editing large AnnData .h5ad and .zarr stores"
 readme = "README.md"
 requires-python = ">=3.12"

src/adata/commands/split.py

@@ -3,7 +3,7 @@
 Requested in issue #2, modelled on cellgeni/scraft's `split_h5ad` but
 streaming: the source is never loaded into memory, only the column being split
 on is read, and each output is produced by the same subset machinery used by
-`adata subset`.
+`adata subset`, fanned out so every matrix is read once for all groups.
 """
 
 from __future__ import annotations
@@ -17,7 +17,7 @@
 from rich.console import Console
 
 from adata.core.select import group_indices
-from adata.core.subset import subset_h5ad
+from adata.core.subset import split_h5ad
 from adata.storage import detect_backend, open_store
 
 
@@ -112,19 +112,22 @@ def split_store(
 
     output_dir.mkdir(parents=True, exist_ok=True)
 
-    for label, out_path, _ in planned:
-        indices = np.sort(groups[label])
-        subset_h5ad(
-            file=file,
-            output=out_path,
-            obs_file=None,
-            var_file=None,
-            chunk_rows=chunk_rows,
-            console=console,
-            obs_indices=indices if axis == "obs" else None,
-            var_indices=indices if axis == "var" else None,
-            zarr_format=zarr_format,
-        )
+    # One call for every group, so each matrix is read once rather than once
+    # per group.
+    split_h5ad(
+        file,
+        [
+            (
+                out_path,
+                np.sort(groups[label]) if axis == "obs" else None,
+                np.sort(groups[label]) if axis == "var" else None,
+            )
+            for label, out_path, _ in planned
+        ],
+        chunk_rows=chunk_rows,
+        console=console,
+        zarr_format=zarr_format,
+    )
 
     if manifest:
         _write_manifest(file, output_dir, column, planned, console)

src/adata/core/concat.py

@@ -389,6 +389,45 @@ def _matrix_kind(obj: Any) -> str:
     return "other"
 
 
+def _widen_to_hold(dtype: np.dtype, largest: int) -> np.dtype:
+    """`dtype`, or the narrowest of int32/int64 wider than it that holds `largest`."""
+    dtype = np.dtype(dtype)
+    if largest <= np.iinfo(dtype).max:
+        return dtype
+    for candidate in (np.int32, np.int64):
+        if largest <= np.iinfo(candidate).max:
+            return np.dtype(np.result_type(dtype, candidate))
+    return np.dtype(np.int64)
+
+
+def _concat_index_dtypes(
+    sources: List[Any], n_minor: int
+) -> Tuple[np.dtype, np.dtype]:
+    """The `indices` and `indptr` dtypes for a concatenation of `sources`.
+
+    The widest the inputs used, so int32 inputs give an int32 output rather
+    than the int64 every output used to get. That is widened only when the
+    result could not be addressed otherwise: `indices` must reach the
+    output's minor dimension, and `indptr` its total nonzero count, which is
+    known before writing as the sum of the inputs'. Both are checked against
+    the chosen dtype's own limit, not int32's, since the spec allows
+    narrower index arrays and an int16 `indptr` would otherwise wrap.
+    """
+    indices = np.result_type(*[s["indices"].dtype for s in sources])
+    indptr = np.result_type(*[s["indptr"].dtype for s in sources])
+    nnz = sum(int(s["data"].shape[0]) for s in sources)
+    return _widen_to_hold(indices, n_minor), _widen_to_hold(indptr, nnz)
+
+
+def _write_indptr(
+    group: Any, indptr: List[int], dtype: np.dtype, template: Any
+) -> None:
+    from adata.core.subset import _sparse_dataset
+
+    values = np.asarray(indptr, dtype=np.int64).astype(dtype, copy=False)
+    _sparse_dataset(group, "indptr", dtype, values.size, template)[:] = values
+
+
 def _concat_sparse(
     dst_parent: Any,
     name: str,
@@ -403,7 +442,7 @@ def _concat_sparse(
     Absent columns need no fill: a sparse matrix's zeros are implicit, so
     dropped entries simply do not appear in the output.
     """
-    from adata.core.subset import _append, _growable
+    from adata.core.subset import _append, _growable_like
 
     group = dst_parent.create_group(name)
     spec.set_encoding(group, spec.CSR_MATRIX)
@@ -412,8 +451,11 @@ def _concat_sparse(
     set_shape_attr(group, (n_rows, n_cols))
 
     dtype = np.result_type(*[s["data"].dtype for s in sources])
-    out_data = _growable(group, "data", dtype)
-    out_indices = _growable(group, "indices", np.int64)
+    index_dtype, pointer_dtype = _concat_index_dtypes(sources, n_cols)
+    out_data = _growable_like(group, "data", dtype, sources[0]["data"])
+    out_indices = _growable_like(
+        group, "indices", index_dtype, sources[0]["indices"]
+    )
     indptr = [0]
     nnz = 0
 
@@ -449,7 +491,7 @@ def _concat_sparse(
                 _append(out_indices, np.concatenate(kept_idx))
                 _append(out_data, np.concatenate(kept_data).astype(dtype))
 
-    create_dataset(group, "indptr", data=np.asarray(indptr, dtype=np.int64))
+    _write_indptr(group, indptr, pointer_dtype, sources[0]["indptr"])
 
 
 def _concat_dense(
@@ -509,16 +551,19 @@ def _concat_csc(
     input's row offset added. Because every input's column is already sorted
     and the offsets increase, the result is sorted without a re-sort.
     """
-    from adata.core.subset import _append, _growable
+    from adata.core.subset import _append, _growable_like
     from adata.elements.write import set_shape_attr
 
     group = dst_parent.create_group(name)
     spec.set_encoding(group, spec.CSC_MATRIX)
     set_shape_attr(group, (n_rows, n_cols))
 
     dtype = np.result_type(*[s["data"].dtype for s in sources])
-    out_data = _growable(group, "data", dtype)
-    out_indices = _growable(group, "indices", np.int64)
+    index_dtype, pointer_dtype = _concat_index_dtypes(sources, n_rows)
+    out_data = _growable_like(group, "data", dtype, sources[0]["data"])
+    out_indices = _growable_like(
+        group, "indices", index_dtype, sources[0]["indices"]
+    )
     indptr = [0]
     nnz = 0
 
@@ -547,7 +592,7 @@ def _concat_csc(
             nnz += int(sum(len(r) for r in rows))
         indptr.append(nnz)
 
-    create_dataset(group, "indptr", data=np.asarray(indptr, dtype=np.int64))
+    _write_indptr(group, indptr, pointer_dtype, sources[0]["indptr"])
 
 
 def check_matrix_encodings(roots: List[Any], console: Console) -> None:

src/adata/core/convert.py

@@ -31,6 +31,7 @@
 
 from adata.elements import spec
 from adata.elements.write import set_shape_attr
+from adata.core.subset import _growable_like, _sparse_dataset
 from adata.storage import (
     copy_attrs,
     create_dataset,
@@ -298,41 +299,6 @@ def _new_sparse_group(
     return group
 
 
-def _sparse_dataset(
-    group: Any, name: str, dtype: np.dtype, n: int, template: Any
-) -> Any:
-    """A 1-D dataset of `n` elements laid out like `template`.
-
-    Forwarding compression and chunking matters more here than anywhere
-    else: the point of a dtype change is usually to make the file smaller,
-    and creating the output with a fixed 65,536-element chunk made a
-    2,400-nonzero matrix allocate 786 KB of mostly empty chunk -- eight
-    times the source, from a conversion asked for to halve it.
-    """
-    from adata.core.subset import _clamp_chunks
-
-    backend = "zarr" if is_zarr_group(group) else "hdf5"
-    kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group)
-    kw = _clamp_chunks(kw, max(1, n))
-    if "chunks" not in kw and n:
-        kw["chunks"] = (min(n, 1 << 16),)
-    return create_dataset(group, name, shape=(n,), dtype=dtype, **kw)
-
-
-def _growable_like(group: Any, name: str, dtype: np.dtype, template: Any) -> Any:
-    """Like `_sparse_dataset`, but extensible for a size not yet known."""
-    backend = "zarr" if is_zarr_group(group) else "hdf5"
-    kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group)
-    kw.pop("shards", None)
-    chunks = kw.pop("chunks", None)
-    step = int(chunks[0]) if chunks else 1 << 16
-    if is_zarr_group(group):
-        return group.create_array(name, shape=(0,), dtype=dtype, chunks=(step,), **kw)
-    return group.create_dataset(
-        name, shape=(0,), maxshape=(None,), dtype=dtype, chunks=(step,), **kw
-    )
-
-
 def _write_sparse_arrays(
     group: Any,
     data: np.ndarray,