reifydb_engine/function/math/aggregate/
sum.rs1use std::collections::HashMap;
5
6use reifydb_core::value::column::ColumnData;
7use reifydb_type::Value;
8
9use crate::function::{AggregateFunction, AggregateFunctionContext};
10
11pub struct Sum {
12 pub sums: HashMap<Vec<Value>, Value>,
13}
14
15impl Sum {
16 pub fn new() -> Self {
17 Self {
18 sums: HashMap::new(),
19 }
20 }
21}
22
23impl AggregateFunction for Sum {
24 fn aggregate(&mut self, ctx: AggregateFunctionContext) -> crate::Result<()> {
25 let column = ctx.column;
26 let groups = &ctx.groups;
27
28 match &column.data() {
29 ColumnData::Float8(container) => {
30 for (group, indices) in groups.iter() {
31 let sum: f64 = indices
32 .iter()
33 .filter(|&i| container.is_defined(*i))
34 .filter_map(|&i| container.get(i))
35 .sum();
36
37 self.sums.insert(group.clone(), Value::float8(sum));
38 }
39 Ok(())
40 }
41 ColumnData::Float4(container) => {
42 for (group, indices) in groups.iter() {
43 let sum: f32 = indices
44 .iter()
45 .filter(|&i| container.is_defined(*i))
46 .filter_map(|&i| container.get(i))
47 .sum();
48
49 self.sums.insert(group.clone(), Value::float4(sum));
50 }
51 Ok(())
52 }
53 ColumnData::Int2(container) => {
54 for (group, indices) in groups.iter() {
55 let sum: i16 = indices.iter().filter_map(|&i| container.get(i)).sum();
56
57 self.sums.insert(group.clone(), Value::Int2(sum));
58 }
59 Ok(())
60 }
61 ColumnData::Int4(container) => {
62 for (group, indices) in groups.iter() {
63 let sum: i32 = indices
64 .iter()
65 .filter(|&i| container.is_defined(*i))
66 .filter_map(|&i| container.get(i))
67 .sum();
68 self.sums.insert(group.clone(), Value::Int4(sum));
69 }
70 Ok(())
71 }
72 ColumnData::Int8(container) => {
73 for (group, indices) in groups.iter() {
74 let sum: i64 = indices
75 .iter()
76 .filter(|&i| container.is_defined(*i))
77 .filter_map(|&i| container.get(i))
78 .sum();
79
80 self.sums.insert(group.clone(), Value::Int8(sum));
81 }
82 Ok(())
83 }
84 _ => unimplemented!("{}", column.get_type()),
85 }
86 }
87
88 fn finalize(&mut self) -> crate::Result<(Vec<Vec<Value>>, ColumnData)> {
89 let mut keys = Vec::with_capacity(self.sums.len());
90 let mut data = ColumnData::undefined(0);
91
92 for (key, sum) in std::mem::take(&mut self.sums) {
93 keys.push(key);
94 data.push_value(sum);
95 }
96
97 Ok((keys, data))
98 }
99}