Stream Matrix Market files: a codec behind import/export sparse

docs/COMMANDS.md

@@ -326,7 +326,7 @@ or a file of names. `--force` overwrites an existing store.
 |---|---|---|
 | `dataframe` | CSV | any dataframe path (`obs`, `var`, `raw/var`, …) |
 | `array` | `.npy` | any dense array |
-| `sparse` | `.mtx` | CSR/CSC groups |
+| `sparse` | `.mtx`, or `.mtx.gz` | CSR/CSC groups |
 | `dict` | JSON | any group or scalar |
 | `image` | PNG/JPEG/TIFF | 2D or 3D arrays |
 
@@ -342,6 +342,10 @@ Omit `--output` to write to stdout. Writing to a file streams in chunks;
 stdout is not seekable, so `export array` holds the array in memory on that
 path.
 
+`export sparse` gzips its output when the name ends in `.gz`. Whole-number
+values are written as integers, and other floats with enough digits to read
+back exactly.
+
 ## `import`
 
 Write data into a store at any path. All subcommands need `--output`/`-o` or
@@ -358,6 +362,16 @@ adata import image     data.h5ad uns/spatial/hires tissue.png --inplace
 Dimensions are validated against the existing obs/var where the path implies
 an axis.
 
+For `sparse`: the file is read in two streaming passes, so its size is not a
+limit, and it may be gzipped (`counts.mtx.gz`). The first pass checks the
+whole file before the store is touched. `integer` values stay integers
+(int32, or int64 when a value needs it), `real` ones become float64, and
+`pattern` entries become ones. Duplicate entries are summed, as scipy does.
+Only `coordinate general` files are read: `symmetric`, `skew-symmetric`,
+`hermitian` and `complex` files, and dense `array` ones, are refused. An
+input whose entries are not in row order is sorted through scratch space
+beside the store.
+
 For `dataframe`: columns that parse as integers or floats become numeric
 arrays; string columns with few distinct values become categoricals. Use
 `--categorical col1,col2` to force specific columns, or

src/adata/cli.py

@@ -1117,7 +1117,7 @@ def export_sparse(
         "--output",
         "-o",
         writable=True,
-        help="Output .mtx file path (defaults to stdout)",
+        help="Output .mtx file path; .mtx.gz gzips it (defaults to stdout)",
     ),
     head: Optional[int] = typer.Option(
         None, "--head", "-n", help="Output only the first n entries of mtx file"
@@ -1407,7 +1407,7 @@ def import_sparse(
         ..., help="Object path to create/replace (e.g., 'X', 'layers/counts')"
     ),
     input_file: Path = typer.Argument(
-        ..., help="Input .mtx file", exists=True, readable=True
+        ..., help="Input .mtx or .mtx.gz file", exists=True, readable=True
     ),
     output: Optional[Path] = typer.Option(
         None,

src/adata/commands/import_data.py

@@ -60,7 +60,9 @@ def import_object(
     console: Console,
 ) -> None:
     target_file = _prepare_target_path(file, output_file, inplace, console)
-    ext = input_file.suffix.lower()
+    suffixes = [x.lower() for x in input_file.suffixes]
+    # A gzipped matrix is still a matrix: `counts.mtx.gz` imports as `.mtx`.
+    ext = suffixes[-2] if len(suffixes) > 1 and suffixes[-1] == ".gz" else input_file.suffix.lower()
 
     if ext not in EXTENSION_FORMAT:
         raise ValueError(

src/adata/core/csr_build.py

@@ -11,6 +11,10 @@
    `BucketSorter.add`), kept on disk in a scratch group (`scratch_group`);
 3. read one bucket at a time, sort it, and append it to the output
    (`BucketSorter.runs`).
+
+Input that already arrives in row order -- as Cell Ranger and STARsolo write
+their MTX files -- skips the buckets (`grouped_runs`). Either way the runs
+go through `DuplicateSummer` and into a `CsrAppender`.
 """
 
 from __future__ import annotations
@@ -51,6 +55,23 @@ def count_keys(counts: np.ndarray, keys: np.ndarray) -> None:
         counts[unique] += seen
 
 
+def sum_keys(sums: np.ndarray, keys: np.ndarray, weights: np.ndarray) -> None:
+    """Add `weights[i]` to `sums[keys[i]]`, over the span the block covers.
+
+    As `count_keys`, and for the same reason: cost set by the block, not by
+    the length of `sums`.
+    """
+    if keys.size == 0:
+        return
+    lo, hi = int(keys.min()), int(keys.max())
+    span = hi - lo + 1
+    if span <= 4 * keys.size:
+        sums[lo : hi + 1] += np.bincount(keys - lo, weights=weights, minlength=span)
+    else:
+        unique, inverse = np.unique(keys, return_inverse=True)
+        sums[unique] += np.bincount(inverse, weights=weights, minlength=unique.size)
+
+
 def bucket_bounds(
     indptr: np.ndarray, *, chunk: int, bucket_entries: Optional[int] = None
 ) -> np.ndarray:
@@ -176,3 +197,132 @@ def runs(self) -> Iterator[Tuple[np.ndarray, np.ndarray, np.ndarray]]:
             values = np.asarray(vals[...])
             order = np.lexsort((minor, major))
             yield major[order], minor[order], values[order]
+
+
+Triple = Tuple[np.ndarray, np.ndarray, np.ndarray]
+
+#: The largest nnz an int32 `indptr` can address.
+INT32_MAX = int(np.iinfo(np.int32).max)
+
+
+def grouped_runs(blocks: Iterator[Triple], *, sort: bool) -> Iterator[Triple]:
+    """Runs in CSR order, from blocks already grouped by row.
+
+    The fast path, for input whose rows arrive in order -- as Cell Ranger and
+    STARsolo write them. A row may straddle two blocks, so the last row of
+    each block is held back and joined to the next. With `sort`, entries are
+    put in column order within each row; without, they are taken to be in it
+    already.
+    """
+    carry: Optional[Triple] = None
+    for rows, cols, values in blocks:
+        if carry is not None:
+            rows = np.concatenate((carry[0], rows))
+            cols = np.concatenate((carry[1], cols))
+            values = np.concatenate((carry[2], values))
+            carry = None
+        if rows.size == 0:
+            continue
+        cut = int(np.searchsorted(rows, rows[-1], side="left"))
+        carry = (rows[cut:], cols[cut:], values[cut:])
+        if cut:
+            yield _ordered(rows[:cut], cols[:cut], values[:cut], sort)
+    if carry is not None and carry[0].size:
+        yield _ordered(*carry, sort)
+
+
+def _ordered(rows: np.ndarray, cols: np.ndarray, values: np.ndarray, sort: bool) -> Triple:
+    if not sort:
+        return rows, cols, values
+    order = np.lexsort((cols, rows))
+    return rows[order], cols[order], values[order]
+
+
+class DuplicateSummer:
+    """Merge repeated (row, column) entries by adding their values.
+
+    As scipy does when it builds CSR from coordinates, and so as scanpy's
+    readers give them. Each run must hold whole rows, sorted, which both
+    `grouped_runs` and `BucketSorter.runs` guarantee.
+    """
+
+    def __init__(self) -> None:
+        self.summed = 0
+
+    def __call__(self, runs: Iterator[Triple]) -> Iterator[Triple]:
+        for rows, cols, values in runs:
+            same = (rows[1:] == rows[:-1]) & (cols[1:] == cols[:-1])
+            if same.any():
+                starts = np.flatnonzero(np.concatenate(([True], ~same)))
+                self.summed += int(rows.size - starts.size)
+                values = np.add.reduceat(values, starts)
+                rows, cols = rows[starts], cols[starts]
+            yield rows, cols, values
+
+
+class CsrAppender:
+    """A CSR matrix written row by row: entries must arrive in row order.
+
+    `data` and `indices` grow as entries are appended; `indptr` is written
+    by `finish` from the rows seen, as int32 when nnz allows, as anndata and
+    scipy would write it.
+    """
+
+    def __init__(
+        self,
+        parent: Any,
+        name: str,
+        *,
+        shape: Tuple[int, int],
+        data_dtype: Any,
+        index_dtype: Any = np.int32,
+        compression: Optional[str] = None,
+    ) -> None:
+        from adata.core.subset import _growable_like
+        from adata.elements import spec
+        from adata.elements.write import set_shape_attr
+
+        self.group = parent.create_group(name)
+        spec.set_encoding(self.group, spec.CSR_MATRIX)
+        set_shape_attr(self.group, shape)
+        self.shape = shape
+        self.data_dtype = np.dtype(data_dtype)
+        self.index_dtype = np.dtype(index_dtype)
+        self.compression = compression
+        self.data = _growable_like(
+            self.group, "data", self.data_dtype, None, compression=compression
+        )
+        self.indices = _growable_like(
+            self.group, "indices", self.index_dtype, None, compression=compression
+        )
+        self.counts = np.zeros(shape[0], dtype=np.int64)
+        self.nnz = 0
+
+    def append(self, rows: np.ndarray, cols: np.ndarray, values: np.ndarray) -> None:
+        from adata.core import subset
+
+        if rows.size == 0:
+            return
+        if self.data_dtype.kind in "iu":
+            info = np.iinfo(self.data_dtype)
+            if values.min() < info.min or values.max() > info.max:
+                raise ValueError(
+                    f"A value does not fit {self.data_dtype}: summing duplicate "
+                    f"entries reached {int(values.max()):,}."
+                )
+        count_keys(self.counts, rows)
+        subset._append(self.data, values.astype(self.data_dtype, copy=False))
+        subset._append(self.indices, cols.astype(self.index_dtype, copy=False))
+        self.nnz += int(rows.size)
+
+    def finish(self) -> int:
+        from adata.core.subset import _sparse_dataset
+
+        dtype = np.int32 if self.nnz <= INT32_MAX else np.int64
+        indptr = np.concatenate(([0], np.cumsum(self.counts))).astype(dtype)
+        dataset = _sparse_dataset(
+            self.group, "indptr", np.dtype(dtype), indptr.size, None,
+            compression=self.compression,
+        )
+        dataset[:] = indptr
+        return self.nnz