datafusion_functions/datetime/
make_time.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::any::Any;
19use std::sync::Arc;
20
21use arrow::array::builder::PrimitiveBuilder;
22use arrow::array::cast::AsArray;
23use arrow::array::types::Int32Type;
24use arrow::array::{Array, PrimitiveArray};
25use arrow::datatypes::DataType::Time32;
26use arrow::datatypes::{DataType, Time32SecondType, TimeUnit};
27use chrono::prelude::*;
28
29use datafusion_common::types::{NativeType, logical_int32, logical_string};
30use datafusion_common::{Result, ScalarValue, exec_err, utils::take_function_args};
31use datafusion_expr::{
32    ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
33};
34use datafusion_expr_common::signature::{Coercion, TypeSignatureClass};
35use datafusion_macros::user_doc;
36
37#[user_doc(
38    doc_section(label = "Time and Date Functions"),
39    description = "Make a time from hour/minute/second component parts.",
40    syntax_example = "make_time(hour, minute, second)",
41    sql_example = r#"```sql
42> select make_time(13, 23, 1);
43+-------------------------------------------+
44| make_time(Int64(13),Int64(23),Int64(1))   |
45+-------------------------------------------+
46| 13:23:01                                  |
47+-------------------------------------------+
48> select make_time('23', '01', '31');
49+-----------------------------------------------+
50| make_time(Utf8("23"),Utf8("01"),Utf8("31"))   |
51+-----------------------------------------------+
52| 23:01:31                                      |
53+-----------------------------------------------+
54```
55
56Additional examples can be found [here](https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/builtin_functions/date_time.rs)
57"#,
58    argument(
59        name = "hour",
60        description = "Hour to use when making the time. Can be a constant, column or function, and any combination of arithmetic operators."
61    ),
62    argument(
63        name = "minute",
64        description = "Minute to use when making the time. Can be a constant, column or function, and any combination of arithmetic operators."
65    ),
66    argument(
67        name = "second",
68        description = "Second to use when making the time. Can be a constant, column or function, and any combination of arithmetic operators."
69    )
70)]
71#[derive(Debug, PartialEq, Eq, Hash)]
72pub struct MakeTimeFunc {
73    signature: Signature,
74}
75
76impl Default for MakeTimeFunc {
77    fn default() -> Self {
78        Self::new()
79    }
80}
81
82impl MakeTimeFunc {
83    pub fn new() -> Self {
84        let int = Coercion::new_implicit(
85            TypeSignatureClass::Native(logical_int32()),
86            vec![
87                TypeSignatureClass::Integer,
88                TypeSignatureClass::Native(logical_string()),
89            ],
90            NativeType::Int32,
91        );
92        Self {
93            signature: Signature::coercible(vec![int; 3], Volatility::Immutable),
94        }
95    }
96}
97
98impl ScalarUDFImpl for MakeTimeFunc {
99    fn as_any(&self) -> &dyn Any {
100        self
101    }
102
103    fn name(&self) -> &str {
104        "make_time"
105    }
106
107    fn signature(&self) -> &Signature {
108        &self.signature
109    }
110
111    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
112        Ok(Time32(TimeUnit::Second))
113    }
114
115    fn invoke_with_args(
116        &self,
117        args: datafusion_expr::ScalarFunctionArgs,
118    ) -> Result<ColumnarValue> {
119        let [hours, minutes, seconds] = take_function_args(self.name(), args.args)?;
120
121        match (hours, minutes, seconds) {
122            (ColumnarValue::Scalar(h), _, _) if h.is_null() => {
123                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(None)))
124            }
125            (_, ColumnarValue::Scalar(m), _) if m.is_null() => {
126                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(None)))
127            }
128            (_, _, ColumnarValue::Scalar(s)) if s.is_null() => {
129                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(None)))
130            }
131            (
132                ColumnarValue::Scalar(ScalarValue::Int32(Some(hours))),
133                ColumnarValue::Scalar(ScalarValue::Int32(Some(minutes))),
134                ColumnarValue::Scalar(ScalarValue::Int32(Some(seconds))),
135            ) => {
136                let mut value = 0;
137                make_time_inner(hours, minutes, seconds, |seconds: i32| value = seconds)?;
138                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(Some(
139                    value,
140                ))))
141            }
142            (hours, minutes, seconds) => {
143                let len = args.number_rows;
144                let hours = hours.into_array(len)?;
145                let minutes = minutes.into_array(len)?;
146                let seconds = seconds.into_array(len)?;
147
148                let hours = hours.as_primitive::<Int32Type>();
149                let minutes = minutes.as_primitive::<Int32Type>();
150                let seconds = seconds.as_primitive::<Int32Type>();
151
152                let mut builder: PrimitiveBuilder<Time32SecondType> =
153                    PrimitiveArray::builder(len);
154
155                for i in 0..len {
156                    // match postgresql behaviour which returns null for any null input
157                    if hours.is_null(i) || minutes.is_null(i) || seconds.is_null(i) {
158                        builder.append_null();
159                    } else {
160                        make_time_inner(
161                            hours.value(i),
162                            minutes.value(i),
163                            seconds.value(i),
164                            |seconds: i32| builder.append_value(seconds),
165                        )?;
166                    }
167                }
168
169                Ok(ColumnarValue::Array(Arc::new(builder.finish())))
170            }
171        }
172    }
173
174    fn documentation(&self) -> Option<&Documentation> {
175        self.doc()
176    }
177}
178
179/// Converts the hour/minute/second fields to an `i32` representing the seconds from
180/// midnight and invokes `time_consumer_fn` with the value
181fn make_time_inner<F: FnMut(i32)>(
182    hour: i32,
183    minute: i32,
184    second: i32,
185    mut time_consumer_fn: F,
186) -> Result<()> {
187    let h = match hour {
188        0..=24 => hour as u32,
189        _ => return exec_err!("Hour value '{hour:?}' is out of range"),
190    };
191    let m = match minute {
192        0..=60 => minute as u32,
193        _ => return exec_err!("Minute value '{minute:?}' is out of range"),
194    };
195    let s = match second {
196        0..=60 => second as u32,
197        _ => return exec_err!("Second value '{second:?}' is out of range"),
198    };
199
200    if let Some(time) = NaiveTime::from_hms_opt(h, m, s) {
201        time_consumer_fn(time.num_seconds_from_midnight() as i32);
202        Ok(())
203    } else {
204        exec_err!("Unable to parse time from {hour}, {minute}, {second}")
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use crate::datetime::make_time::MakeTimeFunc;
211    use arrow::array::{Array, Int32Array, Time32SecondArray};
212    use arrow::datatypes::TimeUnit::Second;
213    use arrow::datatypes::{DataType, Field};
214    use datafusion_common::DataFusionError;
215    use datafusion_common::config::ConfigOptions;
216    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
217    use std::sync::Arc;
218
219    fn invoke_make_time_with_args(
220        args: Vec<ColumnarValue>,
221        number_rows: usize,
222    ) -> Result<ColumnarValue, DataFusionError> {
223        let arg_fields = args
224            .iter()
225            .map(|arg| Field::new("a", arg.data_type(), true).into())
226            .collect::<Vec<_>>();
227        let args = datafusion_expr::ScalarFunctionArgs {
228            args,
229            arg_fields,
230            number_rows,
231            return_field: Field::new("f", DataType::Time32(Second), true).into(),
232            config_options: Arc::new(ConfigOptions::default()),
233        };
234
235        MakeTimeFunc::new().invoke_with_args(args)
236    }
237
238    #[test]
239    fn test_make_time() {
240        let hours = Arc::new((4..8).map(Some).collect::<Int32Array>());
241        let minutes = Arc::new((1..5).map(Some).collect::<Int32Array>());
242        let seconds = Arc::new((11..15).map(Some).collect::<Int32Array>());
243        let batch_len = hours.len();
244        let res = invoke_make_time_with_args(
245            vec![
246                ColumnarValue::Array(hours),
247                ColumnarValue::Array(minutes),
248                ColumnarValue::Array(seconds),
249            ],
250            batch_len,
251        )
252        .unwrap();
253
254        if let ColumnarValue::Array(array) = res {
255            assert_eq!(array.len(), 4);
256
257            let mut builder = Time32SecondArray::builder(4);
258            builder.append_value(14_471);
259            builder.append_value(18_132);
260            builder.append_value(21_793);
261            builder.append_value(25_454);
262            assert_eq!(&builder.finish() as &dyn Array, array.as_ref());
263        } else {
264            panic!("Expected a columnar array")
265        }
266    }
267}