graphrecords_query/operations/traversal/
edges.rs1use crate::{
2 Definite, EdgeDirection, EntityReference, EvaluateOperand, Explain, IndexDomain, Indexed,
3 Multiple, Operand, OrderState, QueryResult, Single, Unit, Unordered,
4 execution::EvaluationCache,
5 operands::EdgesOperand,
6 operations::{Apply, KeyedStream, LaneKernel, Operation, OperationContext, Prepare},
7 optimizer::{OperationInputs, OptimizerHints, PlanIdentity, PlanInputs},
8 registry::operation_manifest,
9 traits::Edges,
10};
11use graphrecords_core::{GraphRecord, graphrecord::NodeIndex};
12use graphrecords_utils::aliases::GrHashSet;
13use std::iter::empty;
14
15#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
16#[operation(scope = Lane)]
17#[explain(label = "Edges")]
18#[plan(optimizer_hints(empty = if_any))]
19pub struct EdgesOperation {
20 #[explain(label)]
21 direction: EdgeDirection,
22}
23
24impl Prepare for EdgesOperation {
25 type Prepared<'a> = EdgeDirection;
26
27 fn prepare<'a>(
28 &'a self,
29 _graphrecord: &'a GraphRecord,
30 _cache: &'a EvaluationCache<'a>,
31 ) -> QueryResult<Self::Prepared<'a>> {
32 Ok(self.direction)
33 }
34}
35
36impl<O: OrderState> LaneKernel<Indexed<NodeIndex, Unit>, Multiple<O>> for EdgesOperation {
37 type Output = EdgesOperand<Unordered>;
38
39 fn execute<'a>(
40 graphrecord: &'a GraphRecord,
41 values: KeyedStream<'a, NodeIndex, Unit, Multiple<O>>,
42 direction: Self::Prepared<'a>,
43 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
44 let mut edges = GrHashSet::default();
45
46 for (node, membership) in values {
47 membership?;
48 edges.extend(direction.edges_for_node(graphrecord, node));
49 }
50
51 Ok(Box::new(edges.into_iter().map(|edge| (edge, Ok(())))))
52 }
53}
54
55impl LaneKernel<Indexed<NodeIndex, Unit>, Single> for EdgesOperation {
56 type Output = EdgesOperand<Unordered>;
57
58 fn execute<'a>(
59 graphrecord: &'a GraphRecord,
60 value: KeyedStream<'a, NodeIndex, Unit, Single>,
61 direction: Self::Prepared<'a>,
62 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
63 let Some((node, membership)) = value else {
64 return Ok(Box::new(empty()));
65 };
66 membership?;
67
68 let edges: GrHashSet<_> = direction.edges_for_node(graphrecord, node).collect();
69
70 Ok(Box::new(edges.into_iter().map(|edge| (edge, Ok(())))))
71 }
72}
73
74impl LaneKernel<Indexed<NodeIndex, Unit>, Definite> for EdgesOperation {
75 type Output = EdgesOperand<Unordered>;
76
77 fn execute<'a>(
78 graphrecord: &'a GraphRecord,
79 value: KeyedStream<'a, NodeIndex, Unit, Definite>,
80 direction: Self::Prepared<'a>,
81 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
82 let (node, membership) = value;
83 membership?;
84
85 let edges: GrHashSet<_> = direction.edges_for_node(graphrecord, node).collect();
86
87 Ok(Box::new(edges.into_iter().map(|edge| (edge, Ok(())))))
88 }
89}
90
91impl<I: IndexDomain, O: OrderState> LaneKernel<Indexed<I, EntityReference<NodeIndex>>, Multiple<O>>
92 for EdgesOperation
93{
94 type Output = EdgesOperand<Unordered>;
95
96 fn execute<'a>(
97 graphrecord: &'a GraphRecord,
98 values: KeyedStream<'a, I, EntityReference<NodeIndex>, Multiple<O>>,
99 direction: Self::Prepared<'a>,
100 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
101 let mut edges = GrHashSet::default();
102
103 for value in values {
104 let node = value.1?;
105 edges.extend(direction.edges_for_node(graphrecord, node));
106 }
107
108 Ok(Box::new(edges.into_iter().map(|edge| (edge, Ok(())))))
109 }
110}
111
112impl<I: IndexDomain> LaneKernel<Indexed<I, EntityReference<NodeIndex>>, Single> for EdgesOperation {
113 type Output = EdgesOperand<Unordered>;
114
115 fn execute<'a>(
116 graphrecord: &'a GraphRecord,
117 value: KeyedStream<'a, I, EntityReference<NodeIndex>, Single>,
118 direction: Self::Prepared<'a>,
119 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
120 let Some(value) = value else {
121 return Ok(Box::new(empty()));
122 };
123 let node = value.1?;
124 let edges: GrHashSet<_> = direction.edges_for_node(graphrecord, node).collect();
125
126 Ok(Box::new(edges.into_iter().map(|edge| (edge, Ok(())))))
127 }
128}
129
130impl<I: IndexDomain> LaneKernel<Indexed<I, EntityReference<NodeIndex>>, Definite>
131 for EdgesOperation
132{
133 type Output = EdgesOperand<Unordered>;
134
135 fn execute<'a>(
136 graphrecord: &'a GraphRecord,
137 value: KeyedStream<'a, I, EntityReference<NodeIndex>, Definite>,
138 direction: Self::Prepared<'a>,
139 ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
140 let node = value.1?;
141 let edges: GrHashSet<_> = direction.edges_for_node(graphrecord, node).collect();
142
143 Ok(Box::new(edges.into_iter().map(|edge| (edge, Ok(())))))
144 }
145}
146
147impl<O: Apply<EdgesOperation>> Edges for O {
148 type ReturnOperand = O::Output;
149
150 fn edges(&self, direction: EdgeDirection) -> Self::ReturnOperand {
151 Self::ReturnOperand::new(OperationContext::new(
152 self.clone(),
153 EdgesOperation { direction },
154 ))
155 }
156}
157
158operation_manifest! {
159 EdgesOperation {
160 method: Edges::edges;
161 scope: lane;
162
163 kernel {
164 parameters: <O: OrderState>;
165 field: direction: EdgeDirection;
166 input: (Indexed<NodeIndex, Unit>, Multiple<O>);
167 output: EdgesOperand<Unordered>;
168 }
169 kernel {
170 parameters: <>;
171 field: direction: EdgeDirection;
172 input: (Indexed<NodeIndex, Unit>, Single);
173 output: EdgesOperand<Unordered>;
174 }
175 kernel {
176 parameters: <>;
177 field: direction: EdgeDirection;
178 input: (Indexed<NodeIndex, Unit>, Definite);
179 output: EdgesOperand<Unordered>;
180 }
181 kernel {
182 parameters: <I: IndexDomain, O: OrderState>;
183 field: direction: EdgeDirection;
184 input: (Indexed<I, EntityReference<NodeIndex>>, Multiple<O>);
185 output: EdgesOperand<Unordered>;
186 }
187 kernel {
188 parameters: <I: IndexDomain>;
189 field: direction: EdgeDirection;
190 input: (Indexed<I, EntityReference<NodeIndex>>, Single);
191 output: EdgesOperand<Unordered>;
192 }
193 kernel {
194 parameters: <I: IndexDomain>;
195 field: direction: EdgeDirection;
196 input: (Indexed<I, EntityReference<NodeIndex>>, Definite);
197 output: EdgesOperand<Unordered>;
198 }
199 }
200}