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    ///
92    /// # Panics
93    /// If `batch_membership` is not one entry per column;
94    /// [`try_register_batch_membership`](Self::try_register_batch_membership)
95    /// returns that as an error instead.
96    pub fn register_batch_membership<T>(&mut self, batch_membership: &[T])
97    where
98        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
99    {
100        self.try_register_batch_membership(batch_membership)
101            .unwrap_or_else(|e| panic!("{e}"));
102    }
103
104    /// [`register_batch_membership`](Self::register_batch_membership) that
105    /// returns a membership of the wrong length as an error. The layers
106    /// share their columns, so a refused membership is registered on none.
107    pub fn try_register_batch_membership<T>(&mut self, batch_membership: &[T]) -> anyhow::Result<()>
108    where
109        T: Sync + Send + std::hash::Hash + Eq + Clone + ToString,
110    {
111        for layer in self.stack.iter_mut() {
112            layer.try_register_batch_membership(batch_membership)?;
113        }
114        Ok(())
115    }
116
117    /// Get the row names combined across all the types. We will
118    /// append additional data type index: `format!("{}/{}", x, d)`
119    ///
120    pub fn row_names(&self) -> anyhow::Result<Vec<Box<str>>> {
121        Ok(self
122            .stack
123            .iter()
124            .enumerate()
125            .map(|(d, x)| {
126                x.row_names().map(|x| {
127                    x.into_iter()
128                        .map(|y| format!("{}/{}", y, d).into_boxed_str())
129                        .collect::<Vec<_>>()
130                })
131            })
132            .collect::<anyhow::Result<Vec<_>>>()?
133            .into_iter()
134            .flatten()
135            .collect())
136    }
137}
data_beans/sparse_io_stack.rs

data_beans/
sparse_io_stack.rs
