corium_query/
aggregate.rs

1//! Aggregate evaluation for `:find` specifications.
2//!
3//! The engine is deterministic: `sample` returns the first `n` distinct
4//! values in sort order rather than a random selection, and `rand` is not
5//! provided. `median` returns the lower middle element for even counts.
6
7use corium_core::{TotalF64, Value};
8
9use crate::QueryError;
10use crate::builtins::compare;
11
12/// An aggregate outcome, distinguishing scalar, ordered-collection, and
13/// set-valued results for output conversion.
14#[derive(Clone, Debug, PartialEq)]
15pub enum AggOut {
16    /// A scalar value.
17    Value(Value),
18    /// An ordered collection (n-arity `min`/`max`, `sample`).
19    Coll(Vec<Value>),
20    /// A set of distinct values (`distinct`).
21    Set(Vec<Value>),
22}
23
24fn distinct_sorted(values: &[Value]) -> Result<Vec<Value>, QueryError> {
25    let mut out: Vec<Value> = Vec::new();
26    for value in values {
27        if !out.iter().any(|v| v == value) {
28            out.push(value.clone());
29        }
30    }
31    sort_values(&mut out)?;
32    Ok(out)
33}
34
35fn sort_values(values: &mut [Value]) -> Result<(), QueryError> {
36    // Validate comparability first so sort_by can stay total.
37    for window in values.windows(2) {
38        compare(&window[0], &window[1])?;
39    }
40    values.sort_by(|left, right| compare(left, right).unwrap_or(std::cmp::Ordering::Equal));
41    Ok(())
42}
43
44#[allow(clippy::cast_precision_loss)]
45fn numeric_sum(values: &[Value]) -> Result<Value, QueryError> {
46    let mut long_sum: i64 = 0;
47    let mut double_sum: f64 = 0.0;
48    let mut any_double = false;
49    for value in values {
50        match value {
51            Value::Long(v) => {
52                long_sum = long_sum
53                    .checked_add(*v)
54                    .ok_or_else(|| QueryError::Type("sum overflowed".into()))?;
55            }
56            Value::Double(v) => {
57                any_double = true;
58                double_sum += v.0;
59            }
60            _ => return Err(QueryError::Type("sum requires numbers".into())),
61        }
62    }
63    if any_double {
64        Ok(Value::Double(TotalF64(double_sum + long_sum as f64)))
65    } else {
66        Ok(Value::Long(long_sum))
67    }
68}
69
70#[allow(clippy::cast_precision_loss)]
71fn as_f64(value: &Value) -> Result<f64, QueryError> {
72    match value {
73        Value::Long(v) => Ok(*v as f64),
74        Value::Double(v) => Ok(v.0),
75        _ => Err(QueryError::Type("aggregate requires numbers".into())),
76    }
77}
78
79/// Applies aggregate `op` (with optional constant argument `n`) to the bag
80/// of `values` collected for one group.
81///
82/// # Errors
83/// Returns [`QueryError`] for unknown operations or unsupported operand
84/// types.
85#[allow(clippy::cast_precision_loss)]
86pub fn apply(op: &str, n: Option<i64>, values: &[Value]) -> Result<AggOut, QueryError> {
87    let take = |n: Option<i64>| -> Result<usize, QueryError> {
88        n.and_then(|n| usize::try_from(n).ok())
89            .ok_or_else(|| QueryError::Arity(format!("({op} n ?x) requires a positive count")))
90    };
91    match (op, n) {
92        ("count", None) => Ok(AggOut::Value(Value::Long(
93            i64::try_from(values.len()).unwrap_or(i64::MAX),
94        ))),
95        ("count-distinct", None) => Ok(AggOut::Value(Value::Long(
96            i64::try_from(distinct_sorted(values)?.len()).unwrap_or(i64::MAX),
97        ))),
98        ("sum", None) => numeric_sum(values).map(AggOut::Value),
99        ("avg", None) => {
100            if values.is_empty() {
101                return Err(QueryError::Type("avg of no values".into()));
102            }
103            let sum: f64 = values.iter().map(as_f64).sum::<Result<f64, _>>()?;
104            Ok(AggOut::Value(Value::Double(TotalF64(
105                sum / values.len() as f64,
106            ))))
107        }
108        ("median", None) => {
109            if values.is_empty() {
110                return Err(QueryError::Type("median of no values".into()));
111            }
112            let mut sorted = values.to_vec();
113            sort_values(&mut sorted)?;
114            Ok(AggOut::Value(sorted[(sorted.len() - 1) / 2].clone()))
115        }
116        ("min" | "max", None) => {
117            if values.is_empty() {
118                return Err(QueryError::Type(format!("{op} of no values")));
119            }
120            let mut best = values[0].clone();
121            for value in &values[1..] {
122                let ordering = compare(value, &best)?;
123                if (op == "min" && ordering.is_lt()) || (op == "max" && ordering.is_gt()) {
124                    best = value.clone();
125                }
126            }
127            Ok(AggOut::Value(best))
128        }
129        ("min" | "sample", Some(_)) => {
130            let count = take(n)?;
131            let sorted = distinct_sorted(values)?;
132            Ok(AggOut::Coll(sorted.into_iter().take(count).collect()))
133        }
134        ("max", Some(_)) => {
135            let count = take(n)?;
136            let sorted = distinct_sorted(values)?;
137            let skip = sorted.len().saturating_sub(count);
138            Ok(AggOut::Coll(sorted.into_iter().skip(skip).collect()))
139        }
140        ("distinct", None) => Ok(AggOut::Set(distinct_sorted(values)?)),
141        ("variance" | "stddev", None) => {
142            if values.is_empty() {
143                return Err(QueryError::Type(format!("{op} of no values")));
144            }
145            let nums: Vec<f64> = values.iter().map(as_f64).collect::<Result<_, _>>()?;
146            let mean = nums.iter().sum::<f64>() / nums.len() as f64;
147            let variance = nums.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / nums.len() as f64;
148            let out = if op == "variance" {
149                variance
150            } else {
151                variance.sqrt()
152            };
153            Ok(AggOut::Value(Value::Double(TotalF64(out))))
154        }
155        _ => Err(QueryError::Unsupported(format!(
156            "unknown aggregate ({op} …)"
157        ))),
158    }
159}
corium_query/aggregate.rs

corium_query/
aggregate.rs
