trs_dataframe/dataframe/column_store/
key_index.rs

1use data_value::DataValue;
2use halfbrown::HashMap;
3use ndarray::ArrayView1;
4use serde::{Deserialize, Serialize};
5
6use crate::{error::Error, Key};
7
8/// [`KeyIndex`] is used to store the keys for the [`super::ColumnFrame`]
9/// The keys are stored in the order they are added - the order is preserved
10/// The keys are stored in the [`Vec`] and the indexes are stored in the [`HashMap`]
11/// The indexes are used to access the data in the [`super::ColumnFrame`] by the column [`Key`]
12/// NOTE: The keys are unique - if the key is already present, it will be removed
13#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
14pub struct KeyIndex {
15    pub keys: Vec<Key>,
16    indexes: HashMap<String, usize>,
17    pub alias: HashMap<String, String>,
18}
19
20impl KeyIndex {
21    pub fn new(keys: Vec<Key>) -> Self {
22        let mut indexes = HashMap::with_capacity(keys.len());
23        let mut removed = 0;
24        let mut actual_keys = Vec::with_capacity(keys.len());
25        for (idx, key) in keys.into_iter().enumerate() {
26            if indexes.contains_key(key.name()) {
27                removed += 1;
28            } else {
29                indexes.insert(key.name().to_string(), idx.saturating_sub(removed));
30                actual_keys.push(key)
31            }
32        }
33        Self {
34            keys: actual_keys,
35            indexes,
36            alias: HashMap::new(),
37        }
38    }
39
40    pub fn len(&self) -> usize {
41        self.keys.len()
42    }
43    pub fn is_empty(&self) -> bool {
44        self.keys.is_empty()
45    }
46
47    pub fn get_column_index(&self, key: &Key) -> Option<usize> {
48        self.get_column_index_by_name(key.name())
49    }
50
51    pub fn get_column_index_by_name(&self, key: &str) -> Option<usize> {
52        if let Some(f) = self.indexes.get(key) {
53            Some(*f)
54        } else {
55            self.alias
56                .get(key)
57                .and_then(|alias| self.indexes.get(alias).copied())
58        }
59    }
60
61    pub fn get_keys(&self) -> &[Key] {
62        &self.keys
63    }
64    pub fn get_keys_mut(&mut self) -> &mut [Key] {
65        &mut self.keys
66    }
67
68    pub fn get_key(&self, idx: usize) -> Option<Key> {
69        self.keys.get(idx).cloned()
70    }
71    pub fn get_complement_keys(&self, keys: &[Key]) -> Vec<Key> {
72        self.keys
73            .iter()
74            .filter(|key| !keys.contains(key))
75            .cloned()
76            .collect()
77    }
78
79    pub fn select(&self, keys: &[Key]) -> KeyIndex {
80        let mut new_keys = Vec::with_capacity(keys.len());
81        let mut new_indexes = HashMap::with_capacity(keys.len());
82
83        for key in keys.iter() {
84            if let Some(idx) = self.indexes.get(key.name()) {
85                new_indexes.insert(key.name().to_string(), *idx);
86                new_keys.push(key.to_owned());
87            } else if let Some(alias_key) = self.alias.get(key.name()) {
88                if let Some(idx) = self.indexes.get(alias_key) {
89                    new_indexes.insert(key.name().to_string(), *idx);
90                    new_keys.push(key.to_owned());
91                }
92            }
93        }
94        Self {
95            keys: new_keys,
96            indexes: new_indexes,
97            alias: HashMap::new(),
98        }
99    }
100
101    pub fn indexes(&self) -> Vec<usize> {
102        self.indexes.values().copied().collect()
103    }
104
105    pub fn store_key(&mut self, key: Key) {
106        if self.indexes.contains_key(key.name()) {
107            return;
108        }
109        self.keys.push(key.clone());
110        self.indexes
111            .insert(key.name().to_string(), self.keys.len() - 1);
112    }
113
114    pub fn remove_key(&mut self, key: &Key) -> Option<(Key, usize)> {
115        let idx = self.indexes.remove(key.name())?;
116        let current = self.keys.remove(idx);
117        Some((current, idx))
118    }
119
120    pub fn get_as_candidate(&self, row: ArrayView1<DataValue>) -> HashMap<Key, DataValue> {
121        let mut result = HashMap::with_capacity(self.keys.len());
122        for (key, idx) in self.indexes.iter() {
123            result.insert(key.into(), row[*idx].clone());
124        }
125        result
126    }
127
128    pub fn to_vec_row(&self, candidate: HashMap<Key, DataValue>) -> Vec<DataValue> {
129        self.keys
130            .iter()
131            .map(|key| candidate.get(key).cloned().unwrap_or_default())
132            .collect()
133    }
134
135    pub fn check_order_of_indexes(&self, other: &Self) -> Result<(), Error> {
136        for (self_key, other_key) in self.keys.iter().zip(other.keys.iter()) {
137            if self_key != other_key {
138                return Err(Error::IndexOutOfOrder(
139                    self.keys.clone(),
140                    other.keys.clone(),
141                ));
142            }
143        }
144        Ok(())
145    }
146
147    pub fn rename_key(&mut self, key: &str, new_key: Key) -> Result<(), Error> {
148        if let Some(idx) = self.indexes.remove(key) {
149            self.indexes.insert(new_key.to_string(), idx);
150            self.keys[idx] = new_key;
151            Ok(())
152        } else {
153            Err(Error::NotFound(key.into()))
154        }
155    }
156
157    pub fn add_alias(&mut self, key: &str, alias: &str) -> Result<(), Error> {
158        if !self.indexes.contains_key(key) {
159            return Err(Error::NotFound(key.into()));
160        }
161        self.alias.insert(alias.to_string(), key.to_string());
162        Ok(())
163    }
164}
165
166impl From<Vec<Key>> for KeyIndex {
167    fn from(keys: Vec<Key>) -> Self {
168        Self::new(keys)
169    }
170}
171
172#[cfg(test)]
173mod test {
174    use crate::DataType;
175
176    use super::*;
177    use rstest::*;
178
179    #[rstest]
180    fn test_alias() {
181        let key = Key::new("a", DataType::U32);
182        let mut key_index = KeyIndex::new(vec![key.clone()]);
183        assert_eq!(key_index.add_alias(key.name(), "alias"), Ok(()));
184        assert!(key_index.add_alias("c", "alias").is_err());
185        assert_eq!(key_index.alias.get("alias"), Some(&key.name().to_string()));
186        assert_eq!(key_index.get_column_index(&key), Some(0));
187        assert_eq!(
188            key_index.get_column_index(&Key::new("alias", DataType::U32)),
189            Some(0)
190        );
191    }
192
193    #[rstest]
194    fn test_rename() {
195        let key = Key::new("a", DataType::U32);
196        let mut key_index = KeyIndex::new(vec![key.clone(), Key::new("b", DataType::U32)]);
197        assert_eq!(key_index.rename_key("a", "new_key".into()), Ok(()));
198        assert_eq!(
199            key_index.get_column_index(&Key::new("new_key", DataType::U32)),
200            Some(0)
201        );
202        assert!(key_index.rename_key("c", "alias".into()).is_err());
203    }
204
205    #[rstest]
206    #[case(
207        vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)],
208        vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)]
209    )]
210    fn test_key_index_new(#[case] keys: Vec<Key>, #[case] expected: Vec<Key>) {
211        let key_index = KeyIndex::new(keys);
212        assert_eq!(key_index.keys, expected);
213        assert_eq!(key_index.get_keys(), expected);
214    }
215
216    #[rstest]
217    #[case(
218        vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)],
219        vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)],
220        Ok(())
221    )]
222    #[case(
223        vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)],
224        vec![Key::new("b", DataType::U32), Key::new("a", DataType::U32)],
225        Err(Error::IndexOutOfOrder(
226            vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)],
227            vec![Key::new("b", DataType::U32), Key::new("a", DataType::U32)]
228        ))
229    )]
230    fn test_key_index_check_order_of_indexes(
231        #[case] keys: Vec<Key>,
232        #[case] other_keys: Vec<Key>,
233        #[case] expected: Result<(), Error>,
234    ) {
235        let key_index: KeyIndex = keys.into();
236        let other_key_index = KeyIndex::new(other_keys);
237        assert_eq!(key_index.check_order_of_indexes(&other_key_index), expected);
238    }
239
240    #[rstest]
241    #[case(
242        KeyIndex::new(vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)]),
243        (Key::new("a", DataType::U32), 0)
244    )]
245    #[case(
246        KeyIndex::new(vec![Key::new("a", DataType::U32), Key::new("b", DataType::U32)]),
247        (Key::new("b", DataType::U32), 1)
248    )]
249    fn test_key_index_remove_key(#[case] mut key_index: KeyIndex, #[case] expected: (Key, usize)) {
250        let key = expected.0.clone();
251        assert_eq!(key_index.remove_key(&key), Some(expected));
252    }
253}
trs_dataframe/dataframe/column_store/key_index.rs

trs_dataframe/dataframe/column_store/
key_index.rs
