data_beans/sparse_io/
helpers.rs

1use ndarray::prelude::*;
2use rustc_hash::FxHashMap as HashMap;
3
4use super::DMatrix;
5
6pub fn build_name2index_map(_names: &[Box<str>]) -> HashMap<Box<str>, usize> {
7    _names
8        .iter()
9        .enumerate()
10        .map(|(r, name)| (name.clone(), r))
11        .collect()
12}
13
14pub fn take_subset_indices_names(
15    new_indices: &[usize],
16    ntot: usize,
17    old_names: Vec<Box<str>>,
18) -> (HashMap<u64, u64>, Vec<Box<str>>) {
19    let mut old2new: HashMap<u64, u64> = Default::default();
20    let mut new2old = vec![];
21    debug_assert!(ntot == old_names.len());
22    let mut k = 0_u64;
23    for idx in new_indices.iter() {
24        if *idx < ntot {
25            old2new.insert(*idx as u64, k);
26            new2old.push(*idx);
27            k += 1;
28        }
29    }
30
31    let new_names = new2old
32        .iter()
33        .map(|&i| old_names[i].clone())
34        .collect::<Vec<Box<str>>>();
35
36    (old2new, new_names)
37}
38
39pub fn take_subset_indices_names_if_needed(
40    new_indices: Option<&Vec<usize>>,
41    ntot: Option<usize>,
42    old_names: Vec<Box<str>>,
43) -> (HashMap<u64, u64>, Vec<Box<str>>) {
44    let ntot = ntot.unwrap_or(old_names.len());
45    if let Some(new_indices) = new_indices {
46        take_subset_indices_names(new_indices, ntot, old_names)
47    } else {
48        let names = old_names;
49        let identity = (0..(ntot as u64))
50            .zip(0..(ntot as u64))
51            .collect::<HashMap<u64, u64>>();
52        (identity, names)
53    }
54}
55
56pub fn ndarray_to_triplets(array: &Array2<f32>) -> Vec<(u64, u64, f32)> {
57    let eps = 1e-6;
58    array
59        .indexed_iter()
60        .filter(|(_, &elem)| elem.abs() > eps)
61        .map(|((row, col), &value)| (row as u64, col as u64, value))
62        .collect::<Vec<(u64, u64, f32)>>()
63}
64
65pub fn dmatrix_to_triplets(matrix: &DMatrix<f32>) -> Vec<(u64, u64, f32)> {
66    let (nrow, _) = matrix.shape();
67    let eps = 1e-6;
68    matrix
69        .iter() // column-major
70        .enumerate()
71        .filter(|(_, &elem)| elem.abs() > eps)
72        .map(|(idx, &value)| {
73            let row = idx % nrow;
74            let col = idx / nrow;
75            (row as u64, col as u64, value)
76        })
77        .collect::<Vec<(u64, u64, f32)>>()
78}
79
80/// Remove a backend at `path`, whether it is a file (`.h5`, `.zarr.zip`) or a
81/// directory (`.zarr`). Nothing happens when the path does not exist.
82pub fn remove_backend_path(path: &str) -> anyhow::Result<()> {
83    let p = std::path::Path::new(path);
84    if p.exists() {
85        if p.is_file() {
86            std::fs::remove_file(p)?;
87        } else {
88            std::fs::remove_dir_all(p)?;
89        }
90    }
91    Ok(())
92}
93
94/// Whether a preload of `nnz` entries fits the budget.
95///
96/// Preloading costs 12 bytes per non-zero (a `u64` index and an `f32` value),
97/// there was no size check anywhere in front of it, and no consumer ever
98/// releases it — so at imaging scale a `--preload-data` was an OOM order, not a
99/// request. The budget turns that into a logged skip: every read path already
100/// handles the not-preloaded state, it is just slower.
101///
102/// `LEGUME_PRELOAD_BUDGET_BYTES` overrides the default, following the
103/// `LEGUME_ZARR_CACHE_CAP` precedent for memory knobs.
104pub fn preload_within_budget(nnz: usize, what: &str) -> bool {
105    const BYTES_PER_NNZ: usize = 12;
106    const DEFAULT_BUDGET_BYTES: usize = 8 << 30;
107    let budget = std::env::var("LEGUME_PRELOAD_BUDGET_BYTES")
108        .ok()
109        .and_then(|v| v.parse::<usize>().ok())
110        .unwrap_or(DEFAULT_BUDGET_BYTES);
111    let cost = nnz.saturating_mul(BYTES_PER_NNZ);
112    if cost > budget {
113        log::warn!(
114            "skipping {what} preload: {cost} bytes ({nnz} nnz x {BYTES_PER_NNZ}) exceeds the \
115             {budget}-byte budget (LEGUME_PRELOAD_BUDGET_BYTES to raise); reads stay on the \
116             streaming path"
117        );
118        false
119    } else {
120        true
121    }
122}
123
124/// Bytes of `(u64, u64, f32)` triplets a streaming-write slab may hold, and the
125/// padded size of one such triplet. One definition, because the two streaming
126/// pipelines (the subset trait method and the handlers' column-selection
127/// writer) each carried their own copy and two memory ceilings drift apart.
128pub const SLAB_BUDGET_BYTES: usize = 256 << 20;
129/// `(u64, u64, f32)` padded.
130pub const TRIPLET_BYTES: usize = 24;
data_beans/sparse_io/helpers.rs

data_beans/sparse_io/
helpers.rs
