reifydb_engine/function/math/scalar/
avg.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 reifydb_core::value::column::ColumnData;
5
6use crate::function::{ScalarFunction, ScalarFunctionContext};
7
8pub struct Avg {}
9
10impl Avg {
11	pub fn new() -> Self {
12		Self {}
13	}
14}
15
16impl ScalarFunction for Avg {
17	fn scalar(&self, ctx: ScalarFunctionContext) -> crate::Result<ColumnData> {
18		let columns = ctx.columns;
19		let row_count = ctx.row_count;
20
21		let mut sum = vec![0.0f64; row_count];
22		let mut count = vec![0u32; row_count];
23
24		for col in columns.iter() {
25			match &col.data() {
26				ColumnData::Int1(container) => {
27					for i in 0..row_count {
28						if let Some(value) = container.get(i) {
29							sum[i] += *value as f64;
30							count[i] += 1;
31						}
32					}
33				}
34				ColumnData::Int2(container) => {
35					for i in 0..row_count {
36						if let Some(value) = container.get(i) {
37							sum[i] += *value as f64;
38							count[i] += 1;
39						}
40					}
41				}
42				ColumnData::Int4(container) => {
43					for i in 0..row_count {
44						if let Some(value) = container.get(i) {
45							sum[i] += *value as f64;
46							count[i] += 1;
47						}
48					}
49				}
50				ColumnData::Float4(container) => {
51					for i in 0..row_count {
52						if let Some(value) = container.get(i) {
53							sum[i] += *value as f64;
54							count[i] += 1;
55						}
56					}
57				}
58				ColumnData::Float8(container) => {
59					for i in 0..row_count {
60						if let Some(value) = container.get(i) {
61							sum[i] += *value;
62							count[i] += 1;
63						}
64					}
65				}
66				data => unimplemented!("{data:?}"),
67			}
68		}
69
70		let mut data = Vec::with_capacity(row_count);
71		let mut valids = Vec::with_capacity(row_count);
72
73		for i in 0..row_count {
74			if count[i] > 0 {
75				data.push(sum[i] / count[i] as f64);
76				valids.push(true);
77			} else {
78				data.push(0.0);
79				valids.push(false);
80			}
81		}
82
83		Ok(ColumnData::float8_with_bitvec(data, valids))
84	}
85}