Skip to main content

cairn_knowledge_graph/operations/
select_vertex.rs

1use once_cell::sync::Lazy;
2
3use graphblas_sparse_linear_algebra::operators::binary_operator::Second;
4use graphblas_sparse_linear_algebra::util::ElementIndex;
5use graphblas_sparse_linear_algebra::value_types::sparse_vector::{
6    FromVectorElementList, SetVectorElement, SparseVector, VectorElement, VectorElementList,
7};
8
9use crate::error::GraphComputingError;
10
11use crate::graph::edge::{EdgeType, EdgeTypeIndex};
12use crate::graph::graph::{Graph, GraphTrait};
13use crate::graph::vertex::{VertexIndex, VertexKeyAndIndexConversion, VertexKeyRef};
14use crate::operations::selection::vertex_selection::VertexSelection;
15
16use super::select_edge_type::EdgeTypeSelectorTrait;
17
18static SECOND_BINARY_OPERATOR: Lazy<Second<bool, bool, bool>> =
19    Lazy::new(|| Second::<bool, bool, bool>::new());
20
21pub trait SelectVertex {
22    fn select_vertex_by_key(
23        &self,
24        vertex_key: &VertexKeyRef,
25    ) -> Result<VertexSelection, GraphComputingError>;
26    fn select_vertex_by_index(
27        &self,
28        vertex_index: VertexIndex,
29    ) -> Result<VertexSelection, GraphComputingError>;
30
31    fn select_vertices_by_key(
32        &self,
33        vertex_keys: Vec<&VertexKeyRef>,
34    ) -> Result<VertexSelection, GraphComputingError>;
35    fn select_vertices_by_index(
36        &self,
37        vertex_indices: Vec<VertexIndex>,
38    ) -> Result<VertexSelection, GraphComputingError>;
39
40    fn select_vertices_connected_to_vertex_by_key(
41        &self,
42        edge_type: EdgeType,
43        to_vertex_key: &VertexKeyRef,
44    ) -> Result<VertexSelection, GraphComputingError>;
45    fn select_vertices_connected_to_vertex_by_index(
46        &self,
47        edge_type: EdgeTypeIndex,
48        to_vertex_index: &VertexIndex,
49    ) -> Result<VertexSelection, GraphComputingError>;
50
51    fn select_vertices_connected_from_vertex_by_key(
52        &self,
53        edge_type: EdgeType,
54        from_vertex_key: &VertexKeyRef,
55    ) -> Result<VertexSelection, GraphComputingError>;
56    fn select_vertices_connected_from_vertex_by_index(
57        &self,
58        edge_type: EdgeTypeIndex,
59        from_vertex_index: &VertexIndex,
60    ) -> Result<VertexSelection, GraphComputingError>;
61}
62
63impl SelectVertex for Graph {
64    fn select_vertex_by_key(
65        &self,
66        vertex_key: &VertexKeyRef,
67    ) -> Result<VertexSelection, GraphComputingError> {
68        let mut vertex_mask =
69            SparseVector::<bool>::new(self.graphblas_context_ref(), &self.vertex_capacity()?)?;
70        match self.vertex_key_ref_to_vertex_index_ref(vertex_key) {
71            Ok(index) => vertex_mask.set_element(VectorElement::new(index.index(), true))?,
72            Err(_) => (),
73        };
74        VertexSelection::new(self, vertex_mask)
75    }
76    fn select_vertex_by_index(
77        &self,
78        vertex_index: VertexIndex,
79    ) -> Result<VertexSelection, GraphComputingError> {
80        let mut vertex_mask =
81            SparseVector::<bool>::new(self.graphblas_context_ref(), &self.vertex_capacity()?)?;
82        vertex_mask.set_element(VectorElement::new(vertex_index.index(), true))?;
83        VertexSelection::new(self, vertex_mask)
84    }
85
86    fn select_vertices_by_key(
87        &self,
88        vertex_keys: Vec<&VertexKeyRef>,
89    ) -> Result<VertexSelection, GraphComputingError> {
90        let mut vertex_indices: Vec<ElementIndex> = vec![0; vertex_keys.len()]; // TODO: review id pre-allocation actually improves performance
91        let mut mask_values: Vec<bool> = vec![true; vertex_keys.len()];
92
93        // TODO: consider to parallelize
94        (0..vertex_keys.len())
95            .into_iter()
96            .for_each(|vertex_key_index| {
97                match self.vertex_key_ref_to_vertex_index_ref(vertex_keys[vertex_key_index]) {
98                    Ok(vertex_index) => vertex_indices[vertex_key_index] = vertex_index.index(),
99                    Err(_) => mask_values[vertex_key_index] = false,
100                }
101            });
102
103        let mask_elements = VectorElementList::from_vectors(vertex_indices, mask_values)?;
104        let selection_mask = SparseVector::from_element_list(
105            self.graphblas_context_ref(),
106            &self.vertex_capacity()?,
107            &mask_elements,
108            &*SECOND_BINARY_OPERATOR,
109        )?;
110        VertexSelection::new(self, selection_mask)
111    }
112    fn select_vertices_by_index(
113        &self,
114        vertex_indices: Vec<VertexIndex>,
115    ) -> Result<VertexSelection, GraphComputingError> {
116        let mask_values: Vec<bool> = vec![true; vertex_indices.len()];
117
118        // TODO: is this efficient? https://stackoverflow.com/questions/48308759/how-do-i-convert-a-vect-to-a-vecu-without-copying-the-vector
119        // Or should the VertexIndex be refactored to a type alias of usize for better performance?
120        let element_indices: Vec<ElementIndex> = vertex_indices
121            .into_iter()
122            .map(|index| index.index())
123            .collect();
124
125        let mask_elements = VectorElementList::<bool>::from_vectors(element_indices, mask_values)?;
126        let selection_mask = SparseVector::from_element_list(
127            self.graphblas_context_ref(),
128            &self.vertex_capacity()?,
129            &mask_elements,
130            &*SECOND_BINARY_OPERATOR,
131        )?;
132        VertexSelection::new(self, selection_mask)
133    }
134
135    fn select_vertices_connected_to_vertex_by_key<'g>(
136        &'g self,
137        edge_type: EdgeType,
138        to_vertex_key: &VertexKeyRef,
139    ) -> Result<VertexSelection<'g>, GraphComputingError> {
140        let edge_selection = self.select_edge_type(edge_type)?;
141        edge_selection.select_vertices_connected_to_vertex(to_vertex_key)
142    }
143    fn select_vertices_connected_to_vertex_by_index<'g>(
144        &'g self,
145        edge_type: EdgeTypeIndex,
146        to_vertex_index: &VertexIndex,
147    ) -> Result<VertexSelection<'g>, GraphComputingError> {
148        let edge_selection = self.select_edge_type_by_index(edge_type)?;
149        edge_selection.select_vertices_connected_to_vertex_by_index(to_vertex_index)
150    }
151
152    fn select_vertices_connected_from_vertex_by_key<'g>(
153        &'g self,
154        edge_type: EdgeType,
155        to_vertex_key: &VertexKeyRef,
156    ) -> Result<VertexSelection<'g>, GraphComputingError> {
157        let edge_selection = self.select_edge_type(edge_type)?;
158        edge_selection.select_vertices_connected_from_vertex(to_vertex_key)
159    }
160    fn select_vertices_connected_from_vertex_by_index<'g>(
161        &'g self,
162        edge_type: EdgeTypeIndex,
163        to_vertex_index: &VertexIndex,
164    ) -> Result<VertexSelection<'g>, GraphComputingError> {
165        let edge_selection = self.select_edge_type_by_index(edge_type)?;
166        edge_selection.select_vertices_connected_from_vertex_by_index(to_vertex_index)
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    use crate::graph::vertex::VertexValue;
175    use crate::tests::standard_graph_for_testing::standard_graph_for_testing;
176
177    #[test]
178    fn select_vertex_by_key() {
179        let graph = standard_graph_for_testing();
180
181        let selection = graph.select_vertex_by_key("1").unwrap();
182        let selected_values = selection.vertex_values_ref().unwrap();
183        assert_eq!(selected_values, vec![&VertexValue::UnsignedInteger8Bit(1)])
184    }
185
186    #[test]
187    fn select_vertex_by_index() {
188        let graph = standard_graph_for_testing();
189
190        let selection = graph.select_vertex_by_key("1").unwrap();
191        let indices = selection.vertex_indices_ref().unwrap();
192
193        let selection = graph.select_vertex_by_index(indices[0]).unwrap();
194        let selected_values = selection.vertex_values_ref().unwrap();
195        assert_eq!(selected_values, vec![&VertexValue::UnsignedInteger8Bit(1)])
196    }
197
198    #[test]
199    fn select_vertices_by_key() {
200        let graph = standard_graph_for_testing();
201
202        let selection = graph.select_vertices_by_key(vec!["1", "2"]).unwrap();
203        let selected_values = selection.vertex_values_ref().unwrap();
204        assert_eq!(selected_values, vec![&1u8.into(), &2u8.into()])
205    }
206
207    #[test]
208    fn select_vertices_by_index() {
209        let graph = standard_graph_for_testing();
210
211        let selection = graph.select_vertices_by_key(vec!["1", "2"]).unwrap();
212        let indices = selection.vertex_indices_ref().unwrap();
213
214        let selection = graph.select_vertices_by_index(indices).unwrap();
215        let selected_values = selection.vertex_values_ref().unwrap();
216        assert_eq!(selected_values, vec![&1u8.into(), &2u8.into()])
217    }
218
219    #[test]
220    fn test_select_vertices_connected_to_vertex() {
221        let graph = standard_graph_for_testing();
222
223        let selection_vertices_smaller_than_minus_one = graph
224            .select_vertices_connected_to_vertex_by_key(String::from("smaller_than"), &"-1")
225            .unwrap();
226        let vertices_smaller_than_minus_one = selection_vertices_smaller_than_minus_one
227            .vertex_values_ref()
228            .unwrap();
229
230        assert_eq!(
231            vertices_smaller_than_minus_one,
232            vec!(&VertexValue::FloatingPoint32Bit(-1.1))
233        );
234    }
235
236    #[test]
237    fn test_select_vertices_connected_from_vertex() {
238        let graph = standard_graph_for_testing();
239
240        let selection_vertices_larger_than_1_dot_2 = graph
241            .select_vertices_connected_to_vertex_by_key(String::from("larger_than"), &"1.2")
242            .unwrap();
243        let vertices_larger_than_one_dot_two = selection_vertices_larger_than_1_dot_2
244            .vertex_values_ref()
245            .unwrap();
246
247        assert_eq!(
248            vertices_larger_than_one_dot_two,
249            vec!(&VertexValue::UnsignedInteger8Bit(2))
250        );
251    }
252}