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
15#[cfg(feature = "ndarray")]
16pub fn ndarray_to_triplets(array: &Array2<f32>) -> Vec<(u64, u64, f32)> {
17    let eps = 1e-6;
18    array
19        .indexed_iter()
20        .filter(|(_, &elem)| elem.abs() > eps)
21        .map(|((row, col), &value)| (row as u64, col as u64, value))
22        .collect::<Vec<(u64, u64, f32)>>()
23}
24
25pub fn dmatrix_to_triplets(matrix: &DMatrix<f32>) -> Vec<(u64, u64, f32)> {
26    let (nrow, _) = matrix.shape();
27    let eps = 1e-6;
28    matrix
29        .iter() // column-major
30        .enumerate()
31        .filter(|(_, &elem)| elem.abs() > eps)
32        .map(|(idx, &value)| {
33            let row = idx % nrow;
34            let col = idx / nrow;
35            (row as u64, col as u64, value)
36        })
37        .collect::<Vec<(u64, u64, f32)>>()
38}
39
40/// Remove a backend at `path`, whether it is a file (`.h5`, `.zarr.zip`) or a
41/// directory (`.zarr`). Nothing happens when the path does not exist.
42pub fn remove_backend_path(path: &str) -> anyhow::Result<()> {
43    let p = std::path::Path::new(path);
44    if p.exists() {
45        if p.is_file() {
46            std::fs::remove_file(p)?;
47        } else {
48            std::fs::remove_dir_all(p)?;
49        }
50    }
51    Ok(())
52}
53
54/// Bytes reserved by live preloads in this process; see [`reserve_preload`].
55static RESERVED: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
56
57/// Preloading costs 12 bytes per non-zero (a `u64` index and an `f32` value);
58/// so does an entry of a transpose band.
59pub(crate) const PRELOAD_BYTES_PER_NNZ: usize = 12;
60
61/// A preload's share of the process-wide preload budget, given back when it
62/// is dropped. Keep it next to the preloaded arrays and drop it with them,
63/// so a later preload in the same process sees the memory as free again.
64///
65/// Cloning reserves the same amount again: a clone carries its own copy of
66/// the arrays.
67#[derive(Debug)]
68pub struct PreloadReservation {
69    bytes: usize,
70    counter: &'static std::sync::atomic::AtomicUsize,
71}
72
73impl PreloadReservation {
74    /// Bytes this reservation holds against the process-wide budget: 0 for
75    /// a preload admitted under a `LEGUME_PRELOAD_BUDGET_BYTES` override,
76    /// which is per preload and reserves nothing.
77    #[allow(dead_code)] // library API; the bin has no caller
78    pub fn bytes(&self) -> usize {
79        self.bytes
80    }
81}
82
83impl Clone for PreloadReservation {
84    fn clone(&self) -> Self {
85        self.counter
86            .fetch_add(self.bytes, std::sync::atomic::Ordering::Relaxed);
87        Self {
88            bytes: self.bytes,
89            counter: self.counter,
90        }
91    }
92}
93
94impl Drop for PreloadReservation {
95    fn drop(&mut self) {
96        self.counter
97            .fetch_sub(self.bytes, std::sync::atomic::Ordering::Relaxed);
98    }
99}
100
101/// Bytes currently held by live [`PreloadReservation`]s in this process.
102#[allow(dead_code)] // library API; the bin has no caller
103pub fn preload_reserved_bytes() -> usize {
104    RESERVED.load(std::sync::atomic::Ordering::Relaxed)
105}
106
107/// Reserve room for a preload of `nnz` entries, or `None` when it does not
108/// fit the budget.
109///
110/// An unchecked preload at imaging scale is an OOM order, not a request. The
111/// budget turns that into a logged skip: every read path already handles the
112/// not-preloaded state, it is just slower.
113///
114/// The budget is half the memory available when the process first asks,
115/// shared by every live preload of the process (8 GiB when the system does
116/// not say). The reservation is given back when the returned guard drops.
117/// `LEGUME_PRELOAD_BUDGET_BYTES` overrides it per preload, following the
118/// `LEGUME_ZARR_CACHE_CAP` precedent for memory knobs.
119pub fn reserve_preload(nnz: usize, what: &str) -> Option<PreloadReservation> {
120    let cost = nnz.saturating_mul(PRELOAD_BYTES_PER_NNZ);
121    if let Some(budget) = preload_budget_override() {
122        return fits_or_warn(cost, budget, nnz, what).then_some(PreloadReservation {
123            bytes: 0,
124            counter: &RESERVED,
125        });
126    }
127    reserve_against(&RESERVED, cost, auto_preload_budget_bytes(), nnz, what)
128}
129
130/// The per-preload budget `LEGUME_PRELOAD_BUDGET_BYTES` sets, when it is set
131/// to a number.
132fn preload_budget_override() -> Option<usize> {
133    std::env::var("LEGUME_PRELOAD_BUDGET_BYTES")
134        .ok()
135        .and_then(|v| v.parse::<usize>().ok())
136}
137
138fn reserve_against(
139    counter: &'static std::sync::atomic::AtomicUsize,
140    cost: usize,
141    budget: usize,
142    nnz: usize,
143    what: &str,
144) -> Option<PreloadReservation> {
145    use std::sync::atomic::Ordering;
146    let reserved = counter.fetch_add(cost, Ordering::Relaxed);
147    if fits_or_warn(reserved.saturating_add(cost), budget, nnz, what) {
148        Some(PreloadReservation {
149            bytes: cost,
150            counter,
151        })
152    } else {
153        counter.fetch_sub(cost, Ordering::Relaxed);
154        None
155    }
156}
157
158/// Bytes of a transpose band granted out of `room`: all of `wanted` when
159/// there is that much, what there is otherwise, and never less than
160/// [`SLAB_BUDGET_BYTES`] (or `wanted`, when smaller) so a transpose always
161/// makes progress.
162fn transpose_grant(wanted: usize, room: usize) -> usize {
163    wanted.min(room).max(wanted.min(SLAB_BUDGET_BYTES))
164}
165
166/// Reserve up to `wanted` bytes for the row band of a CSC → CSR transpose
167/// (see [`transpose_grant`]). The room is what the preload budget has left,
168/// and no more than half the memory available now: the budget was sized when
169/// the process first asked, and the process has allocated since. Returns
170/// the bytes granted and the reservation that holds them against the budget
171/// until it drops. Under a `LEGUME_PRELOAD_BUDGET_BYTES` override, which is
172/// per preload, the grant is cut from the override alone and nothing is
173/// reserved. The memory available is read once per call, and a grant shows
174/// there only once its band is allocated, so the half is a bound per caller,
175/// not across transposes that start together.
176pub(crate) fn reserve_transpose_band(wanted: usize) -> (usize, Option<PreloadReservation>) {
177    use std::sync::atomic::Ordering;
178    if let Some(budget) = preload_budget_override() {
179        return (transpose_grant(wanted, budget), None);
180    }
181    let budget = auto_preload_budget_bytes();
182    let free_now = available_memory_bytes().map_or(usize::MAX, |avail| avail / 2);
183    let mut reserved = RESERVED.load(Ordering::Relaxed);
184    loop {
185        let room = budget.saturating_sub(reserved).min(free_now);
186        let granted = transpose_grant(wanted, room);
187        match RESERVED.compare_exchange_weak(
188            reserved,
189            reserved.saturating_add(granted),
190            Ordering::Relaxed,
191            Ordering::Relaxed,
192        ) {
193            Ok(_) => {
194                return (
195                    granted,
196                    Some(PreloadReservation {
197                        bytes: granted,
198                        counter: &RESERVED,
199                    }),
200                )
201            }
202            Err(now) => reserved = now,
203        }
204    }
205}
206
207/// Whether a preload of `nnz` entries fits the budget. Unlike
208/// [`reserve_preload`], the reservation is held for the life of the process;
209/// prefer [`reserve_preload`] and keep its guard with the preloaded arrays.
210#[allow(dead_code)] // library API; the bin has no caller
211pub fn preload_within_budget(nnz: usize, what: &str) -> bool {
212    reserve_preload(nnz, what).map(std::mem::forget).is_some()
213}
214
215fn fits_or_warn(cost: usize, budget: usize, nnz: usize, what: &str) -> bool {
216    if cost > budget {
217        log::warn!(
218            "skipping {what} preload of {nnz} nnz: {cost} bytes would exceed the \
219             {budget}-byte budget (LEGUME_PRELOAD_BUDGET_BYTES to raise); reads stay on the \
220             streaming path"
221        );
222        false
223    } else {
224        true
225    }
226}
227
228/// Half the memory available at the first call, or 8 GiB when unknown.
229fn auto_preload_budget_bytes() -> usize {
230    static BUDGET: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
231    *BUDGET.get_or_init(|| {
232        const FALLBACK_BYTES: usize = 8 << 30;
233        available_memory_bytes().map_or(FALLBACK_BYTES, |avail| avail / 2)
234    })
235}
236
237/// Memory available to the process right now; `None` when the system does
238/// not say.
239fn available_memory_bytes() -> Option<usize> {
240    let mut sys = sysinfo::System::new();
241    sys.refresh_memory();
242    usize::try_from(sys.available_memory())
243        .ok()
244        .filter(|&avail| avail > 0)
245}
246
247/// Whether loaders preload columns without being asked. On unless
248/// `LEGUME_AUTO_PRELOAD` is `0`, `false` or `off`; the preload itself still
249/// answers to [`preload_within_budget`], so data that does not fit in memory
250/// stays on the streaming path.
251#[allow(dead_code)] // called by library loaders; the bin has none
252pub fn auto_preload_enabled() -> bool {
253    !matches!(
254        std::env::var("LEGUME_AUTO_PRELOAD").as_deref(),
255        Ok("0") | Ok("false") | Ok("off")
256    )
257}
258
259/// Bytes of `(u64, u64, f32)` triplets a streaming-write slab may hold, and the
260/// padded size of one such triplet. One definition, because the two streaming
261/// pipelines (the subset trait method and the handlers' column-selection
262/// writer) each carried their own copy and two memory ceilings drift apart.
263pub const SLAB_BUDGET_BYTES: usize = 256 << 20;
264/// `(u64, u64, f32)` padded.
265pub const TRIPLET_BYTES: usize = 24;
266
267#[cfg(test)]
268mod preload_reservation_tests {
269    use super::*;
270    use std::sync::atomic::{AtomicUsize, Ordering};
271
272    #[test]
273    fn dropping_a_reservation_gives_it_back() {
274        static COUNTER: AtomicUsize = AtomicUsize::new(0);
275        let a = reserve_against(&COUNTER, 60, 100, 5, "test").expect("fits");
276        assert_eq!(COUNTER.load(Ordering::Relaxed), 60);
277        // A second preload that would overflow the budget is refused and
278        // leaves nothing behind.
279        assert!(reserve_against(&COUNTER, 60, 100, 5, "test").is_none());
280        assert_eq!(COUNTER.load(Ordering::Relaxed), 60);
281        drop(a);
282        assert_eq!(COUNTER.load(Ordering::Relaxed), 0);
283        // With the first one gone, the same preload fits again.
284        let b = reserve_against(&COUNTER, 60, 100, 5, "test").expect("fits after release");
285        assert_eq!(b.bytes(), 60);
286        drop(b);
287        assert_eq!(COUNTER.load(Ordering::Relaxed), 0);
288    }
289
290    #[test]
291    fn a_transpose_band_takes_the_room_there_is_but_never_under_a_slab() {
292        let slab = SLAB_BUDGET_BYTES;
293        // Room for all of it.
294        assert_eq!(transpose_grant(4 * slab, 8 * slab), 4 * slab);
295        // Less room than wanted: the room.
296        assert_eq!(transpose_grant(4 * slab, 2 * slab), 2 * slab);
297        // No room to speak of: still one slab, or all of a smaller want.
298        assert_eq!(transpose_grant(4 * slab, 0), slab);
299        assert_eq!(transpose_grant(slab / 2, 0), slab / 2);
300    }
301
302    #[test]
303    fn a_transpose_band_under_one_slab_is_granted_whole() {
304        assert_eq!(reserve_transpose_band(1024).0, 1024);
305    }
306
307    #[test]
308    fn cloning_a_reservation_reserves_again() {
309        static COUNTER: AtomicUsize = AtomicUsize::new(0);
310        let a = reserve_against(&COUNTER, 30, 100, 5, "test").expect("fits");
311        let b = a.clone();
312        assert_eq!(COUNTER.load(Ordering::Relaxed), 60);
313        drop(a);
314        assert_eq!(COUNTER.load(Ordering::Relaxed), 30);
315        drop(b);
316        assert_eq!(COUNTER.load(Ordering::Relaxed), 0);
317    }
318}
data_beans/sparse_io/helpers.rs

data_beans/sparse_io/
helpers.rs
