data_beans/sparse_io_vector/

read.rs

#![allow(dead_code)]

use super::*;
use crate::sparse_io::index_audit;

/// Map a backend-local row through `l2g` then `g2c`: `Ok(None)` for a row
/// outside the intersection, an error naming the backend, the column and the
/// row when either lookup is out of range.
#[inline]
fn remap_row(
    l2g: &[usize],
    g2c: &[Option<usize>],
    row: u64,
    didx: usize,
    col: usize,
) -> anyhow::Result<Option<usize>> {
    let Some(&g) = l2g.get(row as usize) else {
        anyhow::bail!(
            "backend {didx}, column {col}: row index {row} (0x{row:016x}) is outside \
             the backend's {} rows",
            l2g.len()
        );
    };
    g2c.get(g).copied().ok_or_else(|| {
        anyhow::anyhow!(
            "backend {didx}, column {col}: row index {row} maps to global row {g}, \
             outside the global row table"
        )
    })
}

impl SparseIoVec {
    /// True when some row is observed by more than one backend, so a column
    /// that several backends contribute to can receive two values at one row.
    fn backends_share_rows(&self) -> bool {
        self.row_count_by_global.iter().any(|&n| n > 1)
    }

    /// The factor reads divide global column `glob` by: 1 when no scales
    /// are set, and IEEE division by 1 is exact, so such reads are the
    /// counts as stored.
    #[inline]
    fn scale_of(&self, glob: usize) -> f32 {
        self.column_scales.as_ref().map_or(1.0, |s| s[glob])
    }

    ////////////////////
    // access columns //
    ////////////////////

    /// Stream every nonzero of the selected `cells` (global column
    /// indices) through `f(row, col, val)` **without** materializing the
    /// full triplet vector. `row` is a compact-row index and `col` runs
    /// over `0..ncol` in the iteration order of `cells` — exactly the
    /// indices [`Self::columns_triplets`] would emit.
    ///
    /// Columns are processed in slabs of at most `chunk_cols`, so the
    /// only transient allocation is one backend slab's worth of
    /// `(u64, u64, f32)` triplets: peak memory is bounded by `chunk_cols`,
    /// not by the total nnz. Callers can therefore build a compact edge
    /// list (e.g. 12-byte triplets) directly and never pay for the wide
    /// intermediate. Returns the `(nrow, ncol)` dimensions.
    ///
    /// Each `(row, col)` is emitted at most once: where a row canonicalizer
    /// folds several rows of one backend together, or several backends of a
    /// [`ColumnAlignment::Union`] column observe the same row, the values are
    /// summed first.
    pub fn for_each_triplet<I, F>(
        &self,
        cells: I,
        chunk_cols: usize,
        mut f: F,
    ) -> anyhow::Result<(usize, usize)>
    where
        I: Iterator<Item = usize>,
        F: FnMut(u64, u64, f32),
    {
        let nrow = self.num_rows();
        let cells: Vec<usize> = cells.collect();
        let ncol = cells.len();
        let chunk = chunk_cols.max(1);
        let g2c = self.global_to_compact_row.as_slice();
        let scale: Vec<f32> = cells.iter().map(|&g| self.scale_of(g)).collect();

        // Backends of one column can only collide on a row they both observe.
        let union_shares_rows =
            self.column_alignment == ColumnAlignment::Union && self.backends_share_rows();
        let mut pending: Vec<(u64, u64, f32)> = Vec::new();

        for slab_start in (0..ncol).step_by(chunk) {
            let slab_end = (slab_start + chunk).min(ncol);

            // Group this slab's cells by backend, tracking (local_col,
            // out_col) where out_col is the index into the *full* `cells`
            // sequence. Under `ColumnAlignment::Union` one global cell can
            // contribute entries to multiple backend groups (one per
            // backend that observed it); all those reads target the same
            // out_col.
            let mut backend_groups: HashMap<usize, Vec<(usize, usize)>> = HashMap::default();
            // Per slab column: do several backends feed it rows they share?
            let mut multi_source = vec![false; slab_end - slab_start];
            for (k, &glob) in cells[slab_start..slab_end].iter().enumerate() {
                let out_col = slab_start + k;
                let sources = self.col_to_data[glob].as_slice();
                multi_source[k] = union_shares_rows && sources.len() > 1;
                for source in sources {
                    backend_groups
                        .entry(source.backend as usize)
                        .or_default()
                        .push((source.local_col as usize, out_col));
                }
            }

            for (&didx, group) in &backend_groups {
                let local_cols: Vec<usize> = group.iter().map(|&(loc, _)| loc).collect();
                let (_, _, group_triplets) =
                    self.data_vec[didx].read_triplets_by_columns(local_cols)?;

                let l2g = self.data_local_to_global_row[didx].as_slice();
                // A canonicalizer that collapsed >=2 local rows of this
                // dataset to one global row repeats (row, col) within the
                // backend; a multi-source column can repeat it across
                // backends. Those entries wait in `pending` to be summed;
                // every entry of an out_col lives in this slab, so per-slab
                // accumulation is exact. The rest pass straight through.
                let intra_merges = self.data_has_intra_row_merges[didx];
                for (i, j, v) in group_triplets {
                    if let Some(c) = g2c[l2g[i as usize]] {
                        let out_col = group[j as usize].1;
                        if intra_merges || multi_source[out_col - slab_start] {
                            pending.push((c as u64, out_col as u64, v));
                        } else {
                            f(c as u64, out_col as u64, v / scale[out_col]);
                        }
                    }
                }
            }

            // Divided after the duplicates are summed: one factor per column.
            crate::sparse_util::sum_duplicate_triplets(&mut pending);
            for (r, c, v) in pending.drain(..) {
                f(r, c, v / scale[c as usize]);
            }
        }

        Ok((nrow, ncol))
    }

    /// Collect all nonzeros of the selected `cells` into one triplet
    /// vector. Thin wrapper over [`Self::for_each_triplet`] with a single
    /// slab spanning every column (identical one-pass behavior); prefer
    /// `for_each_triplet` when the result is consumed once, to avoid the
    /// full-width intermediate.
    #[allow(clippy::type_complexity)]
    pub fn columns_triplets<I>(
        &self,
        cells: I,
    ) -> anyhow::Result<((usize, usize), Vec<(u64, u64, f32)>)>
    where
        I: Iterator<Item = usize>,
    {
        let cells: Vec<usize> = cells.collect();
        let one_slab = cells.len().max(1);
        let mut triplets = Vec::new();
        let dims = self.for_each_triplet(cells.into_iter(), one_slab, |r, c, v| {
            triplets.push((r, c, v));
        })?;
        Ok((dims, triplets))
    }

    #[cfg(feature = "ndarray")]
    pub fn read_columns_ndarray<I>(&self, cells: I) -> anyhow::Result<ndarray::Array2<f32>>
    where
        I: Iterator<Item = usize>,
    {
        let ((nrow, ncol), triplets) = self.columns_triplets(cells)?;
        ndarray::Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    pub fn read_columns_dmatrix<I>(&self, cells: I) -> anyhow::Result<nalgebra::DMatrix<f32>>
    where
        I: Iterator<Item = usize>,
    {
        let ((nrow, ncol), triplets) = self.columns_triplets(cells)?;
        DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    /// Direct-slice CSC read: bypasses the `triplet → COO → CSC` roundtrip
    /// when the underlying backends have preloaded column arrays.
    ///
    /// For each cell, we slice `(indices, values)` straight out of the
    /// backend's preloaded `by_column_indices` / `by_column_data`, remap
    /// row indices through `l2g` then `g2c` once, drop entries that fall
    /// outside the row intersection, and assemble final CSC arrays in one
    /// pass per column. Backends that aren't preloaded fall back to
    /// per-column triplet reads, which still avoids the global triplet
    /// vec and the column-major sort inside `CscMatrix::from(&coo)`.
    pub fn read_columns_csc<I>(&self, cells: I) -> anyhow::Result<CscMatrix<f32>>
    where
        I: Iterator<Item = usize>,
    {
        let cells: Vec<usize> = cells.collect();
        let nrow = self.num_rows();
        let ncol = cells.len();

        // Group cells by backend, carrying the output column index so we
        // can scatter directly into the final per-column buckets. Under
        // `ColumnAlignment::Union` one global cell can be observed by
        // multiple backends — the loop visits every `(didx, loc)` for
        // the same `out_col`, so the bucket accumulates contributions
        // from all observing backends.
        let mut backend_groups: Vec<Vec<(usize, usize)>> =
            (0..self.data_vec.len()).map(|_| Vec::new()).collect();
        for (out_col, &glob) in cells.iter().enumerate() {
            for source in self.col_to_data[glob].iter() {
                backend_groups[source.backend as usize].push((source.local_col as usize, out_col));
            }
        }

        // Per-output-column buckets of (compact_row, value).
        let mut buckets: Vec<Vec<(u32, f32)>> = (0..ncol).map(|_| Vec::new()).collect();
        let g2c = self.global_to_compact_row.as_slice();

        for (didx, group) in backend_groups.iter().enumerate() {
            if group.is_empty() {
                continue;
            }
            let l2g = self.data_local_to_global_row[didx].as_slice();

            if let Some((indptr, indices, values)) = self.data_vec[didx].csc_column_arrays() {
                // Fast path: zero-copy slicing into preloaded arrays. Every
                // index is checked at use, so a value that went bad after the
                // preload audit fails here, naming its backend, column and row.
                for &(loc, out_col) in group {
                    if loc + 1 >= indptr.len() {
                        continue;
                    }
                    index_audit::check_slot(
                        "preloaded column",
                        loc,
                        indptr[loc],
                        indptr[loc + 1],
                        indices.len().min(values.len()),
                    )?;
                    let s = indptr[loc] as usize;
                    let e = indptr[loc + 1] as usize;
                    let bucket = &mut buckets[out_col];
                    bucket.reserve(e - s);
                    for k in s..e {
                        if let Some(c) = remap_row(l2g, g2c, indices[k], didx, loc)? {
                            bucket.push((c as u32, values[k]));
                        }
                    }
                }
            } else {
                // Cold path: route through `read_triplets_by_columns` so the
                // backend can coalesce abutting indptr ranges into a single
                // zarr/hdf5 retrieval (zarr uses `coalesce_and_emit` + chunk
                // LRU cache). For a contiguous block of N cells this becomes
                // ONE retrieval instead of N — the dominant win on cold reads
                // from `.zarr.zip` over a slow disk.
                // The indptr gives each column's size before any entry is
                // read, so its bucket grows once.
                for &(loc, out_col) in group {
                    if let Some(nnz) = self.data_vec[didx].column_nnz(loc) {
                        buckets[out_col].reserve(nnz as usize);
                    }
                }
                let cols: Vec<usize> = group.iter().map(|&(loc, _)| loc).collect();
                let (_, _, trip) = self.data_vec[didx].read_triplets_by_columns(cols)?;
                for (i, jj, v) in trip {
                    let &(loc, out_col) = group.get(jj as usize).ok_or_else(|| {
                        anyhow::anyhow!(
                            "backend {didx}: read returned output column {jj} of {}",
                            group.len()
                        )
                    })?;
                    if let Some(c) = remap_row(l2g, g2c, i, didx, loc)? {
                        buckets[out_col].push((c as u32, v));
                    }
                }
            }
        }

        // Assemble final CSC arrays in one pass.
        let total_nnz: usize = buckets.iter().map(|b| b.len()).sum();
        let mut col_offsets: Vec<usize> = Vec::with_capacity(ncol + 1);
        let mut row_indices: Vec<usize> = Vec::with_capacity(total_nnz);
        let mut values: Vec<f32> = Vec::with_capacity(total_nnz);
        col_offsets.push(0);

        for (out_col, bucket) in buckets.iter_mut().enumerate() {
            // Canonical CSC requires within-column row indices sorted
            // ascending AND unique. On-disk indices are sorted by local
            // row, but two sources can land on the same compact row:
            //   (a) a row canonicalizer that maps two local rows in the
            //       same backend to one global row;
            //   (b) `ColumnAlignment::Union` where multiple backends
            //       contribute to one output column.
            // The composition `l2g[g2c[..]]` is monotonic in the
            // single-source, no-canonicalizer case (the historical
            // fast path) — detect it cheaply and skip the rebuild;
            // otherwise sort and fold duplicates by summing.
            let strictly_sorted_unique = bucket.windows(2).all(|w| w[0].0 < w[1].0);
            if !strictly_sorted_unique {
                bucket.sort_by_key(|&(r, _)| r);
                // Compact duplicates in-place: sum values for equal rows.
                let mut write = 0usize;
                let mut read = 0usize;
                while read < bucket.len() {
                    let (r, mut v) = bucket[read];
                    read += 1;
                    while read < bucket.len() && bucket[read].0 == r {
                        v += bucket[read].1;
                        read += 1;
                    }
                    bucket[write] = (r, v);
                    write += 1;
                }
                bucket.truncate(write);
            }
            // Divided after the duplicates are summed: one factor per column.
            let scale = self.scale_of(cells[out_col]);
            for &(r, v) in bucket.iter() {
                row_indices.push(r as usize);
                values.push(v / scale);
            }
            col_offsets.push(row_indices.len());
        }

        CscMatrix::try_from_csc_data(nrow, ncol, col_offsets, row_indices, values)
            .map_err(|e| anyhow::anyhow!("CSC construction failed: {:?}", e))
    }

    pub fn read_columns_csr<I>(&self, cells: I) -> anyhow::Result<CsrMatrix<f32>>
    where
        I: Iterator<Item = usize>,
    {
        let ((nrow, ncol), triplets) = self.columns_triplets(cells)?;
        nalgebra_sparse::CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    #[cfg(feature = "tensor")]
    pub fn read_columns_tensor<I>(&self, cells: I) -> anyhow::Result<Tensor>
    where
        I: Iterator<Item = usize>,
    {
        let ((nrow, ncol), triplets) = self.columns_triplets(cells)?;
        Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    /// Build (shape, triplets) for the requested compact rows across
    /// all backends. Output column index is the SparseIoVec-global
    /// column (concatenation of backends in push order); output row
    /// index is the position in `rows`.
    ///
    /// Each `(row, col)` appears at most once, with the same sums
    /// [`Self::for_each_triplet`] reports for it.
    #[allow(clippy::type_complexity)]
    pub fn rows_triplets<I>(
        &self,
        rows: I,
    ) -> anyhow::Result<((usize, usize), Vec<(u64, u64, f32)>)>
    where
        I: Iterator<Item = usize>,
    {
        let rows_compact: Vec<usize> = rows.collect();
        let nrow_out = rows_compact.len();
        let ncol_out = self.num_columns();

        let n_compact = self.cached_num_rows;
        let compact_to_global = self.compact_to_global_row.as_slice();

        // Two backends can only land on one (row, col) when a Union column
        // spans them and they observe the same row.
        let mut has_duplicates =
            self.column_alignment == ColumnAlignment::Union && self.backends_share_rows();

        let mut triplets: Vec<(u64, u64, f32)> = Vec::new();
        let mut local_to_out: Vec<usize> = Vec::with_capacity(rows_compact.len());
        for didx in 0..self.data_vec.len() {
            let g2l = self.data_global_to_local_row[didx].as_slice();

            // Every local row folded onto a requested row is read, each
            // mapped to that row's output position.
            local_to_out.clear();
            let mut local_rows: Vec<usize> = Vec::with_capacity(rows_compact.len());
            for (out_row, &c) in rows_compact.iter().enumerate() {
                if c >= n_compact {
                    continue;
                }
                let g = compact_to_global[c];
                let first = g2l.partition_point(|&(global, _)| global < g);
                for &(_, l) in g2l[first..].iter().take_while(|&&(global, _)| global == g) {
                    local_rows.push(l);
                    local_to_out.push(out_row);
                }
            }
            if local_rows.is_empty() {
                continue;
            }
            has_duplicates |= self.data_has_intra_row_merges[didx];

            let (_, _, group_triplets) = self.data_vec[didx].read_triplets_by_rows(local_rows)?;
            let cols_map = self
                .data_to_cols
                .get(&didx)
                .ok_or_else(|| anyhow::anyhow!("missing data_to_cols entry for didx {}", didx))?;
            triplets.reserve(group_triplets.len());
            for (i, j, v) in group_triplets {
                // `usize::MAX` marks a cell dropped by `mask_columns`.
                let mapped = cols_map[j as usize];
                if mapped == usize::MAX {
                    continue;
                }
                let out_row = local_to_out[i as usize] as u64;
                triplets.push((out_row, mapped as u64, v));
            }
        }

        if has_duplicates {
            crate::sparse_util::sum_duplicate_triplets(&mut triplets);
        }
        // Divided after the duplicates are summed: one factor per column.
        if let Some(scales) = self.column_scales.as_ref() {
            for t in triplets.iter_mut() {
                t.2 /= scales[t.1 as usize];
            }
        }

        Ok(((nrow_out, ncol_out), triplets))
    }

    #[cfg(feature = "ndarray")]
    pub fn read_rows_ndarray<I>(&self, rows: I) -> anyhow::Result<ndarray::Array2<f32>>
    where
        I: Iterator<Item = usize>,
    {
        let ((nrow, ncol), triplets) = self.rows_triplets(rows)?;
        ndarray::Array2::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    pub fn read_rows_dmatrix<I>(&self, rows: I) -> anyhow::Result<nalgebra::DMatrix<f32>>
    where
        I: Iterator<Item = usize>,
    {
        let ((nrow, ncol), triplets) = self.rows_triplets(rows)?;
        DMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    pub fn read_rows_csc<I>(&self, rows: I) -> anyhow::Result<CscMatrix<f32>>
    where
        I: Iterator<Item = usize>,
    {
        let ((nrow, ncol), triplets) = self.rows_triplets(rows)?;
        nalgebra_sparse::CscMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    pub fn read_rows_csr<I>(&self, rows: I) -> anyhow::Result<CsrMatrix<f32>>
    where
        I: Iterator<Item = usize>,
    {
        let ((nrow, ncol), triplets) = self.rows_triplets(rows)?;
        nalgebra_sparse::CsrMatrix::<f32>::from_nonzero_triplets(nrow, ncol, &triplets)
    }

    #[cfg(feature = "tensor")]
    pub fn read_rows_tensor<I>(&self, rows: I) -> anyhow::Result<Tensor>
    where
        I: Iterator<Item = usize>,
    {
        let ((nrow, ncol), triplets) = self.rows_triplets(rows)?;
        Tensor::from_nonzero_triplets(nrow, ncol, &triplets)
    }
}