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        // Columns are dense ids, so the reverse map is a plain vector;
28        // `ncolumns_per_group` can leave a column out of every group.
29        let mut col_to_group = vec![usize::MAX; column_to_group.len()];
30        for (g, cols) in group_to_cols.iter().enumerate() {
31            for &j in cols {
32                col_to_group[j] = g;
33            }
34        }
35
36        self.derived.group_keys = Some(group_keys);
37        self.derived.group_to_cols = Some(group_to_cols);
38        self.derived.col_to_group = Some(col_to_group);
39    }
40
41    /// Take a vector of columns where each vector corresponds to a set
42    pub fn take_grouped_columns(&self) -> Option<&Vec<Vec<usize>>> {
43        self.derived.group_to_cols.as_ref()
44    }
45
46    /// Get the group keys in the same order as group indices
47    pub fn group_keys(&self) -> Option<&Vec<Box<str>>> {
48        self.derived.group_keys.as_ref()
49    }
50
51    /// Get a mapping from group keys to their column indices
52    pub fn group_key_to_cols(&self) -> Option<HashMap<Box<str>, Vec<usize>>> {
53        if let (Some(keys), Some(cols)) = (&self.derived.group_keys, &self.derived.group_to_cols) {
54            Some(
55                keys.iter()
56                    .zip(cols.iter())
57                    .map(|(k, c)| (k.clone(), c.clone()))
58                    .collect(),
59            )
60        } else {
61            None
62        }
63    }
64
65    /// Take a vector of backend file and corresponding column indices
66    pub fn take_backend_columns(&self) -> Vec<(Box<str>, Vec<usize>)> {
67        self.data_to_cols
68            .iter()
69            .filter_map(|(&didx, cols)| {
70                if let Some(arc_data) = self.data_vec.get(didx) {
71                    let k = arc_data.get_backend_file_name();
72                    // Drop sentinels left by `mask_columns` (masked-out cells).
73                    let kept: Vec<usize> =
74                        cols.iter().copied().filter(|&c| c != usize::MAX).collect();
75                    Some((Box::<str>::from(k), kept))
76                } else {
77                    None
78                }
79            })
80            .collect()
81    }
82
83    /// Recall the `cells` group assignment; Note that this can be
84    /// differ from the original vector used in `assign_groups` as we
85    /// can have different number of columns and groups.
86    pub fn get_group_membership<I>(&self, cells: I) -> anyhow::Result<Vec<usize>>
87    where
88        I: Iterator<Item = usize>,
89    {
90        let cell_to_group = self
91            .derived
92            .col_to_group
93            .as_ref()
94            .expect("groups were not assigned");
95
96        cells
97            .map(|j| {
98                cell_to_group
99                    .get(j)
100                    .copied()
101                    .filter(|&g| g != usize::MAX)
102                    .ok_or_else(|| anyhow::anyhow!("missing group membership"))
103            })
104            .collect()
105    }
106
107    /// number of groups
108    pub fn num_groups(&self) -> usize {
109        self.derived
110            .group_to_cols
111            .as_ref()
112            .map(|x| x.len())
113            .unwrap_or(0)
114    }
115}
data_beans/sparse_io_vector/groups.rs

data_beans/sparse_io_vector/
groups.rs
