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;
16use std::sync::Arc;
17
18use common_error::ext::BoxedError;
19use common_query::error;
20use common_time::{Date, Timestamp};
21use datafusion_common::DataFusionError;
22use datafusion_common::arrow::array::{Array, AsArray, StringViewBuilder};
23use datafusion_common::arrow::datatypes::{ArrowTimestampType, DataType, Date32Type, TimeUnit};
24use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature};
25use snafu::ResultExt;
26
27use crate::function::{Function, extract_args, find_function_context};
28use crate::helper;
29use crate::helper::with_match_timestamp_types;
30
31/// A function that formats timestamp/date/datetime into string by the format
32#[derive(Clone, Debug)]
33pub(crate) struct DateFormatFunction {
34    signature: Signature,
35}
36
37impl Default for DateFormatFunction {
38    fn default() -> Self {
39        Self {
40            signature: helper::one_of_sigs2(
41                vec![
42                    DataType::Date32,
43                    DataType::Timestamp(TimeUnit::Second, None),
44                    DataType::Timestamp(TimeUnit::Millisecond, None),
45                    DataType::Timestamp(TimeUnit::Microsecond, None),
46                    DataType::Timestamp(TimeUnit::Nanosecond, None),
47                ],
48                vec![DataType::Utf8],
49            ),
50        }
51    }
52}
53
54impl Function for DateFormatFunction {
55    fn name(&self) -> &str {
56        "date_format"
57    }
58
59    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
60        Ok(DataType::Utf8View)
61    }
62
63    fn signature(&self) -> &Signature {
64        &self.signature
65    }
66
67    fn invoke_with_args(
68        &self,
69        args: ScalarFunctionArgs,
70    ) -> datafusion_common::Result<ColumnarValue> {
71        let ctx = find_function_context(&args)?;
72        let timezone = &ctx.query_ctx.timezone();
73
74        let [left, arg1] = extract_args(self.name(), &args)?;
75        let formats = arg1.as_string::<i32>();
76
77        let size = left.len();
78        let left_datatype = left.data_type();
79        let mut builder = StringViewBuilder::with_capacity(size);
80
81        match left_datatype {
82            DataType::Timestamp(_, _) => {
83                with_match_timestamp_types!(left_datatype, |$S| {
84                    let array = left.as_primitive::<$S>();
85                    for (date, format) in array.iter().zip(formats.iter()) {
86                        let result = match (date, format) {
87                            (Some(date), Some(format)) => {
88                                let ts = Timestamp::new(date, $S::UNIT.into());
89                                let x = ts.as_formatted_string(&format, Some(timezone))
90                                    .map_err(|e| DataFusionError::Execution(format!(
91                                        "cannot format {ts:?} as '{format}': {e}"
92                                    )))?;
93                                Some(x)
94                            }
95                            _ => None
96                        };
97                        builder.append_option(result.as_deref());
98                    }
99                })?;
100            }
101            DataType::Date32 => {
102                let left = left.as_primitive::<Date32Type>();
103                for i in 0..size {
104                    let date = left.is_valid(i).then(|| Date::from(left.value(i)));
105                    let format = formats.is_valid(i).then(|| formats.value(i));
106
107                    let result = match (date, format) {
108                        (Some(date), Some(fmt)) => date
109                            .as_formatted_string(fmt, Some(timezone))
110                            .map_err(BoxedError::new)
111                            .context(error::ExecuteSnafu)?,
112                        _ => None,
113                    };
114
115                    builder.append_option(result.as_deref());
116                }
117            }
118            x => {
119                return Err(DataFusionError::Execution(format!(
120                    "unsupported input data type {x}"
121                )));
122            }
123        }
124
125        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
126    }
127}
128
129impl fmt::Display for DateFormatFunction {
130    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
131        write!(f, "DATE_FORMAT")
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use std::sync::Arc;
138
139    use arrow_schema::Field;
140    use datafusion_common::arrow::array::{Date32Array, StringArray, TimestampSecondArray};
141    use datafusion_common::config::ConfigOptions;
142    use datafusion_expr::{TypeSignature, Volatility};
143
144    use super::{DateFormatFunction, *};
145    use crate::function::FunctionContext;
146
147    #[test]
148    fn test_date_format_misc() {
149        let f = DateFormatFunction::default();
150        assert_eq!("date_format", f.name());
151        assert_eq!(
152            DataType::Utf8View,
153            f.return_type(&[DataType::Timestamp(TimeUnit::Microsecond, None)])
154                .unwrap()
155        );
156        assert_eq!(
157            DataType::Utf8View,
158            f.return_type(&[DataType::Timestamp(TimeUnit::Second, None)])
159                .unwrap()
160        );
161        assert_eq!(
162            DataType::Utf8View,
163            f.return_type(&[DataType::Date32]).unwrap()
164        );
165        assert!(matches!(f.signature(),
166                         Signature {
167                             type_signature: TypeSignature::OneOf(sigs),
168                             volatility: Volatility::Immutable
169                         } if  sigs.len() == 5));
170    }
171
172    #[test]
173    fn test_timestamp_date_format() {
174        let f = DateFormatFunction::default();
175
176        let times = vec![Some(123), None, Some(42), None];
177        let formats = vec![
178            "%Y-%m-%d %T.%3f",
179            "%Y-%m-%d %T.%3f",
180            "%Y-%m-%d %T.%3f",
181            "%Y-%m-%d %T.%3f",
182        ];
183        let results = [
184            Some("1970-01-01 00:02:03.000"),
185            None,
186            Some("1970-01-01 00:00:42.000"),
187            None,
188        ];
189
190        let mut config_options = ConfigOptions::default();
191        config_options.extensions.insert(FunctionContext::default());
192        let config_options = Arc::new(config_options);
193
194        let args = ScalarFunctionArgs {
195            args: vec![
196                ColumnarValue::Array(Arc::new(TimestampSecondArray::from(times))),
197                ColumnarValue::Array(Arc::new(StringArray::from_iter_values(formats))),
198            ],
199            arg_fields: vec![],
200            number_rows: 4,
201            return_field: Arc::new(Field::new("x", DataType::Utf8View, false)),
202            config_options,
203        };
204        let result = f
205            .invoke_with_args(args)
206            .and_then(|x| x.to_array(4))
207            .unwrap();
208        let vector = result.as_string_view();
209
210        assert_eq!(4, vector.len());
211        for (actual, expect) in vector.iter().zip(results) {
212            assert_eq!(actual, expect);
213        }
214    }
215
216    #[test]
217    fn test_date_date_format() {
218        let f = DateFormatFunction::default();
219
220        let dates = vec![Some(123), None, Some(42), None];
221        let formats = vec![
222            "%Y-%m-%d %T.%3f",
223            "%Y-%m-%d %T.%3f",
224            "%Y-%m-%d %T.%3f",
225            "%Y-%m-%d %T.%3f",
226        ];
227        let results = [
228            Some("1970-05-04 00:00:00.000"),
229            None,
230            Some("1970-02-12 00:00:00.000"),
231            None,
232        ];
233
234        let mut config_options = ConfigOptions::default();
235        config_options.extensions.insert(FunctionContext::default());
236        let config_options = Arc::new(config_options);
237
238        let args = ScalarFunctionArgs {
239            args: vec![
240                ColumnarValue::Array(Arc::new(Date32Array::from(dates))),
241                ColumnarValue::Array(Arc::new(StringArray::from_iter_values(formats))),
242            ],
243            arg_fields: vec![],
244            number_rows: 4,
245            return_field: Arc::new(Field::new("x", DataType::Utf8View, false)),
246            config_options,
247        };
248        let result = f
249            .invoke_with_args(args)
250            .and_then(|x| x.to_array(4))
251            .unwrap();
252        let vector = result.as_string_view();
253
254        assert_eq!(4, vector.len());
255        for (actual, expect) in vector.iter().zip(results) {
256            assert_eq!(actual, expect);
257        }
258    }
259}
common_function/scalars/date/date_format.rs

common_function/scalars/date/
date_format.rs
