datafusion 0.2.15

// Copyright 2018 Grove Enterprises LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Count() aggregate function

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) => {
                //println!("Counting array elements: {}", array.len());
                self.count += array.len();
                Ok(())
            }
            Value::Scalar(_) => {
                //println!("Counting scalar value");
                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::*;
    //use arrow::array::*;

    #[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!(),
        }
    }
}