Skip to main content

data_beans/sparse_io/
traits.rs

1#![allow(dead_code, unused_imports)]
2
3#[cfg(feature = "tensor")]
4pub use legume_numeric::candle_core::Tensor;
5pub use nalgebra::DMatrix;
6pub use nalgebra_sparse::{csc::CscMatrix, csr::CsrMatrix};
7pub use ndarray::prelude::*;
8
9pub const MAX_ROW_NAME_IDX: usize = 3;
10pub const MAX_COLUMN_NAME_IDX: usize = 10;
11pub const COLUMN_SEP: &str = "@";
12pub const ROW_SEP: &str = "_";
13
14use super::helpers::*;
15
16use crate::sparse_data_visitors::styled_progress_bar;
17use clap::ValueEnum;
18use indicatif::ParallelProgressIterator;
19use legume_numeric::matrix::mtx_io::*;
20use legume_numeric::matrix::traits::*;
21use log::info;
22use rayon::prelude::*;
23use rustc_hash::FxHashMap as HashMap;
24use std::ops::Range;
25use std::sync::{Arc, Mutex};
26
27#[cfg(test)]
28mod tests;
29
30#[derive(ValueEnum, Clone, Debug, PartialEq)]
31#[clap(rename_all = "lowercase")]
32pub enum SparseIoBackend {
33    Zarr,
34    HDF5,
35}
36
37/// Identifies one of the six 1-D datasets inside a sparse backend.
38/// Used by the streaming write API so we don't have to add six separate
39/// abstract methods per dtype × (csc|csr) × (data|indices|indptr).
40#[derive(Clone, Copy, Debug, PartialEq, Eq)]
41pub enum CsKey {
42    CscData,
43    CscIndices,
44    CscIndptr,
45    CsrData,
46    CsrIndices,
47    CsrIndptr,
48}
49
50/// Entries per slab handed to the backend while a sorted triplet vector is
51/// streamed out as CSC or CSR. Bounds the staging buffers; the triplets
52/// themselves are the only full-size structure alive at that point.
53const SLAB_NNZ: usize = 1 << 20;
54
55/// End of the slab that starts at triplet `start` of a vector sorted on the
56/// major axis `major` (column for CSC, row for CSR): at least `slab_nnz`
57/// entries, or all that remain, and never splitting a major index. Returns
58/// `(end, band_end)` -- the exclusive triplet index and the exclusive major
59/// bound -- so consecutive slabs tile `0..n_major` with no gap, empty
60/// columns or rows included; the last slab runs to `n_major`.
61fn slab_end(
62    triplets: &[(u64, u64, f32)],
63    start: usize,
64    slab_nnz: usize,
65    n_major: usize,
66    major: impl Fn(&(u64, u64, f32)) -> u64,
67) -> (usize, u64) {
68    debug_assert!(slab_nnz > 0);
69    let nnz = triplets.len();
70    let mut end = (start + slab_nnz).min(nnz);
71    while end < nnz && major(&triplets[end]) == major(&triplets[end - 1]) {
72        end += 1;
73    }
74    let band_end = if end == nnz {
75        n_major as u64
76    } else {
77        major(&triplets[end])
78    };
79    (end, band_end)
80}
81
82pub trait SparseIo: Sync + Send {
83    type IndexIter: IntoIterator<Item = usize> + FromIterator<usize>;
84
85    ////////////////////////////
86    // default implementation //
87    ////////////////////////////
88
89    /// Read columns within the range and return dense `ndarray::Array2`
90    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
91    ///
92    fn read_columns_ndarray(&self, columns: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
93        let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
94        Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
95    }
96
97    #[cfg(feature = "tensor")]
98    /// Read columns within the range and return dense `candle_core::Tensor`
99    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
100    ///
101    fn read_columns_tensor(&self, columns: Self::IndexIter) -> anyhow::Result<Tensor> {
102        let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
103        Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
104    }
105
106    /// Read columns within the range and return dense `nalgebrea::DMatrix`
107    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
108    ///
109    fn read_columns_dmatrix(&self, columns: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
110        let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
111        DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
112    }
113
114    /// Read columns within the range and return sparse `CsrMatrix`
115    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
116    ///
117    fn read_columns_csr(&self, columns: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
118        let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
119        CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
120    }
121
122    /// Read columns within the range and return sparse `CsrMatrix`
123    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
124    ///
125    fn read_columns_csc(&self, columns: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
126        let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
127        CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
128    }
129
130    /// Zero-copy view of preloaded column-major CSC arrays as
131    /// `(indptr, indices, data)`. Returns `None` when the backend has
132    /// not preloaded columns or doesn't support direct array access.
133    /// Callers (e.g. `SparseIoVec::read_columns_csc`) use this to skip
134    /// the triplet roundtrip when columns are already in memory.
135    fn csc_column_arrays(&self) -> Option<(&[u64], &[u64], &[f32])> {
136        None
137    }
138
139    /// Read rows within the range and return dense `ndarray::Array2`
140    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
141    ///
142    fn read_rows_ndarray(&self, rows: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
143        let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
144        Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
145    }
146
147    #[cfg(feature = "tensor")]
148    /// Read rows within the range and return dense `candle_core::Tensor`
149    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
150    ///
151    fn read_rows_tensor(&self, rows: Self::IndexIter) -> anyhow::Result<Tensor> {
152        let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
153        Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
154    }
155
156    /// Read rows within the range and return dense `nalgebra::DMatrix`
157    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
158    ///
159    fn read_rows_dmatrix(&self, rows: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
160        let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
161        DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
162    }
163
164    /// Read rows within the range and return sparse `CsrMatrix`
165    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
166    ///
167    fn read_rows_csr(&self, rows: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
168        let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
169        CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
170    }
171
172    /// Read rows within the range and return sparse `CscMatrix`
173    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
174    ///
175    fn read_rows_csc(&self, rows: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
176        let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
177        CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
178    }
179
180    /////////////////////////////
181    // `mtx` related functions //
182    /////////////////////////////
183
184    /// Read an mtx file once and populate the backend: the column (CSC) index
185    /// always, the row (CSR) index as well when `index_by_row`. Both are
186    /// streamed out of the same triplet vector, so the file is inflated once
187    /// and the triplets are the only full-size structure alive.
188    /// * `mtx_file`: mtx file to be read into the backend
189    fn import_mtx_file(&mut self, mtx_file: &str, index_by_row: bool) -> anyhow::Result<()> {
190        let (mut mtx_triplets, mtx_shape) = read_mtx_triplets(mtx_file)?;
191        info!("read mtx file: {}", mtx_file);
192        if mtx_triplets.is_empty() {
193            return Err(anyhow::anyhow!("No data in mtx file"));
194        }
195        self.record_mtx_shape(Some(mtx_shape))?;
196        info!("recording the column index");
197        self.record_triplets_by_col(&mut mtx_triplets)?;
198        if index_by_row {
199            info!("recording the row index");
200            self.record_triplets_by_row(&mut mtx_triplets)?;
201        }
202        Ok(())
203    }
204
205    /////////////////////////////////
206    // `dmatrix` related functions //
207    /////////////////////////////////
208
209    /// Add dmatrix to zarr backend by row (CSR format)
210    /// * `array` - 2D array to be added to the backend
211    fn import_dmatrix_by_row(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
212        let (nrow, ncol) = matrix.shape();
213        let mut mtx_triplets = dmatrix_to_triplets(matrix);
214        let mtx_shape = (nrow, ncol, mtx_triplets.len());
215        self.record_mtx_shape(Some(mtx_shape))?;
216        self.record_triplets_by_row(&mut mtx_triplets)
217    }
218
219    /// Add dmatrix to zarr backend by column (CSC format)
220    /// * `array` - 2D array to be added to the backend
221    fn import_dmatrix_by_col(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
222        let (nrow, ncol) = matrix.shape();
223        let mut mtx_triplets = dmatrix_to_triplets(matrix);
224        let mtx_shape = (nrow, ncol, mtx_triplets.len());
225        self.record_mtx_shape(Some(mtx_shape))?;
226        self.record_triplets_by_col(&mut mtx_triplets)
227    }
228
229    /////////////////////////////////
230    // `ndarray` related functions //
231    /////////////////////////////////
232
233    /// Add ndarray to zarr backend by row (CSR format)
234    /// * `array` - 2D array to be added to the backend
235    fn import_ndarray_by_row(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
236        let nrow = array.shape()[0];
237        let ncol = array.shape()[1];
238
239        // dbg!("importing ndarray by row...");
240        let mut mtx_triplets = ndarray_to_triplets(array);
241
242        let nnz = mtx_triplets.len();
243        let mtx_shape = (nrow, ncol, nnz);
244        self.record_mtx_shape(Some(mtx_shape))?;
245
246        // dbg!(format!("populated: {} elements", mtx_triplets.len()));
247
248        self.record_triplets_by_row(&mut mtx_triplets)
249    }
250
251    /// Add ndarray to zarr backend by column (CSC format)
252    /// * `array` - 2D array to be added to the backend
253    fn import_ndarray_by_col(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
254        let nrow = array.shape()[0];
255        let ncol = array.shape()[1];
256
257        // dbg!("importing ndarray by column...");
258        let mut mtx_triplets = ndarray_to_triplets(array);
259
260        let nnz = mtx_triplets.len();
261        let mtx_shape = (nrow, ncol, nnz);
262        self.record_mtx_shape(Some(mtx_shape))?;
263
264        // dbg!(format!("populated: {} elements", mtx_triplets.len()));
265
266        self.record_triplets_by_col(&mut mtx_triplets)
267    }
268
269    //////////////////////
270    // backend-specific //
271    //////////////////////
272
273    /// Read rows within the range and return a vector of triplets (row, column, value)
274    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
275    ///
276    #[allow(clippy::type_complexity)]
277    fn read_triplets_by_rows(
278        &self,
279        rows: Self::IndexIter,
280    ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
281
282    /// Read columns within the range and return a vector of triplets (row, col, value)
283    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
284    ///
285    #[allow(clippy::type_complexity)]
286    fn read_triplets_by_columns(
287        &self,
288        columns: Self::IndexIter,
289    ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
290
291    /// Read columns within the range and return a vector of triplets (row, col, value)
292    /// * `col` : usize
293    ///
294    #[allow(clippy::type_complexity)]
295    fn read_triplets_by_single_column(
296        &self,
297        col: usize,
298    ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;
299
300    /// Export the data to a mtx file. This will take time.
301    /// * `mtx_file`: mtx file to be written
302    fn to_mtx_file(&self, mtx_file: &str) -> anyhow::Result<()>;
303
304    /// Number of rows in the underlying data matrix
305    fn num_rows(&self) -> Option<usize>;
306
307    /// Number of columns in the underlying data matrix
308    fn num_columns(&self) -> Option<usize>;
309
310    /// Number of non-zero elements
311    fn num_non_zeros(&self) -> Option<usize>;
312
313    /// Re-open handles on the CURRENT backend path after its contents were
314    /// replaced from outside (a finished temp file renamed into place). The
315    /// zarr store is path-addressed so this is a cache refresh; hdf5 holds an
316    /// open file handle that would otherwise point at the deleted inode.
317    fn reopen_backend(&mut self) -> anyhow::Result<()>;
318
319    /// Maintained by [`append_csc_slab`](Self::append_csc_slab) — never call
320    /// this yourself: padding the cursor masks exactly the under-append the
321    /// finalize audit exists to catch.
322    #[doc(hidden)]
323    /// Advance the streaming-write cursor by `n` entries. Called by
324    /// [`append_csc_slab`](Self::append_csc_slab); backends keep the count so
325    /// [`finalize_streaming_csc`](Self::finalize_streaming_csc) can audit the
326    /// declared nnz against what was actually appended — the one violation the
327    /// written indptr cannot reveal, because an over-declared tail leaves it
328    /// perfectly monotone with the phantom hiding between the last written
329    /// pointer and the sentinel.
330    fn note_streamed_nnz(&mut self, n: u64);
331
332    /// Entries appended so far in this streaming build.
333    #[doc(hidden)]
334    fn streamed_nnz(&self) -> u64;
335
336    /// Zero the cursor. Called by [`begin_streaming_csc`](Self::begin_streaming_csc).
337    #[doc(hidden)]
338    fn reset_streamed_nnz(&mut self);
339
340    /// The resident by-column indptr, loaded at `open()`. Empty when the
341    /// backend carries no `/by_column/indptr` array — `read_column_indptr`
342    /// silently does nothing on that failure, and the accessors below must
343    /// report that absence rather than read zeros out of it.
344    fn column_indptr(&self) -> &[u64];
345
346    /// Exact nnz of one column, from the resident indptr — no I/O.
347    ///
348    /// `None` for an out-of-range column or when the indptr is absent. This is
349    /// what lets a streaming writer declare a column subset's total nnz up
350    /// front without a counting pass over the data.
351    fn column_nnz(&self, col: usize) -> Option<u64> {
352        let indptr = self.column_indptr();
353        let hi = *indptr.get(col + 1)?;
354        let lo = *indptr.get(col)?;
355        hi.checked_sub(lo)
356    }
357
358    /// Set row names for the matrix
359    /// * `row_name_file`: a file each line contains row name words
360    fn register_row_names_file(&mut self, row_name_file: &str);
361
362    /// Set column names for the matrix
363    /// * `column_name_file`: a file each line contains column name words
364    fn register_column_names_file(&mut self, column_name_file: &str);
365
366    /// Set row names for the matrix
367    /// * `rows`: a vector of row names
368    fn register_row_names_vec(&mut self, rows: &[Box<str>]);
369
370    /// Set column names for the matrix
371    /// * `columns`: a vector of column names
372    fn register_column_names_vec(&mut self, columns: &[Box<str>]);
373
374    /// Add arbitrary names (a vector of strings)
375    /// * `group_name`: group name
376    /// * `name_file`: a file each line contains name words
377    /// * `name_columns`: range of columns to be used for name
378    /// * `name_sep`: separator for name columns
379    fn register_names_file(
380        &mut self,
381        key: &str,
382        name_file: &str,
383        name_columns: Range<usize>,
384        name_sep: &str,
385    ) -> anyhow::Result<()>;
386
387    /// Add arbitrary names (a vector of strings)
388    /// * `group_name`: group name
389    /// * `names`: a file each line contains name words
390    fn register_names_vec(&mut self, key: &str, names: &[Box<str>]) -> anyhow::Result<()>;
391
392    fn row_names(&self) -> anyhow::Result<Vec<Box<str>>>;
393
394    fn column_names(&self) -> anyhow::Result<Vec<Box<str>>>;
395
396    /// Get back the registered names
397    /// * `key`: key for the registered names
398    fn retrieve_registered_names(&self, key: &str) -> anyhow::Result<Vec<Box<str>>>;
399
400    /////////////////////////////
401    // major structural change //
402    /////////////////////////////
403
404    /// Select the columns of the data and create a new backend file
405    /// * `columns`: columns to be subsetted
406    /// * `rows`: if something, subset the rows
407    fn subset_columns_rows(
408        &mut self,
409        columns: Option<&Vec<usize>>,
410        rows: Option<&Vec<usize>>,
411    ) -> anyhow::Result<()> {
412        let ncol_data = self
413            .num_columns()
414            .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
415        let nrow_data = self
416            .num_rows()
417            .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
418
419        // An empty selection is refused, not honoured: this method DESTROYS the
420        // original backend, and writing a zero-column husk over real data is
421        // almost certainly a caller mistake rather than an intention.
422        // Empty and duplicated selections are refused, not honoured: this
423        // method DESTROYS the original, and a duplicate collapses in the
424        // old→new map while still counting toward the new shape — slabs then
425        // land at wrong offsets and the finalize audit rejects the build with a
426        // message about nnz tiling that names nothing the caller did.
427        let distinct = |sel: &[usize], what: &str| -> anyhow::Result<()> {
428            anyhow::ensure!(!sel.is_empty(), "subset: empty {what} selection");
429            let mut seen = sel.to_vec();
430            seen.sort_unstable();
431            seen.dedup();
432            anyhow::ensure!(
433                seen.len() == sel.len(),
434                "subset: the {what} selection repeats an index ({} of {} are distinct)",
435                seen.len(),
436                sel.len()
437            );
438            Ok(())
439        };
440        if let Some(cols) = columns {
441            distinct(cols, "column")?;
442        }
443        if let Some(rs) = rows {
444            distinct(rs, "row")?;
445        }
446
447        //////////////////////////////////////////////////////
448        // 0. Create a mapping from old to new columns/rows //
449        //////////////////////////////////////////////////////
450
451        let (old2new_cols, new_col_names) =
452            take_subset_indices_names_if_needed(columns, Some(ncol_data), self.column_names()?);
453        let (old2new_rows, new_row_names) =
454            take_subset_indices_names_if_needed(rows, Some(nrow_data), self.row_names()?);
455        let (new_ncol, new_nrow) = (new_col_names.len(), new_row_names.len());
456        anyhow::ensure!(new_ncol > 0, "subset: no column survived the selection");
457        anyhow::ensure!(new_nrow > 0, "subset: no row survived the selection");
458
459        // Old columns in NEW order — the selection's order is the output order.
460        let mut cols_new_order: Vec<(u64, u64)> =
461            old2new_cols.iter().map(|(&o, &n)| (n, o)).collect();
462        cols_new_order.sort_unstable();
463
464        // Dense old-row → new-row map, and whether it preserves order. A
465        // monotone map keeps within-column rows ascending after renumbering, so
466        // no per-column sort is needed; a reordering map costs one small sort
467        // per column.
468        let mut row_map: Vec<Option<u64>> = vec![None; nrow_data];
469        for (&old, &new) in &old2new_rows {
470            row_map[old as usize] = Some(new);
471        }
472        let monotone_rows = row_map.iter().flatten().is_sorted_by(|a, b| a < b);
473
474        ///////////////////////////////////////////////////////
475        // 1. Exact per-new-column nnz, without materialising //
476        ///////////////////////////////////////////////////////
477
478        // No row filter: straight off the resident indptr, zero I/O. With one:
479        // a counting pass — reads every selected column once and keeps counts,
480        // never entries.
481        let full_rows = rows.is_none();
482        let per_col_nnz: Vec<u64> = if full_rows {
483            cols_new_order
484                .iter()
485                .map(|&(_, old)| {
486                    self.column_nnz(old as usize)
487                        .ok_or_else(|| anyhow::anyhow!("subset: no indptr for column {old}"))
488                })
489                .collect::<anyhow::Result<_>>()?
490        } else {
491            // Block reads, never one column at a time: a single-column read
492            // pays the cached-subset machinery per call, which measured two
493            // orders of magnitude slower at imaging scale. Counts only, never
494            // entries.
495            let mut counts = vec![0u64; cols_new_order.len()];
496            let coarse = legume_numeric::matrix::utils::generate_minibatch_intervals(
497                cols_new_order.len(),
498                0,
499                Some(8192),
500            );
501            for (lb, ub) in coarse {
502                let old_cols: Vec<usize> = cols_new_order[lb..ub]
503                    .iter()
504                    .map(|&(_, o)| o as usize)
505                    .collect();
506                let (_, _, triplets) =
507                    self.read_triplets_by_columns(old_cols.into_iter().collect())?;
508                for (i, c_local, _) in triplets {
509                    if row_map[i as usize].is_some() {
510                        counts[lb + c_local as usize] += 1;
511                    }
512                }
513            }
514            counts
515        };
516        let new_nnz: u64 = per_col_nnz.iter().sum();
517
518        ///////////////////////////////////////////////////////////
519        // 2. Stream the survivors into a TEMPORARY sibling file //
520        ///////////////////////////////////////////////////////////
521
522        // Written beside the original (same filesystem, so the final rename is
523        // atomic) and swapped in only when complete. The old implementation
524        // deleted the original FIRST and rewrote it from a RAM buffer, so a
525        // crash mid-write lost the data outright — and that buffer held every
526        // surviving triplet, which is the memory wall this replaces.
527        // CONTRACT of the swap: the original is untouched until the temporary
528        // sibling is complete and finalized; a failure mid-stream leaves the
529        // original intact plus a `{path}.subset_tmp` leftover (cleaned on the
530        // next attempt); the unrecoverable window is only remove→rename below.
531        // A zip-archived backend is refused up front — streaming to a sibling
532        // DIRECTORY and renaming it over the `.zip` name would silently change
533        // the on-disk format under the old extension, and the old code's
534        // "store is read-only" failure was at least loud.
535        let final_path = self.get_backend_file_name().to_string();
536        anyhow::ensure!(
537            !final_path.ends_with(".zip"),
538            "subset: {final_path} is a zip archive; convert it to a directory \
539             backend first (data-beans convert)"
540        );
541        let temp_path = format!("{final_path}.subset_tmp");
542        if std::path::Path::new(&temp_path).exists() {
543            crate::sparse_io::remove_backend_path(&temp_path)?;
544        }
545
546        {
547            let backend_kind = self.backend_type();
548            let mut out = crate::sparse_io::create_sparse_streaming_empty(
549                Some(&temp_path),
550                Some(&backend_kind),
551            )?;
552            out.begin_streaming_csc((new_nrow, new_ncol, new_nnz as usize))?;
553
554            // Blocks bounded by bytes of surviving triplets, from the measured
555            // per-column counts — not by a fixed column count.
556            let blocks = legume_numeric::matrix::utils::byte_budget_intervals(
557                &per_col_nnz,
558                crate::sparse_io::SLAB_BUDGET_BYTES,
559                crate::sparse_io::TRIPLET_BYTES,
560            );
561
562            let mut nnz_offset = 0u64;
563            for (lb, ub) in blocks {
564                // ONE block read per slab (see the counting pass above for why).
565                // The read returns LOCAL column ids in the requested order, rows
566                // ascending within each column.
567                let old_cols: Vec<usize> = cols_new_order[lb..ub]
568                    .iter()
569                    .map(|&(_, o)| o as usize)
570                    .collect();
571                let (_, _, triplets) =
572                    self.read_triplets_by_columns(old_cols.into_iter().collect())?;
573
574                let n_block = ub - lb;
575                let mut per_col: Vec<Vec<(u64, f32)>> = vec![Vec::new(); n_block];
576                for (i, c_local, x) in triplets {
577                    if let Some(new_row) = row_map[i as usize] {
578                        per_col[c_local as usize].push((new_row, x));
579                    }
580                }
581                let mut local_colptr = Vec::with_capacity(n_block);
582                let mut row_indices = Vec::new();
583                let mut values = Vec::new();
584                for entries in &mut per_col {
585                    if !monotone_rows {
586                        // Renumbering scrambled this column's order; restore the
587                        // ascending-rows invariant the writer enforces.
588                        entries.sort_unstable_by_key(|&(r, _)| r);
589                    }
590                    local_colptr.push(row_indices.len() as u64);
591                    for &(r, x) in entries.iter() {
592                        row_indices.push(r);
593                        values.push(x);
594                    }
595                }
596                out.append_csc_slab(lb as u64, nnz_offset, &local_colptr, &row_indices, &values)?;
597                nnz_offset += values.len() as u64;
598            }
599
600            out.finalize_streaming_csc()?;
601            out.build_csr_from_csc_streaming()?;
602            out.register_row_names_vec(&new_row_names);
603            out.register_column_names_vec(&new_col_names);
604        }
605
606        ////////////////////////////////////
607        // 3. Swap the finished file in  //
608        ////////////////////////////////////
609
610        self.remove_backend_file()?;
611        std::fs::rename(&temp_path, &final_path)?;
612        self.reopen_backend()?;
613        self.clean_preloaded_columns();
614        self.clean_preloaded_rows();
615        info!("registered new data to {}", self.get_backend_file_name());
616        Ok(())
617    }
618
619    /// Reposition rows in a new order specified by `remap`
620    /// * `row_names_order` - a vector of row names in the new order
621    fn reorder_rows(&mut self, row_names_order: &[Box<str>]) -> anyhow::Result<()> {
622        let new_col_names = self.column_names()?.clone();
623        let name2new = build_name2index_map(row_names_order);
624
625        let block_size = 100;
626
627        let old2new: HashMap<u64, u64> = self
628            .row_names()?
629            .into_par_iter()
630            .enumerate()
631            .filter_map(|(idx_old, name)| {
632                name2new
633                    .get(&name)
634                    .map(|&idx_new| (idx_old as u64, idx_new as u64))
635            })
636            .collect();
637
638        if let Some(ncol) = self.num_columns() {
639            /////////////////////////////////////////////////////
640            // 1. triplets after filtering and reordering rows //
641            /////////////////////////////////////////////////////
642
643            let arc_triplets = Arc::new(Mutex::new(vec![]));
644
645            let nblock = ncol.div_ceil(block_size);
646
647            info!("remapping triplets ...");
648
649            (0..nblock)
650                .into_par_iter()
651                .progress_with(styled_progress_bar(nblock as u64, "blocks"))
652                .map(|b| {
653                    let lb = (b * block_size) as u64;
654                    let ub = ((b + 1) * block_size).min(ncol) as u64;
655                    (lb, ub)
656                })
657                .for_each(|(lb, ub)| {
658                    let (_, _, _triplets_b) = self
659                        .read_triplets_by_columns(((lb as usize)..(ub as usize)).collect())
660                        .unwrap();
661
662                    let _triplets_b = _triplets_b.into_iter().filter_map(|(i, j_loc, x)| {
663                        let j_glob = j_loc + lb;
664                        old2new.get(&i).map(|&i_new| (i_new, j_glob, x))
665                    });
666
667                    {
668                        let mut triplets = arc_triplets.lock().unwrap();
669                        triplets.extend(_triplets_b);
670                    }
671                });
672
673            /////////////////////////////////////
674            // 2. Remove previous backend file //
675            /////////////////////////////////////
676            self.remove_backend_file()?;
677
678            ///////////////////////////////
679            // 3. populate a new backend //
680            ///////////////////////////////
681            self.initialize_backend()?;
682
683            // populate data from mtx triplets
684            {
685                let mut row_col_val_triplets =
686                    arc_triplets.lock().expect("failed to lock triplets");
687
688                let nnz = row_col_val_triplets.len();
689                debug_assert!(row_col_val_triplets.len() <= nnz); // subset
690                let new_nrow = row_names_order.len();
691                let mtx_shape = (new_nrow, ncol, nnz);
692
693                info!("sorting triplets ...");
694
695                self.record_mtx_shape(Some(mtx_shape))?;
696                self.record_triplets_by_col(&mut row_col_val_triplets)?;
697                self.record_triplets_by_row(&mut row_col_val_triplets)?;
698            }
699            self.read_column_indptr()?;
700            self.read_row_indptr()?;
701
702            self.register_row_names_vec(row_names_order);
703            self.register_column_names_vec(&new_col_names);
704            info!("registered new data to {}", self.get_backend_file_name());
705        }
706
707        self.clean_preloaded_columns();
708        self.clean_preloaded_rows();
709        Ok(())
710    }
711    // fn reorder_rows(&mut self, row_names_order: &[Box<str>]) -> anyhow::Result<()>;
712
713    /// Remove backend file
714    fn remove_backend_file(&self) -> anyhow::Result<()>;
715
716    /// Initialize backend
717    fn initialize_backend(&mut self) -> anyhow::Result<()>;
718
719    fn record_mtx_shape(&mut self, mtx_shape: Option<(usize, usize, usize)>) -> anyhow::Result<()>;
720
721    /// Stream the triplets out as CSR slabs; the row-major twin of
722    /// [`record_triplets_by_col`](Self::record_triplets_by_col).
723    fn record_triplets_by_row(
724        &mut self,
725        row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
726    ) -> anyhow::Result<()> {
727        let nrow = self.num_rows().expect("should have `nrow`");
728        let ncol = self.num_columns().expect("should have `ncol`");
729        let nnz = row_col_val_triplets.len();
730
731        if nnz == 0 {
732            let csr_rowptr = vec![0u64; nrow + 1];
733            return self.record_csr_dataset_backend(&[], &[], &csr_rowptr);
734        }
735
736        // One in-place pass on the full key. A stable sort would allocate a
737        // scratch copy of the whole vector, and duplicate coordinates carry no
738        // meaning in coordinate format, so stability buys nothing.
739        row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (row, col));
740
741        self.begin_streaming_csr((nrow, ncol, nnz))?;
742
743        let mut local_rowptr: Vec<u64> = Vec::new();
744        let mut cols: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
745        let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
746
747        let mut start = 0_usize;
748        let mut row_offset = 0_u64;
749        while (row_offset as usize) < nrow {
750            let (end, band_end_row) =
751                slab_end(row_col_val_triplets, start, SLAB_NNZ, nrow, |t| t.0);
752
753            local_rowptr.clear();
754            cols.clear();
755            vals.clear();
756            let mut i = start;
757            for row in row_offset..band_end_row {
758                local_rowptr.push((i - start) as u64);
759                while i < end && row_col_val_triplets[i].0 == row {
760                    cols.push(row_col_val_triplets[i].1);
761                    vals.push(row_col_val_triplets[i].2);
762                    i += 1;
763                }
764            }
765            debug_assert_eq!(i, end, "every entry of the band belongs to one of its rows");
766
767            self.append_csr_slab(row_offset, start as u64, &local_rowptr, &cols, &vals)?;
768            start = end;
769            row_offset = band_end_row;
770        }
771
772        self.finalize_streaming_csr()
773    }
774
775    /// Stream the triplets out as CSC slabs.
776    ///
777    /// After the one in-place sort the triplet vector is the only full-size
778    /// structure alive; the slab staging buffers are bounded by
779    /// [`SLAB_NNZ`], and the backend's streaming audits check the column
780    /// tiling and the appended count on the way through.
781    fn record_triplets_by_col(
782        &mut self,
783        row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
784    ) -> anyhow::Result<()> {
785        let nrow = self.num_rows().expect("should have `nrow`");
786        let ncol = self.num_columns().expect("should have `ncol`");
787        let nnz = row_col_val_triplets.len();
788
789        if nnz == 0 {
790            let csc_colptr = vec![0u64; ncol + 1];
791            return self.record_csc_dataset_backend(&[], &[], &csc_colptr);
792        }
793
794        // See `record_triplets_by_row` for why this is one unstable pass.
795        row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (col, row));
796
797        self.begin_streaming_csc((nrow, ncol, nnz))?;
798
799        let mut local_colptr: Vec<u64> = Vec::new();
800        let mut rows: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
801        let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);
802
803        let mut start = 0_usize;
804        let mut col_offset = 0_u64;
805        while (col_offset as usize) < ncol {
806            let (end, band_end_col) =
807                slab_end(row_col_val_triplets, start, SLAB_NNZ, ncol, |t| t.1);
808
809            local_colptr.clear();
810            rows.clear();
811            vals.clear();
812            let mut i = start;
813            for col in col_offset..band_end_col {
814                local_colptr.push((i - start) as u64);
815                while i < end && row_col_val_triplets[i].1 == col {
816                    rows.push(row_col_val_triplets[i].0);
817                    vals.push(row_col_val_triplets[i].2);
818                    i += 1;
819                }
820            }
821            debug_assert_eq!(
822                i, end,
823                "every entry of the band belongs to one of its columns"
824            );
825
826            self.append_csc_slab(col_offset, start as u64, &local_colptr, &rows, &vals)?;
827            start = end;
828            col_offset = band_end_col;
829        }
830
831        self.finalize_streaming_csc()
832    }
833
834    /// CSR data structure in Zarr backend
835    ///
836    /// ```text
837    ///     └── by_row
838    ///         ├── data
839    ///         ├── indices (column indices)
840    ///         └── isndptr (row pointers)
841    /// ```
842    fn record_csr_dataset_backend(
843        &mut self,
844        csr_cols: &[u64],
845        csr_vals: &[f32],
846        csr_rowptr: &[u64],
847    ) -> anyhow::Result<()>;
848
849    /// Helper function to add CSC dataset to HDF5 backend
850    ///
851    /// ```text
852    /// Helper function to record the CSC dataset
853    ///     ├── by_column
854    ///     │   ├── data
855    ///     │   ├── indices (row indices)
856    ///     │   └── indptr (column pointers)
857    /// ```
858    fn record_csc_dataset_backend(
859        &mut self,
860        csc_rows: &[u64],
861        csc_vals: &[f32],
862        csc_colptr: &[u64],
863    ) -> anyhow::Result<()>;
864
865    /// Create a fixed-size 1-D backend dataset of `len` elements for the
866    /// given CSC/CSR slot. No data is written yet.
867    fn cs_create(&mut self, key: CsKey, len: usize) -> anyhow::Result<()>;
868
869    /// Write a `u64` slab at `offset` in the specified dataset.
870    /// Used for CSC/CSR `indices` and `indptr`.
871    fn cs_write_u64(&mut self, key: CsKey, offset: u64, data: &[u64]) -> anyhow::Result<()>;
872
873    /// Write an `f32` slab at `offset` in the specified dataset.
874    /// Used for CSC/CSR `data`.
875    fn cs_write_f32(&mut self, key: CsKey, offset: u64, data: &[f32]) -> anyhow::Result<()>;
876
877    /// Begin a streaming CSC build for a sparse matrix of known shape.
878    /// Pre-creates `/by_column/{data, indices, indptr}` at their final sizes
879    /// so subsequent [`append_csc_slab`](Self::append_csc_slab) calls write
880    /// into disjoint hyperslabs without further allocation.
881    fn begin_streaming_csc(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
882        // A reused handle must not inherit a previous build's cursor: the
883        // finalize audit compares appended-vs-declared, and a stale count turns
884        // a correct build into a false accusation.
885        self.reset_streamed_nnz();
886        let (_, ncol, nnz) = shape;
887        self.record_mtx_shape(Some(shape))?;
888        self.cs_create(CsKey::CscData, nnz)?;
889        self.cs_create(CsKey::CscIndices, nnz)?;
890        self.cs_create(CsKey::CscIndptr, ncol + 1)?;
891        Ok(())
892    }
893
894    /// Append one contiguous CSC column band.
895    ///
896    /// * `col_offset` — global column index where this band starts
897    /// * `nnz_offset` — global nnz offset where this band's values land
898    /// * `local_colptr` — length `batch_ncol`, values in `[0, batch_nnz]`,
899    ///   will be shifted by `nnz_offset` before writing
900    /// * `row_indices` — length `batch_nnz`
901    /// * `values`      — length `batch_nnz`
902    fn append_csc_slab(
903        &mut self,
904        col_offset: u64,
905        nnz_offset: u64,
906        local_colptr: &[u64],
907        row_indices: &[u64],
908        values: &[f32],
909    ) -> anyhow::Result<()> {
910        // These checks exist because a violation does NOT fail loudly on its
911        // own: unwritten regions read back as the zarr fill value, so a bad
912        // slab yields a backend that opens and reads cleanly while carrying
913        // poisoned or duplicated entries. Cheap (one pass over the slab, no
914        // I/O) next to the compressed writes below.
915        anyhow::ensure!(
916            row_indices.len() == values.len(),
917            "append_csc_slab: {} row indices vs {} values",
918            row_indices.len(),
919            values.len()
920        );
921        anyhow::ensure!(
922            local_colptr.first().copied() == Some(0) || local_colptr.is_empty(),
923            "append_csc_slab: local_colptr must start at 0"
924        );
925        anyhow::ensure!(
926            local_colptr.windows(2).all(|w| w[0] <= w[1]),
927            "append_csc_slab: local_colptr must be monotone non-decreasing"
928        );
929        if let Some(&last) = local_colptr.last() {
930            anyhow::ensure!(
931                last <= values.len() as u64,
932                "append_csc_slab: colptr claims {last} entries, slab holds {}",
933                values.len()
934            );
935        }
936        if let Some(nrow) = self.num_rows() {
937            if let Some(&bad) = row_indices.iter().find(|&&r| r >= nrow as u64) {
938                anyhow::bail!("append_csc_slab: row index {bad} outside the {nrow}-row matrix");
939            }
940        }
941        // Ascending rows within each column: readers document it as an
942        // invariant, and the h5ad export hands the arrays to scipy as-is.
943        for (c, &start) in local_colptr.iter().enumerate() {
944            let end = local_colptr
945                .get(c + 1)
946                .copied()
947                .unwrap_or(values.len() as u64) as usize;
948            anyhow::ensure!(
949                row_indices[start as usize..end]
950                    .windows(2)
951                    .all(|w| w[0] < w[1]),
952                "append_csc_slab: rows within column {} of this band must be \
953                 strictly ascending — repeated rows usually mean duplicate \
954                 (row, col) coordinates in the source (an MTX with repeated \
955                 entries, or a union remap folding rows together)",
956                col_offset as usize + c
957            );
958        }
959
960        let shifted: Vec<u64> = local_colptr.iter().map(|&p| p + nnz_offset).collect();
961        self.cs_write_u64(CsKey::CscIndptr, col_offset, &shifted)?;
962        self.cs_write_u64(CsKey::CscIndices, nnz_offset, row_indices)?;
963        self.cs_write_f32(CsKey::CscData, nnz_offset, values)?;
964        self.note_streamed_nnz(values.len() as u64);
965        Ok(())
966    }
967
968    /// Finalize CSC streaming by writing the final indptr sentinel at
969    /// position `ncol`, equal to the total nnz.
970    fn finalize_streaming_csc(&mut self) -> anyhow::Result<()> {
971        let ncol = self
972            .num_columns()
973            .ok_or_else(|| anyhow::anyhow!("ncol not set before finalize_streaming_csc"))?;
974        let nnz = self
975            .num_non_zeros()
976            .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csc"))?;
977        self.cs_write_u64(CsKey::CscIndptr, ncol as u64, &[nnz as u64])?;
978        self.read_column_indptr()?;
979
980        // The tiling check the per-slab guards cannot do. A gap or overlap in
981        // the nnz offsets, or an over-declared total, leaves the WRITTEN
982        // indptr non-monotone or short of the declared nnz — and unwritten
983        // indptr slots read back as the fill value 0, so column j would claim
984        // the whole array prefix. `indptr[ncol] - indptr[0]` equals the
985        // declaration by construction, which is why the old debug_assert on it
986        // could never fire; the shape of the vector between the endpoints is
987        // what carries the truth.
988        let indptr = self.column_indptr();
989        anyhow::ensure!(
990            indptr.len() == ncol + 1,
991            "finalize_streaming_csc: indptr has {} entries, expected {}",
992            indptr.len(),
993            ncol + 1
994        );
995        anyhow::ensure!(
996            indptr.first().copied() == Some(0),
997            "finalize_streaming_csc: indptr[0] = {:?}, expected 0 — the first \
998             slab was never appended",
999            indptr.first()
1000        );
1001        if let Some(w) = indptr.windows(2).position(|w| w[0] > w[1]) {
1002            anyhow::bail!(
1003                "finalize_streaming_csc: indptr decreases at column {w} — slabs \
1004                 were appended with a gap or overlap in their nnz offsets"
1005            );
1006        }
1007        // The appended count is the ground truth the indptr cannot carry: an
1008        // over-declared nnz leaves the written indptr perfectly monotone with
1009        // the phantom tail hiding between the last written pointer and the
1010        // sentinel — and the sentinel itself was written by this function, so
1011        // comparing against it can only ever agree.
1012        let appended = self.streamed_nnz();
1013        anyhow::ensure!(
1014            appended == nnz as u64,
1015            "finalize_streaming_csc: {appended} entries appended but {nnz} \
1016             declared — the difference reads back as fill values wearing real \
1017             entries' positions"
1018        );
1019        Ok(())
1020    }
1021
1022    /// Begin a streaming CSR build for a sparse matrix of known shape; the
1023    /// row-major twin of [`begin_streaming_csc`](Self::begin_streaming_csc).
1024    fn begin_streaming_csr(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
1025        self.reset_streamed_nnz();
1026        let (nrow, _, nnz) = shape;
1027        self.record_mtx_shape(Some(shape))?;
1028        self.cs_create(CsKey::CsrData, nnz)?;
1029        self.cs_create(CsKey::CsrIndices, nnz)?;
1030        self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1031        Ok(())
1032    }
1033
1034    /// Append one contiguous CSR row band; the row-major twin of
1035    /// [`append_csc_slab`](Self::append_csc_slab), with the same audits.
1036    ///
1037    /// * `row_offset` — global row index where this band starts
1038    /// * `nnz_offset` — global nnz offset where this band's values land
1039    /// * `local_rowptr` — length `batch_nrow`, values in `[0, batch_nnz]`,
1040    ///   will be shifted by `nnz_offset` before writing
1041    /// * `col_indices` — length `batch_nnz`
1042    /// * `values`      — length `batch_nnz`
1043    fn append_csr_slab(
1044        &mut self,
1045        row_offset: u64,
1046        nnz_offset: u64,
1047        local_rowptr: &[u64],
1048        col_indices: &[u64],
1049        values: &[f32],
1050    ) -> anyhow::Result<()> {
1051        anyhow::ensure!(
1052            col_indices.len() == values.len(),
1053            "append_csr_slab: {} column indices vs {} values",
1054            col_indices.len(),
1055            values.len()
1056        );
1057        anyhow::ensure!(
1058            local_rowptr.first().copied() == Some(0) || local_rowptr.is_empty(),
1059            "append_csr_slab: local_rowptr must start at 0"
1060        );
1061        anyhow::ensure!(
1062            local_rowptr.windows(2).all(|w| w[0] <= w[1]),
1063            "append_csr_slab: local_rowptr must be monotone non-decreasing"
1064        );
1065        if let Some(&last) = local_rowptr.last() {
1066            anyhow::ensure!(
1067                last <= values.len() as u64,
1068                "append_csr_slab: rowptr claims {last} entries, slab holds {}",
1069                values.len()
1070            );
1071        }
1072        if let Some(ncol) = self.num_columns() {
1073            if let Some(&bad) = col_indices.iter().find(|&&c| c >= ncol as u64) {
1074                anyhow::bail!(
1075                    "append_csr_slab: column index {bad} outside the {ncol}-column matrix"
1076                );
1077            }
1078        }
1079        for (r, &start) in local_rowptr.iter().enumerate() {
1080            let end = local_rowptr
1081                .get(r + 1)
1082                .copied()
1083                .unwrap_or(values.len() as u64) as usize;
1084            anyhow::ensure!(
1085                col_indices[start as usize..end]
1086                    .windows(2)
1087                    .all(|w| w[0] < w[1]),
1088                "append_csr_slab: columns within row {} of this band must be \
1089                 strictly ascending — repeated columns usually mean duplicate \
1090                 (row, col) coordinates in the source",
1091                row_offset as usize + r
1092            );
1093        }
1094
1095        let shifted: Vec<u64> = local_rowptr.iter().map(|&p| p + nnz_offset).collect();
1096        self.cs_write_u64(CsKey::CsrIndptr, row_offset, &shifted)?;
1097        self.cs_write_u64(CsKey::CsrIndices, nnz_offset, col_indices)?;
1098        self.cs_write_f32(CsKey::CsrData, nnz_offset, values)?;
1099        self.note_streamed_nnz(values.len() as u64);
1100        Ok(())
1101    }
1102
1103    /// Finalize CSR streaming: write the indptr sentinel at position `nrow`,
1104    /// load the row index, and check the appended count against the declared
1105    /// nnz -- the one violation the written indptr cannot reveal.
1106    fn finalize_streaming_csr(&mut self) -> anyhow::Result<()> {
1107        let nrow = self
1108            .num_rows()
1109            .ok_or_else(|| anyhow::anyhow!("nrow not set before finalize_streaming_csr"))?;
1110        let nnz = self
1111            .num_non_zeros()
1112            .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csr"))?;
1113        self.cs_write_u64(CsKey::CsrIndptr, nrow as u64, &[nnz as u64])?;
1114        self.read_row_indptr()?;
1115
1116        let appended = self.streamed_nnz();
1117        anyhow::ensure!(
1118            appended == nnz as u64,
1119            "finalize_streaming_csr: {appended} entries appended but {nnz} \
1120             declared — the slabs did not cover the matrix"
1121        );
1122        Ok(())
1123    }
1124
1125    /// Build `/by_row/{data, indices, indptr}` by transposing the already-
1126    /// written CSC data on disk. Uses two passes over CSC with bounded
1127    /// auxiliary memory (~`24 B × nrow` plus one row-band worth of CSR).
1128    fn build_csr_from_csc_streaming(&mut self) -> anyhow::Result<()> {
1129        let nrow = self
1130            .num_rows()
1131            .ok_or_else(|| anyhow::anyhow!("nrow not set before build_csr_from_csc_streaming"))?;
1132        let ncol = self
1133            .num_columns()
1134            .ok_or_else(|| anyhow::anyhow!("ncol not set before build_csr_from_csc_streaming"))?;
1135        let nnz = self
1136            .num_non_zeros()
1137            .ok_or_else(|| anyhow::anyhow!("nnz not set before build_csr_from_csc_streaming"))?;
1138
1139        if nnz == 0 {
1140            self.cs_create(CsKey::CsrData, 0)?;
1141            self.cs_create(CsKey::CsrIndices, 0)?;
1142            self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1143            let zeros = vec![0u64; nrow + 1];
1144            self.cs_write_u64(CsKey::CsrIndptr, 0, &zeros)?;
1145            self.read_row_indptr()?;
1146            return Ok(());
1147        }
1148
1149        const COL_BLOCK: usize = 1024;
1150        let n_col_blocks = ncol.div_ceil(COL_BLOCK);
1151        let bar1 = styled_progress_bar(n_col_blocks as u64, "transpose count");
1152        let mut row_counts = vec![0u64; nrow];
1153        let mut col_lo = 0usize;
1154        while col_lo < ncol {
1155            let col_hi = (col_lo + COL_BLOCK).min(ncol);
1156            let cols: Self::IndexIter = (col_lo..col_hi).collect();
1157            let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1158            for (row_i, _, _) in &triplets {
1159                row_counts[*row_i as usize] += 1;
1160            }
1161            col_lo = col_hi;
1162            bar1.inc(1);
1163        }
1164        bar1.finish_and_clear();
1165
1166        let mut rowptr = vec![0u64; nrow + 1];
1167        let mut acc = 0u64;
1168        for i in 0..nrow {
1169            rowptr[i] = acc;
1170            acc += row_counts[i];
1171        }
1172        rowptr[nrow] = acc;
1173        debug_assert_eq!(acc, nnz as u64);
1174
1175        self.cs_create(CsKey::CsrData, nnz)?;
1176        self.cs_create(CsKey::CsrIndices, nnz)?;
1177        self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
1178        self.cs_write_u64(CsKey::CsrIndptr, 0, &rowptr)?;
1179
1180        // Per-band buffers carry 12 B/nnz (u64 col + f32 val); cap aggregate
1181        // at ~256 MB so the row-banded scatter stays within a fixed budget
1182        // regardless of nnz.
1183        const TRANSPOSE_BAND_BYTES: usize = 256 * 1024 * 1024;
1184        let avg_density = nnz.div_ceil(nrow.max(1));
1185        let band_rows = (TRANSPOSE_BAND_BYTES / (12 * avg_density.max(1)))
1186            .max(1)
1187            .min(nrow);
1188        let n_bands = nrow.div_ceil(band_rows);
1189
1190        let bar2 = styled_progress_bar(n_bands as u64, "transpose scatter");
1191        let mut band_lo = 0usize;
1192        while band_lo < nrow {
1193            let band_hi = (band_lo + band_rows).min(nrow);
1194            let band_nnz_start = rowptr[band_lo];
1195            let band_nnz_end = rowptr[band_hi];
1196            let band_nnz = (band_nnz_end - band_nnz_start) as usize;
1197
1198            if band_nnz == 0 {
1199                band_lo = band_hi;
1200                bar2.inc(1);
1201                continue;
1202            }
1203
1204            let mut out_indices = vec![0u64; band_nnz];
1205            let mut out_values = vec![0f32; band_nnz];
1206            let mut cursor = vec![0u64; band_hi - band_lo];
1207
1208            let mut col_lo = 0usize;
1209            while col_lo < ncol {
1210                let col_hi = (col_lo + COL_BLOCK).min(ncol);
1211                let cols: Self::IndexIter = (col_lo..col_hi).collect();
1212                let (_, _, triplets) = self.read_triplets_by_columns(cols)?;
1213                for &(row_i, col_j_local, x) in &triplets {
1214                    let row_i_us = row_i as usize;
1215                    if row_i_us >= band_lo && row_i_us < band_hi {
1216                        let band_idx = row_i_us - band_lo;
1217                        // col_j_local is already the global column index because
1218                        // read_triplets_by_columns returns columns in the passed order
1219                        // (0..batch for standalone call). We passed col_lo..col_hi,
1220                        // which returns local indices 0..(col_hi - col_lo) — so add col_lo.
1221                        let col_j_global = col_j_local + col_lo as u64;
1222                        let offset_in_band =
1223                            (rowptr[band_lo + band_idx] - band_nnz_start) + cursor[band_idx];
1224                        out_indices[offset_in_band as usize] = col_j_global;
1225                        out_values[offset_in_band as usize] = x;
1226                        cursor[band_idx] += 1;
1227                    }
1228                }
1229                col_lo = col_hi;
1230            }
1231
1232            self.cs_write_u64(CsKey::CsrIndices, band_nnz_start, &out_indices)?;
1233            self.cs_write_f32(CsKey::CsrData, band_nnz_start, &out_values)?;
1234
1235            band_lo = band_hi;
1236            bar2.inc(1);
1237        }
1238        bar2.finish_and_clear();
1239
1240        self.read_row_indptr()?;
1241        Ok(())
1242    }
1243
1244    /// preload row index pointers
1245    fn read_row_indptr(&mut self) -> anyhow::Result<()>;
1246
1247    /// preload column index pointers
1248    fn read_column_indptr(&mut self) -> anyhow::Result<()>;
1249
1250    /// preload all the columns for faster processing
1251    fn preload_columns(&mut self) -> anyhow::Result<()>;
1252
1253    /// unload the memory
1254    fn clean_preloaded_columns(&mut self);
1255
1256    /// preload all the rows for faster processing
1257    fn preload_rows(&mut self) -> anyhow::Result<()>;
1258
1259    /// unload the row memory
1260    fn clean_preloaded_rows(&mut self);
1261
1262    /// backend file name
1263    fn get_backend_file_name(&self) -> &str;
1264
1265    /// backend file type
1266    fn backend_type(&self) -> SparseIoBackend;
1267}