graphrecords_query/operations/traversal/
endpoint.rs

1use crate::{
2    EdgeEndpointRole, EntityReference, Explain, IndexDomain, Indexed, Operand, QueryResult, Unit,
3    element::{Pipeline, Preserving},
4    execution::EvaluationCache,
5    operations::{Apply, ElementKernel, ElementPipeline, Operation, OperationContext, Prepare},
6    optimizer::{
7        Count, CountKind, Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs,
8        Stats,
9    },
10    registry::operation_manifest,
11    traits::{Select, SourceNode, TargetNode, ViaSourceNode, ViaTargetNode},
12};
13use graphrecords_core::{
14    GraphRecord,
15    graphrecord::{EdgeIndex, NodeIndex},
16};
17
18fn endpoint_estimate(input: Estimate, stats: &Stats) -> Estimate {
19    let node_count = stats.get::<Count>(&CountKind::Nodes);
20    let distinct = input
21        .distinct
22        .map_or(node_count, |distinct| node_count.min(distinct));
23    let distinct = input
24        .elements
25        .map_or(distinct, |elements| distinct.min(elements));
26
27    Estimate {
28        distinct: Some(distinct),
29        selectivity: None,
30        ..input
31    }
32}
33
34#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
35#[operation(scope = Element)]
36#[explain(label = "Endpoint")]
37#[plan(optimizer_hints(empty = if_any))]
38pub struct EndpointOperation {
39    #[explain(label)]
40    role: EdgeEndpointRole,
41}
42
43impl Prepare for EndpointOperation {
44    type Prepared<'a> = EdgeEndpointRole;
45
46    fn prepare<'a>(
47        &'a self,
48        _graphrecord: &'a GraphRecord,
49        _cache: &'a EvaluationCache<'a>,
50    ) -> QueryResult<Self::Prepared<'a>> {
51        Ok(self.role)
52    }
53}
54
55impl ElementKernel<Indexed<EdgeIndex, Unit>> for EndpointOperation {
56    type Emission = Preserving;
57    type OutShape = Indexed<EdgeIndex, EntityReference<NodeIndex>>;
58
59    fn pipeline<'a>(
60        graphrecord: &'a GraphRecord,
61        prepared: Self::Prepared<'a>,
62    ) -> QueryResult<ElementPipeline<'a, Indexed<EdgeIndex, Unit>, Self>> {
63        Ok(Pipeline::keyed(move |edge, membership: QueryResult<_>| {
64            membership.map(|()| {
65                let (source, target) = graphrecord.edge_endpoints(edge).expect("Edge must exist");
66
67                match prepared {
68                    EdgeEndpointRole::Source => source,
69                    EdgeEndpointRole::Target => target,
70                }
71            })
72        }))
73    }
74
75    fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
76        endpoint_estimate(input, stats)
77    }
78}
79
80impl<I: IndexDomain> ElementKernel<Indexed<I, EntityReference<EdgeIndex>>> for EndpointOperation {
81    type Emission = Preserving;
82    type OutShape = Indexed<I, EntityReference<NodeIndex>>;
83
84    fn pipeline<'a>(
85        graphrecord: &'a GraphRecord,
86        prepared: Self::Prepared<'a>,
87    ) -> QueryResult<ElementPipeline<'a, Indexed<I, EntityReference<EdgeIndex>>, Self>> {
88        Ok(Pipeline::unkeyed(move |edge: QueryResult<_>| {
89            edge.map(|edge| {
90                let (source, target) = graphrecord.edge_endpoints(edge).expect("Edge must exist");
91
92                match prepared {
93                    EdgeEndpointRole::Source => source,
94                    EdgeEndpointRole::Target => target,
95                }
96            })
97        }))
98    }
99
100    fn estimate(&self, input: Estimate, stats: &Stats) -> Estimate {
101        endpoint_estimate(input, stats)
102    }
103}
104
105impl<O: Apply<EndpointOperation>> ViaSourceNode for O {
106    type ReturnOperand = O::Output;
107
108    fn via_source_node(&self) -> Self::ReturnOperand {
109        Self::ReturnOperand::new(OperationContext::new(
110            self.clone(),
111            EndpointOperation {
112                role: EdgeEndpointRole::Source,
113            },
114        ))
115    }
116}
117
118impl<O: Apply<EndpointOperation>> ViaTargetNode for O {
119    type ReturnOperand = O::Output;
120
121    fn via_target_node(&self) -> Self::ReturnOperand {
122        Self::ReturnOperand::new(OperationContext::new(
123            self.clone(),
124            EndpointOperation {
125                role: EdgeEndpointRole::Target,
126            },
127        ))
128    }
129}
130
131impl<O> SourceNode for O
132where
133    O: ViaSourceNode,
134    O::ReturnOperand: Select,
135{
136    type ReturnOperand = <O::ReturnOperand as Select>::ReturnOperand;
137
138    fn source_node(&self) -> Self::ReturnOperand {
139        self.via_source_node().select()
140    }
141}
142
143impl<O> TargetNode for O
144where
145    O: ViaTargetNode,
146    O::ReturnOperand: Select,
147{
148    type ReturnOperand = <O::ReturnOperand as Select>::ReturnOperand;
149
150    fn target_node(&self) -> Self::ReturnOperand {
151        self.via_target_node().select()
152    }
153}
154
155pub(super) mod via_source_node {
156    use super::{
157        EdgeIndex, EndpointOperation, EntityReference, Indexed, NodeIndex, Preserving, Unit,
158        ViaSourceNode, operation_manifest,
159    };
160
161    operation_manifest! {
162        EndpointOperation {
163            method: ViaSourceNode::via_source_node;
164            scope: element;
165
166            kernel {
167                parameters: <>;
168                input: Indexed<EdgeIndex, Unit>;
169                output: Indexed<EdgeIndex, EntityReference<NodeIndex>>;
170                emission: Preserving;
171            }
172            kernel {
173                parameters: <I: IndexDomain>;
174                input: Indexed<I, EntityReference<EdgeIndex>>;
175                output: Indexed<I, EntityReference<NodeIndex>>;
176                emission: Preserving;
177            }
178        }
179    }
180}
181
182pub(super) mod via_target_node {
183    use super::{
184        EdgeIndex, EndpointOperation, EntityReference, Indexed, NodeIndex, Preserving, Unit,
185        ViaTargetNode, operation_manifest,
186    };
187
188    operation_manifest! {
189        EndpointOperation {
190            method: ViaTargetNode::via_target_node;
191            scope: element;
192
193            kernel {
194                parameters: <>;
195                input: Indexed<EdgeIndex, Unit>;
196                output: Indexed<EdgeIndex, EntityReference<NodeIndex>>;
197                emission: Preserving;
198            }
199            kernel {
200                parameters: <I: IndexDomain>;
201                input: Indexed<I, EntityReference<EdgeIndex>>;
202                output: Indexed<I, EntityReference<NodeIndex>>;
203                emission: Preserving;
204            }
205        }
206    }
207}
graphrecords_query/operations/traversal/endpoint.rs

graphrecords_query/operations/traversal/
endpoint.rs
