common_function/aggrs/vector/
sum.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use common_macro::{as_aggr_func_creator, AggrFuncTypeStore};
18use common_query::error::{CreateAccumulatorSnafu, Error, InvalidFuncArgsSnafu};
19use common_query::logical_plan::{
20    create_aggregate_function, Accumulator, AggregateFunctionCreator,
21};
22use common_query::prelude::AccumulatorCreatorFunction;
23use datafusion_expr::AggregateUDF;
24use datatypes::prelude::{ConcreteDataType, Value, *};
25use datatypes::vectors::VectorRef;
26use nalgebra::{Const, DVectorView, Dyn, OVector};
27use snafu::ensure;
28
29use crate::scalars::vector::impl_conv::{as_veclit, as_veclit_if_const, veclit_to_binlit};
30
31/// The accumulator for the `vec_sum` aggregate function.
32#[derive(Debug, Default)]
33pub struct VectorSum {
34    sum: Option<OVector<f32, Dyn>>,
35    has_null: bool,
36}
37
38#[as_aggr_func_creator]
39#[derive(Debug, Default, AggrFuncTypeStore)]
40pub struct VectorSumCreator {}
41
42impl AggregateFunctionCreator for VectorSumCreator {
43    fn creator(&self) -> AccumulatorCreatorFunction {
44        let creator: AccumulatorCreatorFunction = Arc::new(move |types: &[ConcreteDataType]| {
45            ensure!(
46                types.len() == 1,
47                InvalidFuncArgsSnafu {
48                    err_msg: format!(
49                        "The length of the args is not correct, expect exactly one, have: {}",
50                        types.len()
51                    )
52                }
53            );
54            let input_type = &types[0];
55            match input_type {
56                ConcreteDataType::String(_) | ConcreteDataType::Binary(_) => {
57                    Ok(Box::new(VectorSum::default()))
58                }
59                _ => {
60                    let err_msg = format!(
61                        "\"VEC_SUM\" aggregate function not support data type {:?}",
62                        input_type.logical_type_id(),
63                    );
64                    CreateAccumulatorSnafu { err_msg }.fail()?
65                }
66            }
67        });
68        creator
69    }
70
71    fn output_type(&self) -> common_query::error::Result<ConcreteDataType> {
72        Ok(ConcreteDataType::binary_datatype())
73    }
74
75    fn state_types(&self) -> common_query::error::Result<Vec<ConcreteDataType>> {
76        Ok(vec![self.output_type()?])
77    }
78}
79
80impl VectorSum {
81    /// Create a new `AggregateUDF` for the `vec_sum` aggregate function.
82    pub fn uadf_impl() -> AggregateUDF {
83        create_aggregate_function(
84            "vec_sum".to_string(),
85            1,
86            Arc::new(VectorSumCreator::default()),
87        )
88        .into()
89    }
90
91    fn inner(&mut self, len: usize) -> &mut OVector<f32, Dyn> {
92        self.sum
93            .get_or_insert_with(|| OVector::zeros_generic(Dyn(len), Const::<1>))
94    }
95
96    fn update(&mut self, values: &[VectorRef], is_update: bool) -> Result<(), Error> {
97        if values.is_empty() || self.has_null {
98            return Ok(());
99        };
100        let column = &values[0];
101        let len = column.len();
102
103        match as_veclit_if_const(column)? {
104            Some(column) => {
105                let vec_column = DVectorView::from_slice(&column, column.len()).scale(len as f32);
106                *self.inner(vec_column.len()) += vec_column;
107            }
108            None => {
109                for i in 0..len {
110                    let Some(arg0) = as_veclit(column.get_ref(i))? else {
111                        if is_update {
112                            self.has_null = true;
113                            self.sum = None;
114                        }
115                        return Ok(());
116                    };
117                    let vec_column = DVectorView::from_slice(&arg0, arg0.len());
118                    *self.inner(vec_column.len()) += vec_column;
119                }
120            }
121        }
122        Ok(())
123    }
124}
125
126impl Accumulator for VectorSum {
127    fn state(&self) -> common_query::error::Result<Vec<Value>> {
128        self.evaluate().map(|v| vec![v])
129    }
130
131    fn update_batch(&mut self, values: &[VectorRef]) -> common_query::error::Result<()> {
132        self.update(values, true)
133    }
134
135    fn merge_batch(&mut self, states: &[VectorRef]) -> common_query::error::Result<()> {
136        self.update(states, false)
137    }
138
139    fn evaluate(&self) -> common_query::error::Result<Value> {
140        match &self.sum {
141            None => Ok(Value::Null),
142            Some(vector) => Ok(Value::from(veclit_to_binlit(vector.as_slice()))),
143        }
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use std::sync::Arc;
150
151    use datatypes::vectors::{ConstantVector, StringVector};
152
153    use super::*;
154
155    #[test]
156    fn test_update_batch() {
157        // test update empty batch, expect not updating anything
158        let mut vec_sum = VectorSum::default();
159        vec_sum.update_batch(&[]).unwrap();
160        assert!(vec_sum.sum.is_none());
161        assert!(!vec_sum.has_null);
162        assert_eq!(Value::Null, vec_sum.evaluate().unwrap());
163
164        // test update one not-null value
165        let mut vec_sum = VectorSum::default();
166        let v: Vec<VectorRef> = vec![Arc::new(StringVector::from(vec![Some(
167            "[1.0,2.0,3.0]".to_string(),
168        )]))];
169        vec_sum.update_batch(&v).unwrap();
170        assert_eq!(
171            Value::from(veclit_to_binlit(&[1.0, 2.0, 3.0])),
172            vec_sum.evaluate().unwrap()
173        );
174
175        // test update one null value
176        let mut vec_sum = VectorSum::default();
177        let v: Vec<VectorRef> = vec![Arc::new(StringVector::from(vec![Option::<String>::None]))];
178        vec_sum.update_batch(&v).unwrap();
179        assert_eq!(Value::Null, vec_sum.evaluate().unwrap());
180
181        // test update no null-value batch
182        let mut vec_sum = VectorSum::default();
183        let v: Vec<VectorRef> = vec![Arc::new(StringVector::from(vec![
184            Some("[1.0,2.0,3.0]".to_string()),
185            Some("[4.0,5.0,6.0]".to_string()),
186            Some("[7.0,8.0,9.0]".to_string()),
187        ]))];
188        vec_sum.update_batch(&v).unwrap();
189        assert_eq!(
190            Value::from(veclit_to_binlit(&[12.0, 15.0, 18.0])),
191            vec_sum.evaluate().unwrap()
192        );
193
194        // test update null-value batch
195        let mut vec_sum = VectorSum::default();
196        let v: Vec<VectorRef> = vec![Arc::new(StringVector::from(vec![
197            Some("[1.0,2.0,3.0]".to_string()),
198            None,
199            Some("[7.0,8.0,9.0]".to_string()),
200        ]))];
201        vec_sum.update_batch(&v).unwrap();
202        assert_eq!(Value::Null, vec_sum.evaluate().unwrap());
203
204        // test update with constant vector
205        let mut vec_sum = VectorSum::default();
206        let v: Vec<VectorRef> = vec![Arc::new(ConstantVector::new(
207            Arc::new(StringVector::from_vec(vec!["[1.0,2.0,3.0]".to_string()])),
208            4,
209        ))];
210        vec_sum.update_batch(&v).unwrap();
211        assert_eq!(
212            Value::from(veclit_to_binlit(&[4.0, 8.0, 12.0])),
213            vec_sum.evaluate().unwrap()
214        );
215    }
216}
common_function/aggrs/vector/sum.rs

common_function/aggrs/vector/
sum.rs
