spatialrust_search/
staging.rs

1//! Shared parallel index-range staging for per-point algorithms.
2//!
3//! Aligns worker counts and chunk boundaries with [`SpatialTensor`](spatialrust_core::SpatialTensor)
4//! defaults so feature estimation and filtering reuse the same staging policy.
5
6use std::ops::Range;
7
8use spatialrust_core::DEFAULT_SPATIAL_TENSOR_CHUNK_SIZE;
9
10/// Minimum point count before multi-threaded index staging is used.
11pub const PARALLEL_STAGING_MIN_POINTS: usize = 4_096;
12
13/// Returns the number of worker threads to use for `point_count` points.
14#[must_use]
15pub fn parallel_worker_count(point_count: usize) -> usize {
16    parallel_worker_count_with_chunk(point_count, DEFAULT_SPATIAL_TENSOR_CHUNK_SIZE)
17}
18
19/// Returns the number of worker threads for an explicit staging chunk size.
20#[must_use]
21pub fn parallel_worker_count_with_chunk(point_count: usize, chunk_size: usize) -> usize {
22    if point_count < PARALLEL_STAGING_MIN_POINTS || chunk_size == 0 {
23        return 1;
24    }
25    let available = std::thread::available_parallelism().map_or(1, |count| count.get());
26    let useful = (point_count / chunk_size).max(1);
27    available.min(useful)
28}
29
30/// Returns disjoint half-open index ranges covering `[0, point_count)`.
31pub fn parallel_index_ranges(
32    point_count: usize,
33    worker_count: usize,
34) -> impl Iterator<Item = Range<usize>> {
35    let worker_count = worker_count.max(1);
36    let chunk_size = point_count.div_ceil(worker_count);
37    (0..point_count).step_by(chunk_size).map(move |start| {
38        let end = (start + chunk_size).min(point_count);
39        start..end
40    })
41}
42
43/// Runs `work` over disjoint index ranges using [`parallel_worker_count`].
44pub fn parallel_index_for_each<W>(point_count: usize, work: W)
45where
46    W: Fn(Range<usize>) + Send + Sync,
47{
48    let worker_count = parallel_worker_count(point_count);
49    if worker_count == 1 {
50        work(0..point_count);
51        return;
52    }
53
54    std::thread::scope(|scope| {
55        for range in parallel_index_ranges(point_count, worker_count) {
56            scope.spawn(|| work(range));
57        }
58    });
59}
60
61#[cfg(test)]
62mod tests {
63    use super::{parallel_index_ranges, parallel_worker_count, PARALLEL_STAGING_MIN_POINTS};
64
65    #[test]
66    fn single_worker_below_threshold() {
67        assert_eq!(parallel_worker_count(PARALLEL_STAGING_MIN_POINTS - 1), 1);
68    }
69
70    #[test]
71    fn ranges_cover_all_indices() {
72        let point_count = 50_000;
73        let workers = parallel_worker_count(point_count);
74        assert!(workers > 1);
75        let ranges: Vec<_> = parallel_index_ranges(point_count, workers).collect();
76        let covered: usize = ranges.iter().map(|range| range.len()).sum();
77        assert_eq!(covered, point_count);
78        for window in ranges.windows(2) {
79            assert_eq!(window[0].end, window[1].start);
80        }
81    }
82}
spatialrust_search/staging.rs

spatialrust_search/
staging.rs
