1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
use std::rc::Rc;
use super::super::errors::*;
use super::super::types::*;
use arrow::datatypes::*;
pub struct CountFunction {
count: usize,
}
impl CountFunction {
pub fn new() -> Self {
CountFunction { count: 0 }
}
}
impl AggregateFunction for CountFunction {
fn name(&self) -> String {
"COUNT".to_string()
}
fn args(&self) -> Vec<Field> {
vec![Field::new("arg", DataType::UInt32, true)]
}
fn return_type(&self) -> DataType {
DataType::UInt64
}
fn execute(&mut self, args: &[Value]) -> Result<()> {
assert_eq!(1, args.len());
match args[0] {
Value::Column(ref array) => {
self.count += array.len();
Ok(())
}
Value::Scalar(_) => {
self.count += 1;
Ok(())
}
}
}
fn finish(&self) -> Result<Value> {
Ok(Value::Scalar(Rc::new(ScalarValue::UInt64(
self.count as u64,
))))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_count() {
let mut count = CountFunction::new();
assert_eq!(DataType::UInt64, count.return_type());
let values: Vec<f64> = vec![12.0, 22.0, 32.0, 6.0, 58.1];
count
.execute(&vec![Value::Column(Rc::new(Array::from(values)))])
.unwrap();
let result = count.finish().unwrap();
match result {
Value::Scalar(ref v) => assert_eq!(v.get_u64().unwrap(), 5),
_ => panic!(),
}
}
}