use serde_json::Value;
use crate::{
JsonPrimitive,
parser::{
aggregators_helper::{Accumulator, AggregateImpl},
analyzer::{AnalysisContext, AnalyzerError},
ast::Function,
},
};
pub struct CountImpl;
impl AggregateImpl for CountImpl {
fn name(&self) -> &'static str {
"count"
}
fn infer_type(
&self,
fun: &Function,
_ctx: &AnalysisContext,
) -> Result<(JsonPrimitive, bool), AnalyzerError> {
if fun.args.len() == 1 {
Ok((JsonPrimitive::Int, false))
} else {
Err(AnalyzerError::FunctionArgMismatch {
name: fun.name.clone(),
expected: "COUNT(*|expr)".into(),
got: vec![],
})
}
}
fn create_accumulator(&self) -> Box<dyn Accumulator> {
Box::new(CountAcc {
cnt: 0,
is_star: false,
})
}
}
struct CountAcc {
cnt: i64,
is_star: bool,
}
impl Accumulator for CountAcc {
fn update(&mut self, args: &[Value]) -> Result<(), AnalyzerError> {
match args {
[] => {
self.is_star = true;
self.cnt += 1;
}
[v] => {
if matches!(v, Value::Null) {
} else {
self.cnt += 1;
}
}
_ => {
return Err(AnalyzerError::FunctionArgMismatch {
name: "COUNT".into(),
expected: "COUNT(*|expr)".into(),
got: vec![],
});
}
}
Ok(())
}
fn finalize(&self) -> Value {
Value::Number(serde_json::Number::from(self.cnt))
}
}
#[cfg(test)]
mod tests {
use super::CountImpl;
use crate::parser::aggregators_helper::AggregateImpl;
use serde_json::json;
#[test]
fn accumulator_counts_star_and_non_null_values_only() {
let mut acc = CountImpl.create_accumulator();
acc.update(&[]).unwrap();
acc.update(&[json!("Ada")]).unwrap();
acc.update(&[json!(null)]).unwrap();
assert_eq!(acc.finalize(), json!(2));
}
#[test]
fn accumulator_rejects_multiple_arguments() {
let mut acc = CountImpl.create_accumulator();
assert!(acc.update(&[json!(1), json!(2)]).is_err());
}
#[test]
fn aggregate_impl_default_does_not_allow_folding() {
assert!(!CountImpl.allow_fold());
}
}