Skip to main content

graphrecords_query/operations/aggregation/
count.rs

1use crate::{
2    Bare, BareValueDomain, Definite, EvaluateOperand, Explain, IndexDomain, Indexed, Multiple,
3    Operand, OrderState, QueryResult, Single, ValueDomain,
4    execution::EvaluationCache,
5    operands::DefiniteBareValueOperand,
6    operations::{
7        Apply, BareStream, KeyedStream, LaneKernel, Operation, OperationContext, Prepare,
8    },
9    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
10    registry::operation_manifest,
11    traits::Count,
12};
13use graphrecords_core::{GraphRecord, graphrecord::GraphRecordValue};
14
15#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
16#[operation(scope = Lane)]
17#[explain(label = "Count")]
18pub struct CountOperation;
19
20impl Prepare for CountOperation {
21    type Prepared<'a> = ();
22
23    fn prepare<'a>(
24        &'a self,
25        _graphrecord: &'a GraphRecord,
26        _cache: &'a EvaluationCache<'a>,
27    ) -> QueryResult<Self::Prepared<'a>> {
28        Ok(())
29    }
30}
31
32impl<I: IndexDomain, V: ValueDomain, O: OrderState> LaneKernel<Indexed<I, V>, Multiple<O>>
33    for CountOperation
34{
35    type Output = DefiniteBareValueOperand;
36
37    fn execute<'a>(
38        _graphrecord: &'a GraphRecord,
39        mut values: KeyedStream<'a, I, V, Multiple<O>>,
40        _prepared: Self::Prepared<'a>,
41    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
42        let count = values.try_fold(0_i64, |count, (_, item)| item.map(|_| count + 1));
43
44        Ok(count.map(GraphRecordValue::Int))
45    }
46
47    fn estimate(&self, _input: Estimate, _stats: &Stats) -> Estimate {
48        Estimate::singleton()
49    }
50}
51
52impl<V: BareValueDomain, O: OrderState> LaneKernel<Bare<V>, Multiple<O>> for CountOperation {
53    type Output = DefiniteBareValueOperand;
54
55    fn execute<'a>(
56        _graphrecord: &'a GraphRecord,
57        mut values: BareStream<'a, V, Multiple<O>>,
58        _prepared: Self::Prepared<'a>,
59    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
60        let count = values.try_fold(0_i64, |count, item| item.map(|_| count + 1));
61
62        Ok(count.map(GraphRecordValue::Int))
63    }
64
65    fn estimate(&self, _input: Estimate, _stats: &Stats) -> Estimate {
66        Estimate::singleton()
67    }
68}
69
70impl<I: IndexDomain, V: ValueDomain> LaneKernel<Indexed<I, V>, Single> for CountOperation {
71    type Output = DefiniteBareValueOperand;
72
73    fn execute<'a>(
74        _graphrecord: &'a GraphRecord,
75        value: KeyedStream<'a, I, V, Single>,
76        _prepared: Self::Prepared<'a>,
77    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
78        let count = match value {
79            Some((_, item)) => item.map(|_| 1_i64),
80            None => Ok(0_i64),
81        };
82
83        Ok(count.map(GraphRecordValue::Int))
84    }
85
86    fn estimate(&self, _input: Estimate, _stats: &Stats) -> Estimate {
87        Estimate::singleton()
88    }
89}
90
91impl<V: BareValueDomain> LaneKernel<Bare<V>, Single> for CountOperation {
92    type Output = DefiniteBareValueOperand;
93
94    fn execute<'a>(
95        _graphrecord: &'a GraphRecord,
96        value: BareStream<'a, V, Single>,
97        _prepared: Self::Prepared<'a>,
98    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
99        let count = match value {
100            Some(item) => item.map(|_| 1_i64),
101            None => Ok(0_i64),
102        };
103
104        Ok(count.map(GraphRecordValue::Int))
105    }
106
107    fn estimate(&self, _input: Estimate, _stats: &Stats) -> Estimate {
108        Estimate::singleton()
109    }
110}
111
112impl<I: IndexDomain, V: ValueDomain> LaneKernel<Indexed<I, V>, Definite> for CountOperation {
113    type Output = DefiniteBareValueOperand;
114
115    fn execute<'a>(
116        _graphrecord: &'a GraphRecord,
117        value: KeyedStream<'a, I, V, Definite>,
118        _prepared: Self::Prepared<'a>,
119    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
120        Ok(value.1.map(|_| GraphRecordValue::Int(1)))
121    }
122
123    fn estimate(&self, _input: Estimate, _stats: &Stats) -> Estimate {
124        Estimate::singleton()
125    }
126}
127
128impl<V: BareValueDomain> LaneKernel<Bare<V>, Definite> for CountOperation {
129    type Output = DefiniteBareValueOperand;
130
131    fn execute<'a>(
132        _graphrecord: &'a GraphRecord,
133        value: BareStream<'a, V, Definite>,
134        _prepared: Self::Prepared<'a>,
135    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
136        Ok(value.map(|_| GraphRecordValue::Int(1)))
137    }
138
139    fn estimate(&self, _input: Estimate, _stats: &Stats) -> Estimate {
140        Estimate::singleton()
141    }
142}
143
144impl<O: Apply<CountOperation>> Count for O {
145    type ReturnOperand = O::Output;
146
147    fn count(&self) -> Self::ReturnOperand {
148        Self::ReturnOperand::new(OperationContext::new(self.clone(), CountOperation))
149    }
150}
151
152operation_manifest! {
153    CountOperation {
154        method: Count::count;
155        scope: lane;
156
157        kernel {
158            parameters: <
159                I: IndexDomain,
160                V: ValueDomain,
161                O: OrderState,
162            >;
163            input: (Indexed<I, V>, Multiple<O>);
164            output: DefiniteBareValueOperand;
165        }
166
167        kernel {
168            parameters: <
169                V: BareValueDomain,
170                O: OrderState,
171            >;
172            input: (Bare<V>, Multiple<O>);
173            output: DefiniteBareValueOperand;
174        }
175
176        kernel {
177            parameters: <
178                I: IndexDomain,
179                V: ValueDomain,
180            >;
181            input: (Indexed<I, V>, Single);
182            output: DefiniteBareValueOperand;
183        }
184
185        kernel {
186            parameters: <V: BareValueDomain>;
187            input: (Bare<V>, Single);
188            output: DefiniteBareValueOperand;
189        }
190
191        kernel {
192            parameters: <
193                I: IndexDomain,
194                V: ValueDomain,
195            >;
196            input: (Indexed<I, V>, Definite);
197            output: DefiniteBareValueOperand;
198        }
199
200        kernel {
201            parameters: <V: BareValueDomain>;
202            input: (Bare<V>, Definite);
203            output: DefiniteBareValueOperand;
204        }
205    }
206}