graphrecords_query/operations/grouping/
keys.rs

1use super::reject_key_failures;
2use crate::{
3    Arity, ElementShape, EvaluateOperand, Explain, IndexDomain, Labeled, Operand, QueryResult,
4    Unordered,
5    execution::EvaluationCache,
6    index::GroupKey,
7    operands::{ElementsOperand, OperandHandle, Partition},
8    operations::{Apply, GroupKernel, Operation, OperationContext, Prepare},
9    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
10    registry::operation_manifest,
11    traits::Keys,
12};
13use graphrecords_core::GraphRecord;
14
15#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
16#[operation(scope = Group)]
17#[explain(label = "Keys")]
18#[plan(optimizer_hints(empty = if_any))]
19pub struct KeysOperation;
20
21impl Prepare for KeysOperation {
22    type Prepared<'a> = ();
23
24    fn prepare<'a>(
25        &'a self,
26        _graphrecord: &'a GraphRecord,
27        _cache: &'a EvaluationCache<'a>,
28    ) -> QueryResult<Self::Prepared<'a>> {
29        Ok(())
30    }
31}
32
33impl<M: IndexDomain, K: GroupKey, S: ElementShape, C: Arity> GroupKernel<M, K, OperandHandle<S, C>>
34    for KeysOperation
35{
36    type Output = ElementsOperand<K, Unordered>;
37
38    fn execute<'a>(
39        graphrecord: &'a GraphRecord,
40        partition: Partition<'a, M, K, OperandHandle<S, C>>,
41        _prepared: Self::Prepared<'a>,
42    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
43        let (buckets, key_failures) = partition.into_parts();
44
45        reject_key_failures::<M>(key_failures, Self::LABEL)?;
46
47        let elements: Vec<_> = buckets
48            .into_iter()
49            .map(|(key, _, payload)| {
50                let index = K::resolve_key(Self::LABEL, graphrecord, &key)?;
51
52                Ok((index, payload.map(|_| ())))
53            })
54            .collect::<QueryResult<_>>()?;
55
56        Ok(Box::new(elements.into_iter()))
57    }
58
59    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
60        Estimate {
61            elements: input.elements,
62            distinct: input.elements,
63            selectivity: None,
64            per_group: None,
65        }
66    }
67}
68
69impl<O: Apply<KeysOperation>> Keys for O {
70    type ReturnOperand = O::Output;
71
72    fn keys(&self) -> Self::ReturnOperand {
73        Self::ReturnOperand::new(OperationContext::new(self.clone(), KeysOperation))
74    }
75}
76
77operation_manifest! {
78    KeysOperation {
79        method: Keys::keys;
80        scope: group;
81
82        kernel {
83            group: <M: IndexDomain, K: GroupKey>;
84            parameters: <S: ElementShape, C: Arity>;
85            input: OperandHandle<S, C>;
86            output: ElementsOperand<K, Unordered>;
87        }
88    }
89}
graphrecords_query/operations/grouping/keys.rs

graphrecords_query/operations/grouping/
keys.rs
