raphtory_api/core/entities/properties/
props.rs

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