Skip to main content

graphrecords_query/operations/traversal/
edges.rs

1use 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}