From aef5dcb8a8d8a8976c79884f861fbc51c2a000c5 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sat, 22 Jun 2024 19:32:41 +0800 Subject: [PATCH] Support to unparse ScalarValue::TimestampMillisecond to String (#11046) * wip Signed-off-by: Kevin Su * add a test Signed-off-by: Kevin Su --------- Signed-off-by: Kevin Su --- datafusion/sql/src/unparser/expr.rs | 58 ++++++++++++++++++++++++----- 1 file changed, 49 insertions(+), 9 deletions(-) diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 7db6956e3244..0aefefc1f007 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -19,7 +19,9 @@ use arrow::util::display::array_value_to_string; use core::fmt; use std::{fmt::Display, vec}; -use arrow_array::{Date32Array, Date64Array, TimestampNanosecondArray}; +use arrow_array::{ + Date32Array, Date64Array, TimestampMillisecondArray, TimestampNanosecondArray, +}; use arrow_schema::DataType; use sqlparser::ast::Value::SingleQuotedString; use sqlparser::ast::{ @@ -647,6 +649,15 @@ impl Unparser<'_> { } } + fn timestamp_string_to_sql(&self, ts: String) -> Result { + Ok(ast::Expr::Cast { + kind: ast::CastKind::Cast, + expr: Box::new(ast::Expr::Value(SingleQuotedString(ts))), + data_type: ast::DataType::Timestamp(None, TimezoneInfo::None), + format: None, + }) + } + /// DataFusion ScalarValues sometimes require a ast::Expr to construct. /// For example ScalarValue::Date32(d) corresponds to the ast::Expr CAST('datestr' as DATE) fn scalar_to_sql(&self, v: &ScalarValue) -> Result { @@ -808,8 +819,31 @@ impl Unparser<'_> { ScalarValue::TimestampSecond(None, _) => { Ok(ast::Expr::Value(ast::Value::Null)) } - ScalarValue::TimestampMillisecond(Some(_ts), _) => { - not_impl_err!("Unsupported scalar: {v:?}") + ScalarValue::TimestampMillisecond(Some(_ts), tz) => { + let result = if let Some(tz) = tz { + v.to_array()? + .as_any() + .downcast_ref::() + .ok_or(internal_datafusion_err!( + "Unable to downcast to TimestampMillisecond from TimestampMillisecond scalar" + ))? + .value_as_datetime_with_tz(0, tz.parse()?) + .ok_or(internal_datafusion_err!( + "Unable to convert TimestampMillisecond to DateTime" + ))?.to_string() + } else { + v.to_array()? + .as_any() + .downcast_ref::() + .ok_or(internal_datafusion_err!( + "Unable to downcast to TimestampMillisecond from TimestampMillisecond scalar" + ))? + .value_as_datetime(0) + .ok_or(internal_datafusion_err!( + "Unable to convert TimestampMillisecond to NaiveDateTime" + ))?.to_string() + }; + self.timestamp_string_to_sql(result) } ScalarValue::TimestampMillisecond(None, _) => { Ok(ast::Expr::Value(ast::Value::Null)) @@ -844,12 +878,7 @@ impl Unparser<'_> { "Unable to convert TimestampNanosecond to NaiveDateTime" ))?.to_string() }; - Ok(ast::Expr::Cast { - kind: ast::CastKind::Cast, - expr: Box::new(ast::Expr::Value(SingleQuotedString(result))), - data_type: ast::DataType::Timestamp(None, TimezoneInfo::None), - format: None, - }) + self.timestamp_string_to_sql(result) } ScalarValue::TimestampNanosecond(None, _) => { Ok(ast::Expr::Value(ast::Value::Null)) @@ -1180,6 +1209,17 @@ mod tests { Expr::Literal(ScalarValue::Date32(Some(-1))), r#"CAST('1969-12-31' AS DATE)"#, ), + ( + Expr::Literal(ScalarValue::TimestampMillisecond(Some(10001), None)), + r#"CAST('1970-01-01 00:00:10.001' AS TIMESTAMP)"#, + ), + ( + Expr::Literal(ScalarValue::TimestampMillisecond( + Some(10001), + Some("+08:00".into()), + )), + r#"CAST('1970-01-01 08:00:10.001 +08:00' AS TIMESTAMP)"#, + ), ( Expr::Literal(ScalarValue::TimestampNanosecond(Some(10001), None)), r#"CAST('1970-01-01 00:00:00.000010001' AS TIMESTAMP)"#,