feat: Allow Date32 as argument to date_trunc

datafusion/core/tests/sql/timestamp.rs

@@ -846,3 +846,34 @@ async fn test_current_date() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn date_trunc_date32_test() -> Result<()> {
+    let ctx = SessionContext::new();
+
+    let sql = "select date_trunc('month', cast('2023-02-28' as date)) as dt";
+    let results = execute_to_batches(&ctx, sql).await;
+
+    let expected = vec![
+        "+---------------------+",
+        "| dt                  |",
+        "+---------------------+",
+        "| 2023-02-01 00:00:00 |",
+        "+---------------------+",
+    ];
+    assert_batches_eq!(expected, &results);
+
+    let sql = "with w as (select cast('2023-02-28' as date) d) select date_trunc('month', d) as dt from w";
+    let results = execute_to_batches(&ctx, sql).await;
+
+    let expected = vec![
+        "+---------------------+",
+        "| dt                  |",
+        "+---------------------+",
+        "| 2023-02-01 00:00:00 |",
+        "+---------------------+",
+    ];
+    assert_batches_eq!(expected, &results);
+
+    Ok(())
+}

datafusion/physical-expr/src/datetime_expressions.rs

@@ -452,6 +452,12 @@ pub fn date_trunc(args: &[ColumnarValue]) -> Result<ColumnarValue> {
         };
 
     let f = |x: Option<i64>| x.map(|x| date_trunc_single(granularity, x)).transpose();
+    let date_to_timestamp = |x: Option<i32>| {
+        x.map(|x| {
+            let nanoseconds_in_day = 86_400_000_000_000_i64;
+            x as i64 * nanoseconds_in_day
+        })
+    };
 
     Ok(match array {
         ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(v, tz_opt)) => {
@@ -460,22 +466,39 @@ pub fn date_trunc(args: &[ColumnarValue]) -> Result<ColumnarValue> {
                 tz_opt.clone(),
             ))
         }
-        ColumnarValue::Scalar(ScalarValue::Date32(_)) => {
-            return Err(DataFusionError::Execution(
-                "`date_trunc` does not accept Date32 type, it's a stub".to_string(),
-            ));
-        }
+        ColumnarValue::Scalar(ScalarValue::Date32(d)) => ColumnarValue::Scalar(
+            ScalarValue::TimestampNanosecond((f)(date_to_timestamp(*d))?, None),
+        ),
         ColumnarValue::Array(array) => {
-            let array = array
-                .as_any()
-                .downcast_ref::<TimestampNanosecondArray>()
-                .unwrap();
-            let array = array
-                .iter()
-                .map(f)
-                .collect::<Result<TimestampNanosecondArray>>()?;
+            let array = match array.data_type() {
+                DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+                    let array = array
+                        .as_any()
+                        .downcast_ref::<TimestampNanosecondArray>()
+                        .unwrap();
+                    array
+                        .iter()
+                        .map(f)
+                        .collect::<Result<TimestampNanosecondArray>>()
+                }
+                DataType::Date32 => {
+                    let array = array.as_any().downcast_ref::<Date32Array>().unwrap();
+                    let array = array
+                        .iter()
+                        .map(date_to_timestamp)
+                        .collect::<TimestampNanosecondArray>();
+                    array
+                        .iter()
+                        .map(f)
+                        .collect::<Result<TimestampNanosecondArray>>()
+                }
+                t => Err(DataFusionError::Execution(format!(
+                    "`date_trunc` does not accept type {} as a second argument",
+                    t
+                ))),
+            };
 
-            ColumnarValue::Array(Arc::new(array))
+            ColumnarValue::Array(Arc::new(array?))
         }
         _ => {
             return Err(DataFusionError::Execution(