feat(spark): add trunc, date_trunc and time_trunc functions

datafusion/functions/src/datetime/date_trunc.rs

@@ -34,14 +34,16 @@ use arrow::array::types::{
 use arrow::array::{Array, ArrayRef, PrimitiveArray};
 use arrow::datatypes::DataType::{self, Time32, Time64, Timestamp};
 use arrow::datatypes::TimeUnit::{self, Microsecond, Millisecond, Nanosecond, Second};
+use arrow::datatypes::{Field, FieldRef};
 use datafusion_common::cast::as_primitive_array;
 use datafusion_common::types::{NativeType, logical_date, logical_string};
 use datafusion_common::{
     DataFusionError, Result, ScalarValue, exec_datafusion_err, exec_err,
 };
 use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
 use datafusion_expr::{
-    ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature, Volatility,
+    ColumnarValue, Documentation, ReturnFieldArgs, ScalarUDFImpl, Signature,
+    TypeSignature, Volatility,
 };
 use datafusion_expr_common::signature::{Coercion, TypeSignatureClass};
 use datafusion_macros::user_doc;
@@ -221,6 +223,7 @@ impl ScalarUDFImpl for DateTruncFunc {
         &self.signature
     }
 
+    // keep return_type implementation for information schema generation
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
         if arg_types[1].is_null() {
             Ok(Timestamp(Nanosecond, None))
@@ -229,6 +232,21 @@ impl ScalarUDFImpl for DateTruncFunc {
         }
     }
 
+    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
+        let data_types = args
+            .arg_fields
+            .iter()
+            .map(|f| f.data_type())
+            .cloned()
+            .collect::<Vec<_>>();
+        let return_type = self.return_type(&data_types)?;
+        Ok(Arc::new(Field::new(
+            self.name(),
+            return_type,
+            args.arg_fields[1].is_nullable(),
+        )))
+    }
+
     fn invoke_with_args(
         &self,
         args: datafusion_expr::ScalarFunctionArgs,

datafusion/spark/src/function/datetime/date_trunc.rs

@@ -0,0 +1,172 @@
+// 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::datatypes::{DataType, Field, FieldRef, TimeUnit};
+use datafusion_common::types::{NativeType, logical_string};
+use datafusion_common::utils::take_function_args;
+use datafusion_common::{Result, ScalarValue, internal_err, plan_err};
+use datafusion_expr::expr::ScalarFunction;
+use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext};
+use datafusion_expr::{
+    Coercion, ColumnarValue, Expr, ExprSchemable, ReturnFieldArgs, ScalarFunctionArgs,
+    ScalarUDFImpl, Signature, TypeSignatureClass, Volatility,
+};
+
+/// Spark date_trunc supports extra format aliases.
+/// It also handles timestamps with timezones by converting to session timezone first.
+/// <https://spark.apache.org/docs/latest/api/sql/index.html#date_trunc>
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkDateTrunc {
+    signature: Signature,
+}
+
+impl Default for SparkDateTrunc {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl SparkDateTrunc {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::coercible(
+                vec![
+                    Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
+                    Coercion::new_implicit(
+                        TypeSignatureClass::Timestamp,
+                        vec![TypeSignatureClass::Native(logical_string())],
+                        NativeType::Timestamp(TimeUnit::Microsecond, None),
+                    ),
+                ],
+                Volatility::Immutable,
+            ),
+        }
+    }
+}
+
+impl ScalarUDFImpl for SparkDateTrunc {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "date_trunc"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        internal_err!("return_field_from_args should be used instead")
+    }
+
+    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
+        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
+
+        Ok(Arc::new(Field::new(
+            self.name(),
+            args.arg_fields[1].data_type().clone(),
+            nullable,
+        )))
+    }
+
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        internal_err!(
+            "spark date_trunc should have been simplified to standard date_trunc"
+        )
+    }
+
+    fn simplify(
+        &self,
+        args: Vec<Expr>,
+        info: &SimplifyContext,
+    ) -> Result<ExprSimplifyResult> {
+        let [fmt_expr, ts_expr] = take_function_args(self.name(), args)?;
+
+        let fmt = match fmt_expr.as_literal() {
+            Some(ScalarValue::Utf8(Some(v)))
+            | Some(ScalarValue::Utf8View(Some(v)))
+            | Some(ScalarValue::LargeUtf8(Some(v))) => v.to_lowercase(),
+            _ => {
+                return plan_err!(
+                    "First argument of `DATE_TRUNC` must be non-null scalar Utf8"
+                );
+            }
+        };
+
+        // Map Spark-specific fmt aliases to datafusion ones
+        let fmt = match fmt.as_str() {
+            "yy" | "yyyy" => "year",
+            "mm" | "mon" => "month",
+            "dd" => "day",
+            other => other,
+        };
+
+        let session_tz = info.config_options().execution.time_zone.clone();
+        let ts_type = ts_expr.get_type(info.schema())?;
+
+        // Spark interprets timestamps in the session timezone before truncating,
+        // then returns a timestamp at microsecond precision.
+        // See: https://github.com/apache/spark/blob/f310f4fcc95580a6824bc7d22b76006f79b8804a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala#L492
+        //
+        // For sub-second truncations (second, millisecond, microsecond), timezone
+        // adjustment is unnecessary since timezone offsets are whole seconds.
+        let ts_expr = match (&ts_type, fmt) {
+            // Sub-second truncations don't need timezone adjustment
+            (_, "second" | "millisecond" | "microsecond") => ts_expr,
+
+            // convert to session timezone, strip timezone and convert back to original timezone
+            (DataType::Timestamp(unit, tz), _) => {
+                let ts_expr = match &session_tz {
+                    Some(session_tz) => ts_expr.cast_to(
+                        &DataType::Timestamp(
+                            TimeUnit::Microsecond,
+                            Some(Arc::from(session_tz.as_str())),
+                        ),
+                        info.schema(),
+                    )?,
+                    None => ts_expr,
+                };
+                Expr::ScalarFunction(ScalarFunction::new_udf(
+                    datafusion_functions::datetime::to_local_time(),
+                    vec![ts_expr],
+                ))
+                .cast_to(&DataType::Timestamp(*unit, tz.clone()), info.schema())?
+            }
+
+            _ => {
+                return plan_err!(
+                    "Second argument of `DATE_TRUNC` must be Timestamp, got {}",
+                    ts_type
+                );
+            }
+        };
+
+        let fmt_expr = Expr::Literal(ScalarValue::new_utf8(fmt), None);
+
+        Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction(
+            ScalarFunction::new_udf(
+                datafusion_functions::datetime::date_trunc(),
+                vec![fmt_expr, ts_expr],
+            ),
+        )))
+    }
+}

datafusion/spark/src/function/datetime/mod.rs

@@ -18,26 +18,32 @@
 pub mod date_add;
 pub mod date_part;
 pub mod date_sub;
+pub mod date_trunc;
 pub mod extract;
 pub mod last_day;
 pub mod make_dt_interval;
 pub mod make_interval;
 pub mod next_day;
+pub mod time_trunc;
+pub mod trunc;
 
 use datafusion_expr::ScalarUDF;
 use datafusion_functions::make_udf_function;
 use std::sync::Arc;
 
 make_udf_function!(date_add::SparkDateAdd, date_add);
+make_udf_function!(date_part::SparkDatePart, date_part);
 make_udf_function!(date_sub::SparkDateSub, date_sub);
+make_udf_function!(date_trunc::SparkDateTrunc, date_trunc);
 make_udf_function!(extract::SparkHour, hour);
 make_udf_function!(extract::SparkMinute, minute);
 make_udf_function!(extract::SparkSecond, second);
 make_udf_function!(last_day::SparkLastDay, last_day);
 make_udf_function!(make_dt_interval::SparkMakeDtInterval, make_dt_interval);
 make_udf_function!(make_interval::SparkMakeInterval, make_interval);
 make_udf_function!(next_day::SparkNextDay, next_day);
-make_udf_function!(date_part::SparkDatePart, date_part);
+make_udf_function!(time_trunc::SparkTimeTrunc, time_trunc);
+make_udf_function!(trunc::SparkTrunc, trunc);
 
 pub mod expr_fn {
     use datafusion_functions::export_functions;
@@ -85,24 +91,43 @@ pub mod expr_fn {
         "Returns the first date which is later than start_date and named as indicated. The function returns NULL if at least one of the input parameters is NULL.",
         arg1 arg2
     ));
+    export_functions!((
+        date_trunc,
+        "Truncates a timestamp `ts` to the unit specified by the format `fmt`.",
+        fmt ts
+    ));
+    export_functions!((
+        time_trunc,
+        "Truncates a time `t` to the unit specified by the format `fmt`.",
+        fmt t
+    ));
+    export_functions!((
+        trunc,
+        "Truncates a date `dt` to the unit specified by the format `fmt`.",
+        dt fmt
+    ));
     export_functions!((
         date_part,
         "Extracts a part of the date or time from a date, time, or timestamp expression.",
         arg1 arg2
+
     ));
 }
 
 pub fn functions() -> Vec<Arc<ScalarUDF>> {
     vec![
         date_add(),
+        date_part(),
         date_sub(),
+        date_trunc(),
         hour(),
-        minute(),
-        second(),
         last_day(),
         make_dt_interval(),
         make_interval(),
+        minute(),
         next_day(),
-        date_part(),
+        second(),
+        time_trunc(),
+        trunc(),
     ]
 }