reifydb_engine/function/math/scalar/
avg.rs1use 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}