data_beans/sparse_io/
helpers.rs

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

data_beans/sparse_io/
helpers.rs
