common_function/scalars/date/
date_add.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::fmt;
16
17use common_query::error::{ArrowComputeSnafu, IntoVectorSnafu, InvalidFuncArgsSnafu, Result};
18use common_query::prelude::Signature;
19use datatypes::arrow::compute::kernels::numeric;
20use datatypes::prelude::ConcreteDataType;
21use datatypes::vectors::{Helper, VectorRef};
22use snafu::{ensure, ResultExt};
23
24use crate::function::{Function, FunctionContext};
25use crate::helper;
26
27/// A function adds an interval value to Timestamp, Date, and return the result.
28/// The implementation of datetime type is based on Date64 which is incorrect so this function
29/// doesn't support the datetime type.
30#[derive(Clone, Debug, Default)]
31pub struct DateAddFunction;
32
33const NAME: &str = "date_add";
34
35impl Function for DateAddFunction {
36    fn name(&self) -> &str {
37        NAME
38    }
39
40    fn return_type(&self, input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
41        Ok(input_types[0].clone())
42    }
43
44    fn signature(&self) -> Signature {
45        helper::one_of_sigs2(
46            vec![
47                ConcreteDataType::date_datatype(),
48                ConcreteDataType::timestamp_second_datatype(),
49                ConcreteDataType::timestamp_millisecond_datatype(),
50                ConcreteDataType::timestamp_microsecond_datatype(),
51                ConcreteDataType::timestamp_nanosecond_datatype(),
52            ],
53            vec![
54                ConcreteDataType::interval_month_day_nano_datatype(),
55                ConcreteDataType::interval_year_month_datatype(),
56                ConcreteDataType::interval_day_time_datatype(),
57            ],
58        )
59    }
60
61    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
62        ensure!(
63            columns.len() == 2,
64            InvalidFuncArgsSnafu {
65                err_msg: format!(
66                    "The length of the args is not correct, expect 2, have: {}",
67                    columns.len()
68                ),
69            }
70        );
71
72        let left = columns[0].to_arrow_array();
73        let right = columns[1].to_arrow_array();
74
75        let result = numeric::add(&left, &right).context(ArrowComputeSnafu)?;
76        let arrow_type = result.data_type().clone();
77        Helper::try_into_vector(result).context(IntoVectorSnafu {
78            data_type: arrow_type,
79        })
80    }
81}
82
83impl fmt::Display for DateAddFunction {
84    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
85        write!(f, "DATE_ADD")
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use std::sync::Arc;
92
93    use common_query::prelude::{TypeSignature, Volatility};
94    use datatypes::arrow::datatypes::IntervalDayTime;
95    use datatypes::prelude::ConcreteDataType;
96    use datatypes::value::Value;
97    use datatypes::vectors::{
98        DateVector, IntervalDayTimeVector, IntervalYearMonthVector, TimestampSecondVector,
99    };
100
101    use super::{DateAddFunction, *};
102
103    #[test]
104    fn test_date_add_misc() {
105        let f = DateAddFunction;
106        assert_eq!("date_add", f.name());
107        assert_eq!(
108            ConcreteDataType::timestamp_microsecond_datatype(),
109            f.return_type(&[ConcreteDataType::timestamp_microsecond_datatype()])
110                .unwrap()
111        );
112        assert_eq!(
113            ConcreteDataType::timestamp_second_datatype(),
114            f.return_type(&[ConcreteDataType::timestamp_second_datatype()])
115                .unwrap()
116        );
117        assert_eq!(
118            ConcreteDataType::date_datatype(),
119            f.return_type(&[ConcreteDataType::date_datatype()]).unwrap()
120        );
121        assert!(
122            matches!(f.signature(),
123                         Signature {
124                             type_signature: TypeSignature::OneOf(sigs),
125                             volatility: Volatility::Immutable
126                         } if  sigs.len() == 15),
127            "{:?}",
128            f.signature()
129        );
130    }
131
132    #[test]
133    fn test_timestamp_date_add() {
134        let f = DateAddFunction;
135
136        let times = vec![Some(123), None, Some(42), None];
137        // Intervals in milliseconds
138        let intervals = vec![
139            IntervalDayTime::new(0, 1000),
140            IntervalDayTime::new(0, 2000),
141            IntervalDayTime::new(0, 3000),
142            IntervalDayTime::new(0, 1000),
143        ];
144        let results = [Some(124), None, Some(45), None];
145
146        let time_vector = TimestampSecondVector::from(times.clone());
147        let interval_vector = IntervalDayTimeVector::from_vec(intervals);
148        let args: Vec<VectorRef> = vec![Arc::new(time_vector), Arc::new(interval_vector)];
149        let vector = f.eval(&FunctionContext::default(), &args).unwrap();
150
151        assert_eq!(4, vector.len());
152        for (i, _t) in times.iter().enumerate() {
153            let v = vector.get(i);
154            let result = results.get(i).unwrap();
155
156            if result.is_none() {
157                assert_eq!(Value::Null, v);
158                continue;
159            }
160            match v {
161                Value::Timestamp(ts) => {
162                    assert_eq!(ts.value(), result.unwrap());
163                }
164                _ => unreachable!(),
165            }
166        }
167    }
168
169    #[test]
170    fn test_date_date_add() {
171        let f = DateAddFunction;
172
173        let dates = vec![Some(123), None, Some(42), None];
174        // Intervals in months
175        let intervals = vec![1, 2, 3, 1];
176        let results = [Some(154), None, Some(131), None];
177
178        let date_vector = DateVector::from(dates.clone());
179        let interval_vector = IntervalYearMonthVector::from_vec(intervals);
180        let args: Vec<VectorRef> = vec![Arc::new(date_vector), Arc::new(interval_vector)];
181        let vector = f.eval(&FunctionContext::default(), &args).unwrap();
182
183        assert_eq!(4, vector.len());
184        for (i, _t) in dates.iter().enumerate() {
185            let v = vector.get(i);
186            let result = results.get(i).unwrap();
187
188            if result.is_none() {
189                assert_eq!(Value::Null, v);
190                continue;
191            }
192            match v {
193                Value::Date(date) => {
194                    assert_eq!(date.val(), result.unwrap());
195                }
196                _ => unreachable!(),
197            }
198        }
199    }
200}
common_function/scalars/date/date_add.rs

common_function/scalars/date/
date_add.rs
