graphrecords_query/operations/traversal/
via_neighbors.rs

1use crate::{
2    EdgeDirection, EntityReference, ExpandedChild, ExpandedIndex, Explain, IndexDomain, Indexed,
3    Operand, QueryResult, Unit, Unordered,
4    element::{Expanding, Pipeline},
5    execution::EvaluationCache,
6    operations::{Apply, ElementKernel, ElementPipeline, Operation, OperationContext, Prepare},
7    optimizer::{OperationInputs, OptimizerHints, PlanIdentity, PlanInputs},
8    registry::operation_manifest,
9    traits::ViaNeighbors,
10};
11use graphrecords_core::{GraphRecord, graphrecord::NodeIndex};
12use graphrecords_utils::aliases::GrHashSet;
13
14#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
15#[operation(scope = Element)]
16#[explain(label = "ViaNeighbors")]
17#[plan(optimizer_hints(empty = if_any))]
18pub struct ViaNeighborsOperation {
19    #[explain(label)]
20    direction: EdgeDirection,
21}
22
23impl Prepare for ViaNeighborsOperation {
24    type Prepared<'a> = EdgeDirection;
25
26    fn prepare<'a>(
27        &'a self,
28        _graphrecord: &'a GraphRecord,
29        _cache: &'a EvaluationCache<'a>,
30    ) -> QueryResult<Self::Prepared<'a>> {
31        Ok(self.direction)
32    }
33}
34
35impl ElementKernel<Indexed<NodeIndex, Unit>> for ViaNeighborsOperation {
36    type Emission = Expanding<Unordered>;
37    type OutShape = Indexed<ExpandedIndex<NodeIndex, NodeIndex>, EntityReference<NodeIndex>>;
38
39    fn pipeline<'a>(
40        graphrecord: &'a GraphRecord,
41        prepared: Self::Prepared<'a>,
42    ) -> QueryResult<ElementPipeline<'a, Indexed<NodeIndex, Unit>, Self>> {
43        Ok(Pipeline::keyed(move |parent_index, ()| {
44            let neighbors: GrHashSet<_> = prepared
45                .neighbors_for_node(graphrecord, parent_index)
46                .collect();
47
48            Ok(neighbors
49                .into_iter()
50                .map(|neighbor| ExpandedChild::success(neighbor, neighbor))
51                .collect())
52        }))
53    }
54}
55
56impl<I: IndexDomain> ElementKernel<Indexed<I, EntityReference<NodeIndex>>>
57    for ViaNeighborsOperation
58{
59    type Emission = Expanding<Unordered>;
60    type OutShape = Indexed<ExpandedIndex<I, NodeIndex>, EntityReference<NodeIndex>>;
61
62    fn pipeline<'a>(
63        graphrecord: &'a GraphRecord,
64        prepared: Self::Prepared<'a>,
65    ) -> QueryResult<ElementPipeline<'a, Indexed<I, EntityReference<NodeIndex>>, Self>> {
66        Ok(Pipeline::unkeyed(move |node| {
67            let neighbors: GrHashSet<_> = prepared.neighbors_for_node(graphrecord, node).collect();
68
69            Ok(neighbors
70                .into_iter()
71                .map(|neighbor| ExpandedChild::success(neighbor, neighbor))
72                .collect())
73        }))
74    }
75}
76
77impl<O: Apply<ViaNeighborsOperation>> ViaNeighbors for O {
78    type ReturnOperand = O::Output;
79
80    fn via_neighbors(&self, direction: EdgeDirection) -> Self::ReturnOperand {
81        Self::ReturnOperand::new(OperationContext::new(
82            self.clone(),
83            ViaNeighborsOperation { direction },
84        ))
85    }
86}
87
88operation_manifest! {
89    ViaNeighborsOperation {
90        method: ViaNeighbors::via_neighbors;
91        scope: element;
92
93        kernel {
94            parameters: <>;
95            field: direction: EdgeDirection;
96            input: Indexed<NodeIndex, Unit>;
97            output: Indexed<ExpandedIndex<NodeIndex, NodeIndex>, EntityReference<NodeIndex>>;
98            emission: Expanding<Unordered>;
99        }
100        kernel {
101            parameters: <I: IndexDomain>;
102            field: direction: EdgeDirection;
103            input: Indexed<I, EntityReference<NodeIndex>>;
104            output: Indexed<ExpandedIndex<I, NodeIndex>, EntityReference<NodeIndex>>;
105            emission: Expanding<Unordered>;
106        }
107    }
108}
graphrecords_query/operations/traversal/via_neighbors.rs

graphrecords_query/operations/traversal/
via_neighbors.rs
