data_beans/sparse_io/

traits.rs

#![allow(dead_code, unused_imports)]

#[cfg(feature = "tensor")]
pub use legume_numeric::candle_core::Tensor;
pub use nalgebra::DMatrix;
pub use nalgebra_sparse::{csc::CscMatrix, csr::CsrMatrix};
#[cfg(feature = "ndarray")]
pub use ndarray::prelude::*;

pub const MAX_ROW_NAME_IDX: usize = 3;
pub const MAX_COLUMN_NAME_IDX: usize = 10;
pub const COLUMN_SEP: &str = "@";
pub const ROW_SEP: &str = "_";

use super::helpers::*;
use super::index_audit;
use super::meta::Metadata;
use anyhow::Context;

use crate::column_subset::{distinct_within, stream_csc_selection, RowMap};
use crate::sparse_data_visitors::styled_progress_bar;
use clap::ValueEnum;
use indicatif::ParallelProgressIterator;
use legume_numeric::matrix::mtx_io::*;
use legume_numeric::matrix::traits::*;
use legume_numeric::matrix::utils::byte_budget_intervals;
use log::info;
use rayon::prelude::*;
use rustc_hash::FxHashMap as HashMap;
use std::ops::Range;

#[cfg(test)]
mod tests;

#[derive(ValueEnum, Clone, Debug, PartialEq)]
#[clap(rename_all = "lowercase")]
pub enum SparseIoBackend {
    Zarr,
    HDF5,
}

/// Identifies one of the six 1-D datasets inside a sparse backend.
/// Used by the streaming write API so we don't have to add six separate
/// abstract methods per dtype × (csc|csr) × (data|indices|indptr).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CsKey {
    CscData,
    CscIndices,
    CscIndptr,
    CsrData,
    CsrIndices,
    CsrIndptr,
}

/// Entries per slab handed to the backend while a sorted triplet vector is
/// streamed out as CSC or CSR. Bounds the staging buffers; the triplets
/// themselves are the only full-size structure alive at that point.
const SLAB_NNZ: usize = 1 << 20;

/// End of the slab that starts at triplet `start` of a vector sorted on the
/// major axis `major` (column for CSC, row for CSR): at least `slab_nnz`
/// entries, or all that remain, and never splitting a major index. Returns
/// `(end, band_end)` -- the exclusive triplet index and the exclusive major
/// bound -- so consecutive slabs tile `0..n_major` with no gap, empty
/// columns or rows included; the last slab runs to `n_major`.
fn slab_end(
    triplets: &[(u64, u64, f32)],
    start: usize,
    slab_nnz: usize,
    n_major: usize,
    major: impl Fn(&(u64, u64, f32)) -> u64,
) -> (usize, u64) {
    debug_assert!(slab_nnz > 0);
    let nnz = triplets.len();
    let mut end = (start + slab_nnz).min(nnz);
    while end < nnz && major(&triplets[end]) == major(&triplets[end - 1]) {
        end += 1;
    }
    let band_end = if end == nnz {
        n_major as u64
    } else {
        major(&triplets[end])
    };
    (end, band_end)
}

/// Every triplet must sit inside the declared `nrow × ncol` matrix before
/// the streaming writers see it: one that does not is never consumed by the
/// band walk, and would otherwise surface only as a count mismatch at the
/// finalize audit, with no trace of which entry was wrong.
fn check_triplet_bounds(
    label: &str,
    triplets: &[(u64, u64, f32)],
    nrow: usize,
    ncol: usize,
) -> anyhow::Result<()> {
    let (nrow_u, ncol_u) = (nrow as u64, ncol as u64);
    if let Some(k) = triplets
        .par_iter()
        .position_first(|&(r, c, _)| r >= nrow_u || c >= ncol_u)
    {
        let (r, c, v) = triplets[k];
        anyhow::bail!(
            "{label}: triplet {k} = (row {r} (0x{r:016x}), column {c} (0x{c:016x}), {v}) \
             is outside the {nrow} x {ncol} matrix"
        );
    }
    Ok(())
}

/// Replace the content of `data` with what `build` writes into a fresh
/// backend of the same kind.
///
/// The replacement is written beside the original (same filesystem, so the
/// final rename is atomic) and swapped in only when complete. CONTRACT of the
/// swap: the original is untouched until the temporary sibling is complete; a
/// failure in `build` leaves the original intact plus a
/// `.data-beans-stage/{name}.{stage_tag}` leftover (cleaned on the next
/// attempt); the unrecoverable window is only remove→rename.
///
/// A zip-archived backend is refused up front — streaming to a sibling
/// DIRECTORY and renaming it over the `.zip` name would silently change the
/// on-disk format under the old extension.
fn replace_backend_content<S>(
    data: &mut S,
    stage_tag: &str,
    build: impl FnOnce(&S, &mut dyn SparseIo<IndexIter = Vec<usize>>) -> anyhow::Result<()>,
) -> anyhow::Result<()>
where
    S: SparseIo + ?Sized,
{
    let final_path = data.get_backend_file_name().to_string();
    anyhow::ensure!(
        !final_path.ends_with(".zip"),
        "{final_path} is a zip archive; convert it to a directory \
         backend first (data-beans convert)"
    );
    // In the hidden stage beside the output (same disk, so the rename
    // stays atomic), where a sync client does not see it come and go.
    let name = std::path::Path::new(&final_path)
        .file_name()
        .map(|n| n.to_string_lossy().into_owned())
        .unwrap_or_default();
    let temp_path = crate::zarr_io::in_stage(&final_path, &format!("{name}.{stage_tag}"));
    crate::zarr_io::make_stage(&temp_path)?;
    if std::path::Path::new(&temp_path).exists() {
        crate::sparse_io::remove_backend_path(&temp_path)?;
    }

    {
        let backend_kind = data.backend_type();
        let mut out =
            crate::sparse_io::create_sparse_streaming_empty(Some(&temp_path), Some(&backend_kind))?;
        build(data, &mut *out)?;
        out.set_metadata(&data.metadata())?;
    }

    data.remove_backend_file()?;
    std::fs::rename(&temp_path, &final_path)?;
    crate::zarr_io::leave_stage(&temp_path);
    data.reopen_backend()?;
    info!("registered new data to {}", data.get_backend_file_name());
    Ok(())
}

/// [`SparseIo::build_csr_from_csc_streaming`] with an explicit cap of `band_bytes` on the entries one row band
/// holds in memory (12 B each; a single row larger than the cap is still
/// one band). Every scatter pass re-reads the whole CSC, so the cap
/// trades memory against passes; the result does not depend on it.
pub(crate) fn build_csr_from_csc_banded<S>(data: &mut S, band_bytes: usize) -> anyhow::Result<()>
where
    S: SparseIo + ?Sized,
{
    let nrow = data
        .num_rows()
        .ok_or_else(|| anyhow::anyhow!("nrow not set before build_csr_from_csc_streaming"))?;
    let ncol = data
        .num_columns()
        .ok_or_else(|| anyhow::anyhow!("ncol not set before build_csr_from_csc_streaming"))?;
    let nnz = data
        .num_non_zeros()
        .ok_or_else(|| anyhow::anyhow!("nnz not set before build_csr_from_csc_streaming"))?;

    if nnz == 0 {
        data.cs_create(CsKey::CsrData, 0)?;
        data.cs_create(CsKey::CsrIndices, 0)?;
        data.cs_create(CsKey::CsrIndptr, nrow + 1)?;
        let zeros = vec![0u64; nrow + 1];
        data.cs_write_u64(CsKey::CsrIndptr, 0, &zeros)?;
        data.read_row_indptr()?;
        return Ok(());
    }

    const COL_BLOCK: usize = 1024;
    let n_col_blocks = ncol.div_ceil(COL_BLOCK);
    let bar1 = styled_progress_bar(n_col_blocks as u64, "transpose count");
    let mut row_counts = vec![0u64; nrow];
    let mut col_lo = 0usize;
    while col_lo < ncol {
        let col_hi = (col_lo + COL_BLOCK).min(ncol);
        let cols: S::IndexIter = (col_lo..col_hi).collect();
        let (_, _, triplets) = data.read_triplets_by_columns(cols)?;
        for t in &triplets {
            let slot = row_counts.get_mut(t.0 as usize).ok_or_else(|| {
                anyhow::anyhow!(
                    "build_csr_from_csc_streaming: row index {} (0x{:016x}) outside \
                     the {nrow}-row matrix",
                    t.0,
                    t.0
                )
            })?;
            *slot += 1;
        }
        col_lo = col_hi;
        bar1.inc(1);
    }
    bar1.finish_and_clear();

    let mut rowptr = vec![0u64; nrow + 1];
    let mut acc = 0u64;
    for i in 0..nrow {
        rowptr[i] = acc;
        acc += row_counts[i];
    }
    rowptr[nrow] = acc;
    anyhow::ensure!(
        acc == nnz as u64,
        "build_csr_from_csc_streaming: counted {acc} entries in the CSC arrays, \
         but {nnz} are declared"
    );

    data.cs_create(CsKey::CsrData, nnz)?;
    data.cs_create(CsKey::CsrIndices, nnz)?;
    data.cs_create(CsKey::CsrIndptr, nrow + 1)?;
    data.cs_write_u64(CsKey::CsrIndptr, 0, &rowptr)?;

    // Row bands cut where the cumulative nnz reaches the cap, so a band
    // of dense rows holds no more than a band of sparse ones. A band of
    // empty rows has nothing to scatter.
    let bands: Vec<(usize, usize)> =
        byte_budget_intervals(&row_counts, band_bytes, PRELOAD_BYTES_PER_NNZ)
            .into_iter()
            .filter(|&(lo, hi)| rowptr[hi] > rowptr[lo])
            .collect();

    let bar2 = styled_progress_bar(bands.len() as u64, "transpose scatter");
    for (band_lo, band_hi) in bands {
        let band_nnz_start = rowptr[band_lo];
        let band_nnz = (rowptr[band_hi] - band_nnz_start) as usize;

        let mut out_indices = vec![0u64; band_nnz];
        let mut out_values = vec![0f32; band_nnz];
        let mut cursor = vec![0u64; band_hi - band_lo];

        let mut col_lo = 0usize;
        while col_lo < ncol {
            let col_hi = (col_lo + COL_BLOCK).min(ncol);
            let cols: S::IndexIter = (col_lo..col_hi).collect();
            let (_, _, triplets) = data.read_triplets_by_columns(cols)?;
            for &(row_i, col_j_local, x) in &triplets {
                let row_i_us = row_i as usize;
                if row_i_us >= band_lo && row_i_us < band_hi {
                    let band_idx = row_i_us - band_lo;
                    // col_j_local is already the global column index because
                    // read_triplets_by_columns returns columns in the passed order
                    // (0..batch for standalone call). We passed col_lo..col_hi,
                    // which returns local indices 0..(col_hi - col_lo) — so add col_lo.
                    let col_j_global = col_j_local + col_lo as u64;
                    let offset_in_band =
                        (rowptr[band_lo + band_idx] - band_nnz_start) + cursor[band_idx];
                    out_indices[offset_in_band as usize] = col_j_global;
                    out_values[offset_in_band as usize] = x;
                    cursor[band_idx] += 1;
                }
            }
            col_lo = col_hi;
        }

        data.cs_write_u64(CsKey::CsrIndices, band_nnz_start, &out_indices)?;
        data.cs_write_f32(CsKey::CsrData, band_nnz_start, &out_values)?;

        bar2.inc(1);
    }
    bar2.finish_and_clear();

    data.read_row_indptr()?;
    Ok(())
}

pub trait SparseIo: Sync + Send {
    type IndexIter: IntoIterator<Item = usize> + FromIterator<usize>;

    ////////////////////////////
    // default implementation //
    ////////////////////////////

    #[cfg(feature = "ndarray")]
    /// Read columns within the range and return dense `ndarray::Array2`
    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    fn read_columns_ndarray(&self, columns: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
        let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
        Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    #[cfg(feature = "tensor")]
    /// Read columns within the range and return dense `candle_core::Tensor`
    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    fn read_columns_tensor(&self, columns: Self::IndexIter) -> anyhow::Result<Tensor> {
        let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
        Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    /// Read columns within the range and return dense `nalgebrea::DMatrix`
    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    fn read_columns_dmatrix(&self, columns: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
        let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
        DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    /// Read columns within the range and return sparse `CsrMatrix`
    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    fn read_columns_csr(&self, columns: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
        let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
        CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    /// Read columns within the range and return sparse `CsrMatrix`
    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    fn read_columns_csc(&self, columns: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
        let (nrow, ncol, triplets) = self.read_triplets_by_columns(columns)?;
        CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    /// Zero-copy view of preloaded column-major CSC arrays as
    /// `(indptr, indices, data)`. Returns `None` when the backend has
    /// not preloaded columns or doesn't support direct array access.
    /// Callers (e.g. `SparseIoVec::read_columns_csc`) use this to skip
    /// the triplet roundtrip when columns are already in memory.
    fn csc_column_arrays(&self) -> Option<(&[u64], &[u64], &[f32])> {
        None
    }

    #[cfg(feature = "ndarray")]
    /// Read rows within the range and return dense `ndarray::Array2`
    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    fn read_rows_ndarray(&self, rows: Self::IndexIter) -> anyhow::Result<Array2<f32>> {
        let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
        Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    #[cfg(feature = "tensor")]
    /// Read rows within the range and return dense `candle_core::Tensor`
    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    fn read_rows_tensor(&self, rows: Self::IndexIter) -> anyhow::Result<Tensor> {
        let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
        Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    /// Read rows within the range and return dense `nalgebra::DMatrix`
    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    fn read_rows_dmatrix(&self, rows: Self::IndexIter) -> anyhow::Result<DMatrix<f32>> {
        let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
        DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    /// Read rows within the range and return sparse `CsrMatrix`
    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    fn read_rows_csr(&self, rows: Self::IndexIter) -> anyhow::Result<CsrMatrix<f32>> {
        let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
        CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    /// Read rows within the range and return sparse `CscMatrix`
    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    fn read_rows_csc(&self, rows: Self::IndexIter) -> anyhow::Result<CscMatrix<f32>> {
        let (nrow, ncol, triplets) = self.read_triplets_by_rows(rows)?;
        CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    /////////////////////////////
    // `mtx` related functions //
    /////////////////////////////

    /// Read an mtx file once and populate the backend: the column (CSC) index
    /// always, the row (CSR) index as well when `index_by_row`. Both are
    /// streamed out of the same triplet vector, so the file is inflated once
    /// and the triplets are the only full-size structure alive.
    /// * `mtx_file`: mtx file to be read into the backend
    fn import_mtx_file(&mut self, mtx_file: &str, index_by_row: bool) -> anyhow::Result<()> {
        let (mut mtx_triplets, mtx_shape) = read_mtx_triplets(mtx_file)?;
        info!("read mtx file: {}", mtx_file);
        if mtx_triplets.is_empty() {
            return Err(anyhow::anyhow!("No data in mtx file"));
        }
        self.record_mtx_shape(Some(mtx_shape))?;
        info!("recording the column index");
        self.record_triplets_by_col(&mut mtx_triplets)?;
        if index_by_row {
            info!("recording the row index");
            self.record_triplets_by_row(&mut mtx_triplets)?;
        }
        Ok(())
    }

    /////////////////////////////////
    // `dmatrix` related functions //
    /////////////////////////////////

    /// Add dmatrix to zarr backend by row (CSR format)
    /// * `array` - 2D array to be added to the backend
    fn import_dmatrix_by_row(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
        let (nrow, ncol) = matrix.shape();
        let mut mtx_triplets = dmatrix_to_triplets(matrix);
        let mtx_shape = (nrow, ncol, mtx_triplets.len());
        self.record_mtx_shape(Some(mtx_shape))?;
        self.record_triplets_by_row(&mut mtx_triplets)
    }

    /// Add dmatrix to zarr backend by column (CSC format)
    /// * `array` - 2D array to be added to the backend
    fn import_dmatrix_by_col(&mut self, matrix: &DMatrix<f32>) -> anyhow::Result<()> {
        let (nrow, ncol) = matrix.shape();
        let mut mtx_triplets = dmatrix_to_triplets(matrix);
        let mtx_shape = (nrow, ncol, mtx_triplets.len());
        self.record_mtx_shape(Some(mtx_shape))?;
        self.record_triplets_by_col(&mut mtx_triplets)
    }

    /////////////////////////////////
    // `ndarray` related functions //
    /////////////////////////////////

    #[cfg(feature = "ndarray")]
    /// Add ndarray to zarr backend by row (CSR format)
    /// * `array` - 2D array to be added to the backend
    fn import_ndarray_by_row(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
        let nrow = array.shape()[0];
        let ncol = array.shape()[1];

        // dbg!("importing ndarray by row...");
        let mut mtx_triplets = ndarray_to_triplets(array);

        let nnz = mtx_triplets.len();
        let mtx_shape = (nrow, ncol, nnz);
        self.record_mtx_shape(Some(mtx_shape))?;

        // dbg!(format!("populated: {} elements", mtx_triplets.len()));

        self.record_triplets_by_row(&mut mtx_triplets)
    }

    #[cfg(feature = "ndarray")]
    /// Add ndarray to zarr backend by column (CSC format)
    /// * `array` - 2D array to be added to the backend
    fn import_ndarray_by_col(&mut self, array: &Array2<f32>) -> anyhow::Result<()> {
        let nrow = array.shape()[0];
        let ncol = array.shape()[1];

        // dbg!("importing ndarray by column...");
        let mut mtx_triplets = ndarray_to_triplets(array);

        let nnz = mtx_triplets.len();
        let mtx_shape = (nrow, ncol, nnz);
        self.record_mtx_shape(Some(mtx_shape))?;

        // dbg!(format!("populated: {} elements", mtx_triplets.len()));

        self.record_triplets_by_col(&mut mtx_triplets)
    }

    //////////////////////
    // backend-specific //
    //////////////////////

    /// Read rows within the range and return a vector of triplets (row, column, value)
    /// * `rows` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    #[allow(clippy::type_complexity)]
    fn read_triplets_by_rows(
        &self,
        rows: Self::IndexIter,
    ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;

    /// Read columns within the range and return a vector of triplets (row, col, value)
    /// * `columns` : range e.g., 0..3 -> [0, 1, 2] or vec![0, 1, 2]
    ///
    #[allow(clippy::type_complexity)]
    fn read_triplets_by_columns(
        &self,
        columns: Self::IndexIter,
    ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;

    /// Read columns within the range and return a vector of triplets (row, col, value)
    /// * `col` : usize
    ///
    #[allow(clippy::type_complexity)]
    fn read_triplets_by_single_column(
        &self,
        col: usize,
    ) -> anyhow::Result<(usize, usize, Vec<(u64, u64, f32)>)>;

    /// Export the data to a mtx file. This will take time.
    /// * `mtx_file`: mtx file to be written
    fn to_mtx_file(&self, mtx_file: &str) -> anyhow::Result<()>;

    /// Number of rows in the underlying data matrix
    fn num_rows(&self) -> Option<usize>;

    /// Number of columns in the underlying data matrix
    fn num_columns(&self) -> Option<usize>;

    /// Number of non-zero elements
    fn num_non_zeros(&self) -> Option<usize>;

    /// Re-open handles on the CURRENT backend path after its contents were
    /// replaced from outside (a finished temp file renamed into place). The
    /// zarr store is path-addressed so this is a cache refresh; hdf5 holds an
    /// open file handle that would otherwise point at the deleted inode.
    fn reopen_backend(&mut self) -> anyhow::Result<()>;

    /// Maintained by [`append_csc_slab`](Self::append_csc_slab) — never call
    /// this yourself: padding the cursor masks exactly the under-append the
    /// finalize audit exists to catch.
    #[doc(hidden)]
    /// Advance the streaming-write cursor by `n` entries. Called by
    /// [`append_csc_slab`](Self::append_csc_slab); backends keep the count so
    /// [`finalize_streaming_csc`](Self::finalize_streaming_csc) can audit the
    /// declared nnz against what was actually appended — the one violation the
    /// written indptr cannot reveal, because an over-declared tail leaves it
    /// perfectly monotone with the phantom hiding between the last written
    /// pointer and the sentinel.
    fn note_streamed_nnz(&mut self, n: u64);

    /// Entries appended so far in this streaming build.
    #[doc(hidden)]
    fn streamed_nnz(&self) -> u64;

    /// Zero the cursor. Called by [`begin_streaming_csc`](Self::begin_streaming_csc).
    #[doc(hidden)]
    fn reset_streamed_nnz(&mut self);

    /// The resident by-column indptr, loaded at `open()`. Empty when the
    /// backend carries no `/by_column/indptr` array — `read_column_indptr`
    /// silently does nothing on that failure, and the accessors below must
    /// report that absence rather than read zeros out of it.
    fn column_indptr(&self) -> &[u64];

    /// The resident by-row indptr; the row-major twin of
    /// [`column_indptr`](Self::column_indptr), empty when absent.
    fn row_indptr(&self) -> &[u64];

    /// Exact nnz of one column, from the resident indptr — no I/O.
    ///
    /// `None` for an out-of-range column or when the indptr is absent. This is
    /// what lets a streaming writer declare a column subset's total nnz up
    /// front without a counting pass over the data.
    fn column_nnz(&self, col: usize) -> Option<u64> {
        let indptr = self.column_indptr();
        let hi = *indptr.get(col + 1)?;
        let lo = *indptr.get(col)?;
        hi.checked_sub(lo)
    }

    /// Set row names for the matrix
    /// * `row_name_file`: a file each line contains row name words
    fn register_row_names_file(&mut self, row_name_file: &str);

    /// Set column names for the matrix
    /// * `column_name_file`: a file each line contains column name words
    fn register_column_names_file(&mut self, column_name_file: &str);

    /// Set row names for the matrix
    /// * `rows`: a vector of row names
    fn register_row_names_vec(&mut self, rows: &[Box<str>]);

    /// Set column names for the matrix
    /// * `columns`: a vector of column names
    fn register_column_names_vec(&mut self, columns: &[Box<str>]);

    /// Add arbitrary names (a vector of strings)
    /// * `group_name`: group name
    /// * `name_file`: a file each line contains name words
    /// * `name_columns`: range of columns to be used for name
    /// * `name_sep`: separator for name columns
    fn register_names_file(
        &mut self,
        key: &str,
        name_file: &str,
        name_columns: Range<usize>,
        name_sep: &str,
    ) -> anyhow::Result<()>;

    /// Add arbitrary names (a vector of strings)
    /// * `group_name`: group name
    /// * `names`: a file each line contains name words
    fn register_names_vec(&mut self, key: &str, names: &[Box<str>]) -> anyhow::Result<()>;

    fn row_names(&self) -> anyhow::Result<Vec<Box<str>>>;

    fn column_names(&self) -> anyhow::Result<Vec<Box<str>>>;

    /// Get back the registered names
    /// * `key`: key for the registered names
    fn retrieve_registered_names(&self, key: &str) -> anyhow::Result<Vec<Box<str>>>;

    //////////////
    // metadata //
    //////////////

    /// The strings stored with the data to say what it is (see
    /// [`meta`](crate::sparse_io::meta) for the keys programs share); empty
    /// when none were set.
    fn metadata(&self) -> Metadata;

    /// Store `meta` as the data's metadata, replacing what was there. It
    /// needs a writable backend: one being created (set it before zipping),
    /// or an unzipped zarr store. A zipped zarr store, or an HDF5 file opened
    /// to read, refuses it.
    fn set_metadata(&mut self, meta: &Metadata) -> anyhow::Result<()>;

    /// One metadata value, e.g. `meta(meta::SAMPLE)`.
    fn meta(&self, key: &str) -> Option<String> {
        self.metadata().remove(key)
    }

    /// Set one metadata value, keeping the others.
    fn set_meta(&mut self, key: &str, value: &str) -> anyhow::Result<()> {
        let mut meta = self.metadata();
        meta.insert(key.to_string(), value.to_string());
        self.set_metadata(&meta)
    }

    /////////////////////////////
    // major structural change //
    /////////////////////////////

    /// Keep the columns `columns` and the rows `rows` of the data, each in
    /// the order listed (`None` keeps that side whole), replacing the
    /// backend's content. The result is streamed into a sibling and swapped
    /// in once complete, so a refused selection leaves the data as it was.
    ///
    /// # Errors
    /// - a selection is empty: this method replaces the original, and an
    ///   empty matrix over real data is not what a caller means;
    /// - an index is past the data's rows or columns;
    /// - a selection lists an index twice;
    /// - the backend is a zip archive, which cannot be rewritten in place
    ///   (convert it to a directory backend first).
    fn subset_columns_rows(
        &mut self,
        columns: Option<&Vec<usize>>,
        rows: Option<&Vec<usize>>,
    ) -> anyhow::Result<()> {
        let ncol_data = self
            .num_columns()
            .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;
        let nrow_data = self
            .num_rows()
            .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;

        if let Some(cols) = columns {
            anyhow::ensure!(!cols.is_empty(), "subset: empty column selection");
            distinct_within(cols, ncol_data, "column")?;
        }
        let row_map = match rows {
            Some(rs) => {
                anyhow::ensure!(!rs.is_empty(), "subset: empty row selection");
                Some(RowMap::keeping(rs, nrow_data)?)
            }
            None => None,
        };

        // The selection's order is the output order, for names and data alike.
        let take_names = |sel: Option<&Vec<usize>>,
                          names: Vec<Box<str>>,
                          what: &str|
         -> anyhow::Result<Vec<Box<str>>> {
            let Some(sel) = sel else {
                return Ok(names);
            };
            sel.iter()
                .map(|&i| {
                    names.get(i).cloned().ok_or_else(|| {
                        anyhow::anyhow!("subset: {what} {i} has no name ({} named)", names.len())
                    })
                })
                .collect()
        };
        let new_col_names = take_names(columns, self.column_names()?, "column")?;
        let new_row_names = take_names(rows, self.row_names()?, "row")?;

        let every_column: Vec<usize>;
        let old_cols: &[usize] = match columns {
            Some(cols) => cols,
            None => {
                every_column = (0..ncol_data).collect();
                &every_column
            }
        };
        replace_backend_content(self, "subset_tmp", |data, out| {
            stream_csc_selection(data, old_cols, row_map.as_ref(), out)?;
            out.register_row_names_vec(&new_row_names);
            out.register_column_names_vec(&new_col_names);
            Ok(())
        })
    }

    /// Reposition rows in the order of `row_names_order`: row `i` of the
    /// result is the row named `row_names_order[i]`, all-zero when the data
    /// has no such row; rows whose name is not listed are dropped. The result
    /// is streamed into a sibling and swapped in once complete, like
    /// [`subset_columns_rows`](Self::subset_columns_rows).
    ///
    /// # Errors
    /// If `row_names_order` lists a name twice, or two rows of the data carry
    /// one listed name (see [`RowMap::by_names`]).
    fn reorder_rows(&mut self, row_names_order: &[Box<str>]) -> anyhow::Result<()> {
        let Some(ncol) = self.num_columns() else {
            self.clean_preloaded_columns();
            self.clean_preloaded_rows();
            return Ok(());
        };
        let nrow_data = self
            .num_rows()
            .ok_or_else(|| anyhow::anyhow!("missing shape information"))?;

        let row_map = RowMap::by_names(&self.row_names()?, row_names_order, nrow_data)?;

        let new_col_names = self.column_names()?;
        let old_cols: Vec<usize> = (0..ncol).collect();

        info!("remapping triplets ...");
        replace_backend_content(self, "reorder_tmp", |data, out| {
            stream_csc_selection(data, &old_cols, Some(&row_map), out)?;
            out.register_row_names_vec(row_names_order);
            out.register_column_names_vec(&new_col_names);
            Ok(())
        })
    }

    /// Remove backend file
    fn remove_backend_file(&self) -> anyhow::Result<()>;

    /// Initialize backend
    fn initialize_backend(&mut self) -> anyhow::Result<()>;

    fn record_mtx_shape(&mut self, mtx_shape: Option<(usize, usize, usize)>) -> anyhow::Result<()>;

    /// Stream the triplets out as CSR slabs; the row-major twin of
    /// [`record_triplets_by_col`](Self::record_triplets_by_col).
    fn record_triplets_by_row(
        &mut self,
        row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
    ) -> anyhow::Result<()> {
        let nrow = self.num_rows().expect("should have `nrow`");
        let ncol = self.num_columns().expect("should have `ncol`");
        let nnz = row_col_val_triplets.len();

        if nnz == 0 {
            let csr_rowptr = vec![0u64; nrow + 1];
            return self.record_csr_dataset_backend(&[], &[], &csr_rowptr);
        }

        check_triplet_bounds("record_triplets_by_row", row_col_val_triplets, nrow, ncol)?;

        // One in-place pass on the full key. A stable sort would allocate a
        // scratch copy of the whole vector, and duplicate coordinates carry no
        // meaning in coordinate format, so stability buys nothing.
        row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (row, col));

        self.begin_streaming_csr((nrow, ncol, nnz))?;

        let mut local_rowptr: Vec<u64> = Vec::new();
        let mut cols: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
        let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);

        let mut start = 0_usize;
        let mut row_offset = 0_u64;
        while (row_offset as usize) < nrow {
            let (end, band_end_row) =
                slab_end(row_col_val_triplets, start, SLAB_NNZ, nrow, |t| t.0);

            local_rowptr.clear();
            cols.clear();
            vals.clear();
            let mut i = start;
            for row in row_offset..band_end_row {
                local_rowptr.push((i - start) as u64);
                while i < end && row_col_val_triplets[i].0 == row {
                    cols.push(row_col_val_triplets[i].1);
                    vals.push(row_col_val_triplets[i].2);
                    i += 1;
                }
            }
            if i != end {
                anyhow::bail!(
                    "record_triplets_by_row: band walk left triplet {:?} at {i}",
                    row_col_val_triplets[i]
                );
            }

            self.append_csr_slab(row_offset, start as u64, &local_rowptr, &cols, &vals)?;
            start = end;
            row_offset = band_end_row;
        }

        self.finalize_streaming_csr()
    }

    /// Stream the triplets out as CSC slabs.
    ///
    /// After the one in-place sort the triplet vector is the only full-size
    /// structure alive; the slab staging buffers are bounded by
    /// [`SLAB_NNZ`], and the backend's streaming audits check the column
    /// tiling and the appended count on the way through.
    fn record_triplets_by_col(
        &mut self,
        row_col_val_triplets: &mut Vec<(u64, u64, f32)>,
    ) -> anyhow::Result<()> {
        let nrow = self.num_rows().expect("should have `nrow`");
        let ncol = self.num_columns().expect("should have `ncol`");
        let nnz = row_col_val_triplets.len();

        if nnz == 0 {
            let csc_colptr = vec![0u64; ncol + 1];
            return self.record_csc_dataset_backend(&[], &[], &csc_colptr);
        }

        check_triplet_bounds("record_triplets_by_col", row_col_val_triplets, nrow, ncol)?;

        // See `record_triplets_by_row` for why this is one unstable pass.
        row_col_val_triplets.par_sort_unstable_by_key(|&(row, col, _)| (col, row));

        self.begin_streaming_csc((nrow, ncol, nnz))?;

        let mut local_colptr: Vec<u64> = Vec::new();
        let mut rows: Vec<u64> = Vec::with_capacity(SLAB_NNZ);
        let mut vals: Vec<f32> = Vec::with_capacity(SLAB_NNZ);

        let mut start = 0_usize;
        let mut col_offset = 0_u64;
        while (col_offset as usize) < ncol {
            let (end, band_end_col) =
                slab_end(row_col_val_triplets, start, SLAB_NNZ, ncol, |t| t.1);

            local_colptr.clear();
            rows.clear();
            vals.clear();
            let mut i = start;
            for col in col_offset..band_end_col {
                local_colptr.push((i - start) as u64);
                while i < end && row_col_val_triplets[i].1 == col {
                    rows.push(row_col_val_triplets[i].0);
                    vals.push(row_col_val_triplets[i].2);
                    i += 1;
                }
            }
            if i != end {
                anyhow::bail!(
                    "record_triplets_by_col: band walk left triplet {:?} at {i}",
                    row_col_val_triplets[i]
                );
            }

            self.append_csc_slab(col_offset, start as u64, &local_colptr, &rows, &vals)?;
            start = end;
            col_offset = band_end_col;
        }

        self.finalize_streaming_csc()
    }

    /// CSR data structure in Zarr backend
    ///
    /// ```text
    ///     └── by_row
    ///         ├── data
    ///         ├── indices (column indices)
    ///         └── isndptr (row pointers)
    /// ```
    fn record_csr_dataset_backend(
        &mut self,
        csr_cols: &[u64],
        csr_vals: &[f32],
        csr_rowptr: &[u64],
    ) -> anyhow::Result<()>;

    /// Helper function to add CSC dataset to HDF5 backend
    ///
    /// ```text
    /// Helper function to record the CSC dataset
    ///     ├── by_column
    ///     │   ├── data
    ///     │   ├── indices (row indices)
    ///     │   └── indptr (column pointers)
    /// ```
    fn record_csc_dataset_backend(
        &mut self,
        csc_rows: &[u64],
        csc_vals: &[f32],
        csc_colptr: &[u64],
    ) -> anyhow::Result<()>;

    /// Create a fixed-size 1-D backend dataset of `len` elements for the
    /// given CSC/CSR slot. No data is written yet.
    fn cs_create(&mut self, key: CsKey, len: usize) -> anyhow::Result<()>;

    /// Write a `u64` slab at `offset` in the specified dataset.
    /// Used for CSC/CSR `indices` and `indptr`.
    fn cs_write_u64(&mut self, key: CsKey, offset: u64, data: &[u64]) -> anyhow::Result<()>;

    /// Write an `f32` slab at `offset` in the specified dataset.
    /// Used for CSC/CSR `data`.
    fn cs_write_f32(&mut self, key: CsKey, offset: u64, data: &[f32]) -> anyhow::Result<()>;

    /// Begin a streaming CSC build for a sparse matrix of known shape.
    /// Pre-creates `/by_column/{data, indices, indptr}` at their final sizes
    /// so subsequent [`append_csc_slab`](Self::append_csc_slab) calls write
    /// into disjoint hyperslabs without further allocation.
    fn begin_streaming_csc(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
        // A reused handle must not inherit a previous build's cursor: the
        // finalize audit compares appended-vs-declared, and a stale count turns
        // a correct build into a false accusation.
        self.reset_streamed_nnz();
        let (_, ncol, nnz) = shape;
        self.record_mtx_shape(Some(shape))?;
        self.cs_create(CsKey::CscData, nnz)?;
        self.cs_create(CsKey::CscIndices, nnz)?;
        self.cs_create(CsKey::CscIndptr, ncol + 1)?;
        Ok(())
    }

    /// Append one contiguous CSC column band.
    ///
    /// * `col_offset` — global column index where this band starts
    /// * `nnz_offset` — global nnz offset where this band's values land
    /// * `local_colptr` — length `batch_ncol`, values in `[0, batch_nnz]`,
    ///   will be shifted by `nnz_offset` before writing
    /// * `row_indices` — length `batch_nnz`
    /// * `values`      — length `batch_nnz`
    fn append_csc_slab(
        &mut self,
        col_offset: u64,
        nnz_offset: u64,
        local_colptr: &[u64],
        row_indices: &[u64],
        values: &[f32],
    ) -> anyhow::Result<()> {
        // These checks exist because a violation does NOT fail loudly on its
        // own: unwritten regions read back as the zarr fill value, so a bad
        // slab yields a backend that opens and reads cleanly while carrying
        // poisoned or duplicated entries. Cheap (one pass over the slab, no
        // I/O) next to the compressed writes below.
        anyhow::ensure!(
            row_indices.len() == values.len(),
            "append_csc_slab: {} row indices vs {} values",
            row_indices.len(),
            values.len()
        );
        anyhow::ensure!(
            local_colptr.first().copied() == Some(0) || local_colptr.is_empty(),
            "append_csc_slab: local_colptr must start at 0"
        );
        anyhow::ensure!(
            local_colptr.windows(2).all(|w| w[0] <= w[1]),
            "append_csc_slab: local_colptr must be monotone non-decreasing"
        );
        if let Some(&last) = local_colptr.last() {
            anyhow::ensure!(
                last <= values.len() as u64,
                "append_csc_slab: colptr claims {last} entries, slab holds {}",
                values.len()
            );
        }
        if let Some(nrow) = self.num_rows() {
            if let Some(&bad) = row_indices.iter().find(|&&r| r >= nrow as u64) {
                anyhow::bail!("append_csc_slab: row index {bad} outside the {nrow}-row matrix");
            }
        }
        // Ascending rows within each column: readers document it as an
        // invariant, and the h5ad export hands the arrays to scipy as-is.
        for (c, &start) in local_colptr.iter().enumerate() {
            let end = local_colptr
                .get(c + 1)
                .copied()
                .unwrap_or(values.len() as u64) as usize;
            anyhow::ensure!(
                row_indices[start as usize..end]
                    .windows(2)
                    .all(|w| w[0] < w[1]),
                "append_csc_slab: rows within column {} of this band must be \
                 strictly ascending — repeated rows usually mean duplicate \
                 (row, col) coordinates in the source (an MTX with repeated \
                 entries, or a union remap folding rows together)",
                col_offset as usize + c
            );
        }

        let shifted: Vec<u64> = local_colptr.iter().map(|&p| p + nnz_offset).collect();
        self.cs_write_u64(CsKey::CscIndptr, col_offset, &shifted)?;
        self.cs_write_u64(CsKey::CscIndices, nnz_offset, row_indices)?;
        self.cs_write_f32(CsKey::CscData, nnz_offset, values)?;
        self.note_streamed_nnz(values.len() as u64);
        Ok(())
    }

    /// Finalize CSC streaming by writing the final indptr sentinel at
    /// position `ncol`, equal to the total nnz.
    fn finalize_streaming_csc(&mut self) -> anyhow::Result<()> {
        let ncol = self
            .num_columns()
            .ok_or_else(|| anyhow::anyhow!("ncol not set before finalize_streaming_csc"))?;
        let nnz = self
            .num_non_zeros()
            .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csc"))?;
        self.cs_write_u64(CsKey::CscIndptr, ncol as u64, &[nnz as u64])?;
        self.read_column_indptr()?;

        // The tiling check the per-slab guards cannot do. A gap or overlap in
        // the nnz offsets, or an over-declared total, leaves the WRITTEN
        // indptr non-monotone or short of the declared nnz — and unwritten
        // indptr slots read back as the fill value 0, so column j would claim
        // the whole array prefix. `indptr[ncol] - indptr[0]` equals the
        // declaration by construction, which is why the old debug_assert on it
        // could never fire; the shape of the vector between the endpoints is
        // what carries the truth.
        index_audit::check_indptr("finalize_streaming_csc", self.column_indptr(), ncol, nnz)
            .context(
                "the slabs were appended with a gap or overlap in their nnz offsets, \
             or the first slab was never appended",
            )?;
        // The appended count is the ground truth the indptr cannot carry: an
        // over-declared nnz leaves the written indptr perfectly monotone with
        // the phantom tail hiding between the last written pointer and the
        // sentinel — and the sentinel itself was written by this function, so
        // comparing against it can only ever agree.
        let appended = self.streamed_nnz();
        anyhow::ensure!(
            appended == nnz as u64,
            "finalize_streaming_csc: {appended} entries appended but {nnz} \
             declared — the difference reads back as fill values wearing real \
             entries' positions"
        );
        Ok(())
    }

    /// Begin a streaming CSR build for a sparse matrix of known shape; the
    /// row-major twin of [`begin_streaming_csc`](Self::begin_streaming_csc).
    fn begin_streaming_csr(&mut self, shape: (usize, usize, usize)) -> anyhow::Result<()> {
        self.reset_streamed_nnz();
        let (nrow, _, nnz) = shape;
        self.record_mtx_shape(Some(shape))?;
        self.cs_create(CsKey::CsrData, nnz)?;
        self.cs_create(CsKey::CsrIndices, nnz)?;
        self.cs_create(CsKey::CsrIndptr, nrow + 1)?;
        Ok(())
    }

    /// Append one contiguous CSR row band; the row-major twin of
    /// [`append_csc_slab`](Self::append_csc_slab), with the same audits.
    ///
    /// * `row_offset` — global row index where this band starts
    /// * `nnz_offset` — global nnz offset where this band's values land
    /// * `local_rowptr` — length `batch_nrow`, values in `[0, batch_nnz]`,
    ///   will be shifted by `nnz_offset` before writing
    /// * `col_indices` — length `batch_nnz`
    /// * `values`      — length `batch_nnz`
    fn append_csr_slab(
        &mut self,
        row_offset: u64,
        nnz_offset: u64,
        local_rowptr: &[u64],
        col_indices: &[u64],
        values: &[f32],
    ) -> anyhow::Result<()> {
        anyhow::ensure!(
            col_indices.len() == values.len(),
            "append_csr_slab: {} column indices vs {} values",
            col_indices.len(),
            values.len()
        );
        anyhow::ensure!(
            local_rowptr.first().copied() == Some(0) || local_rowptr.is_empty(),
            "append_csr_slab: local_rowptr must start at 0"
        );
        anyhow::ensure!(
            local_rowptr.windows(2).all(|w| w[0] <= w[1]),
            "append_csr_slab: local_rowptr must be monotone non-decreasing"
        );
        if let Some(&last) = local_rowptr.last() {
            anyhow::ensure!(
                last <= values.len() as u64,
                "append_csr_slab: rowptr claims {last} entries, slab holds {}",
                values.len()
            );
        }
        if let Some(ncol) = self.num_columns() {
            if let Some(&bad) = col_indices.iter().find(|&&c| c >= ncol as u64) {
                anyhow::bail!(
                    "append_csr_slab: column index {bad} outside the {ncol}-column matrix"
                );
            }
        }
        for (r, &start) in local_rowptr.iter().enumerate() {
            let end = local_rowptr
                .get(r + 1)
                .copied()
                .unwrap_or(values.len() as u64) as usize;
            anyhow::ensure!(
                col_indices[start as usize..end]
                    .windows(2)
                    .all(|w| w[0] < w[1]),
                "append_csr_slab: columns within row {} of this band must be \
                 strictly ascending — repeated columns usually mean duplicate \
                 (row, col) coordinates in the source",
                row_offset as usize + r
            );
        }

        let shifted: Vec<u64> = local_rowptr.iter().map(|&p| p + nnz_offset).collect();
        self.cs_write_u64(CsKey::CsrIndptr, row_offset, &shifted)?;
        self.cs_write_u64(CsKey::CsrIndices, nnz_offset, col_indices)?;
        self.cs_write_f32(CsKey::CsrData, nnz_offset, values)?;
        self.note_streamed_nnz(values.len() as u64);
        Ok(())
    }

    /// Finalize CSR streaming: write the indptr sentinel at position `nrow`,
    /// load the row index, and check the appended count against the declared
    /// nnz -- the one violation the written indptr cannot reveal.
    fn finalize_streaming_csr(&mut self) -> anyhow::Result<()> {
        let nrow = self
            .num_rows()
            .ok_or_else(|| anyhow::anyhow!("nrow not set before finalize_streaming_csr"))?;
        let nnz = self
            .num_non_zeros()
            .ok_or_else(|| anyhow::anyhow!("nnz not set before finalize_streaming_csr"))?;
        self.cs_write_u64(CsKey::CsrIndptr, nrow as u64, &[nnz as u64])?;
        self.read_row_indptr()?;

        // The same tiling audit as `finalize_streaming_csc`.
        index_audit::check_indptr("finalize_streaming_csr", self.row_indptr(), nrow, nnz).context(
            "the slabs were appended with a gap or overlap in their nnz offsets, \
             or the first slab was never appended",
        )?;

        let appended = self.streamed_nnz();
        anyhow::ensure!(
            appended == nnz as u64,
            "finalize_streaming_csr: {appended} entries appended but {nnz} \
             declared — the slabs did not cover the matrix"
        );
        Ok(())
    }

    /// Build `/by_row/{data, indices, indptr}` by transposing the already-
    /// written CSC data on disk: one pass over the CSC to count each row,
    /// then one scatter pass per row band. The band is as large as the
    /// process-wide preload budget and the memory available now allow
    /// (never under [`SLAB_BUDGET_BYTES`]), so a matrix that fits is
    /// transposed in a single scatter pass.
    fn build_csr_from_csc_streaming(&mut self) -> anyhow::Result<()> {
        let wanted = self
            .num_non_zeros()
            .unwrap_or(0)
            .saturating_mul(PRELOAD_BYTES_PER_NNZ);
        // The reservation is held until the transpose is done.
        let (band_bytes, _reservation) = reserve_transpose_band(wanted);
        build_csr_from_csc_banded(self, band_bytes)
    }

    /// preload row index pointers
    fn read_row_indptr(&mut self) -> anyhow::Result<()>;

    /// preload column index pointers
    fn read_column_indptr(&mut self) -> anyhow::Result<()>;

    /// preload all the columns for faster processing
    fn preload_columns(&mut self) -> anyhow::Result<()>;

    /// unload the memory
    fn clean_preloaded_columns(&mut self);

    /// preload all the rows for faster processing
    fn preload_rows(&mut self) -> anyhow::Result<()>;

    /// unload the row memory
    fn clean_preloaded_rows(&mut self);

    /// backend file name
    fn get_backend_file_name(&self) -> &str;

    /// backend file type
    fn backend_type(&self) -> SparseIoBackend;
}