data_beans/
sparse_io_stack.rs

1// #![allow(dead_code)]
2
3use crate::sparse_io_vector::*;
4use log::info;
5
6/// `sparse_io_stack` is a stack of `sparse_io_vector`
7pub struct SparseIoStack {
8    pub stack: Vec<SparseIoVec>,
9    column_names: Vec<Box<str>>,
10}
11
12impl Default for SparseIoStack {
13    /// an empty sparse io vector for horizontal data integration
14    fn default() -> Self {
15        Self::new()
16    }
17}
18
19impl SparseIoStack {
20    pub fn new() -> Self {
21        Self {
22            stack: vec![],
23            column_names: vec![],
24        }
25    }
26
27    pub fn push(&mut self, data: SparseIoVec) -> anyhow::Result<()> {
28        if self.stack.is_empty() {
29            self.column_names.extend(data.column_names()?);
30            self.stack.push(data);
31            return Ok(());
32        }
33
34        info!("Checking column names...");
35        if self.column_names != data.column_names()? {
36            return Err(anyhow::anyhow!("column names don't match"));
37        }
38        self.stack.push(data);
39        Ok(())
40    }
41
42    /// number of data types
43    pub fn num_types(&self) -> usize {
44        self.stack.len()
45    }
46
47    /// number of shared columns
48    pub fn num_columns(&self) -> anyhow::Result<usize> {
49        self.stack
50            .iter()
51            .map(|x| x.num_columns())
52            .max()
53            .ok_or(anyhow::anyhow!("can't figure out the max"))
54    }
55
56    /// Get the shared column names.
57    pub fn column_names(&self) -> anyhow::Result<Vec<Box<str>>> {
58        if self.column_names.len() != self.num_columns()? {
59            return Err(anyhow::anyhow!("inconsistent columns"));
60        }
61        Ok(self.column_names.clone())
62    }
63
64    /// Exclude columns (cells) across every layer in the stack in
65    /// lockstep. `keep[global_col]` is `true` to keep, `false` to drop.
66    /// Mirror of [`SparseIoVec::mask_columns`]; the members share one
67    /// synchronized cell axis, so the same mask applies to all. MUST be
68    /// called before batch/group registration.
69    pub fn mask_columns_all(&mut self, keep: &[bool]) -> anyhow::Result<()> {
70        let n = self.num_columns()?;
71        if keep.len() != n {
72            return Err(anyhow::anyhow!(
73                "mask_columns_all: keep.len()={} != num_columns={}",
74                keep.len(),
75                n
76            ));
77        }
78        for layer in self.stack.iter_mut() {
79            layer.mask_columns(keep)?;
80        }
81        // Refresh the cached shared column names from the (now filtered)
82        // first member; all members were masked identically.
83        self.column_names = match self.stack.first() {
84            Some(first) => first.column_names()?,
85            None => vec![],
86        };
87        Ok(())
88    }
89
90    /// Register batch membership for all layers in the stack.
91    pub fn register_batch_membership<T>(&mut self, batch_membership: &[T])
92    where
93        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
94    {
95        for layer in self.stack.iter_mut() {
96            layer.register_batch_membership(batch_membership);
97        }
98    }
99
100    /// Get the row names combined across all the types. We will
101    /// append additional data type index: `format!("{}/{}", x, d)`
102    ///
103    pub fn row_names(&self) -> anyhow::Result<Vec<Box<str>>> {
104        Ok(self
105            .stack
106            .iter()
107            .enumerate()
108            .map(|(d, x)| {
109                x.row_names().map(|x| {
110                    x.into_iter()
111                        .map(|y| format!("{}/{}", y, d).into_boxed_str())
112                        .collect::<Vec<_>>()
113                })
114            })
115            .collect::<anyhow::Result<Vec<_>>>()?
116            .into_iter()
117            .flatten()
118            .collect())
119    }
120}
data_beans/sparse_io_stack.rs

data_beans/
sparse_io_stack.rs
