Skip to main content

graphrecords_query/
selection.rs

1use crate::{
2    Arity, EvaluateOperand, Explanation, IndexDomain, NodesOperand, Operand, QueryResult,
3    ReturnShape, Unordered,
4    execution::EvaluationCache,
5    index::GroupKey,
6    operands::{
7        AllEdges, AllNodes, EdgesOperand, GroupOperand, OperandHandle, Partition, ReturnPartition,
8    },
9    optimizer::{OptimizationReport, Optimizer, Stats},
10};
11use graphrecords_core::GraphRecord;
12use std::fmt::{self, Display, Formatter};
13
14pub trait ReturnOperand<'a>: Clone {
15    type ReturnValue;
16
17    fn evaluate(
18        &'a self,
19        graphrecord: &'a GraphRecord,
20        cache: &'a EvaluationCache<'a>,
21    ) -> QueryResult<Self::ReturnValue>;
22
23    #[allow(unused_variables)]
24    fn optimize(self, optimizer: &Optimizer, stats: &Stats) -> (Self, OptimizationReport)
25    where
26        Self: Sized,
27    {
28        (self, OptimizationReport::default())
29    }
30
31    fn fmt_plan(&self, formatter: &mut Formatter<'_>) -> fmt::Result;
32}
33
34pub trait IntoReturn: Operand {
35    type Return<'a>: 'a
36    where
37        Self: 'a;
38
39    fn into_return<'a>(values: Self::ReturnValue<'a>) -> Self::Return<'a>
40    where
41        Self: 'a;
42}
43
44impl<S: ReturnShape, C: Arity> IntoReturn for OperandHandle<S, C> {
45    type Return<'a>
46        = C::Container<'a, S::ReturnElement<'a>>
47    where
48        Self: 'a;
49
50    fn into_return<'a>(values: Self::ReturnValue<'a>) -> Self::Return<'a>
51    where
52        Self: 'a,
53    {
54        C::map_elements(values, S::into_return_element)
55    }
56}
57
58impl<M: IndexDomain, K: GroupKey, O: IntoReturn> IntoReturn for GroupOperand<M, K, O> {
59    type Return<'a>
60        = ReturnPartition<'a, M, K, O::Return<'a>>
61    where
62        Self: 'a;
63
64    fn into_return<'a>(values: Partition<'a, M, K, O>) -> Self::Return<'a>
65    where
66        Self: 'a,
67    {
68        values.into_return_partition(|payload| payload.map(O::into_return))
69    }
70}
71
72impl<'a, S: ReturnShape, C: Arity> ReturnOperand<'a> for OperandHandle<S, C> {
73    type ReturnValue = <Self as IntoReturn>::Return<'a>;
74
75    fn evaluate(
76        &'a self,
77        graphrecord: &'a GraphRecord,
78        cache: &'a EvaluationCache<'a>,
79    ) -> QueryResult<Self::ReturnValue> {
80        let values = EvaluateOperand::evaluate(self, graphrecord, cache)?;
81
82        Ok(Self::into_return(values))
83    }
84
85    fn optimize(self, optimizer: &Optimizer, stats: &Stats) -> (Self, OptimizationReport) {
86        optimizer.run_reported(stats, &self)
87    }
88
89    fn fmt_plan(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
90        write!(formatter, "{}", Explanation::new(self))
91    }
92}
93
94impl<'a, M: IndexDomain, K: GroupKey, O: IntoReturn> ReturnOperand<'a> for GroupOperand<M, K, O> {
95    type ReturnValue = <Self as IntoReturn>::Return<'a>;
96
97    fn evaluate(
98        &'a self,
99        graphrecord: &'a GraphRecord,
100        cache: &'a EvaluationCache<'a>,
101    ) -> QueryResult<Self::ReturnValue> {
102        let values = EvaluateOperand::evaluate(self, graphrecord, cache)?;
103
104        Ok(Self::into_return(values))
105    }
106
107    fn optimize(self, optimizer: &Optimizer, stats: &Stats) -> (Self, OptimizationReport) {
108        optimizer.run_reported(stats, &self)
109    }
110
111    fn fmt_plan(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
112        write!(formatter, "{}", Explanation::new(self))
113    }
114}
115
116macro_rules! impl_return_operand_for_tuples {
117    ($($T:ident),+) => {
118        impl<'a, $($T: ReturnOperand<'a>),+> ReturnOperand<'a> for ($($T,)+) {
119            type ReturnValue = ($($T::ReturnValue,)+);
120
121            #[allow(non_snake_case)]
122            fn evaluate(&'a self, graphrecord: &'a GraphRecord, cache: &'a EvaluationCache<'a>) -> QueryResult<Self::ReturnValue> {
123                let ($($T,)+) = self;
124
125                $(let $T = $T.evaluate(graphrecord, cache)?;)+
126
127                Ok(($($T,)+))
128            }
129
130            #[allow(non_snake_case)]
131            fn optimize(self, optimizer: &Optimizer, stats: &Stats) -> (Self, OptimizationReport) {
132                let ($($T,)+) = self;
133
134                let mut report = OptimizationReport::default();
135                $(
136                    let ($T, sub_report) = $T.optimize(optimizer, stats);
137                    report.phases.extend(sub_report.phases);
138                )+
139
140                (($($T,)+), report)
141            }
142
143            #[allow(non_snake_case)]
144            fn fmt_plan(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
145                let ($($T,)+) = self;
146
147                let mut index = 0;
148                $(
149                    if index > 0 {
150                        writeln!(formatter)?;
151                    }
152                    writeln!(formatter, "[{index}]")?;
153                    $T.fmt_plan(formatter)?;
154                    index += 1;
155                )+
156                let _ = index;
157
158                Ok(())
159            }
160        }
161    };
162}
163
164impl_return_operand_for_tuples!(R1);
165impl_return_operand_for_tuples!(R1, R2);
166impl_return_operand_for_tuples!(R1, R2, R3);
167impl_return_operand_for_tuples!(R1, R2, R3, R4);
168impl_return_operand_for_tuples!(R1, R2, R3, R4, R5);
169impl_return_operand_for_tuples!(R1, R2, R3, R4, R5, R6);
170impl_return_operand_for_tuples!(R1, R2, R3, R4, R5, R6, R7);
171impl_return_operand_for_tuples!(R1, R2, R3, R4, R5, R6, R7, R8);
172impl_return_operand_for_tuples!(R1, R2, R3, R4, R5, R6, R7, R8, R9);
173impl_return_operand_for_tuples!(R1, R2, R3, R4, R5, R6, R7, R8, R9, R10);
174impl_return_operand_for_tuples!(R1, R2, R3, R4, R5, R6, R7, R8, R9, R10, R11);
175impl_return_operand_for_tuples!(R1, R2, R3, R4, R5, R6, R7, R8, R9, R10, R11, R12);
176impl_return_operand_for_tuples!(R1, R2, R3, R4, R5, R6, R7, R8, R9, R10, R11, R12, R13);
177impl_return_operand_for_tuples!(R1, R2, R3, R4, R5, R6, R7, R8, R9, R10, R11, R12, R13, R14);
178impl_return_operand_for_tuples!(
179    R1, R2, R3, R4, R5, R6, R7, R8, R9, R10, R11, R12, R13, R14, R15
180);
181
182pub struct Selection<'a, R: ReturnOperand<'a>> {
183    graphrecord: &'a GraphRecord,
184    cache: EvaluationCache<'a>,
185    unoptimized_return_operand: R,
186    optimized_return_operand: R,
187    report: OptimizationReport,
188}
189
190impl<'a, R: ReturnOperand<'a>> Selection<'a, R> {
191    pub fn new_node<Q>(graphrecord: &'a GraphRecord, query: Q) -> Self
192    where
193        Q: FnOnce(&NodesOperand<Unordered>) -> R,
194    {
195        Self::new_node_with(graphrecord, Optimizer::shared_builtin(), query)
196    }
197
198    pub fn new_node_with<Q>(graphrecord: &'a GraphRecord, optimizer: &Optimizer, query: Q) -> Self
199    where
200        Q: FnOnce(&NodesOperand<Unordered>) -> R,
201    {
202        let operand = NodesOperand::new(AllNodes);
203        let unoptimized_return_operand = query(&operand);
204        let (optimized_return_operand, report) =
205            Self::optimize(unoptimized_return_operand.clone(), optimizer, graphrecord);
206
207        Self {
208            graphrecord,
209            cache: EvaluationCache::new(graphrecord),
210            unoptimized_return_operand,
211            optimized_return_operand,
212            report,
213        }
214    }
215
216    pub fn new_edge<Q>(graphrecord: &'a GraphRecord, query: Q) -> Self
217    where
218        Q: FnOnce(&EdgesOperand<Unordered>) -> R,
219    {
220        Self::new_edge_with(graphrecord, Optimizer::shared_builtin(), query)
221    }
222
223    pub fn new_edge_with<Q>(graphrecord: &'a GraphRecord, optimizer: &Optimizer, query: Q) -> Self
224    where
225        Q: FnOnce(&EdgesOperand<Unordered>) -> R,
226    {
227        let operand = EdgesOperand::new(AllEdges);
228        let unoptimized_return_operand = query(&operand);
229        let (optimized_return_operand, report) =
230            Self::optimize(unoptimized_return_operand.clone(), optimizer, graphrecord);
231
232        Self {
233            graphrecord,
234            cache: EvaluationCache::new(graphrecord),
235            unoptimized_return_operand,
236            optimized_return_operand,
237            report,
238        }
239    }
240
241    fn optimize(
242        return_operand: R,
243        optimizer: &Optimizer,
244        graphrecord: &GraphRecord,
245    ) -> (R, OptimizationReport) {
246        if optimizer.is_empty() {
247            return (return_operand, OptimizationReport::default());
248        }
249
250        let stats = Stats::new(graphrecord);
251
252        return_operand.optimize(optimizer, &stats)
253    }
254
255    pub fn evaluate(&'a self) -> QueryResult<R::ReturnValue> {
256        self.optimized_return_operand
257            .evaluate(self.graphrecord, &self.cache)
258    }
259
260    pub const fn explain(&'a self) -> QueryExplanation<'a, R> {
261        QueryExplanation {
262            operand: &self.optimized_return_operand,
263            report: Some(&self.report),
264        }
265    }
266
267    pub const fn explain_unoptimized(&'a self) -> QueryExplanation<'a, R> {
268        QueryExplanation {
269            operand: &self.unoptimized_return_operand,
270            report: None,
271        }
272    }
273}
274
275pub struct QueryExplanation<'a, R> {
276    operand: &'a R,
277    report: Option<&'a OptimizationReport>,
278}
279
280impl<'a, R: ReturnOperand<'a>> Display for QueryExplanation<'a, R> {
281    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
282        self.operand.fmt_plan(formatter)?;
283
284        if let Some(report) = self.report.filter(|report| !report.phases.is_empty()) {
285            write!(formatter, "\n\noptimization:\n{}", report.display())?;
286        }
287
288        Ok(())
289    }
290}
291
292pub trait QueryNodes {
293    fn query_nodes<'a, Q, R>(&'a self, query: Q) -> Selection<'a, R>
294    where
295        Q: FnOnce(&NodesOperand<Unordered>) -> R,
296        R: ReturnOperand<'a>;
297
298    fn query_nodes_with<'a, Q, R>(&'a self, optimizer: &Optimizer, query: Q) -> Selection<'a, R>
299    where
300        Q: FnOnce(&NodesOperand<Unordered>) -> R,
301        R: ReturnOperand<'a>;
302}
303
304impl QueryNodes for GraphRecord {
305    fn query_nodes<'a, Q, R>(&'a self, query: Q) -> Selection<'a, R>
306    where
307        Q: FnOnce(&NodesOperand<Unordered>) -> R,
308        R: ReturnOperand<'a>,
309    {
310        Selection::new_node(self, query)
311    }
312
313    fn query_nodes_with<'a, Q, R>(&'a self, optimizer: &Optimizer, query: Q) -> Selection<'a, R>
314    where
315        Q: FnOnce(&NodesOperand<Unordered>) -> R,
316        R: ReturnOperand<'a>,
317    {
318        Selection::new_node_with(self, optimizer, query)
319    }
320}
321
322pub trait QueryEdges {
323    fn query_edges<'a, Q, R>(&'a self, query: Q) -> Selection<'a, R>
324    where
325        Q: FnOnce(&EdgesOperand<Unordered>) -> R,
326        R: ReturnOperand<'a>;
327
328    fn query_edges_with<'a, Q, R>(&'a self, optimizer: &Optimizer, query: Q) -> Selection<'a, R>
329    where
330        Q: FnOnce(&EdgesOperand<Unordered>) -> R,
331        R: ReturnOperand<'a>;
332}
333
334impl QueryEdges for GraphRecord {
335    fn query_edges<'a, Q, R>(&'a self, query: Q) -> Selection<'a, R>
336    where
337        Q: FnOnce(&EdgesOperand<Unordered>) -> R,
338        R: ReturnOperand<'a>,
339    {
340        Selection::new_edge(self, query)
341    }
342
343    fn query_edges_with<'a, Q, R>(&'a self, optimizer: &Optimizer, query: Q) -> Selection<'a, R>
344    where
345        Q: FnOnce(&EdgesOperand<Unordered>) -> R,
346        R: ReturnOperand<'a>,
347    {
348        Selection::new_edge_with(self, optimizer, query)
349    }
350}