akar_function/aggregate/
mod.rs

1use crate::registry::*;
2use crate::scalar::comparison::double_cmp;
3use crate::scalar::{evaluate_scalar, numeric_to_f64};
4use akar_common::types::Value;
5use arrow::array::{
6    ArrayRef, Float32Builder, Float64Builder, Int8Builder, Int16Builder, Int32Builder, Int64Builder, PrimitiveArray,
7    UInt8Builder, UInt16Builder, UInt32Builder, UInt64Builder,
8};
9use arrow::compute;
10use arrow::datatypes::{
11    Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
12};
13use std::sync::Arc;
14
15// ==================== Aggregate ====================
16
17/// State for aggregate function computation over Values.
18#[derive(Debug, Clone)]
19pub enum AggValueState {
20    Count(u64),
21    Sum(Value),
22    Min(Value),
23    Max(Value),
24    Avg {
25        sum: Value,
26        count: u64,
27    },
28    Collect(Vec<Value>),
29    StdDev {
30        sum: f64,
31        sum_sq: f64,
32        count: u64,
33    },
34    Variance {
35        sum: f64,
36        sum_sq: f64,
37        count: u64,
38    },
39    /// Percentile state — collects all non-null values for percentile computation.
40    Percentile {
41        values: Vec<f64>,
42        percentile: f64,
43    },
44    /// COUNT_IF state — counts rows where the condition evaluated to TRUE.
45    CountIf(u64),
46    /// STRING_AGG state — collects string pieces to be concatenated.
47    StringAgg {
48        pieces: Vec<String>,
49        delimiter: String,
50    },
51}
52
53impl AggValueState {
54    /// Create a new initial state for the given aggregate function.
55    pub fn new(func: &AggregateFunction) -> Self {
56        match func {
57            AggregateFunction::Count | AggregateFunction::CountStar => AggValueState::Count(0),
58            AggregateFunction::Sum => AggValueState::Sum(Value::Null),
59            AggregateFunction::Min => AggValueState::Min(Value::Null),
60            AggregateFunction::Max => AggValueState::Max(Value::Null),
61            AggregateFunction::Avg => AggValueState::Avg {
62                sum: Value::Int64(0),
63                count: 0,
64            },
65            AggregateFunction::Collect => AggValueState::Collect(Vec::new()),
66            AggregateFunction::StdDev => AggValueState::StdDev {
67                sum: 0.0,
68                sum_sq: 0.0,
69                count: 0,
70            },
71            AggregateFunction::Variance => AggValueState::Variance {
72                sum: 0.0,
73                sum_sq: 0.0,
74                count: 0,
75            },
76            AggregateFunction::PercentileDisc { percentile } => AggValueState::Percentile {
77                values: Vec::new(),
78                percentile: *percentile,
79            },
80            AggregateFunction::PercentileCont { percentile } => AggValueState::Percentile {
81                values: Vec::new(),
82                percentile: *percentile,
83            },
84            AggregateFunction::CountIf => AggValueState::CountIf(0),
85            AggregateFunction::StringAgg { delimiter } => AggValueState::StringAgg {
86                pieces: Vec::new(),
87                delimiter: delimiter.clone(),
88            },
89        }
90    }
91
92    /// Update the state with a new input value.
93    #[inline(always)]
94    pub fn update(&mut self, val: &Value) {
95        if matches!(val, Value::Null) {
96            // Most aggregates skip NULLs (except COUNT which counts them)
97            return;
98        }
99        match self {
100            AggValueState::Count(n) => *n += 1,
101            AggValueState::Sum(current) => {
102                if matches!(current, Value::Null) {
103                    *current = val.clone();
104                } else {
105                    *current = add_values_for_agg(current.clone(), val.clone());
106                }
107            }
108            AggValueState::Min(current) => {
109                if matches!(current, Value::Null) {
110                    *current = val.clone();
111                } else if let Ok(Value::Bool(true)) = evaluate_scalar(
112                    &ScalarFunction::Comparison { op: ComparisonOp::Lt },
113                    &[val.clone(), current.clone()],
114                ) {
115                    *current = val.clone();
116                }
117            }
118            AggValueState::Max(current) => {
119                if matches!(current, Value::Null) {
120                    *current = val.clone();
121                } else if let Ok(Value::Bool(true)) = evaluate_scalar(
122                    &ScalarFunction::Comparison { op: ComparisonOp::Gt },
123                    &[val.clone(), current.clone()],
124                ) {
125                    *current = val.clone();
126                }
127            }
128            AggValueState::Avg { sum, count } => {
129                if matches!(sum, Value::Int64(0)) {
130                    *sum = val.clone();
131                } else {
132                    *sum = add_values_for_agg(sum.clone(), val.clone());
133                }
134                *count += 1;
135            }
136            AggValueState::Collect(items) => {
137                items.push(val.clone());
138            }
139            AggValueState::StdDev { sum, sum_sq, count } | AggValueState::Variance { sum, sum_sq, count } => {
140                let v = numeric_to_f64(val).unwrap_or(0.0);
141                *sum += v;
142                *sum_sq += v * v;
143                *count += 1;
144            }
145            AggValueState::Percentile { values, .. } => {
146                if let Ok(v) = numeric_to_f64(val) {
147                    values.push(v);
148                }
149            }
150            AggValueState::CountIf(n) => {
151                // Only count if the condition value is TRUE (non-null and Bool(true))
152                if matches!(val, Value::Bool(true)) {
153                    *n += 1;
154                }
155                // NULL and false are ignored (not counted)
156            }
157            AggValueState::StringAgg { pieces, .. } => match val {
158                Value::String(s) => pieces.push(s.clone()),
159                other => pieces.push(format!("{:?}", other)),
160            },
161        }
162    }
163
164    /// Finalize the state into a Value.
165    pub fn finalize(&self) -> Value {
166        match self {
167            AggValueState::Count(n) => Value::Int64(*n as i64),
168            AggValueState::CountIf(n) => Value::Int64(*n as i64),
169            AggValueState::Sum(v) => v.clone(),
170            AggValueState::Min(v) => v.clone(),
171            AggValueState::Max(v) => v.clone(),
172            AggValueState::Avg { sum, count } => {
173                if *count == 0 {
174                    return Value::Null;
175                }
176                match sum {
177                    Value::Int64(s) => Value::Double(*s as f64 / *count as f64),
178                    Value::Double(s) => Value::Double(*s / *count as f64),
179                    _ => sum.clone(),
180                }
181            }
182            AggValueState::Collect(items) => Value::List(items.clone()),
183            AggValueState::StdDev { sum, sum_sq, count } => {
184                if *count == 0 {
185                    return Value::Null;
186                }
187                let n = *count as f64;
188                let variance = (sum_sq - (sum * sum) / n) / n;
189                Value::Double(variance.sqrt())
190            }
191            AggValueState::Variance { sum, sum_sq, count } => {
192                if *count == 0 {
193                    return Value::Null;
194                }
195                let n = *count as f64;
196                let variance = (sum_sq - (sum * sum) / n) / n;
197                Value::Double(variance)
198            }
199            AggValueState::Percentile { values, percentile } => {
200                if values.is_empty() {
201                    return Value::Null;
202                }
203                let mut sorted = values.clone();
204                sorted.sort_by(|a, b| double_cmp(*a, *b));
205                let n = sorted.len();
206                let p = *percentile;
207                // Discrete percentile: pick the value at ceil(p * n) - 1 index
208                let idx = ((p * n as f64).ceil() as usize).saturating_sub(1).min(n - 1);
209                Value::Double(sorted[idx])
210            }
211            AggValueState::StringAgg { pieces, delimiter } => {
212                if pieces.is_empty() {
213                    return Value::Null;
214                }
215                Value::String(pieces.join(delimiter))
216            }
217        }
218    }
219
220    /// Merge another AggValueState into this one (for parallel aggregation).
221    #[inline(always)]
222    pub fn merge(&mut self, other: &Self) {
223        match (self, other) {
224            (AggValueState::Count(a), AggValueState::Count(b)) => *a += b,
225            (AggValueState::Sum(a), AggValueState::Sum(b)) => {
226                if matches!(a, Value::Null) {
227                    *a = b.clone();
228                } else if !matches!(b, Value::Null) {
229                    *a = add_values_for_agg(a.clone(), b.clone());
230                }
231            }
232            (AggValueState::Min(a), AggValueState::Min(b)) => {
233                if !matches!(b, Value::Null) {
234                    if matches!(a, Value::Null) {
235                        *a = b.clone();
236                    } else if let Ok(Value::Bool(true)) = evaluate_scalar(
237                        &ScalarFunction::Comparison { op: ComparisonOp::Lt },
238                        &[b.clone(), a.clone()],
239                    ) {
240                        *a = b.clone();
241                    }
242                }
243            }
244            (AggValueState::Max(a), AggValueState::Max(b)) => {
245                if !matches!(b, Value::Null) {
246                    if matches!(a, Value::Null) {
247                        *a = b.clone();
248                    } else if let Ok(Value::Bool(true)) = evaluate_scalar(
249                        &ScalarFunction::Comparison { op: ComparisonOp::Gt },
250                        &[b.clone(), a.clone()],
251                    ) {
252                        *a = b.clone();
253                    }
254                }
255            }
256            (AggValueState::Avg { sum: s1, count: c1 }, AggValueState::Avg { sum: s2, count: c2 }) => {
257                if *c2 > 0 {
258                    if *c1 == 0 {
259                        *s1 = s2.clone();
260                    } else {
261                        *s1 = add_values_for_agg(s1.clone(), s2.clone());
262                    }
263                    *c1 += c2;
264                }
265            }
266            (AggValueState::Collect(a), AggValueState::Collect(b)) => a.extend(b.iter().cloned()),
267            (
268                AggValueState::StdDev {
269                    sum: s1,
270                    sum_sq: sq1,
271                    count: c1,
272                },
273                AggValueState::StdDev {
274                    sum: s2,
275                    sum_sq: sq2,
276                    count: c2,
277                },
278            ) => {
279                *s1 += s2;
280                *sq1 += sq2;
281                *c1 += c2;
282            }
283            (
284                AggValueState::Variance {
285                    sum: s1,
286                    sum_sq: sq1,
287                    count: c1,
288                },
289                AggValueState::Variance {
290                    sum: s2,
291                    sum_sq: sq2,
292                    count: c2,
293                },
294            ) => {
295                *s1 += s2;
296                *sq1 += sq2;
297                *c1 += c2;
298            }
299            (AggValueState::Percentile { values: a_vals, .. }, AggValueState::Percentile { values: b_vals, .. }) => {
300                a_vals.extend_from_slice(b_vals);
301            }
302            (AggValueState::CountIf(a), AggValueState::CountIf(b)) => *a += b,
303            (AggValueState::StringAgg { pieces: a_pieces, .. }, AggValueState::StringAgg { pieces: b_pieces, .. }) => {
304                a_pieces.extend_from_slice(b_pieces);
305            }
306            _ => {}
307        }
308    }
309}
310
311/// Add two Values for aggregate summation (numeric promotion).
312fn add_values_for_agg(a: Value, b: Value) -> Value {
313    match (&a, &b) {
314        (Value::Int64(x), Value::Int64(y)) => Value::Int64(x + y),
315        (Value::Double(x), Value::Double(y)) => Value::Double(x + y),
316        (Value::Int64(x), Value::Double(y)) => Value::Double(*x as f64 + y),
317        (Value::Double(x), Value::Int64(y)) => Value::Double(x + *y as f64),
318        _ => a, // fallback
319    }
320}
321
322/// Convert a slice of Values to an Arrow ArrayRef for numeric types.
323/// Returns None if values are not numeric or have mixed incompatible types.
324fn values_to_arrow_array(values: &[Value], first: &Value) -> Option<ArrayRef> {
325    match first {
326        Value::Int64(_) => {
327            let mut builder = Int64Builder::new();
328            for v in values {
329                match v {
330                    Value::Int64(n) => builder.append_value(*n),
331                    Value::Null => builder.append_null(),
332                    _ => return None,
333                }
334            }
335            Some(Arc::new(builder.finish()))
336        }
337        Value::Double(_) => {
338            let mut builder = Float64Builder::new();
339            for v in values {
340                match v {
341                    Value::Double(n) => builder.append_value(*n),
342                    Value::Null => builder.append_null(),
343                    _ => return None,
344                }
345            }
346            Some(Arc::new(builder.finish()))
347        }
348        Value::Int32(_) => {
349            let mut builder = Int32Builder::new();
350            for v in values {
351                match v {
352                    Value::Int32(n) => builder.append_value(*n),
353                    Value::Null => builder.append_null(),
354                    _ => return None,
355                }
356            }
357            Some(Arc::new(builder.finish()))
358        }
359        Value::Float(_) => {
360            let mut builder = Float32Builder::new();
361            for v in values {
362                match v {
363                    Value::Float(n) => builder.append_value(*n),
364                    Value::Null => builder.append_null(),
365                    _ => return None,
366                }
367            }
368            Some(Arc::new(builder.finish()))
369        }
370        Value::Int16(_) => {
371            let mut builder = Int16Builder::new();
372            for v in values {
373                match v {
374                    Value::Int16(n) => builder.append_value(*n),
375                    Value::Null => builder.append_null(),
376                    _ => return None,
377                }
378            }
379            Some(Arc::new(builder.finish()))
380        }
381        Value::Int8(_) => {
382            let mut builder = Int8Builder::new();
383            for v in values {
384                match v {
385                    Value::Int8(n) => builder.append_value(*n),
386                    Value::Null => builder.append_null(),
387                    _ => return None,
388                }
389            }
390            Some(Arc::new(builder.finish()))
391        }
392        Value::UInt64(_) => {
393            let mut builder = UInt64Builder::new();
394            for v in values {
395                match v {
396                    Value::UInt64(n) => builder.append_value(*n),
397                    Value::Null => builder.append_null(),
398                    _ => return None,
399                }
400            }
401            Some(Arc::new(builder.finish()))
402        }
403        Value::UInt32(_) => {
404            let mut builder = UInt32Builder::new();
405            for v in values {
406                match v {
407                    Value::UInt32(n) => builder.append_value(*n),
408                    Value::Null => builder.append_null(),
409                    _ => return None,
410                }
411            }
412            Some(Arc::new(builder.finish()))
413        }
414        Value::UInt16(_) => {
415            let mut builder = UInt16Builder::new();
416            for v in values {
417                match v {
418                    Value::UInt16(n) => builder.append_value(*n),
419                    Value::Null => builder.append_null(),
420                    _ => return None,
421                }
422            }
423            Some(Arc::new(builder.finish()))
424        }
425        Value::UInt8(_) => {
426            let mut builder = UInt8Builder::new();
427            for v in values {
428                match v {
429                    Value::UInt8(n) => builder.append_value(*n),
430                    Value::Null => builder.append_null(),
431                    _ => return None,
432                }
433            }
434            Some(Arc::new(builder.finish()))
435        }
436        _ => None,
437    }
438}
439
440/// Try to evaluate aggregate using Arrow compute SIMD kernels.
441/// Returns None if the aggregate type or column type doesn't support SIMD.
442fn try_simd_aggregate(func: &AggregateFunction, args: &[Value]) -> Option<Value> {
443    match func {
444        AggregateFunction::Sum | AggregateFunction::Min | AggregateFunction::Max => {}
445        _ => return None,
446    }
447
448    if args.is_empty() {
449        return None;
450    }
451
452    let first = args.iter().find(|v| !matches!(v, Value::Null))?;
453    let array = values_to_arrow_array(args, first)?;
454
455    macro_rules! simd_dispatch {
456        ($arr:expr, $func:expr, $ty:ty, $variant:ident) => {{
457            let typed: &PrimitiveArray<$ty> = $arr.as_any().downcast_ref()?;
458            match $func {
459                AggregateFunction::Sum => compute::sum(typed).map(Value::$variant),
460                AggregateFunction::Min => compute::min(typed).map(Value::$variant),
461                AggregateFunction::Max => compute::max(typed).map(Value::$variant),
462                _ => return None,
463            }
464        }};
465    }
466
467    match first {
468        Value::Int64(_) => simd_dispatch!(array, func, Int64Type, Int64),
469        Value::Double(_) => simd_dispatch!(array, func, Float64Type, Double),
470        Value::Int32(_) => simd_dispatch!(array, func, Int32Type, Int32),
471        Value::Float(_) => simd_dispatch!(array, func, Float32Type, Float),
472        Value::Int16(_) => simd_dispatch!(array, func, Int16Type, Int16),
473        Value::Int8(_) => simd_dispatch!(array, func, Int8Type, Int8),
474        Value::UInt64(_) => simd_dispatch!(array, func, UInt64Type, UInt64),
475        Value::UInt32(_) => simd_dispatch!(array, func, UInt32Type, UInt32),
476        Value::UInt16(_) => simd_dispatch!(array, func, UInt16Type, UInt16),
477        Value::UInt8(_) => simd_dispatch!(array, func, UInt8Type, UInt8),
478        _ => None,
479    }
480}
481
482/// Evaluate an aggregate function across a slice of Values.
483/// Returns the final aggregate value.
484pub fn evaluate_aggregate(func: &AggregateFunction, args: &[Value]) -> Result<Value, String> {
485    // Try SIMD-accelerated path for numeric aggregates (Sum, Min, Max)
486    if let Some(result) = try_simd_aggregate(func, args) {
487        return Ok(result);
488    }
489
490    let mut state = AggValueState::new(func);
491
492    // COUNT(*) counts all rows regardless of arguments
493    if matches!(func, AggregateFunction::CountStar) {
494        if let AggValueState::Count(n) = &mut state {
495            *n = args.len() as u64;
496        }
497        return Ok(state.finalize());
498    }
499
500    for arg in args {
501        state.update(arg);
502    }
503
504    Ok(state.finalize())
505}
akar_function/aggregate/mod.rs

akar_function/aggregate/
mod.rs
