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 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}