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

common_function/scalars/date/
date_format.rs
