graphrecords_query/operations/aggregation/
variance.rs1use 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}