raphtory_api/core/entities/properties/
meta.rs

1use std::{ops::Deref, sync::Arc};
2
3use parking_lot::RwLock;
4use serde::{Deserialize, Serialize};
5
6use crate::core::{
7    entities::properties::prop::{unify_types, PropError, PropType},
8    storage::{
9        arc_str::ArcStr,
10        dict_mapper::{DictMapper, MaybeNew},
11        locked_vec::ArcReadLockedVec,
12    },
13};
14
15#[derive(Serialize, Deserialize, Debug)]
16pub struct Meta {
17    temporal_prop_mapper: PropMapper,
18    metadata_mapper: PropMapper,
19    layer_mapper: DictMapper,
20    node_type_mapper: DictMapper,
21}
22
23impl Default for Meta {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl Meta {
30    pub fn set_metadata_mapper(&mut self, meta: PropMapper) {
31        self.metadata_mapper = meta;
32    }
33    pub fn set_temporal_prop_meta(&mut self, meta: PropMapper) {
34        self.temporal_prop_mapper = meta;
35    }
36    pub fn metadata_mapper(&self) -> &PropMapper {
37        &self.metadata_mapper
38    }
39
40    pub fn temporal_prop_mapper(&self) -> &PropMapper {
41        &self.temporal_prop_mapper
42    }
43
44    pub fn layer_meta(&self) -> &DictMapper {
45        &self.layer_mapper
46    }
47
48    pub fn node_type_meta(&self) -> &DictMapper {
49        &self.node_type_mapper
50    }
51
52    pub fn new() -> Self {
53        let meta_layer = DictMapper::default();
54        let meta_node_type = DictMapper::default();
55        meta_node_type.get_or_create_id("_default");
56        Self {
57            temporal_prop_mapper: PropMapper::default(),
58            metadata_mapper: PropMapper::default(),
59            layer_mapper: meta_layer,
60            node_type_mapper: meta_node_type, // type 0 is the default type for a node
61        }
62    }
63
64    #[inline]
65    pub fn resolve_prop_id(
66        &self,
67        prop: &str,
68        dtype: PropType,
69        is_static: bool,
70    ) -> Result<MaybeNew<usize>, PropError> {
71        if is_static {
72            self.metadata_mapper.get_or_create_and_validate(prop, dtype)
73        } else {
74            self.temporal_prop_mapper
75                .get_or_create_and_validate(prop, dtype)
76        }
77    }
78
79    pub fn get_prop_id(&self, name: &str, is_static: bool) -> Option<usize> {
80        if is_static {
81            self.metadata_mapper.get_id(name)
82        } else {
83            self.temporal_prop_mapper.get_id(name)
84        }
85    }
86
87    pub fn get_prop_id_and_type(&self, name: &str, is_static: bool) -> Option<(usize, PropType)> {
88        if is_static {
89            self.metadata_mapper.get_id_and_dtype(name)
90        } else {
91            self.temporal_prop_mapper.get_id_and_dtype(name)
92        }
93    }
94
95    #[inline]
96    pub fn get_or_create_layer_id(&self, name: Option<&str>) -> MaybeNew<usize> {
97        self.layer_mapper
98            .get_or_create_id(name.unwrap_or("_default"))
99    }
100
101    #[inline]
102    pub fn get_default_node_type_id(&self) -> usize {
103        0usize
104    }
105
106    #[inline]
107    pub fn get_or_create_node_type_id(&self, node_type: &str) -> MaybeNew<usize> {
108        self.node_type_mapper.get_or_create_id(node_type)
109    }
110
111    #[inline]
112    pub fn get_layer_id(&self, name: &str) -> Option<usize> {
113        self.layer_mapper.get_id(name)
114    }
115
116    #[inline]
117    pub fn get_default_layer_id(&self) -> Option<usize> {
118        self.layer_mapper.get_id("_default")
119    }
120
121    #[inline]
122    pub fn get_node_type_id(&self, node_type: &str) -> Option<usize> {
123        self.node_type_mapper.get_id(node_type)
124    }
125
126    pub fn get_layer_name_by_id(&self, id: usize) -> ArcStr {
127        self.layer_mapper.get_name(id)
128    }
129
130    pub fn get_node_type_name_by_id(&self, id: usize) -> Option<ArcStr> {
131        if id == 0 {
132            None
133        } else {
134            Some(self.node_type_mapper.get_name(id))
135        }
136    }
137
138    pub fn get_all_layers(&self) -> Vec<usize> {
139        self.layer_mapper.get_values()
140    }
141
142    pub fn get_all_node_types(&self) -> Vec<ArcStr> {
143        self.node_type_mapper
144            .get_keys()
145            .iter()
146            .filter_map(|key| {
147                if key != "_default" {
148                    Some(key.clone())
149                } else {
150                    None
151                }
152            })
153            .collect()
154    }
155
156    pub fn get_all_property_names(&self, is_static: bool) -> ArcReadLockedVec<ArcStr> {
157        if is_static {
158            self.metadata_mapper.get_keys()
159        } else {
160            self.temporal_prop_mapper.get_keys()
161        }
162    }
163
164    pub fn get_prop_name(&self, prop_id: usize, is_static: bool) -> ArcStr {
165        if is_static {
166            self.metadata_mapper.get_name(prop_id)
167        } else {
168            self.temporal_prop_mapper.get_name(prop_id)
169        }
170    }
171}
172
173#[derive(Default, Debug, Serialize, Deserialize)]
174pub struct PropMapper {
175    id_mapper: DictMapper,
176    dtypes: Arc<RwLock<Vec<PropType>>>,
177}
178
179impl Deref for PropMapper {
180    type Target = DictMapper;
181
182    #[inline]
183    fn deref(&self) -> &Self::Target {
184        &self.id_mapper
185    }
186}
187
188impl PropMapper {
189    pub fn deep_clone(&self) -> Self {
190        let dtypes = self.dtypes.read_recursive().clone();
191        Self {
192            id_mapper: self.id_mapper.deep_clone(),
193            dtypes: Arc::new(RwLock::new(dtypes)),
194        }
195    }
196
197    pub fn get_id_and_dtype(&self, prop: &str) -> Option<(usize, PropType)> {
198        self.get_id(prop).map(|id| {
199            let existing_dtype = self
200                .get_dtype(id)
201                .expect("Existing id should always have a dtype");
202            (id, existing_dtype)
203        })
204    }
205
206    pub fn get_or_create_and_validate(
207        &self,
208        prop: &str,
209        dtype: PropType,
210    ) -> Result<MaybeNew<usize>, PropError> {
211        let wrapped_id = self.id_mapper.get_or_create_id(prop);
212        let id = wrapped_id.inner();
213        let dtype_read = self.dtypes.read_recursive();
214        if let Some(old_type) = dtype_read.get(id) {
215            let mut unified = false;
216            if unify_types(&dtype, old_type, &mut unified).is_ok() {
217                if !unified {
218                    // means the types were equal, no change needed
219                    return Ok(wrapped_id);
220                }
221            } else {
222                return Err(PropError {
223                    name: prop.to_owned(),
224                    expected: old_type.clone(),
225                    actual: dtype,
226                });
227            }
228        }
229        drop(dtype_read); // drop the read lock and wait for write lock as type did not exist yet
230        let mut dtype_write = self.dtypes.write();
231        match dtype_write.get(id).cloned() {
232            Some(old_type) => {
233                if let Ok(tpe) = unify_types(&dtype, &old_type, &mut false) {
234                    dtype_write[id] = tpe;
235                    Ok(wrapped_id)
236                } else {
237                    Err(PropError {
238                        name: prop.to_owned(),
239                        expected: old_type,
240                        actual: dtype,
241                    })
242                }
243            }
244            None => {
245                // vector not resized yet, resize it and set the dtype and return id
246                dtype_write.resize(id + 1, PropType::Empty);
247                dtype_write[id] = dtype;
248                Ok(wrapped_id)
249            }
250        }
251    }
252
253    pub fn set_id_and_dtype(&self, key: impl Into<ArcStr>, id: usize, dtype: PropType) {
254        let mut dtypes = self.dtypes.write();
255        self.set_id(key, id);
256        if dtypes.len() <= id {
257            dtypes.resize(id + 1, PropType::Empty);
258        }
259        dtypes[id] = dtype;
260    }
261
262    pub fn get_dtype(&self, prop_id: usize) -> Option<PropType> {
263        self.dtypes.read_recursive().get(prop_id).cloned()
264    }
265
266    pub fn dtypes(&self) -> impl Deref<Target = Vec<PropType>> + '_ {
267        self.dtypes.read_recursive()
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274
275    #[test]
276    fn test_get_or_create_and_validate_new_property() {
277        let prop_mapper = PropMapper::default();
278        let result = prop_mapper.get_or_create_and_validate("new_prop", PropType::U8);
279        assert!(result.is_ok());
280        assert_eq!(result.unwrap().inner(), 0);
281        assert_eq!(prop_mapper.get_dtype(0), Some(PropType::U8));
282    }
283
284    #[test]
285    fn test_get_or_create_and_validate_existing_property_same_type() {
286        let prop_mapper = PropMapper::default();
287        prop_mapper
288            .get_or_create_and_validate("existing_prop", PropType::U8)
289            .unwrap();
290        let result = prop_mapper.get_or_create_and_validate("existing_prop", PropType::U8);
291        assert!(result.is_ok());
292        assert_eq!(result.unwrap().inner(), 0);
293        assert_eq!(prop_mapper.get_dtype(0), Some(PropType::U8));
294    }
295
296    #[test]
297    fn test_get_or_create_and_validate_existing_property_different_type() {
298        let prop_mapper = PropMapper::default();
299        prop_mapper
300            .get_or_create_and_validate("existing_prop", PropType::U8)
301            .unwrap();
302        let result = prop_mapper.get_or_create_and_validate("existing_prop", PropType::U16);
303        assert!(result.is_err());
304        if let Err(PropError {
305            name,
306            expected,
307            actual,
308        }) = result
309        {
310            assert_eq!(name, "existing_prop");
311            assert_eq!(expected, PropType::U8);
312            assert_eq!(actual, PropType::U16);
313        } else {
314            panic!("Expected PropertyTypeError");
315        }
316    }
317
318    #[test]
319    fn test_get_or_create_and_validate_unify_types() {
320        let prop_mapper = PropMapper::default();
321        prop_mapper
322            .get_or_create_and_validate("prop", PropType::Empty)
323            .unwrap();
324        let result = prop_mapper.get_or_create_and_validate("prop", PropType::U8);
325        assert!(result.is_ok());
326        assert_eq!(result.unwrap().inner(), 0);
327        assert_eq!(prop_mapper.get_dtype(0), Some(PropType::U8));
328    }
329
330    #[test]
331    fn test_get_or_create_and_validate_resize_vector() {
332        let prop_mapper = PropMapper::default();
333        prop_mapper.set_id_and_dtype("existing_prop", 5, PropType::U8);
334        let result = prop_mapper.get_or_create_and_validate("new_prop", PropType::U16);
335        assert!(result.is_ok());
336        assert_eq!(result.unwrap().inner(), 6);
337        assert_eq!(prop_mapper.get_dtype(6), Some(PropType::U16));
338    }
339
340    #[test]
341    fn test_get_or_create_and_validate_two_independent_properties() {
342        let prop_mapper = PropMapper::default();
343        let result1 = prop_mapper.get_or_create_and_validate("prop1", PropType::U8);
344        let result2 = prop_mapper.get_or_create_and_validate("prop2", PropType::U16);
345        assert!(result1.is_ok());
346        assert!(result2.is_ok());
347        assert_eq!(result1.unwrap().inner(), 0);
348        assert_eq!(result2.unwrap().inner(), 1);
349        assert_eq!(prop_mapper.get_dtype(0), Some(PropType::U8));
350        assert_eq!(prop_mapper.get_dtype(1), Some(PropType::U16));
351    }
352}