graphrecords_query/operations/traversal/
neighbors.rs

1use crate::{
2    Definite, EdgeDirection, EntityReference, EvaluateOperand, Explain, IndexDomain, Indexed,
3    Multiple, Operand, OrderState, QueryResult, Single, Unit, Unordered,
4    execution::EvaluationCache,
5    operands::NodesOperand,
6    operations::{Apply, KeyedStream, LaneKernel, Operation, OperationContext, Prepare},
7    optimizer::{OperationInputs, OptimizerHints, PlanIdentity, PlanInputs},
8    registry::operation_manifest,
9    traits::Neighbors,
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 = "Neighbors")]
18#[plan(optimizer_hints(empty = if_any))]
19pub struct NeighborsOperation {
20    #[explain(label)]
21    direction: EdgeDirection,
22}
23
24impl Prepare for NeighborsOperation {
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 NeighborsOperation {
37    type Output = NodesOperand<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 neighbors = GrHashSet::default();
45
46        for (node, membership) in values {
47            membership?;
48            neighbors.extend(direction.neighbors_for_node(graphrecord, node));
49        }
50
51        Ok(Box::new(neighbors.into_iter().map(|node| (node, Ok(())))))
52    }
53}
54
55impl LaneKernel<Indexed<NodeIndex, Unit>, Single> for NeighborsOperation {
56    type Output = NodesOperand<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 neighbors: GrHashSet<_> = direction.neighbors_for_node(graphrecord, node).collect();
69
70        Ok(Box::new(
71            neighbors.into_iter().map(|neighbor| (neighbor, Ok(()))),
72        ))
73    }
74}
75
76impl LaneKernel<Indexed<NodeIndex, Unit>, Definite> for NeighborsOperation {
77    type Output = NodesOperand<Unordered>;
78
79    fn execute<'a>(
80        graphrecord: &'a GraphRecord,
81        value: KeyedStream<'a, NodeIndex, Unit, Definite>,
82        direction: Self::Prepared<'a>,
83    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
84        let (node, membership) = value;
85        membership?;
86
87        let neighbors: GrHashSet<_> = direction.neighbors_for_node(graphrecord, node).collect();
88
89        Ok(Box::new(
90            neighbors.into_iter().map(|neighbor| (neighbor, Ok(()))),
91        ))
92    }
93}
94
95impl<I: IndexDomain, O: OrderState> LaneKernel<Indexed<I, EntityReference<NodeIndex>>, Multiple<O>>
96    for NeighborsOperation
97{
98    type Output = NodesOperand<Unordered>;
99
100    fn execute<'a>(
101        graphrecord: &'a GraphRecord,
102        values: KeyedStream<'a, I, EntityReference<NodeIndex>, Multiple<O>>,
103        direction: Self::Prepared<'a>,
104    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
105        let mut neighbors = GrHashSet::default();
106
107        for value in values {
108            let node = value.1?;
109            neighbors.extend(direction.neighbors_for_node(graphrecord, node));
110        }
111
112        Ok(Box::new(
113            neighbors.into_iter().map(|neighbor| (neighbor, Ok(()))),
114        ))
115    }
116}
117
118impl<I: IndexDomain> LaneKernel<Indexed<I, EntityReference<NodeIndex>>, Single>
119    for NeighborsOperation
120{
121    type Output = NodesOperand<Unordered>;
122
123    fn execute<'a>(
124        graphrecord: &'a GraphRecord,
125        value: KeyedStream<'a, I, EntityReference<NodeIndex>, Single>,
126        direction: Self::Prepared<'a>,
127    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
128        let Some(value) = value else {
129            return Ok(Box::new(empty()));
130        };
131        let node = value.1?;
132        let neighbors: GrHashSet<_> = direction.neighbors_for_node(graphrecord, node).collect();
133
134        Ok(Box::new(
135            neighbors.into_iter().map(|neighbor| (neighbor, Ok(()))),
136        ))
137    }
138}
139
140impl<I: IndexDomain> LaneKernel<Indexed<I, EntityReference<NodeIndex>>, Definite>
141    for NeighborsOperation
142{
143    type Output = NodesOperand<Unordered>;
144
145    fn execute<'a>(
146        graphrecord: &'a GraphRecord,
147        value: KeyedStream<'a, I, EntityReference<NodeIndex>, Definite>,
148        direction: Self::Prepared<'a>,
149    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
150        let node = value.1?;
151        let neighbors: GrHashSet<_> = direction.neighbors_for_node(graphrecord, node).collect();
152
153        Ok(Box::new(
154            neighbors.into_iter().map(|neighbor| (neighbor, Ok(()))),
155        ))
156    }
157}
158
159impl<O: Apply<NeighborsOperation>> Neighbors for O {
160    type ReturnOperand = O::Output;
161
162    fn neighbors(&self, direction: EdgeDirection) -> Self::ReturnOperand {
163        Self::ReturnOperand::new(OperationContext::new(
164            self.clone(),
165            NeighborsOperation { direction },
166        ))
167    }
168}
169
170operation_manifest! {
171    NeighborsOperation {
172        method: Neighbors::neighbors;
173        scope: lane;
174
175        kernel {
176            parameters: <O: OrderState>;
177            field: direction: EdgeDirection;
178            input: (Indexed<NodeIndex, Unit>, Multiple<O>);
179            output: NodesOperand<Unordered>;
180        }
181        kernel {
182            parameters: <>;
183            field: direction: EdgeDirection;
184            input: (Indexed<NodeIndex, Unit>, Single);
185            output: NodesOperand<Unordered>;
186        }
187        kernel {
188            parameters: <>;
189            field: direction: EdgeDirection;
190            input: (Indexed<NodeIndex, Unit>, Definite);
191            output: NodesOperand<Unordered>;
192        }
193        kernel {
194            parameters: <I: IndexDomain, O: OrderState>;
195            field: direction: EdgeDirection;
196            input: (Indexed<I, EntityReference<NodeIndex>>, Multiple<O>);
197            output: NodesOperand<Unordered>;
198        }
199        kernel {
200            parameters: <I: IndexDomain>;
201            field: direction: EdgeDirection;
202            input: (Indexed<I, EntityReference<NodeIndex>>, Single);
203            output: NodesOperand<Unordered>;
204        }
205        kernel {
206            parameters: <I: IndexDomain>;
207            field: direction: EdgeDirection;
208            input: (Indexed<I, EntityReference<NodeIndex>>, Definite);
209            output: NodesOperand<Unordered>;
210        }
211    }
212}
graphrecords_query/operations/traversal/neighbors.rs

graphrecords_query/operations/traversal/
neighbors.rs
