diff --git a/datafusion/core/tests/sql/timestamp.rs b/datafusion/core/tests/sql/timestamp.rs index 1bde202d4376..912e0b072760 100644 --- a/datafusion/core/tests/sql/timestamp.rs +++ b/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(()) +} diff --git a/datafusion/physical-expr/src/datetime_expressions.rs b/datafusion/physical-expr/src/datetime_expressions.rs index b1f785e14e28..b685330648d6 100644 --- a/datafusion/physical-expr/src/datetime_expressions.rs +++ b/datafusion/physical-expr/src/datetime_expressions.rs @@ -452,6 +452,12 @@ pub fn date_trunc(args: &[ColumnarValue]) -> Result { }; let f = |x: Option| x.map(|x| date_trunc_single(granularity, x)).transpose(); + let date_to_timestamp = |x: Option| { + 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 { 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::() - .unwrap(); - let array = array - .iter() - .map(f) - .collect::>()?; + let array = match array.data_type() { + DataType::Timestamp(TimeUnit::Nanosecond, _) => { + let array = array + .as_any() + .downcast_ref::() + .unwrap(); + array + .iter() + .map(f) + .collect::>() + } + DataType::Date32 => { + let array = array.as_any().downcast_ref::().unwrap(); + let array = array + .iter() + .map(date_to_timestamp) + .collect::(); + array + .iter() + .map(f) + .collect::>() + } + 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(