Skip to content

Commit ce19812

Browse files
authored
Merge pull request #48 from cellgeni/feat/mtx-codec
Stream Matrix Market files: a codec behind import/export sparse
2 parents 09416d1 + 27bf873 commit ce19812

10 files changed

Lines changed: 1528 additions & 208 deletions

File tree

‎docs/COMMANDS.md‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -326,7 +326,7 @@ or a file of names. `--force` overwrites an existing store.
326326
|---|---|---|
327327
| `dataframe` | CSV | any dataframe path (`obs`, `var`, `raw/var`, …) |
328328
| `array` | `.npy` | any dense array |
329-
| `sparse` | `.mtx` | CSR/CSC groups |
329+
| `sparse` | `.mtx`, or `.mtx.gz` | CSR/CSC groups |
330330
| `dict` | JSON | any group or scalar |
331331
| `image` | PNG/JPEG/TIFF | 2D or 3D arrays |
332332

@@ -342,6 +342,10 @@ Omit `--output` to write to stdout. Writing to a file streams in chunks;
342342
stdout is not seekable, so `export array` holds the array in memory on that
343343
path.
344344

345+
`export sparse` gzips its output when the name ends in `.gz`. Whole-number
346+
values are written as integers, and other floats with enough digits to read
347+
back exactly.
348+
345349
## `import`
346350

347351
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
358362
Dimensions are validated against the existing obs/var where the path implies
359363
an axis.
360364

365+
For `sparse`: the file is read in two streaming passes, so its size is not a
366+
limit, and it may be gzipped (`counts.mtx.gz`). The first pass checks the
367+
whole file before the store is touched. `integer` values stay integers
368+
(int32, or int64 when a value needs it), `real` ones become float64, and
369+
`pattern` entries become ones. Duplicate entries are summed, as scipy does.
370+
Only `coordinate general` files are read: `symmetric`, `skew-symmetric`,
371+
`hermitian` and `complex` files, and dense `array` ones, are refused. An
372+
input whose entries are not in row order is sorted through scratch space
373+
beside the store.
374+
361375
For `dataframe`: columns that parse as integers or floats become numeric
362376
arrays; string columns with few distinct values become categoricals. Use
363377
`--categorical col1,col2` to force specific columns, or

‎src/adata/cli.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1117,7 +1117,7 @@ def export_sparse(
11171117
"--output",
11181118
"-o",
11191119
writable=True,
1120-
help="Output .mtx file path (defaults to stdout)",
1120+
help="Output .mtx file path; .mtx.gz gzips it (defaults to stdout)",
11211121
),
11221122
head: Optional[int] = typer.Option(
11231123
None, "--head", "-n", help="Output only the first n entries of mtx file"
@@ -1407,7 +1407,7 @@ def import_sparse(
14071407
..., help="Object path to create/replace (e.g., 'X', 'layers/counts')"
14081408
),
14091409
input_file: Path = typer.Argument(
1410-
..., help="Input .mtx file", exists=True, readable=True
1410+
..., help="Input .mtx or .mtx.gz file", exists=True, readable=True
14111411
),
14121412
output: Optional[Path] = typer.Option(
14131413
None,

‎src/adata/commands/import_data.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,9 @@ def import_object(
6060
console: Console,
6161
) -> None:
6262
target_file = _prepare_target_path(file, output_file, inplace, console)
63-
ext = input_file.suffix.lower()
63+
suffixes = [x.lower() for x in input_file.suffixes]
64+
# A gzipped matrix is still a matrix: `counts.mtx.gz` imports as `.mtx`.
65+
ext = suffixes[-2] if len(suffixes) > 1 and suffixes[-1] == ".gz" else input_file.suffix.lower()
6466

6567
if ext not in EXTENSION_FORMAT:
6668
raise ValueError(

‎src/adata/core/csr_build.py‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,10 @@
1111
`BucketSorter.add`), kept on disk in a scratch group (`scratch_group`);
1212
3. read one bucket at a time, sort it, and append it to the output
1313
(`BucketSorter.runs`).
14+
15+
Input that already arrives in row order -- as Cell Ranger and STARsolo write
16+
their MTX files -- skips the buckets (`grouped_runs`). Either way the runs
17+
go through `DuplicateSummer` and into a `CsrAppender`.
1418
"""
1519

1620
from __future__ import annotations
@@ -51,6 +55,23 @@ def count_keys(counts: np.ndarray, keys: np.ndarray) -> None:
5155
counts[unique] += seen
5256

5357

58+
def sum_keys(sums: np.ndarray, keys: np.ndarray, weights: np.ndarray) -> None:
59+
"""Add `weights[i]` to `sums[keys[i]]`, over the span the block covers.
60+
61+
As `count_keys`, and for the same reason: cost set by the block, not by
62+
the length of `sums`.
63+
"""
64+
if keys.size == 0:
65+
return
66+
lo, hi = int(keys.min()), int(keys.max())
67+
span = hi - lo + 1
68+
if span <= 4 * keys.size:
69+
sums[lo : hi + 1] += np.bincount(keys - lo, weights=weights, minlength=span)
70+
else:
71+
unique, inverse = np.unique(keys, return_inverse=True)
72+
sums[unique] += np.bincount(inverse, weights=weights, minlength=unique.size)
73+
74+
5475
def bucket_bounds(
5576
indptr: np.ndarray, *, chunk: int, bucket_entries: Optional[int] = None
5677
) -> np.ndarray:
@@ -176,3 +197,132 @@ def runs(self) -> Iterator[Tuple[np.ndarray, np.ndarray, np.ndarray]]:
176197
values = np.asarray(vals[...])
177198
order = np.lexsort((minor, major))
178199
yield major[order], minor[order], values[order]
200+
201+
202+
Triple = Tuple[np.ndarray, np.ndarray, np.ndarray]
203+
204+
#: The largest nnz an int32 `indptr` can address.
205+
INT32_MAX = int(np.iinfo(np.int32).max)
206+
207+
208+
def grouped_runs(blocks: Iterator[Triple], *, sort: bool) -> Iterator[Triple]:
209+
"""Runs in CSR order, from blocks already grouped by row.
210+
211+
The fast path, for input whose rows arrive in order -- as Cell Ranger and
212+
STARsolo write them. A row may straddle two blocks, so the last row of
213+
each block is held back and joined to the next. With `sort`, entries are
214+
put in column order within each row; without, they are taken to be in it
215+
already.
216+
"""
217+
carry: Optional[Triple] = None
218+
for rows, cols, values in blocks:
219+
if carry is not None:
220+
rows = np.concatenate((carry[0], rows))
221+
cols = np.concatenate((carry[1], cols))
222+
values = np.concatenate((carry[2], values))
223+
carry = None
224+
if rows.size == 0:
225+
continue
226+
cut = int(np.searchsorted(rows, rows[-1], side="left"))
227+
carry = (rows[cut:], cols[cut:], values[cut:])
228+
if cut:
229+
yield _ordered(rows[:cut], cols[:cut], values[:cut], sort)
230+
if carry is not None and carry[0].size:
231+
yield _ordered(*carry, sort)
232+
233+
234+
def _ordered(rows: np.ndarray, cols: np.ndarray, values: np.ndarray, sort: bool) -> Triple:
235+
if not sort:
236+
return rows, cols, values
237+
order = np.lexsort((cols, rows))
238+
return rows[order], cols[order], values[order]
239+
240+
241+
class DuplicateSummer:
242+
"""Merge repeated (row, column) entries by adding their values.
243+
244+
As scipy does when it builds CSR from coordinates, and so as scanpy's
245+
readers give them. Each run must hold whole rows, sorted, which both
246+
`grouped_runs` and `BucketSorter.runs` guarantee.
247+
"""
248+
249+
def __init__(self) -> None:
250+
self.summed = 0
251+
252+
def __call__(self, runs: Iterator[Triple]) -> Iterator[Triple]:
253+
for rows, cols, values in runs:
254+
same = (rows[1:] == rows[:-1]) & (cols[1:] == cols[:-1])
255+
if same.any():
256+
starts = np.flatnonzero(np.concatenate(([True], ~same)))
257+
self.summed += int(rows.size - starts.size)
258+
values = np.add.reduceat(values, starts)
259+
rows, cols = rows[starts], cols[starts]
260+
yield rows, cols, values
261+
262+
263+
class CsrAppender:
264+
"""A CSR matrix written row by row: entries must arrive in row order.
265+
266+
`data` and `indices` grow as entries are appended; `indptr` is written
267+
by `finish` from the rows seen, as int32 when nnz allows, as anndata and
268+
scipy would write it.
269+
"""
270+
271+
def __init__(
272+
self,
273+
parent: Any,
274+
name: str,
275+
*,
276+
shape: Tuple[int, int],
277+
data_dtype: Any,
278+
index_dtype: Any = np.int32,
279+
compression: Optional[str] = None,
280+
) -> None:
281+
from adata.core.subset import _growable_like
282+
from adata.elements import spec
283+
from adata.elements.write import set_shape_attr
284+
285+
self.group = parent.create_group(name)
286+
spec.set_encoding(self.group, spec.CSR_MATRIX)
287+
set_shape_attr(self.group, shape)
288+
self.shape = shape
289+
self.data_dtype = np.dtype(data_dtype)
290+
self.index_dtype = np.dtype(index_dtype)
291+
self.compression = compression
292+
self.data = _growable_like(
293+
self.group, "data", self.data_dtype, None, compression=compression
294+
)
295+
self.indices = _growable_like(
296+
self.group, "indices", self.index_dtype, None, compression=compression
297+
)
298+
self.counts = np.zeros(shape[0], dtype=np.int64)
299+
self.nnz = 0
300+
301+
def append(self, rows: np.ndarray, cols: np.ndarray, values: np.ndarray) -> None:
302+
from adata.core import subset
303+
304+
if rows.size == 0:
305+
return
306+
if self.data_dtype.kind in "iu":
307+
info = np.iinfo(self.data_dtype)
308+
if values.min() < info.min or values.max() > info.max:
309+
raise ValueError(
310+
f"A value does not fit {self.data_dtype}: summing duplicate "
311+
f"entries reached {int(values.max()):,}."
312+
)
313+
count_keys(self.counts, rows)
314+
subset._append(self.data, values.astype(self.data_dtype, copy=False))
315+
subset._append(self.indices, cols.astype(self.index_dtype, copy=False))
316+
self.nnz += int(rows.size)
317+
318+
def finish(self) -> int:
319+
from adata.core.subset import _sparse_dataset
320+
321+
dtype = np.int32 if self.nnz <= INT32_MAX else np.int64
322+
indptr = np.concatenate(([0], np.cumsum(self.counts))).astype(dtype)
323+
dataset = _sparse_dataset(
324+
self.group, "indptr", np.dtype(dtype), indptr.size, None,
325+
compression=self.compression,
326+
)
327+
dataset[:] = indptr
328+
return self.nnz

0 commit comments

Comments
 (0)