Add make_time function

datafusion/functions/src/datetime/make_time.rs

@@ -0,0 +1,270 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::sync::Arc;
+
+use arrow::array::builder::PrimitiveBuilder;
+use arrow::array::cast::AsArray;
+use arrow::array::types::Int32Type;
+use arrow::array::{Array, PrimitiveArray};
+use arrow::datatypes::DataType::Time32;
+use arrow::datatypes::{DataType, Time32SecondType, TimeUnit};
+use chrono::prelude::*;
+
+use datafusion_common::types::{NativeType, logical_int32, logical_string};
+use datafusion_common::{Result, ScalarValue, exec_err, utils::take_function_args};
+use datafusion_expr::{
+    ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
+};
+use datafusion_expr_common::signature::{Coercion, TypeSignatureClass};
+use datafusion_macros::user_doc;
+
+#[user_doc(
+    doc_section(label = "Time and Date Functions"),
+    description = "Make a time from hour/minute/second component parts.",
+    syntax_example = "make_time(hour, minute, second)",
+    sql_example = r#"```sql
+> select make_time(13, 23, 1);
++-------------------------------------------+
+| make_time(Int64(13),Int64(23),Int64(1))   |
++-------------------------------------------+
+| 13:23:01                                  |
++-------------------------------------------+
+> select make_time('23', '01', '31');
++-----------------------------------------------+
+| make_time(Utf8("23"),Utf8("01"),Utf8("31"))   |
++-----------------------------------------------+
+| 23:01:31                                      |
++-----------------------------------------------+
+```
+
+Additional examples can be found [here](https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/builtin_functions/date_time.rs)
+"#,
+    argument(
+        name = "hour",
+        description = "Hour to use when making the time. Can be a constant, column or function, and any combination of arithmetic operators."
+    ),
+    argument(
+        name = "minute",
+        description = "Minute to use when making the time. Can be a constant, column or function, and any combination of arithmetic operators."
+    ),
+    argument(
+        name = "second",
+        description = "Second to use when making the time. Can be a constant, column or function, and any combination of arithmetic operators."
+    )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct MakeTimeFunc {
+    signature: Signature,
+}
+
+impl Default for MakeTimeFunc {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl MakeTimeFunc {
+    pub fn new() -> Self {
+        let int = Coercion::new_implicit(
+            TypeSignatureClass::Native(logical_int32()),
+            vec![
+                TypeSignatureClass::Integer,
+                TypeSignatureClass::Native(logical_string()),
+            ],
+            NativeType::Int32,
+        );
+        Self {
+            signature: Signature::coercible(
+                vec![int.clone(), int.clone(), int.clone()],
+                Volatility::Immutable,
+            ),
+        }
+    }
+}
+
+impl ScalarUDFImpl for MakeTimeFunc {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "make_time"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(Time32(TimeUnit::Second))
+    }
+
+    fn invoke_with_args(
+        &self,
+        args: datafusion_expr::ScalarFunctionArgs,
+    ) -> Result<ColumnarValue> {
+        let [hours, minutes, seconds] = take_function_args(self.name(), args.args)?;
+
+        match (hours, minutes, seconds) {
+            (ColumnarValue::Scalar(h), _, _) if h.is_null() => {
+                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(None)))
+            }
+            (_, ColumnarValue::Scalar(m), _) if m.is_null() => {
+                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(None)))
+            }
+            (_, _, ColumnarValue::Scalar(s)) if s.is_null() => {
+                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(None)))
+            }
+            (
+                ColumnarValue::Scalar(ScalarValue::Int32(Some(hours))),
+                ColumnarValue::Scalar(ScalarValue::Int32(Some(minutes))),
+                ColumnarValue::Scalar(ScalarValue::Int32(Some(seconds))),
+            ) => {
+                let mut value = 0;
+                make_time_inner(hours, minutes, seconds, |seconds: i32| value = seconds)?;
+                Ok(ColumnarValue::Scalar(ScalarValue::Time32Second(Some(
+                    value,
+                ))))
+            }
+            (hours, minutes, seconds) => {
+                let len = args.number_rows;
+                let hours = hours.into_array(len)?;
+                let minutes = minutes.into_array(len)?;
+                let seconds = seconds.into_array(len)?;
+
+                let hours = hours.as_primitive::<Int32Type>();
+                let minutes = minutes.as_primitive::<Int32Type>();
+                let seconds = seconds.as_primitive::<Int32Type>();
+
+                let mut builder: PrimitiveBuilder<Time32SecondType> =
+                    PrimitiveArray::builder(len);
+
+                for i in 0..len {
+                    // match postgresql behaviour which returns null for any null input
+                    if hours.is_null(i) || minutes.is_null(i) || seconds.is_null(i) {
+                        builder.append_null();
+                    } else {
+                        make_time_inner(
+                            hours.value(i),
+                            minutes.value(i),
+                            seconds.value(i),
+                            |seconds: i32| builder.append_value(seconds),
+                        )?;
+                    }
+                }
+
+                Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+            }
+        }
+    }
+
+    fn documentation(&self) -> Option<&Documentation> {
+        self.doc()
+    }
+}
+
+/// Converts the hour/minute/second fields to an `i32` representing the seconds from
+/// midnight and invokes `time_consumer_fn` with the value
+fn make_time_inner<F: FnMut(i32)>(
+    hour: i32,
+    minute: i32,
+    second: i32,
+    mut time_consumer_fn: F,
+) -> Result<()> {
+    let h = match hour {
+        0..=24 => hour as u32,
+        _ => return exec_err!("Hour value '{hour:?}' is out of range"),
+    };
+    let m = match minute {
+        0..=60 => minute as u32,
+        _ => return exec_err!("Minute value '{minute:?}' is out of range"),
+    };
+    let s = match second {
+        0..=60 => second as u32,
+        _ => return exec_err!("Second value '{second:?}' is out of range"),
+    };
+
+    if let Some(time) = NaiveTime::from_hms_opt(h, m, s) {
+        time_consumer_fn(time.num_seconds_from_midnight() as i32);
+        Ok(())
+    } else {
+        exec_err!("Unable to parse time from {hour}, {minute}, {second}")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::datetime::make_time::MakeTimeFunc;
+    use arrow::array::{Array, Int32Array, Time32SecondArray};
+    use arrow::datatypes::TimeUnit::Second;
+    use arrow::datatypes::{DataType, Field};
+    use datafusion_common::DataFusionError;
+    use datafusion_common::config::ConfigOptions;
+    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+    use std::sync::Arc;
+
+    fn invoke_make_time_with_args(
+        args: Vec<ColumnarValue>,
+        number_rows: usize,
+    ) -> Result<ColumnarValue, DataFusionError> {
+        let arg_fields = args
+            .iter()
+            .map(|arg| Field::new("a", arg.data_type(), true).into())
+            .collect::<Vec<_>>();
+        let args = datafusion_expr::ScalarFunctionArgs {
+            args,
+            arg_fields,
+            number_rows,
+            return_field: Field::new("f", DataType::Time32(Second), true).into(),
+            config_options: Arc::new(ConfigOptions::default()),
+        };
+
+        MakeTimeFunc::new().invoke_with_args(args)
+    }
+
+    #[test]
+    fn test_make_time() {
+        let hours = Arc::new((4..8).map(Some).collect::<Int32Array>());
+        let minutes = Arc::new((1..5).map(Some).collect::<Int32Array>());
+        let seconds = Arc::new((11..15).map(Some).collect::<Int32Array>());
+        let batch_len = hours.len();
+        let res = invoke_make_time_with_args(
+            vec![
+                ColumnarValue::Array(hours),
+                ColumnarValue::Array(minutes),
+                ColumnarValue::Array(seconds),
+            ],
+            batch_len,
+        )
+        .unwrap();
+
+        if let ColumnarValue::Array(array) = res {
+            assert_eq!(array.len(), 4);
+
+            let mut builder = Time32SecondArray::builder(4);
+            builder.append_value(14_471);
+            builder.append_value(18_132);
+            builder.append_value(21_793);
+            builder.append_value(25_454);
+            assert_eq!(&builder.finish() as &dyn Array, array.as_ref());
+        } else {
+            panic!("Expected a columnar array")
+        }
+    }
+}

datafusion/functions/src/datetime/mod.rs

@@ -29,6 +29,7 @@ pub mod date_part;
 pub mod date_trunc;
 pub mod from_unixtime;
 pub mod make_date;
+pub mod make_time;
 pub mod now;
 pub mod planner;
 pub mod to_char;
@@ -44,6 +45,7 @@ make_udf_function!(date_bin::DateBinFunc, date_bin);
 make_udf_function!(date_part::DatePartFunc, date_part);
 make_udf_function!(date_trunc::DateTruncFunc, date_trunc);
 make_udf_function!(make_date::MakeDateFunc, make_date);
+make_udf_function!(make_time::MakeTimeFunc, make_time);
 make_udf_function!(from_unixtime::FromUnixtimeFunc, from_unixtime);
 make_udf_function!(to_char::ToCharFunc, to_char);
 make_udf_function!(to_date::ToDateFunc, to_date);
@@ -90,6 +92,10 @@ pub mod expr_fn {
         make_date,
         "make a date from year, month and day component parts",
         year month day
+    ),(
+        make_time,
+        "make a time from hour, minute and second component parts",
+        hour minute second
     ),(
         now,
         "returns the current timestamp in nanoseconds, using the same value for all instances of now() in same statement",
@@ -267,6 +273,7 @@ pub fn functions() -> Vec<Arc<ScalarUDF>> {
         date_trunc(),
         from_unixtime(),
         make_date(),
+        make_time(),
         now(&ConfigOptions::default()),
         to_char(),
         to_date(),