reifydb_engine/function/math/aggregate/
sum.rs

1// Copyright (c) reifydb.com 2025
2// This file is licensed under the AGPL-3.0-or-later, see license.md file
3
4use 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}