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::sync::Arc;
19
20use arrow::array::cast::AsArray;
21use arrow::array::types::Int32Type;
22use arrow::array::{Array, PrimitiveArray};
23use arrow::buffer::NullBuffer;
24use arrow::datatypes::DataType::Time32;
25use arrow::datatypes::{DataType, Time32SecondType, TimeUnit};
26use chrono::prelude::*;
27
28use datafusion_common::types::{NativeType, logical_int32, logical_string};
29use datafusion_common::{Result, ScalarValue, exec_err, utils::take_function_args};
30use datafusion_expr::{
31    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
32    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 name(&self) -> &str {
100        "make_time"
101    }
102
103    fn signature(&self) -> &Signature {
104        &self.signature
105    }
106
107    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
108        Ok(Time32(TimeUnit::Second))
109    }
110
111    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
112        let [hours, minutes, seconds] = take_function_args(self.name(), args.args)?;
113
114        match (hours, minutes, seconds) {
115            (ColumnarValue::Scalar(h), _, _) if h.is_null() => {
116                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(None)))
117            }
118            (_, ColumnarValue::Scalar(m), _) if m.is_null() => {
119                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(None)))
120            }
121            (_, _, ColumnarValue::Scalar(s)) if s.is_null() => {
122                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(None)))
123            }
124            (
125                ColumnarValue::Scalar(ScalarValue::Int32(Some(hours))),
126                ColumnarValue::Scalar(ScalarValue::Int32(Some(minutes))),
127                ColumnarValue::Scalar(ScalarValue::Int32(Some(seconds))),
128            ) => {
129                let mut value = 0;
130                make_time_inner(hours, minutes, seconds, |seconds: i32| value = seconds)?;
131                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(Some(
132                    value,
133                ))))
134            }
135            (hours, minutes, seconds) => {
136                let len = args.number_rows;
137                let hours = hours.into_array(len)?;
138                let minutes = minutes.into_array(len)?;
139                let seconds = seconds.into_array(len)?;
140
141                let hours = hours.as_primitive::<Int32Type>();
142                let minutes = minutes.as_primitive::<Int32Type>();
143                let seconds = seconds.as_primitive::<Int32Type>();
144
145                let nulls = NullBuffer::union_many([
146                    hours.nulls(),
147                    minutes.nulls(),
148                    seconds.nulls(),
149                ]);
150
151                let mut values = Vec::with_capacity(len);
152                for i in 0..len {
153                    // Match Postgres behaviour which returns null for any null input
154                    if nulls.as_ref().is_some_and(|n| n.is_null(i)) {
155                        values.push(0);
156                    } else {
157                        make_time_inner(
158                            hours.value(i),
159                            minutes.value(i),
160                            seconds.value(i),
161                            |seconds: i32| values.push(seconds),
162                        )?;
163                    }
164                }
165
166                Ok(ColumnarValue::Array(Arc::new(PrimitiveArray::<
167                    Time32SecondType,
168                >::new(
169                    values.into(), nulls
170                ))))
171            }
172        }
173    }
174
175    fn documentation(&self) -> Option<&Documentation> {
176        self.doc()
177    }
178}
179
180/// Converts the hour/minute/second fields to an `i32` representing the seconds from
181/// midnight and invokes `time_consumer_fn` with the value
182fn make_time_inner<F: FnMut(i32)>(
183    hour: i32,
184    minute: i32,
185    second: i32,
186    mut time_consumer_fn: F,
187) -> Result<()> {
188    let h = match hour {
189        0..=24 => hour as u32,
190        _ => return exec_err!("Hour value '{hour:?}' is out of range"),
191    };
192    let m = match minute {
193        0..=60 => minute as u32,
194        _ => return exec_err!("Minute value '{minute:?}' is out of range"),
195    };
196    let s = match second {
197        0..=60 => second as u32,
198        _ => return exec_err!("Second value '{second:?}' is out of range"),
199    };
200
201    if let Some(time) = NaiveTime::from_hms_opt(h, m, s) {
202        time_consumer_fn(time.num_seconds_from_midnight() as i32);
203        Ok(())
204    } else {
205        exec_err!("Unable to parse time from {hour}, {minute}, {second}")
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use crate::datetime::make_time::MakeTimeFunc;
212    use arrow::array::{Array, Int32Array, Time32SecondArray};
213    use arrow::datatypes::TimeUnit::Second;
214    use arrow::datatypes::{DataType, Field};
215    use datafusion_common::DataFusionError;
216    use datafusion_common::config::ConfigOptions;
217    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
218    use std::sync::Arc;
219
220    fn invoke_make_time_with_args(
221        args: Vec<ColumnarValue>,
222        number_rows: usize,
223    ) -> Result<ColumnarValue, DataFusionError> {
224        let arg_fields = args
225            .iter()
226            .map(|arg| Field::new("a", arg.data_type(), true).into())
227            .collect::<Vec<_>>();
228        let args = ScalarFunctionArgs {
229            args,
230            arg_fields,
231            number_rows,
232            return_field: Field::new("f", DataType::Time32(Second), true).into(),
233            config_options: Arc::new(ConfigOptions::default()),
234        };
235
236        MakeTimeFunc::new().invoke_with_args(args)
237    }
238
239    #[test]
240    fn test_make_time() {
241        let hours = Arc::new((4..8).map(Some).collect::<Int32Array>());
242        let minutes = Arc::new((1..5).map(Some).collect::<Int32Array>());
243        let seconds = Arc::new((11..15).map(Some).collect::<Int32Array>());
244        let batch_len = hours.len();
245        let res = invoke_make_time_with_args(
246            vec![
247                ColumnarValue::Array(hours),
248                ColumnarValue::Array(minutes),
249                ColumnarValue::Array(seconds),
250            ],
251            batch_len,
252        )
253        .unwrap();
254
255        if let ColumnarValue::Array(array) = res {
256            assert_eq!(array.len(), 4);
257
258            let mut builder = Time32SecondArray::builder(4);
259            builder.append_value(14_471);
260            builder.append_value(18_132);
261            builder.append_value(21_793);
262            builder.append_value(25_454);
263            assert_eq!(&builder.finish() as &dyn Array, array.as_ref());
264        } else {
265            panic!("Expected a columnar array")
266        }
267    }
268}
datafusion_functions/datetime/make_time.rs

datafusion_functions/datetime/
make_time.rs
