raphtory_api/
compute.rs

1use rayon::prelude::*;
2/// Compute cumulative sum in parallel over `num_chunks` chunks
3pub fn par_cum_sum(values: &mut [usize]) {
4    let num_chunks = rayon::current_num_threads();
5    let chunk_size = values.len().div_ceil(num_chunks);
6    let mut chunk_sums = Vec::with_capacity(num_chunks);
7    values
8        .par_chunks_mut(chunk_size)
9        .map(|chunk| {
10            let mut cum_sum = 0;
11            for v in chunk {
12                *v += cum_sum;
13                cum_sum = *v;
14            }
15            cum_sum
16        })
17        .collect_into_vec(&mut chunk_sums);
18
19    let mut cum_sum = 0;
20    for (partial_sum, next_chunk) in chunk_sums
21        .into_iter()
22        .zip(values.chunks_mut(chunk_size).skip(1))
23    {
24        cum_sum += partial_sum;
25        next_chunk.par_iter_mut().for_each(|v| *v += cum_sum);
26    }
27}
28
29#[cfg(test)]
30mod test {
31    use super::par_cum_sum;
32
33    #[test]
34    fn test_cum_sum() {
35        let mut values: Vec<_> = (0..100).collect();
36        par_cum_sum(&mut values);
37        let mut cum_sum = 0;
38        for (index, v) in values.into_iter().enumerate() {
39            cum_sum += index;
40            assert_eq!(v, cum_sum)
41        }
42    }
43}