Skip to main content

graphrecords_query/operations/aggregation/
sum.rs

1use crate::{
2    Bare, BareValueDomain, EvaluateOperand, Explain, IndexDomain, Indexed, Labeled, Multiple,
3    Operand, OrderState, QueryResult, Single,
4    capabilities::ValueAdd,
5    execution::EvaluationCache,
6    operands::OperandHandle,
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::Sum,
13};
14use graphrecords_core::GraphRecord;
15
16#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
17#[operation(scope = Lane)]
18#[explain(label = "Sum")]
19#[plan(optimizer_hints(empty = if_any))]
20pub struct SumOperation;
21
22impl Prepare for SumOperation {
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, V, O> LaneKernel<Indexed<I, V>, Multiple<O>> for SumOperation
35where
36    I: IndexDomain,
37    V: ValueAdd + BareValueDomain,
38    O: OrderState,
39{
40    type Output = OperandHandle<Bare<V>, Single>;
41
42    fn execute<'a>(
43        _graphrecord: &'a GraphRecord,
44        mut values: KeyedStream<'a, I, V, Multiple<O>>,
45        _prepared: Self::Prepared<'a>,
46    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
47        let sum = values.try_fold(None, |sum, (index, value)| {
48            let value = value?;
49
50            match sum {
51                Some(sum) => V::add(Self::LABEL, sum, value)
52                    .map(Some)
53                    .map_err(|failure| failure.at::<I>(&index)),
54                None => Ok(Some(value)),
55            }
56        });
57
58        Ok(sum.transpose())
59    }
60
61    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
62        input.zero_or_one()
63    }
64}
65
66impl<V, O> LaneKernel<Bare<V>, Multiple<O>> for SumOperation
67where
68    V: ValueAdd + BareValueDomain,
69    O: OrderState,
70{
71    type Output = OperandHandle<Bare<V>, Single>;
72
73    fn execute<'a>(
74        _graphrecord: &'a GraphRecord,
75        mut values: BareStream<'a, V, Multiple<O>>,
76        _prepared: Self::Prepared<'a>,
77    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
78        let sum = values.try_fold(None, |sum, value| {
79            let value = value?;
80
81            match sum {
82                Some(sum) => V::add(Self::LABEL, sum, value).map(Some),
83                None => Ok(Some(value)),
84            }
85        });
86
87        Ok(sum.transpose())
88    }
89
90    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
91        input.zero_or_one()
92    }
93}
94
95impl<O: Apply<SumOperation>> Sum for O {
96    type ReturnOperand = O::Output;
97
98    fn sum(&self) -> Self::ReturnOperand {
99        Self::ReturnOperand::new(OperationContext::new(self.clone(), SumOperation))
100    }
101}
102
103operation_manifest! {
104    SumOperation {
105        method: Sum::sum;
106        scope: lane;
107
108        kernel {
109            parameters: <
110                I: IndexDomain,
111                V: ValueAdd + BareValueDomain,
112                O: OrderState,
113            >;
114            input: (Indexed<I, V>, Multiple<O>);
115            output: OperandHandle<Bare<V>, Single>;
116        }
117
118        kernel {
119            parameters: <
120                V: ValueAdd + BareValueDomain,
121                O: OrderState,
122            >;
123            input: (Bare<V>, Multiple<O>);
124            output: OperandHandle<Bare<V>, Single>;
125        }
126    }
127}