Skip to main content

graphrecords_query/operations/grouping/
having.rs

1use crate::{
2    EvaluateOperand, Explain, IndexDomain, Labeled, Mask, Operand, QueryResult,
3    element::Retention,
4    execution::EvaluationCache,
5    index::GroupKey,
6    operands::{BucketChange, GroupOperand, Partition},
7    operations::{Apply, ArgumentSource, GroupKernel, Keyed, Operation, OperationContext, Prepare},
8    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
9    registry::operation_manifest,
10    traits::Having,
11};
12use graphrecords_core::GraphRecord;
13
14#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
15#[operation(scope = Group)]
16#[explain(label = "Having")]
17#[plan(optimizer_hints(empty = if_all))]
18pub struct HavingOperation<P> {
19    #[argument]
20    predicate: P,
21}
22
23impl<P: Prepare> Prepare for HavingOperation<P> {
24    type Prepared<'a>
25        = P::Prepared<'a>
26    where
27        Self: 'a;
28
29    fn prepare<'a>(
30        &'a self,
31        graphrecord: &'a GraphRecord,
32        cache: &'a EvaluationCache<'a>,
33    ) -> QueryResult<Self::Prepared<'a>> {
34        self.predicate.prepare(graphrecord, cache)
35    }
36}
37
38impl<M, K, O, P> GroupKernel<M, K, O> for HavingOperation<P>
39where
40    M: IndexDomain,
41    K: GroupKey,
42    O: Operand,
43    P: ArgumentSource<Keyed<K>, Mask>,
44{
45    type Output = GroupOperand<M, K, O>;
46
47    fn execute<'a>(
48        graphrecord: &'a GraphRecord,
49        partition: Partition<'a, M, K, O>,
50        prepared: Self::Prepared<'a>,
51    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
52        Ok(partition.change_buckets(|bucket| {
53            if bucket.payload().is_err() {
54                return None;
55            }
56
57            let key = match K::resolve_key(Self::LABEL, graphrecord, bucket.key()) {
58                Ok(key) => key,
59                Err(failure) => {
60                    return Some(BucketChange::ReplacePayload(Err(failure)));
61                }
62            };
63            let step = P::resolve(&prepared, &key, Self::LABEL);
64
65            match P::Retention::collapse(step) {
66                None | Some(Ok(false)) => Some(BucketChange::Drop),
67                Some(Ok(true)) => None,
68                Some(Err(failure)) => Some(BucketChange::ReplacePayload(Err(failure))),
69            }
70        }))
71    }
72
73    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
74        Estimate {
75            per_group: input.per_group,
76            ..Estimate::UNKNOWN
77        }
78    }
79}
80
81impl<O, P> Having<P> for O
82where
83    O: Apply<HavingOperation<P>>,
84    HavingOperation<P>: Operation,
85{
86    type ReturnOperand = O::Output;
87
88    fn having(&self, predicate: P) -> Self::ReturnOperand {
89        Self::ReturnOperand::new(OperationContext::new(
90            self.clone(),
91            HavingOperation { predicate },
92        ))
93    }
94}
95
96operation_manifest! {
97    HavingOperation<P> {
98        method: Having<P>::having;
99        scope: group;
100
101        kernel {
102            group: <M: IndexDomain, K: GroupKey>;
103            parameters: <O: Lane>;
104            argument: P: ArgumentSource<Keyed<K>, Mask>;
105            input: O;
106            output: GroupOperand<M, K, O>;
107        }
108    }
109}