common_function/scalars/date/
date_format.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_error::ext::BoxedError;
18use common_query::error::{self, InvalidFuncArgsSnafu, Result, UnsupportedInputDataTypeSnafu};
19use common_query::prelude::Signature;
20use datatypes::prelude::{ConcreteDataType, MutableVector, ScalarVectorBuilder};
21use datatypes::vectors::{StringVectorBuilder, VectorRef};
22use snafu::{ensure, ResultExt};
23
24use crate::function::{Function, FunctionContext};
25use crate::helper;
26
27/// A function that formats timestamp/date/datetime into string by the format
28#[derive(Clone, Debug, Default)]
29pub struct DateFormatFunction;
30
31const NAME: &str = "date_format";
32
33impl Function for DateFormatFunction {
34    fn name(&self) -> &str {
35        NAME
36    }
37
38    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
39        Ok(ConcreteDataType::string_datatype())
40    }
41
42    fn signature(&self) -> Signature {
43        helper::one_of_sigs2(
44            vec![
45                ConcreteDataType::date_datatype(),
46                ConcreteDataType::timestamp_second_datatype(),
47                ConcreteDataType::timestamp_millisecond_datatype(),
48                ConcreteDataType::timestamp_microsecond_datatype(),
49                ConcreteDataType::timestamp_nanosecond_datatype(),
50            ],
51            vec![ConcreteDataType::string_datatype()],
52        )
53    }
54
55    fn eval(&self, func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
56        ensure!(
57            columns.len() == 2,
58            InvalidFuncArgsSnafu {
59                err_msg: format!(
60                    "The length of the args is not correct, expect 2, have: {}",
61                    columns.len()
62                ),
63            }
64        );
65
66        let left = &columns[0];
67        let formats = &columns[1];
68
69        let size = left.len();
70        let left_datatype = columns[0].data_type();
71        let mut results = StringVectorBuilder::with_capacity(size);
72
73        match left_datatype {
74            ConcreteDataType::Timestamp(_) => {
75                for i in 0..size {
76                    let ts = left.get(i).as_timestamp();
77                    let format = formats.get(i).as_string();
78
79                    let result = match (ts, format) {
80                        (Some(ts), Some(fmt)) => Some(
81                            ts.as_formatted_string(&fmt, Some(&func_ctx.query_ctx.timezone()))
82                                .map_err(BoxedError::new)
83                                .context(error::ExecuteSnafu)?,
84                        ),
85                        _ => None,
86                    };
87
88                    results.push(result.as_deref());
89                }
90            }
91            ConcreteDataType::Date(_) => {
92                for i in 0..size {
93                    let date = left.get(i).as_date();
94                    let format = formats.get(i).as_string();
95
96                    let result = match (date, format) {
97                        (Some(date), Some(fmt)) => date
98                            .as_formatted_string(&fmt, Some(&func_ctx.query_ctx.timezone()))
99                            .map_err(BoxedError::new)
100                            .context(error::ExecuteSnafu)?,
101                        _ => None,
102                    };
103
104                    results.push(result.as_deref());
105                }
106            }
107            _ => {
108                return UnsupportedInputDataTypeSnafu {
109                    function: NAME,
110                    datatypes: columns.iter().map(|c| c.data_type()).collect::<Vec<_>>(),
111                }
112                .fail();
113            }
114        }
115
116        Ok(results.to_vector())
117    }
118}
119
120impl fmt::Display for DateFormatFunction {
121    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
122        write!(f, "DATE_FORMAT")
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use std::sync::Arc;
129
130    use common_query::prelude::{TypeSignature, Volatility};
131    use datatypes::prelude::{ConcreteDataType, ScalarVector};
132    use datatypes::value::Value;
133    use datatypes::vectors::{DateVector, StringVector, TimestampSecondVector};
134
135    use super::{DateFormatFunction, *};
136
137    #[test]
138    fn test_date_format_misc() {
139        let f = DateFormatFunction;
140        assert_eq!("date_format", f.name());
141        assert_eq!(
142            ConcreteDataType::string_datatype(),
143            f.return_type(&[ConcreteDataType::timestamp_microsecond_datatype()])
144                .unwrap()
145        );
146        assert_eq!(
147            ConcreteDataType::string_datatype(),
148            f.return_type(&[ConcreteDataType::timestamp_second_datatype()])
149                .unwrap()
150        );
151        assert_eq!(
152            ConcreteDataType::string_datatype(),
153            f.return_type(&[ConcreteDataType::date_datatype()]).unwrap()
154        );
155        assert!(matches!(f.signature(),
156                         Signature {
157                             type_signature: TypeSignature::OneOf(sigs),
158                             volatility: Volatility::Immutable
159                         } if  sigs.len() == 5));
160    }
161
162    #[test]
163    fn test_timestamp_date_format() {
164        let f = DateFormatFunction;
165
166        let times = vec![Some(123), None, Some(42), None];
167        let formats = vec![
168            "%Y-%m-%d %T.%3f",
169            "%Y-%m-%d %T.%3f",
170            "%Y-%m-%d %T.%3f",
171            "%Y-%m-%d %T.%3f",
172        ];
173        let results = [
174            Some("1970-01-01 00:02:03.000"),
175            None,
176            Some("1970-01-01 00:00:42.000"),
177            None,
178        ];
179
180        let time_vector = TimestampSecondVector::from(times.clone());
181        let interval_vector = StringVector::from_vec(formats);
182        let args: Vec<VectorRef> = vec![Arc::new(time_vector), Arc::new(interval_vector)];
183        let vector = f.eval(&FunctionContext::default(), &args).unwrap();
184
185        assert_eq!(4, vector.len());
186        for (i, _t) in times.iter().enumerate() {
187            let v = vector.get(i);
188            let result = results.get(i).unwrap();
189
190            if result.is_none() {
191                assert_eq!(Value::Null, v);
192                continue;
193            }
194            match v {
195                Value::String(s) => {
196                    assert_eq!(s.as_utf8(), result.unwrap());
197                }
198                _ => unreachable!(),
199            }
200        }
201    }
202
203    #[test]
204    fn test_date_date_format() {
205        let f = DateFormatFunction;
206
207        let dates = vec![Some(123), None, Some(42), None];
208        let formats = vec![
209            "%Y-%m-%d %T.%3f",
210            "%Y-%m-%d %T.%3f",
211            "%Y-%m-%d %T.%3f",
212            "%Y-%m-%d %T.%3f",
213        ];
214        let results = [
215            Some("1970-05-04 00:00:00.000"),
216            None,
217            Some("1970-02-12 00:00:00.000"),
218            None,
219        ];
220
221        let date_vector = DateVector::from(dates.clone());
222        let interval_vector = StringVector::from_vec(formats);
223        let args: Vec<VectorRef> = vec![Arc::new(date_vector), Arc::new(interval_vector)];
224        let vector = f.eval(&FunctionContext::default(), &args).unwrap();
225
226        assert_eq!(4, vector.len());
227        for (i, _t) in dates.iter().enumerate() {
228            let v = vector.get(i);
229            let result = results.get(i).unwrap();
230
231            if result.is_none() {
232                assert_eq!(Value::Null, v);
233                continue;
234            }
235            match v {
236                Value::String(s) => {
237                    assert_eq!(s.as_utf8(), result.unwrap());
238                }
239                _ => unreachable!(),
240            }
241        }
242    }
243}