data_beans/sparse_io_vector/
groups.rs

1#![allow(dead_code)]
2
3use super::*;
4
5impl SparseIoVec {
6    /// Assign columns to groups
7    ///
8    /// * `column_to_group` - column to group membership
9    /// * `ncolumns_per_group` - number of columns per group. `None`: assign all the columns to the groups; `Some(x)`: limit the maximum number of columns per group to at most `x`.
10    ///
11    pub fn assign_groups<T>(&mut self, column_to_group: &[T], ncolumns_per_group: Option<usize>)
12    where
13        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
14    {
15        let partitions = partition_by_membership(column_to_group, ncolumns_per_group);
16
17        // Sort by keys to ensure consistent ordering
18        let mut sorted_partitions: Vec<_> = partitions
19            .into_iter()
20            .map(|(k, cols)| (k.to_string().into_boxed_str(), cols))
21            .collect();
22        sorted_partitions.sort_by(|a, b| a.0.cmp(&b.0));
23
24        let (group_keys, group_to_cols): (Vec<Box<str>>, Vec<Vec<usize>>) =
25            sorted_partitions.into_iter().unzip();
26
27        let col_to_group: HashMap<_, _> = group_to_cols
28            .iter()
29            .enumerate()
30            .flat_map(|(g, cols)| cols.iter().map(move |&j| (j, g)))
31            .collect();
32
33        self.derived.group_keys = Some(group_keys);
34        self.derived.group_to_cols = Some(group_to_cols);
35        self.derived.col_to_group = Some(col_to_group);
36    }
37
38    /// Take a vector of columns where each vector corresponds to a set
39    pub fn take_grouped_columns(&self) -> Option<&Vec<Vec<usize>>> {
40        self.derived.group_to_cols.as_ref()
41    }
42
43    /// Get the group keys in the same order as group indices
44    pub fn group_keys(&self) -> Option<&Vec<Box<str>>> {
45        self.derived.group_keys.as_ref()
46    }
47
48    /// Get a mapping from group keys to their column indices
49    pub fn group_key_to_cols(&self) -> Option<HashMap<Box<str>, Vec<usize>>> {
50        if let (Some(keys), Some(cols)) = (&self.derived.group_keys, &self.derived.group_to_cols) {
51            Some(
52                keys.iter()
53                    .zip(cols.iter())
54                    .map(|(k, c)| (k.clone(), c.clone()))
55                    .collect(),
56            )
57        } else {
58            None
59        }
60    }
61
62    /// Take a vector of backend file and corresponding column indices
63    pub fn take_backend_columns(&self) -> Vec<(Box<str>, Vec<usize>)> {
64        self.data_to_cols
65            .iter()
66            .filter_map(|(&didx, cols)| {
67                if let Some(arc_data) = self.data_vec.get(didx) {
68                    let k = arc_data.get_backend_file_name();
69                    // Drop sentinels left by `mask_columns` (masked-out cells).
70                    let kept: Vec<usize> =
71                        cols.iter().copied().filter(|&c| c != usize::MAX).collect();
72                    Some((Box::<str>::from(k), kept))
73                } else {
74                    None
75                }
76            })
77            .collect()
78    }
79
80    /// Recall the `cells` group assignment; Note that this can be
81    /// differ from the original vector used in `assign_groups` as we
82    /// can have different number of columns and groups.
83    pub fn get_group_membership<I>(&self, cells: I) -> anyhow::Result<Vec<usize>>
84    where
85        I: Iterator<Item = usize>,
86    {
87        let cell_to_group = self
88            .derived
89            .col_to_group
90            .as_ref()
91            .expect("groups were not assigned");
92
93        cells
94            .map(|j| {
95                cell_to_group
96                    .get(&j)
97                    .copied()
98                    .ok_or_else(|| anyhow::anyhow!("missing group membership"))
99            })
100            .collect()
101    }
102
103    /// number of groups
104    pub fn num_groups(&self) -> usize {
105        self.derived
106            .group_to_cols
107            .as_ref()
108            .map(|x| x.len())
109            .unwrap_or(0)
110    }
111}
data_beans/sparse_io_vector/groups.rs

data_beans/sparse_io_vector/
groups.rs
