reifydb_engine/function/math/aggregate/
count.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 Count {
12	pub counts: HashMap<Vec<Value>, i64>,
13}
14
15impl Count {
16	pub fn new() -> Self {
17		Self {
18			counts: HashMap::new(),
19		}
20	}
21}
22
23impl AggregateFunction for Count {
24	fn aggregate(&mut self, ctx: AggregateFunctionContext) -> crate::Result<()> {
25		let column = ctx.column;
26		let groups = &ctx.groups;
27
28		// Check if this is count(*) by examining if we have a dummy column
29		let is_count_star = column.name.text() == "dummy" && matches!(column.data(), ColumnData::Int4(_));
30
31		if is_count_star {
32			// For count(*), count all rows including those with undefined values
33			for (group, indices) in groups.iter() {
34				let count = indices.len() as i64;
35				self.counts.insert(group.clone(), count);
36			}
37		} else {
38			// For count(column), only count defined (non-null) values
39			match &column.data() {
40				ColumnData::Bool(container) => {
41					for (group, indices) in groups.iter() {
42						let count = indices.iter().filter(|&i| container.is_defined(*i)).count()
43							as i64;
44						self.counts.insert(group.clone(), count);
45					}
46				}
47				ColumnData::Float8(container) => {
48					for (group, indices) in groups.iter() {
49						let count = indices.iter().filter(|&i| container.is_defined(*i)).count()
50							as i64;
51						self.counts.insert(group.clone(), count);
52					}
53				}
54				ColumnData::Float4(container) => {
55					for (group, indices) in groups.iter() {
56						let count = indices.iter().filter(|&i| container.is_defined(*i)).count()
57							as i64;
58						self.counts.insert(group.clone(), count);
59					}
60				}
61				ColumnData::Int4(container) => {
62					for (group, indices) in groups.iter() {
63						let count = indices.iter().filter(|&i| container.is_defined(*i)).count()
64							as i64;
65						self.counts.insert(group.clone(), count);
66					}
67				}
68				ColumnData::Int8(container) => {
69					for (group, indices) in groups.iter() {
70						let count = indices.iter().filter(|&i| container.is_defined(*i)).count()
71							as i64;
72						self.counts.insert(group.clone(), count);
73					}
74				}
75				ColumnData::Int2(container) => {
76					for (group, indices) in groups.iter() {
77						let count = indices.iter().filter(|&i| container.is_defined(*i)).count()
78							as i64;
79						self.counts.insert(group.clone(), count);
80					}
81				}
82				ColumnData::Int1(container) => {
83					for (group, indices) in groups.iter() {
84						let count = indices.iter().filter(|&i| container.is_defined(*i)).count()
85							as i64;
86						self.counts.insert(group.clone(), count);
87					}
88				}
89				ColumnData::Int16(container) => {
90					for (group, indices) in groups.iter() {
91						let count = indices.iter().filter(|&i| container.is_defined(*i)).count()
92							as i64;
93						self.counts.insert(group.clone(), count);
94					}
95				}
96				ColumnData::Utf8 {
97					container,
98					..
99				} => {
100					for (group, indices) in groups.iter() {
101						let count = indices.iter().filter(|&i| container.is_defined(*i)).count()
102							as i64;
103						self.counts.insert(group.clone(), count);
104					}
105				}
106				ColumnData::Date(container) => {
107					for (group, indices) in groups.iter() {
108						let count = indices.iter().filter(|&i| container.is_defined(*i)).count()
109							as i64;
110						self.counts.insert(group.clone(), count);
111					}
112				}
113				ColumnData::DateTime(container) => {
114					for (group, indices) in groups.iter() {
115						let count = indices.iter().filter(|&i| container.is_defined(*i)).count()
116							as i64;
117						self.counts.insert(group.clone(), count);
118					}
119				}
120				ColumnData::Undefined(_) => {
121					// Undefined columns have no defined values to count
122					for (group, _indices) in groups.iter() {
123						self.counts.insert(group.clone(), 0);
124					}
125				}
126				_ => {
127					// For other column types, use generic is_defined check
128					for (group, indices) in groups.iter() {
129						let count = indices
130							.iter()
131							.filter(|&i| column.data().is_defined(*i))
132							.count() as i64;
133						self.counts.insert(group.clone(), count);
134					}
135				}
136			}
137		}
138		Ok(())
139	}
140
141	fn finalize(&mut self) -> crate::Result<(Vec<Vec<Value>>, ColumnData)> {
142		let mut keys = Vec::with_capacity(self.counts.len());
143		let mut data = ColumnData::int8_with_capacity(self.counts.len());
144
145		for (key, count) in std::mem::take(&mut self.counts) {
146			keys.push(key);
147			data.push_value(Value::Int8(count));
148		}
149
150		Ok((keys, data))
151	}
152}
reifydb_engine/function/math/aggregate/count.rs

reifydb_engine/function/math/aggregate/
count.rs
