data_beans/
sparse_util.rs

1use clap::ValueEnum;
2
3pub struct ValuesIndicesPointers<'a> {
4    pub values: &'a [f32],
5    pub indices: &'a [u64],
6    pub indptr: &'a [u64],
7}
8
9pub struct CooTripletsShape {
10    pub triplets: Vec<(u64, u64, f32)>,
11    pub shape: TripletsShape,
12}
13
14pub struct TripletsShape {
15    pub nrows: usize,
16    pub ncols: usize,
17    pub nnz: usize,
18}
19
20#[derive(ValueEnum, Clone, Copy, Debug, PartialEq)]
21#[clap(rename_all = "lowercase")]
22pub enum IndexPointerType {
23    Column,
24    Row,
25}
26
27pub trait SparseTripletsTraits {
28    /// convert sparse pointers into (row, column, value) triplets
29    fn to_coo(&self, pointer_type: IndexPointerType) -> anyhow::Result<CooTripletsShape>;
30}
31
32/////////////////////
33// implementations //
34/////////////////////
35
36impl<'a> SparseTripletsTraits for ValuesIndicesPointers<'a> {
37    fn to_coo(&self, pointer_type: IndexPointerType) -> anyhow::Result<CooTripletsShape> {
38        use rayon::prelude::*;
39
40        let indices = self.indices;
41        let indptr = self.indptr;
42        let values = self.values;
43
44        let nelem = values.len();
45        if nelem != indices.len() {
46            return Err(anyhow::anyhow!(
47                "`values` and `indices` have different sizes"
48            ));
49        }
50        if indptr.is_empty() {
51            return Err(anyhow::anyhow!("`indptr` is empty"));
52        }
53        let nvectors = indptr.len() - 1;
54
55        // Build the triplet list in parallel without a global lock. `flat_map_iter`
56        // lets each rayon worker produce a contiguous run of triplets from one
57        // compressed vector, and rayon's ordered collect joins them into a single
58        // Vec — no Arc<Mutex>, no per-vector intermediate Vec, no lock contention.
59        let triplets: Vec<(u64, u64, f32)> = (0..nvectors)
60            .into_par_iter()
61            .flat_map_iter(|idx| {
62                let j = idx as u64;
63                let start = indptr[idx] as usize;
64                let end = indptr[idx + 1] as usize;
65                let end = end.min(nelem);
66                let start = start.min(end);
67                indices[start..end]
68                    .iter()
69                    .zip(values[start..end].iter())
70                    .map(move |(&i, &x_ij)| match pointer_type {
71                        IndexPointerType::Column => (i, j, x_ij),
72                        IndexPointerType::Row => (j, i, x_ij),
73                    })
74            })
75            .collect();
76
77        let nnz = triplets.len();
78
79        // Shape is derivable from the compressed layout directly — no need to
80        // scan the whole triplets Vec again.
81        let max_idx_index = indices.par_iter().copied().max().unwrap_or(0) as usize + 1;
82        let (nrows, ncols) = match pointer_type {
83            IndexPointerType::Column => (max_idx_index, nvectors),
84            IndexPointerType::Row => (nvectors, max_idx_index),
85        };
86
87        Ok(CooTripletsShape {
88            triplets,
89            shape: TripletsShape { nrows, ncols, nnz },
90        })
91    }
92}
data_beans/sparse_util.rs

data_beans/
sparse_util.rs
