Skip to main content

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