data_beans/
column_subset.rs

1//! Stream a column selection (and an optional ascending row selection) of a
2//! sparse backend into a fresh backend without materialising the survivors.
3//! Shared by the `split` / `subsample` CLI handlers and by downstream crates
4//! (faba's `qc`) that filter a matrix by cells and features.
5
6/// Stream a column selection (with an optional ASCENDING row filter) into a
7/// fresh backend, without ever materialising the survivors.
8///
9/// The shape `split` and `subsample` share: pick columns, optionally keep an
10/// ascending subset of rows, write a new file. Both used to read every
11/// surviving triplet into one `Vec` and hand it to the sorting triplet writer —
12/// 24 B/nnz of residency that OOMs at imaging scale. Here the exact output nnz
13/// comes from the resident indptr (or one counting pass when rows are
14/// filtered), and survivors flow through the validated streaming writer in
15/// byte-budgeted slabs.
16///
17/// `row_filter`, when given, must be ascending: the renumbering is then
18/// monotone, so within-column row order survives and no per-column sort is
19/// needed. Both callers sample or select ascending; the debug assert keeps the
20/// next caller honest rather than silently writing unsorted columns (which the
21/// streaming writer would refuse anyway).
22pub fn stream_column_selection(
23    data: &dyn crate::sparse_io::SparseIo<IndexIter = Vec<usize>>,
24    selected_columns: &[usize],
25    row_filter: Option<&[usize]>,
26    out_row_names: &[Box<str>],
27    out_col_names: &[Box<str>],
28    file_out: &str,
29    backend_out: &crate::sparse_io::SparseIoBackend,
30) -> anyhow::Result<(usize, usize, usize)> {
31    use crate::sparse_io::*;
32
33    let nrow_full = data
34        .num_rows()
35        .ok_or_else(|| anyhow::anyhow!("backend has no `nrow`"))?;
36
37    // Old-row → new-row, monotone by construction.
38    let row_map: Option<Vec<Option<u64>>> = row_filter.map(|keep| {
39        debug_assert!(
40            keep.windows(2).all(|w| w[0] < w[1]),
41            "row filter must ascend"
42        );
43        let mut map = vec![None; nrow_full];
44        for (new, &old) in keep.iter().enumerate() {
45            map[old] = Some(new as u64);
46        }
47        map
48    });
49    let out_nrow = row_filter.map_or(nrow_full, <[usize]>::len);
50    let out_ncol = selected_columns.len();
51
52    // Exact per-output-column nnz. Free from the indptr when every row
53    // survives; one counting pass — counts, never entries — otherwise.
54    let per_col_nnz: Vec<u64> = match &row_map {
55        None => selected_columns
56            .iter()
57            .map(|&c| {
58                data.column_nnz(c)
59                    .ok_or_else(|| anyhow::anyhow!("no indptr entry for column {c}"))
60            })
61            .collect::<anyhow::Result<_>>()?,
62        Some(map) => {
63            // Block reads here too; counts only, never entries.
64            let mut counts = vec![0u64; out_ncol];
65            let coarse = legume_numeric::matrix::utils::generate_minibatch_intervals(
66                out_ncol,
67                0,
68                Some(8192),
69            );
70            for (lb, ub) in coarse {
71                let (_, _, triplets) =
72                    data.read_triplets_by_columns(selected_columns[lb..ub].to_vec())?;
73                for (r, c_local, _) in triplets {
74                    if map[r as usize].is_some() {
75                        counts[lb + c_local as usize] += 1;
76                    }
77                }
78            }
79            counts
80        }
81    };
82    let nnz: u64 = per_col_nnz.iter().sum();
83
84    let mut out = create_sparse_streaming_empty(Some(file_out), Some(backend_out))?;
85    out.begin_streaming_csc((out_nrow, out_ncol, nnz as usize))?;
86
87    let blocks = legume_numeric::matrix::utils::byte_budget_intervals(
88        &per_col_nnz,
89        crate::sparse_io::SLAB_BUDGET_BYTES,
90        crate::sparse_io::TRIPLET_BYTES,
91    );
92
93    let t_stream = std::time::Instant::now();
94    let mut nnz_offset = 0u64;
95    for (lb, ub) in blocks {
96        // ONE block read per slab, never a read per column: a single-column
97        // read pays the cached-subset machinery per call, and at hundreds of
98        // thousands of columns that made this path slower than the memory wall
99        // it replaced. The block read returns LOCAL column ids for the
100        // requested set, rows ascending within each column — the writer's
101        // invariant already.
102        let (_, _, triplets) = data.read_triplets_by_columns(selected_columns[lb..ub].to_vec())?;
103
104        let n_block = ub - lb;
105        let mut per_col: Vec<Vec<(u64, f32)>> = vec![Vec::new(); n_block];
106        for (r, c_local, x) in triplets {
107            let kept = match &row_map {
108                None => Some(r),
109                Some(map) => map[r as usize],
110            };
111            if let Some(new_r) = kept {
112                per_col[c_local as usize].push((new_r, x));
113            }
114        }
115        let mut local_colptr = Vec::with_capacity(n_block);
116        let mut row_indices = Vec::new();
117        let mut values = Vec::new();
118        for entries in &per_col {
119            local_colptr.push(row_indices.len() as u64);
120            for &(r, x) in entries {
121                row_indices.push(r);
122                values.push(x);
123            }
124        }
125        out.append_csc_slab(lb as u64, nnz_offset, &local_colptr, &row_indices, &values)?;
126        nnz_offset += values.len() as u64;
127    }
128
129    out.finalize_streaming_csc()?;
130    out.build_csr_from_csc_streaming()?;
131    log::info!(
132        "streamed {nnz} entries in {out_ncol} columns ({:.1}s)",
133        t_stream.elapsed().as_secs_f32()
134    );
135    out.register_row_names_vec(out_row_names);
136    out.register_column_names_vec(out_col_names);
137    Ok((out_nrow, out_ncol, nnz as usize))
138}
data_beans/column_subset.rs

data_beans/
column_subset.rs
