graphrecords_query/operations/aggregation/
unique_count.rs

1use crate::{
2    Bare, BareValueDomain, EvaluateOperand, Explain, IndexDomain, Indexed, Multiple, Operand,
3    OrderState, QueryResult,
4    capabilities::ValueEquivalence,
5    execution::EvaluationCache,
6    operands::DefiniteBareValueOperand,
7    operations::{
8        Apply, BareStream, KeyedStream, LaneKernel, Operation, OperationContext, Prepare,
9    },
10    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
11    registry::operation_manifest,
12    traits::UniqueCount,
13};
14use graphrecords_core::{GraphRecord, graphrecord::GraphRecordValue};
15use graphrecords_utils::aliases::GrHashSet;
16
17#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
18#[operation(scope = Lane)]
19#[explain(label = "NUnique")]
20pub struct UniqueCountOperation;
21
22impl Prepare for UniqueCountOperation {
23    type Prepared<'a> = ();
24
25    fn prepare<'a>(
26        &'a self,
27        _graphrecord: &'a GraphRecord,
28        _cache: &'a EvaluationCache<'a>,
29    ) -> QueryResult<Self::Prepared<'a>> {
30        Ok(())
31    }
32}
33
34impl<I: IndexDomain, V: ValueEquivalence, O: OrderState> LaneKernel<Indexed<I, V>, Multiple<O>>
35    for UniqueCountOperation
36{
37    type Output = DefiniteBareValueOperand;
38
39    fn execute<'a>(
40        _graphrecord: &'a GraphRecord,
41        mut values: KeyedStream<'a, I, V, Multiple<O>>,
42        _prepared: Self::Prepared<'a>,
43    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
44        let count = values.try_fold(
45            (GrHashSet::default(), 0),
46            |(mut unique, count), (_, value)| {
47                let value = value?;
48                let inserted = unique.insert(V::equivalence_key(&value));
49
50                Ok((unique, count + i64::from(inserted)))
51            },
52        );
53
54        Ok(count.map(|(_, count)| GraphRecordValue::Int(count)))
55    }
56
57    fn estimate(&self, _input: Estimate, _stats: &Stats) -> Estimate {
58        Estimate::singleton()
59    }
60}
61
62impl<V: ValueEquivalence + BareValueDomain, O: OrderState> LaneKernel<Bare<V>, Multiple<O>>
63    for UniqueCountOperation
64{
65    type Output = DefiniteBareValueOperand;
66
67    fn execute<'a>(
68        _graphrecord: &'a GraphRecord,
69        mut values: BareStream<'a, V, Multiple<O>>,
70        _prepared: Self::Prepared<'a>,
71    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
72        let count = values.try_fold((GrHashSet::default(), 0), |(mut unique, count), value| {
73            let value = value?;
74            let inserted = unique.insert(V::equivalence_key(&value));
75
76            Ok((unique, count + i64::from(inserted)))
77        });
78
79        Ok(count.map(|(_, count)| GraphRecordValue::Int(count)))
80    }
81
82    fn estimate(&self, _input: Estimate, _stats: &Stats) -> Estimate {
83        Estimate::singleton()
84    }
85}
86
87impl<O: Apply<UniqueCountOperation>> UniqueCount for O {
88    type ReturnOperand = O::Output;
89
90    fn n_unique(&self) -> Self::ReturnOperand {
91        Self::ReturnOperand::new(OperationContext::new(self.clone(), UniqueCountOperation))
92    }
93}
94
95operation_manifest! {
96    UniqueCountOperation {
97        method: UniqueCount::n_unique;
98        scope: lane;
99
100        kernel {
101            parameters: <
102                I: IndexDomain,
103                V: ValueEquivalence,
104                O: OrderState,
105            >;
106            input: (Indexed<I, V>, Multiple<O>);
107            output: DefiniteBareValueOperand;
108        }
109
110        kernel {
111            parameters: <
112                V: ValueEquivalence + BareValueDomain,
113                O: OrderState,
114            >;
115            input: (Bare<V>, Multiple<O>);
116            output: DefiniteBareValueOperand;
117        }
118    }
119}
graphrecords_query/operations/aggregation/unique_count.rs

graphrecords_query/operations/aggregation/
unique_count.rs
