spatialrust_search/
parallel_chunked.rs

1//! Parallel dispatch for chunked spatial queries (`parallel` feature).
2//!
3//! Each [`SpatialTensor`] chunk is queried on its own thread; partial results are
4//! merged without locks because [`KdTree`](crate::KdTree) search is read-only.
5
6use spatialrust_core::SpatialTensor;
7
8use crate::chunked::{
9    nearest_k_spatial_tensor, radius_search_spatial_tensor, ChunkedNearestNeighborIndex,
10    ChunkedRadiusSearchIndex,
11};
12use crate::Neighbor;
13
14/// Minimum point count before threaded chunk dispatch is used.
15pub const PARALLEL_CHUNK_QUERY_MIN_POINTS: usize = 4_096;
16
17/// Parallel radius search over [`SpatialTensor`] chunks.
18pub fn radius_search_spatial_tensor_parallel<I: ChunkedRadiusSearchIndex + Sync>(
19    index: &I,
20    x: &[f32],
21    y: &[f32],
22    z: &[f32],
23    tensor: &SpatialTensor<'_>,
24    radius: f32,
25) -> Vec<(usize, Neighbor)> {
26    let mut out = Vec::new();
27    radius_search_spatial_tensor_parallel_into(index, x, y, z, tensor, radius, &mut out);
28    out
29}
30
31/// Parallel radius search, appending `(query_index, neighbor)` pairs to `out`.
32pub fn radius_search_spatial_tensor_parallel_into<I: ChunkedRadiusSearchIndex + Sync>(
33    index: &I,
34    x: &[f32],
35    y: &[f32],
36    z: &[f32],
37    tensor: &SpatialTensor<'_>,
38    radius: f32,
39    out: &mut Vec<(usize, Neighbor)>,
40) {
41    if tensor.is_empty() {
42        return;
43    }
44    if tensor.len() < PARALLEL_CHUNK_QUERY_MIN_POINTS {
45        radius_search_spatial_tensor(index, x, y, z, tensor, radius, out);
46        return;
47    }
48
49    let ranges: Vec<_> = tensor.chunks().map(|chunk| chunk.range()).collect();
50    std::thread::scope(|scope| {
51        let mut handles = Vec::with_capacity(ranges.len());
52        for query_range in ranges {
53            handles.push(scope.spawn(|| {
54                let mut local = Vec::new();
55                index.radius_search_chunk_into(x, y, z, query_range, radius, &mut local);
56                local
57            }));
58        }
59
60        for handle in handles {
61            out.extend(handle.join().expect("chunk radius search thread panicked"));
62        }
63    });
64}
65
66/// Parallel k-NN search over [`SpatialTensor`] chunks.
67pub fn nearest_k_spatial_tensor_parallel<I: ChunkedNearestNeighborIndex + Sync>(
68    index: &I,
69    x: &[f32],
70    y: &[f32],
71    z: &[f32],
72    tensor: &SpatialTensor<'_>,
73    k: usize,
74) -> Vec<(usize, Neighbor)> {
75    let mut out = Vec::new();
76    nearest_k_spatial_tensor_parallel_into(index, x, y, z, tensor, k, &mut out);
77    out
78}
79
80/// Parallel k-NN search, appending `(query_index, neighbor)` pairs to `out`.
81pub fn nearest_k_spatial_tensor_parallel_into<I: ChunkedNearestNeighborIndex + Sync>(
82    index: &I,
83    x: &[f32],
84    y: &[f32],
85    z: &[f32],
86    tensor: &SpatialTensor<'_>,
87    k: usize,
88    out: &mut Vec<(usize, Neighbor)>,
89) {
90    if tensor.is_empty() || k == 0 {
91        return;
92    }
93    if tensor.len() < PARALLEL_CHUNK_QUERY_MIN_POINTS {
94        nearest_k_spatial_tensor(index, x, y, z, tensor, k, out);
95        return;
96    }
97
98    let ranges: Vec<_> = tensor.chunks().map(|chunk| chunk.range()).collect();
99    std::thread::scope(|scope| {
100        let mut handles = Vec::with_capacity(ranges.len());
101        for query_range in ranges {
102            handles.push(scope.spawn(|| {
103                let mut local = Vec::new();
104                index.nearest_k_chunk_into(x, y, z, query_range, k, &mut local);
105                local
106            }));
107        }
108
109        for handle in handles {
110            out.extend(handle.join().expect("chunk k-NN search thread panicked"));
111        }
112    });
113}
114
115#[cfg(test)]
116mod tests {
117    use super::{
118        nearest_k_spatial_tensor_parallel, radius_search_spatial_tensor_parallel,
119        PARALLEL_CHUNK_QUERY_MIN_POINTS,
120    };
121    use crate::chunked::{nearest_k_spatial_tensor, radius_search_spatial_tensor};
122    use crate::kdtree::KdTree;
123    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
124
125    fn grid_cloud(side: usize) -> (Vec<f32>, Vec<f32>, Vec<f32>, spatialrust_core::PointCloud) {
126        let mut x = Vec::with_capacity(side * side);
127        let mut y = Vec::with_capacity(side * side);
128        let mut z = Vec::with_capacity(side * side);
129        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
130        for row in 0..side {
131            for col in 0..side {
132                let px = col as f32 * 0.01;
133                let py = row as f32 * 0.01;
134                x.push(px);
135                y.push(py);
136                z.push(0.0);
137                builder.push_point([px, py, 0.0]).unwrap();
138            }
139        }
140        (x, y, z, builder.build().unwrap())
141    }
142
143    fn sort_neighbors(neighbors: &mut [(usize, crate::Neighbor)]) {
144        neighbors.sort_by_key(|a| (a.0, a.1.index));
145    }
146
147    #[test]
148    fn parallel_radius_matches_sequential_on_large_grid() {
149        let side = 70; // 4900 points > PARALLEL_CHUNK_QUERY_MIN_POINTS
150        assert!(side * side >= PARALLEL_CHUNK_QUERY_MIN_POINTS);
151        let (x, y, z, cloud) = grid_cloud(side);
152        let tree = KdTree::from_slices(&x, &y, &z);
153        let tensor = cloud.spatial_tensor_chunks(512).unwrap();
154
155        let mut sequential = Vec::new();
156        radius_search_spatial_tensor(&tree, &x, &y, &z, &tensor, 0.02, &mut sequential);
157
158        let mut parallel = radius_search_spatial_tensor_parallel(&tree, &x, &y, &z, &tensor, 0.02);
159
160        sort_neighbors(&mut sequential);
161        sort_neighbors(&mut parallel);
162        assert_eq!(parallel, sequential);
163    }
164
165    #[test]
166    fn parallel_knn_matches_sequential_on_large_grid() {
167        let side = 70;
168        let (x, y, z, cloud) = grid_cloud(side);
169        let tree = KdTree::from_slices(&x, &y, &z);
170        let tensor = cloud.spatial_tensor_chunks(512).unwrap();
171
172        let mut sequential = Vec::new();
173        nearest_k_spatial_tensor(&tree, &x, &y, &z, &tensor, 8, &mut sequential);
174
175        let mut parallel = nearest_k_spatial_tensor_parallel(&tree, &x, &y, &z, &tensor, 8);
176
177        sort_neighbors(&mut sequential);
178        sort_neighbors(&mut parallel);
179        assert_eq!(parallel, sequential);
180    }
181
182    #[test]
183    fn parallel_falls_back_below_threshold() {
184        let (x, y, z, cloud) = grid_cloud(10); // 100 points
185        assert!(cloud.len() < PARALLEL_CHUNK_QUERY_MIN_POINTS);
186        let tree = KdTree::from_slices(&x, &y, &z);
187        let tensor = cloud.spatial_tensor_chunks(32).unwrap();
188
189        let mut sequential = Vec::new();
190        radius_search_spatial_tensor(&tree, &x, &y, &z, &tensor, 0.05, &mut sequential);
191
192        let mut parallel = radius_search_spatial_tensor_parallel(&tree, &x, &y, &z, &tensor, 0.05);
193
194        sort_neighbors(&mut sequential);
195        sort_neighbors(&mut parallel);
196        assert_eq!(parallel, sequential);
197    }
198}
spatialrust_search/parallel_chunked.rs

spatialrust_search/
parallel_chunked.rs
