Skip to main content

graphrecords_query/operations/aggregation/
variance.rs

1use crate::{
2    Bare, BareValueDomain, EvaluateOperand, Explain, Failure, IndexDomain, Indexed, Labeled,
3    Multiple, Operand, OrderState, QueryResult,
4    capabilities::ValueScalar,
5    error::aggregation::InvalidVarianceValue,
6    execution::EvaluationCache,
7    operands::BareValueOperand,
8    operations::{
9        Apply, BareStream, KeyedStream, LaneKernel, Operation, OperationContext, Prepare,
10    },
11    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
12    registry::operation_manifest,
13    traits::Variance,
14};
15use graphrecords_core::{GraphRecord, graphrecord::GraphRecordValue};
16
17#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
18#[operation(scope = Lane)]
19#[explain(label = "Var")]
20#[plan(optimizer_hints(empty = if_any))]
21pub struct VarianceOperation;
22
23impl Prepare for VarianceOperation {
24    type Prepared<'a> = ();
25
26    fn prepare<'a>(
27        &'a self,
28        _graphrecord: &'a GraphRecord,
29        _cache: &'a EvaluationCache<'a>,
30    ) -> QueryResult<Self::Prepared<'a>> {
31        Ok(())
32    }
33}
34
35fn update_state(
36    (count, mean, squared_deviation): (usize, f64, f64),
37    value: f64,
38) -> (usize, f64, f64) {
39    let count = count + 1;
40    let difference = value - mean;
41    let mean = mean + difference / count as f64;
42    let updated_difference = value - mean;
43    let squared_deviation = difference.mul_add(updated_difference, squared_deviation);
44
45    (count, mean, squared_deviation)
46}
47
48impl<I, V, O> LaneKernel<Indexed<I, V>, Multiple<O>> for VarianceOperation
49where
50    I: IndexDomain,
51    V: ValueScalar + BareValueDomain,
52    O: OrderState,
53{
54    type Output = BareValueOperand;
55
56    fn execute<'a>(
57        _graphrecord: &'a GraphRecord,
58        mut values: KeyedStream<'a, I, V, Multiple<O>>,
59        _prepared: Self::Prepared<'a>,
60    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
61        let variance = values
62            .try_fold((0, 0.0, 0.0), |state, (index, value)| {
63                let value = V::into_scalar(Self::LABEL, value?)
64                    .map_err(|failure| failure.at::<I>(&index))?;
65                let value = match value {
66                    GraphRecordValue::Int(value) => value as f64,
67                    GraphRecordValue::Float(value) => value,
68                    value => {
69                        return Err(Failure::new_at::<I, _>(
70                            Self::LABEL,
71                            InvalidVarianceValue::new(value),
72                            &index,
73                        ));
74                    }
75                };
76
77                Ok(update_state(state, value))
78            })
79            .map(|(count, _, squared_deviation)| {
80                (count > 1).then(|| GraphRecordValue::Float(squared_deviation / (count - 1) as f64))
81            });
82
83        Ok(variance.transpose())
84    }
85
86    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
87        input.zero_or_one()
88    }
89}
90
91impl<V, O> LaneKernel<Bare<V>, Multiple<O>> for VarianceOperation
92where
93    V: ValueScalar + BareValueDomain,
94    O: OrderState,
95{
96    type Output = BareValueOperand;
97
98    fn execute<'a>(
99        _graphrecord: &'a GraphRecord,
100        mut values: BareStream<'a, V, Multiple<O>>,
101        _prepared: Self::Prepared<'a>,
102    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
103        let variance = values
104            .try_fold((0, 0.0, 0.0), |state, value| {
105                let value = V::into_scalar(Self::LABEL, value?)?;
106                let value = match value {
107                    GraphRecordValue::Int(value) => value as f64,
108                    GraphRecordValue::Float(value) => value,
109                    value => {
110                        return Err(Failure::new(Self::LABEL, InvalidVarianceValue::new(value)));
111                    }
112                };
113
114                Ok(update_state(state, value))
115            })
116            .map(|(count, _, squared_deviation)| {
117                (count > 1).then(|| GraphRecordValue::Float(squared_deviation / (count - 1) as f64))
118            });
119
120        Ok(variance.transpose())
121    }
122
123    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
124        input.zero_or_one()
125    }
126}
127
128impl<O: Apply<VarianceOperation>> Variance for O {
129    type ReturnOperand = O::Output;
130
131    fn var(&self) -> Self::ReturnOperand {
132        Self::ReturnOperand::new(OperationContext::new(self.clone(), VarianceOperation))
133    }
134}
135
136operation_manifest! {
137    VarianceOperation {
138        method: Variance::var;
139        scope: lane;
140
141        kernel {
142            parameters: <
143                I: IndexDomain,
144                V: ValueScalar + BareValueDomain,
145                O: OrderState,
146            >;
147            input: (Indexed<I, V>, Multiple<O>);
148            output: BareValueOperand;
149        }
150
151        kernel {
152            parameters: <
153                V: ValueScalar + BareValueDomain,
154                O: OrderState,
155            >;
156            input: (Bare<V>, Multiple<O>);
157            output: BareValueOperand;
158        }
159    }
160}