Skip to main content

graphrecords_query/operations/
cache.rs

1use crate::{
2    Cache, EvaluateContext, EvaluateOperand, Explain, Failure, Labeled, QueryResult,
3    error::execution::EvaluationCacheGraphRecordMismatch,
4    execution::{CacheSlot, CacheableOperand, EvaluationCache},
5    optimizer::{
6        Estimate, Estimated, MatchInputs, OptimizePlan, OptimizerHints, PlanNode, Session, Stats,
7        Transformed,
8    },
9};
10use graphrecords_core::GraphRecord;
11use std::{
12    any::Any,
13    hash::{Hash, Hasher},
14};
15
16#[derive(MatchInputs, OptimizerHints, Explain)]
17#[explain(label = "Cache")]
18pub struct CacheContext<O: CacheableOperand> {
19    #[input]
20    input: O,
21    slot: CacheSlot,
22}
23
24impl<O: CacheableOperand> CacheContext<O> {
25    #[must_use]
26    pub fn new(input: O) -> Self {
27        Self {
28            input,
29            slot: CacheSlot::new(),
30        }
31    }
32}
33
34impl<O: CacheableOperand> Cache for O {
35    fn cache(&self) -> Self {
36        Self::new(CacheContext::new(self.clone()))
37    }
38}
39
40impl<O: CacheableOperand> PlanNode for CacheContext<O> {
41    fn inputs(&self) -> Vec<&dyn PlanNode> {
42        vec![self.input.as_plan_node()]
43    }
44
45    fn dyn_eq(&self, other: &dyn PlanNode) -> bool {
46        let Some(other) = other.downcast::<Self>() else {
47            return false;
48        };
49
50        self.input.as_plan_node().dyn_eq(other.input.as_plan_node())
51    }
52
53    fn dyn_hash(&self, mut state: &mut dyn Hasher) {
54        self.type_id().hash(&mut state);
55        self.input.as_plan_node().dyn_hash(state);
56    }
57}
58
59impl<O: CacheableOperand> Estimated for CacheContext<O> {
60    fn estimate(&self, stats: &Stats) -> Estimate {
61        self.input.context().estimate(stats)
62    }
63}
64
65impl<O: CacheableOperand> OptimizePlan for CacheContext<O> {
66    type Output = O;
67
68    fn optimize(&self, original: &Self::Output, session: &Session) -> Transformed<Self::Output> {
69        let input = session.optimize(&self.input);
70
71        if !input.is_changed() {
72            return Transformed::unchanged(original.clone());
73        }
74
75        let input = input.into_parts().0;
76
77        Transformed::changed(O::new(Self {
78            input,
79            slot: self.slot.clone(),
80        }))
81    }
82}
83
84impl<O: CacheableOperand> EvaluateContext for CacheContext<O> {
85    type Operand = O;
86
87    fn evaluate<'a>(
88        &'a self,
89        graphrecord: &'a GraphRecord,
90        cache: &'a EvaluationCache<'a>,
91    ) -> QueryResult<<O as EvaluateOperand>::ReturnValue<'a>> {
92        if !cache.is_bound_to(graphrecord) {
93            return Err(Failure::new(
94                Self::LABEL,
95                EvaluationCacheGraphRecordMismatch,
96            ));
97        }
98
99        cache.materialize::<O>(&self.slot, || self.input.evaluate(graphrecord, cache))
100    }
101}