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