Skip to main content

harmos_storage/query/
accumulate.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12/* === Re-Exports === */
13
14/* === Imports === */
15
16// Standard Library Imports
17use std::{cmp::Ordering, num::NonZeroU128};
18
19// Third Party Imports
20use crate::signal::{Series, Summary};
21
22// Local Imports
23use super::scalar::{self, Agg, ExactNumber, ScalarKind, ScalarValue};
24use crate::aggregate;
25use crate::{Aggregate, Error, Result};
26
27/* === Constants === */
28
29/* === Macros === */
30
31/* === Type Aliases === */
32
33/* === Structs === */
34
35#[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
51/* === Enums === */
52
53/* === Traits === */
54
55/* === Implementations === */
56
57impl 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    // Median and percentiles through `Series`, polling cancellation around
151    // the one sort that orders every retained value.
152    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            // The empty sum is zero, as `Summary` answers it.
177            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
209/* === Public Functions === */
210
211// Fallible In-Place Heapsort Polls Original Cancellation During Comparisons.
212pub 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
237/* === Private Functions === */
238
239fn numeric_kind(kind: ScalarKind) -> bool {
240    matches!(kind, ScalarKind::Signed | ScalarKind::Unsigned | ScalarKind::Float | ScalarKind::Exact)
241}
242// The nearest float to a numeric scalar, for statistics that are not exact.
243fn 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/* === Tests === */
277
278#[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        // Floating Statistics Answer Through The Signal Kernels: Population Std.
313        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}