common_function/scalars/date/
date_sub.rs1use 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#[derive(Clone, Debug, Default)]
31pub struct DateSubFunction;
32
33const NAME: &str = "date_sub";
34
35impl Function for DateSubFunction {
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::sub(&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 DateSubFunction {
84 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
85 write!(f, "DATE_SUB")
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::{DateSubFunction, *};
102
103 #[test]
104 fn test_date_sub_misc() {
105 let f = DateSubFunction;
106 assert_eq!("date_sub", 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_sub() {
134 let f = DateSubFunction;
135
136 let times = vec![Some(123), None, Some(42), None];
137 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(122), None, Some(39), 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_sub() {
171 let f = DateSubFunction;
172 let days_per_month = 30;
173
174 let dates = vec![
175 Some(123 * days_per_month),
176 None,
177 Some(42 * days_per_month),
178 None,
179 ];
180 let intervals = vec![1, 2, 3, 1];
182 let results = [Some(3659), None, Some(1168), None];
183
184 let date_vector = DateVector::from(dates.clone());
185 let interval_vector = IntervalYearMonthVector::from_vec(intervals);
186 let args: Vec<VectorRef> = vec![Arc::new(date_vector), Arc::new(interval_vector)];
187 let vector = f.eval(&FunctionContext::default(), &args).unwrap();
188
189 assert_eq!(4, vector.len());
190 for (i, _t) in dates.iter().enumerate() {
191 let v = vector.get(i);
192 let result = results.get(i).unwrap();
193
194 if result.is_none() {
195 assert_eq!(Value::Null, v);
196 continue;
197 }
198 match v {
199 Value::Date(date) => {
200 assert_eq!(date.val(), result.unwrap());
201 }
202 _ => unreachable!(),
203 }
204 }
205 }
206}