graphrecords_query/operations/aggregation/
mean.rs

1use crate::{
2    Bare, BareValueDomain, EvaluateOperand, Explain, Failure, IndexDomain, Indexed, Labeled,
3    Multiple, Operand, OrderState, QueryResult, Single,
4    capabilities::ValueScalar,
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::Mean,
13};
14use graphrecords_core::{GraphRecord, graphrecord::GraphRecordValue};
15use std::ops::{Add, Div};
16
17#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
18#[operation(scope = Lane)]
19#[explain(label = "Mean")]
20#[plan(optimizer_hints(empty = if_any))]
21pub struct MeanOperation;
22
23impl Prepare for MeanOperation {
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
35impl<I, V, O> LaneKernel<Indexed<I, V>, Multiple<O>> for MeanOperation
36where
37    I: IndexDomain,
38    V: ValueScalar + BareValueDomain,
39    O: OrderState,
40{
41    type Output = OperandHandle<Bare<V>, Single>;
42
43    fn execute<'a>(
44        _graphrecord: &'a GraphRecord,
45        mut values: KeyedStream<'a, I, V, Multiple<O>>,
46        _prepared: Self::Prepared<'a>,
47    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
48        let mean = values
49            .try_fold(
50                (None, 0_i64, None),
51                |(sum, count, role): (Option<GraphRecordValue>, _, _), (index, value)| {
52                    let value = value?;
53                    let role = role.or_else(|| Some(value.clone()));
54                    let value = V::into_scalar(Self::LABEL, value)
55                        .map_err(|failure| failure.at::<I>(&index))?;
56                    let sum = match sum {
57                        Some(sum) => sum
58                            .add(value)
59                            .map_err(|error| Failure::new_at::<I, _>(Self::LABEL, error, &index))?,
60                        None => value,
61                    };
62
63                    Ok((Some(sum), count + 1, role))
64                },
65            )
66            .and_then(|(sum, count, role)| {
67                sum.zip(role)
68                    .map(|(sum, role)| {
69                        sum.div(GraphRecordValue::Int(count))
70                            .map(|value| V::from_scalar(&role, value))
71                            .map_err(|error| Failure::new(Self::LABEL, error))
72                    })
73                    .transpose()
74            });
75
76        Ok(mean.transpose())
77    }
78
79    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
80        input.zero_or_one()
81    }
82}
83
84impl<V, O> LaneKernel<Bare<V>, Multiple<O>> for MeanOperation
85where
86    V: ValueScalar + BareValueDomain,
87    O: OrderState,
88{
89    type Output = OperandHandle<Bare<V>, Single>;
90
91    fn execute<'a>(
92        _graphrecord: &'a GraphRecord,
93        mut values: BareStream<'a, V, Multiple<O>>,
94        _prepared: Self::Prepared<'a>,
95    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
96        let mean = values
97            .try_fold(
98                (None, 0_i64, None),
99                |(sum, count, role): (Option<GraphRecordValue>, _, _), value| {
100                    let value = value?;
101                    let role = role.or_else(|| Some(value.clone()));
102                    let value = V::into_scalar(Self::LABEL, value)?;
103                    let sum = match sum {
104                        Some(sum) => sum
105                            .add(value)
106                            .map_err(|error| Failure::new(Self::LABEL, error))?,
107                        None => value,
108                    };
109
110                    Ok((Some(sum), count + 1, role))
111                },
112            )
113            .and_then(|(sum, count, role)| {
114                sum.zip(role)
115                    .map(|(sum, role)| {
116                        sum.div(GraphRecordValue::Int(count))
117                            .map(|value| V::from_scalar(&role, value))
118                            .map_err(|error| Failure::new(Self::LABEL, error))
119                    })
120                    .transpose()
121            });
122
123        Ok(mean.transpose())
124    }
125
126    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
127        input.zero_or_one()
128    }
129}
130
131impl<O: Apply<MeanOperation>> Mean for O {
132    type ReturnOperand = O::Output;
133
134    fn mean(&self) -> Self::ReturnOperand {
135        Self::ReturnOperand::new(OperationContext::new(self.clone(), MeanOperation))
136    }
137}
138
139operation_manifest! {
140    MeanOperation {
141        method: Mean::mean;
142        scope: lane;
143
144        kernel {
145            parameters: <
146                I: IndexDomain,
147                V: ValueScalar + BareValueDomain,
148                O: OrderState,
149            >;
150            input: (Indexed<I, V>, Multiple<O>);
151            output: OperandHandle<Bare<V>, Single>;
152        }
153
154        kernel {
155            parameters: <
156                V: ValueScalar + BareValueDomain,
157                O: OrderState,
158            >;
159            input: (Bare<V>, Multiple<O>);
160            output: OperandHandle<Bare<V>, Single>;
161        }
162    }
163}
graphrecords_query/operations/aggregation/mean.rs

graphrecords_query/operations/aggregation/
mean.rs
