1use std::{cmp::Ordering, num::NonZeroU128};
18
19use crate::signal::{Series, Summary};
21
22use super::scalar::{self, Agg, ExactNumber, ScalarKind, ScalarValue};
24use crate::aggregate;
25use crate::{Aggregate, Error, Result};
26
27#[derive(Clone, Debug)]
36pub struct Accumulator {
37 agg: Agg,
38 rows: u64,
39 count: u64,
40 kind: Option<ScalarKind>,
41 selected: Option<(u64, ScalarValue)>,
42 total: ExactNumber,
43 running: Summary,
44 scratch: Vec<ScalarValue>,
45}
46#[derive(Clone, Debug)]
47pub struct Scratch {
48 remaining: usize,
49}
50
51impl Scratch {
58 pub fn new(slots: usize) -> Self {
59 Self { remaining: slots }
60 }
61}
62impl Accumulator {
63 pub fn new(agg: Agg) -> Self {
64 Self {
65 agg,
66 rows: 0,
67 count: 0,
68 kind: None,
69 selected: None,
70 total: ExactNumber::new(false, 0, NonZeroU128::MIN),
71 running: Summary::new(),
72 scratch: Vec::new(),
73 }
74 }
75
76 pub fn push(&mut self, capture: u64, value: &ScalarValue, scratch: &mut Scratch) -> Result<()> {
77 value.check()?;
78 self.rows = self.rows.checked_add(1).ok_or(invalid("row count overflow"))?;
79 if self.agg == Agg::Last && !matches!(value, ScalarValue::Missing) {
80 let later = self.selected.as_ref().is_none_or(|(previous, _)| capture > *previous);
81 if later {
82 self.selected = Some((capture, value.clone()));
83 }
84 }
85 if matches!(value, ScalarValue::Missing | ScalarValue::Null) {
86 return Ok(());
87 }
88 self.count = self.count.checked_add(1).ok_or(invalid("value count overflow"))?;
89 let Agg::Of(statistic) = self.agg else {
90 return Ok(());
91 };
92 if statistic == Aggregate::Count {
93 return Ok(());
94 }
95 let value_kind = value.kind().expect("checked non-missing scalar has a kind");
96 match self.kind {
97 Some(previous) if previous != value_kind && !(exact_kind(previous) && exact_kind(value_kind)) => {
98 return Err(invalid("mixed scalar aggregate types"));
99 }
100 Some(_) => {}
101 None => self.kind = Some(value_kind),
102 }
103 match statistic {
104 Aggregate::Min | Aggregate::Max => {
105 let replace = match &self.selected {
106 None => true,
107 Some((_, previous)) => {
108 let order = value.compare(previous)?;
109 if statistic == Aggregate::Min { order.is_lt() } else { order.is_gt() }
110 }
111 };
112 if replace {
113 self.selected = Some((capture, value.clone()));
114 }
115 }
116 Aggregate::Sum | Aggregate::Mean => {
117 if let ScalarValue::Float(value) = value {
118 self.running.push(*value);
119 let answered = aggregate::of_summary(&self.running, statistic);
120 if !answered.is_some_and(f64::is_finite) {
121 return Err(invalid("aggregate exceeds finite range"));
122 }
123 } else {
124 let exact = scalar::exact(value).ok_or(invalid("aggregate requires numbers"))?;
125 self.total = self.total.checked_add(exact)?;
126 }
127 }
128 Aggregate::Rms | Aggregate::Std => {
129 self.running.push(approximate(value).ok_or(invalid("aggregate requires numbers"))?);
130 }
131 Aggregate::Median | Aggregate::Percentile(_) => {
132 if !numeric_kind(value_kind) {
133 return Err(invalid("order statistic requires numbers"));
134 }
135 if self.scratch.len() == self.scratch.capacity() {
136 let slots = self.scratch.capacity().max(4).min(scratch.remaining);
137 if slots == 0 {
138 return Err(invalid("order statistic values exceed admitted scratch"));
139 }
140 self.scratch.try_reserve_exact(slots).map_err(|_| invalid("order statistic scratch allocation failed"))?;
141 scratch.remaining -= slots;
142 }
143 self.scratch.push(value.clone());
144 }
145 Aggregate::Count => unreachable!("handled before type checks"),
146 }
147 Ok(())
148 }
149
150 fn ordered(&mut self, statistic: Aggregate, check: &mut impl FnMut() -> Result<()>) -> Result<ScalarValue> {
153 let mut values: Vec<f64> = Vec::new();
154 values.try_reserve_exact(self.scratch.len()).map_err(|_| invalid("order statistic allocation failed"))?;
155 values.extend(self.scratch.iter().map(|value| approximate(value).expect("order statistic values are numeric")));
156 self.scratch = Vec::new();
157 check()?;
158 let mut series = Series::from(values);
159 series.sort();
160 check()?;
161 Ok(ScalarValue::Float(finite(aggregate::of_series(&series, statistic))?))
162 }
163
164 pub fn finish(mut self, check: &mut impl FnMut() -> Result<()>) -> Result<ScalarValue> {
165 let statistic = match self.agg {
166 Agg::Len => return Ok(ScalarValue::Unsigned(self.rows)),
167 Agg::Last => {
168 return Ok(self.selected.map_or(ScalarValue::Missing, |(_, value)| value));
169 }
170 Agg::Of(statistic) => statistic,
171 };
172 let floating = self.kind == Some(ScalarKind::Float);
173 Ok(match statistic {
174 Aggregate::Count => ScalarValue::Unsigned(self.count),
175 Aggregate::Min | Aggregate::Max => self.selected.map_or(ScalarValue::Missing, |(_, value)| value),
176 Aggregate::Sum if self.count == 0 => ScalarValue::Unsigned(0),
178 Aggregate::Sum | Aggregate::Mean if floating => ScalarValue::Float(finite(aggregate::of_summary(&self.running, statistic))?),
179 Aggregate::Sum => integral(self.total, self.kind),
180 Aggregate::Mean if self.count == 0 => ScalarValue::Missing,
181 Aggregate::Mean => integral(self.total.divide(u128::from(self.count))?, self.kind),
182 Aggregate::Rms | Aggregate::Std if self.count == 0 => ScalarValue::Missing,
183 Aggregate::Rms | Aggregate::Std => ScalarValue::Float(finite(aggregate::of_summary(&self.running, statistic))?),
184 Aggregate::Median | Aggregate::Percentile(_) if self.count == 0 => ScalarValue::Missing,
185 Aggregate::Percentile(_) => self.ordered(statistic, check)?,
186 Aggregate::Median if floating => self.ordered(statistic, check)?,
187 Aggregate::Median => {
188 let mut order: Vec<usize> = Vec::new();
189 order.try_reserve_exact(self.scratch.len()).map_err(|_| invalid("median order allocation failed"))?;
190 order.extend(0..self.scratch.len());
191 let scratch = &self.scratch;
192 sort_values(&mut order, |a, b| {
193 check()?;
194 Ok(scratch[a].compare(&scratch[b])?)
195 })?;
196 let lower = &scratch[order[(order.len() - 1) / 2]];
197 let upper = &scratch[order[order.len() / 2]];
198 let midpoint = scalar::exact(lower)
199 .expect("median values are numeric")
200 .checked_add(scalar::exact(upper).expect("median values are numeric"))?
201 .divide(2)?;
202 self.scratch.clear();
203 integral(midpoint, self.kind)
204 }
205 })
206 }
207}
208
209pub fn sort_values<T: Copy>(order: &mut [T], mut compare: impl FnMut(T, T) -> Result<Ordering>) -> Result<()> {
213 fn sift<T: Copy>(order: &mut [T], mut root: usize, end: usize, compare: &mut impl FnMut(T, T) -> Result<Ordering>) -> Result<()> {
214 while root < end / 2 {
215 let mut child = root * 2 + 1;
216 if child + 1 < end && compare(order[child], order[child + 1])?.is_lt() {
217 child += 1;
218 }
219 if !compare(order[root], order[child])?.is_lt() {
220 break;
221 }
222 order.swap(root, child);
223 root = child;
224 }
225 Ok(())
226 }
227 for root in (0..order.len() / 2).rev() {
228 sift(order, root, order.len(), &mut compare)?;
229 }
230 for end in (1..order.len()).rev() {
231 order.swap(0, end);
232 sift(order, 0, end, &mut compare)?;
233 }
234 Ok(())
235}
236
237fn numeric_kind(kind: ScalarKind) -> bool {
240 matches!(kind, ScalarKind::Signed | ScalarKind::Unsigned | ScalarKind::Float | ScalarKind::Exact)
241}
242fn approximate(value: &ScalarValue) -> Option<f64> {
244 if let ScalarValue::Float(value) = value {
245 return Some(*value);
246 }
247 let exact = scalar::exact(value)?;
248 let magnitude = exact.numerator as f64 / exact.denominator.get() as f64;
249 Some(if exact.negative { -magnitude } else { magnitude })
250}
251fn finite(answer: Option<f64>) -> Result<f64> {
252 answer.filter(|value| value.is_finite()).ok_or(invalid("aggregate exceeds finite range"))
253}
254fn exact_kind(kind: ScalarKind) -> bool {
255 matches!(kind, ScalarKind::Signed | ScalarKind::Unsigned | ScalarKind::Exact)
256}
257fn integral(value: ExactNumber, kind: Option<ScalarKind>) -> ScalarValue {
258 if value.denominator.get() == 1 {
259 if kind == Some(ScalarKind::Signed) {
260 let signed = i128::try_from(value.numerator).ok().and_then(|n| i64::try_from(if value.negative { -n } else { n }).ok());
261 if let Some(signed) = signed {
262 return ScalarValue::Signed(signed);
263 }
264 } else if !value.negative
265 && let Ok(unsigned) = u64::try_from(value.numerator)
266 {
267 return ScalarValue::Unsigned(unsigned);
268 }
269 }
270 ScalarValue::Exact(value)
271}
272pub(super) fn invalid(reason: &'static str) -> Error {
273 Error::InvalidInput { field: "storage reduce", reason }
274}
275
276#[cfg(test)]
279mod tests {
280 use super::*;
281
282 fn fold(agg: Agg, values: &[ScalarValue]) -> ScalarValue {
283 let mut scratch = Scratch::new(values.len());
284 let mut accumulator = Accumulator::new(agg);
285 for (capture, value) in values.iter().enumerate() {
286 accumulator.push(capture as u64, value, &mut scratch).unwrap();
287 }
288 accumulator.finish(&mut || Ok(())).unwrap()
289 }
290
291 fn ratio(negative: bool, numerator: u128, denominator: u128) -> ScalarValue {
292 ScalarValue::Exact(ExactNumber::new(negative, numerator, NonZeroU128::new(denominator).unwrap()))
293 }
294
295 #[test]
296 fn public_scalar_fold() {
297 let unsigned = [1u64, 2, 3, 4].map(ScalarValue::Unsigned);
298 assert_eq!(fold(Agg::Of(Aggregate::Sum), &unsigned), ScalarValue::Unsigned(10));
299 assert_eq!(fold(Agg::Of(Aggregate::Mean), &unsigned), ratio(false, 5, 2));
300 assert_eq!(fold(Agg::Of(Aggregate::Median), &unsigned), ratio(false, 5, 2));
301 let signed = [ScalarValue::Signed(-3), ScalarValue::Signed(1)];
302 assert_eq!(fold(Agg::Of(Aggregate::Mean), &signed), ScalarValue::Signed(-1));
303 let text = ["b", "c", "a"].map(|value| ScalarValue::Text(value.into()));
304 assert_eq!(fold(Agg::Of(Aggregate::Min), &text), ScalarValue::Text("a".into()));
305 assert_eq!(fold(Agg::Of(Aggregate::Max), &text), ScalarValue::Text("c".into()));
306 assert_eq!(fold(Agg::Last, &text), ScalarValue::Text("a".into()));
307 let sparse = [ScalarValue::Unsigned(1), ScalarValue::Null, ScalarValue::Missing];
308 assert_eq!(fold(Agg::Len, &sparse), ScalarValue::Unsigned(3));
309 assert_eq!(fold(Agg::Of(Aggregate::Count), &sparse), ScalarValue::Unsigned(1));
310 assert_eq!(fold(Agg::Of(Aggregate::Mean), &[ScalarValue::Null]), ScalarValue::Missing);
311
312 let values = [1.0, 2.0, 4.0, 8.0].map(ScalarValue::Float);
314 let float = |agg| match fold(Agg::Of(agg), &values) {
315 ScalarValue::Float(value) => value,
316 other => panic!("floating statistic answered {other:?}"),
317 };
318 assert!((float(Aggregate::Std) - 7.1875f64.sqrt()).abs() < 1e-12);
319 assert!((float(Aggregate::Rms) - 21.25f64.sqrt()).abs() < 1e-12);
320 assert_eq!(fold(Agg::Of(Aggregate::Median), &values), ScalarValue::Float(3.0));
321 let mut accumulator = Accumulator::new(Agg::Of(Aggregate::Sum));
322 let mut scratch = Scratch::new(0);
323 accumulator.push(0, &ScalarValue::Unsigned(1), &mut scratch).unwrap();
324 assert!(accumulator.push(1, &ScalarValue::Text("x".into()), &mut scratch).is_err());
325 let mut median = Accumulator::new(Agg::Of(Aggregate::Median));
326 assert!(median.push(0, &ScalarValue::Unsigned(1), &mut scratch).is_err());
327 }
328}