cairn_knowledge_graph/operations/
select_vertex.rs1use 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()]; let mut mask_values: Vec<bool> = vec![true; vertex_keys.len()];
92
93 (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 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}