fix: Native engine crashes on literal DateTrunc and TimestampTrunc (apache#3668)

Author: 0lai0

native/spark-expr/src/datetime_funcs/date_trunc.rs

@@ -16,7 +16,9 @@
 // under the License.
 
 use arrow::datatypes::DataType;
-use datafusion::common::{utils::take_function_args, DataFusionError, Result, ScalarValue::Utf8};
+use datafusion::common::{
+    utils::take_function_args, DataFusionError, Result, ScalarValue, ScalarValue::Utf8,
+};
 use datafusion::logical_expr::{
     ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
 };
@@ -76,9 +78,14 @@ impl ScalarUDFImpl for SparkDateTrunc {
                 let result = date_trunc_array_fmt_dyn(&date, &formats)?;
                 Ok(ColumnarValue::Array(result))
             }
+            (ColumnarValue::Scalar(date_scalar), ColumnarValue::Scalar(Utf8(Some(format)))) => {
+                let date_arr = date_scalar.to_array()?;
+                let result = date_trunc_dyn(&date_arr, format)?;
+                let scalar = ScalarValue::try_from_array(&result, 0)?;
+                Ok(ColumnarValue::Scalar(scalar))
+            }
             _ => Err(DataFusionError::Execution(
-                "Invalid input to function DateTrunc. Expected (PrimitiveArray<Date32>, Scalar) or \
-                    (PrimitiveArray<Date32>, StringArray)".to_string(),
+                "Invalid input to function DateTrunc. Expected (Date32, Utf8)".to_string(),
             )),
         }
     }

native/spark-expr/src/datetime_funcs/timestamp_trunc.rs

@@ -18,7 +18,7 @@
 use crate::utils::array_with_timezone;
 use arrow::datatypes::{DataType, Schema, TimeUnit::Microsecond};
 use arrow::record_batch::RecordBatch;
-use datafusion::common::{DataFusionError, ScalarValue::Utf8};
+use datafusion::common::{DataFusionError, ScalarValue, ScalarValue::Utf8};
 use datafusion::logical_expr::ColumnarValue;
 use datafusion::physical_expr::PhysicalExpr;
 use std::hash::Hash;
@@ -130,10 +130,20 @@ impl PhysicalExpr for TimestampTruncExpr {
                 let result = timestamp_trunc_array_fmt_dyn(&ts, &formats)?;
                 Ok(ColumnarValue::Array(result))
             }
+            (ColumnarValue::Scalar(ts_scalar), ColumnarValue::Scalar(Utf8(Some(format)))) => {
+                let ts_arr = ts_scalar.to_array()?;
+                let ts = array_with_timezone(
+                    ts_arr,
+                    tz.clone(),
+                    Some(&DataType::Timestamp(Microsecond, Some(tz.into()))),
+                )?;
+                let result = timestamp_trunc_dyn(&ts, format)?;
+                let scalar = ScalarValue::try_from_array(&result, 0)?;
+                Ok(ColumnarValue::Scalar(scalar))
+            }
             _ => Err(DataFusionError::Execution(
                 "Invalid input to function TimestampTrunc. \
-                    Expected (PrimitiveArray<TimestampMicrosecondType>, Scalar, String) or \
-                    (PrimitiveArray<TimestampMicrosecondType>, StringArray, String)"
+                    Expected (Timestamp, Utf8)"
                     .to_string(),
             )),
         }

spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql

@@ -33,5 +33,5 @@ query
 SELECT trunc(d, 'quarter') FROM test_trunc_date
 
 -- literal arguments
-query ignore(https://github.com/apache/datafusion-comet/issues/3342)
+query
 SELECT trunc(date('2024-06-15'), 'year'), trunc(date('2024-06-15'), 'month'), trunc(date('2024-06-15'), 'quarter')

spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql

@@ -37,5 +37,5 @@ query
 SELECT date_trunc('hour', ts) FROM test_trunc_ts
 
 -- literal arguments
-query ignore(https://github.com/apache/datafusion-comet/issues/3342)
+query
 SELECT date_trunc('year', timestamp('2024-06-15 10:30:45')), date_trunc('month', timestamp('2024-06-15 10:30:45')), date_trunc('day', timestamp('2024-06-15 10:30:45'))