From 068778d43da99ffc9140bad5d98ca5bed7bc0b86 Mon Sep 17 00:00:00 2001 From: Jax Liu Date: Wed, 26 Jun 2024 23:25:32 +0800 Subject: [PATCH] Support to unparse `TimestampSecond` and `TimestampMicrosecond` to String (#11120) * support to unparse TimestampSecond and TimestampMicrosecond * cargo fmt * fmt --- datafusion/sql/src/unparser/expr.rs | 140 +++++++++++++++------------- 1 file changed, 74 insertions(+), 66 deletions(-) diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 241d4148efae9..ad898de5987ac 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -16,18 +16,17 @@ // under the License. use core::fmt; +use std::sync::Arc; use std::{fmt::Display, vec}; use arrow::datatypes::{Decimal128Type, Decimal256Type, DecimalType}; use arrow::util::display::array_value_to_string; use arrow_array::types::{ ArrowTemporalType, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, - Time64NanosecondType, -}; -use arrow_array::{ - Date32Array, Date64Array, PrimitiveArray, TimestampMillisecondArray, - TimestampNanosecondArray, + Time64NanosecondType, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, }; +use arrow_array::{Date32Array, Date64Array, PrimitiveArray}; use arrow_schema::DataType; use sqlparser::ast::Value::SingleQuotedString; use sqlparser::ast::{ @@ -655,6 +654,47 @@ impl Unparser<'_> { } } + fn handle_timestamp( + &self, + v: &ScalarValue, + tz: &Option>, + ) -> Result + where + i64: From, + { + let ts = if let Some(tz) = tz { + v.to_array()? + .as_any() + .downcast_ref::>() + .ok_or(internal_datafusion_err!( + "Failed to downcast type {v:?} to arrow array" + ))? + .value_as_datetime_with_tz(0, tz.parse()?) + .ok_or(internal_datafusion_err!( + "Unable to convert {v:?} to DateTime" + ))? + .to_string() + } else { + v.to_array()? + .as_any() + .downcast_ref::>() + .ok_or(internal_datafusion_err!( + "Failed to downcast type {v:?} to arrow array" + ))? + .value_as_datetime(0) + .ok_or(internal_datafusion_err!( + "Unable to convert {v:?} to DateTime" + ))? + .to_string() + }; + 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, + }) + } + fn handle_time(&self, v: &ScalarValue) -> Result where i64: From, @@ -677,15 +717,6 @@ 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 { @@ -847,72 +878,26 @@ impl Unparser<'_> { self.handle_time::(v) } ScalarValue::Time64Nanosecond(None) => Ok(ast::Expr::Value(ast::Value::Null)), - ScalarValue::TimestampSecond(Some(_ts), _) => { - not_impl_err!("Unsupported scalar: {v:?}") + ScalarValue::TimestampSecond(Some(_ts), tz) => { + self.handle_timestamp::(v, tz) } ScalarValue::TimestampSecond(None, _) => { Ok(ast::Expr::Value(ast::Value::Null)) } 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) + self.handle_timestamp::(v, tz) } ScalarValue::TimestampMillisecond(None, _) => { Ok(ast::Expr::Value(ast::Value::Null)) } - ScalarValue::TimestampMicrosecond(Some(_ts), _) => { - not_impl_err!("Unsupported scalar: {v:?}") + ScalarValue::TimestampMicrosecond(Some(_ts), tz) => { + self.handle_timestamp::(v, tz) } ScalarValue::TimestampMicrosecond(None, _) => { Ok(ast::Expr::Value(ast::Value::Null)) } ScalarValue::TimestampNanosecond(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 TimestampNanosecond from TimestampNanosecond scalar" - ))? - .value_as_datetime_with_tz(0, tz.parse()?) - .ok_or(internal_datafusion_err!( - "Unable to convert TimestampNanosecond to DateTime" - ))?.to_string() - } else { - v.to_array()? - .as_any() - .downcast_ref::() - .ok_or(internal_datafusion_err!( - "Unable to downcast to TimestampNanosecond from TimestampNanosecond scalar" - ))? - .value_as_datetime(0) - .ok_or(internal_datafusion_err!( - "Unable to convert TimestampNanosecond to NaiveDateTime" - ))?.to_string() - }; - self.timestamp_string_to_sql(result) + self.handle_timestamp::(v, tz) } ScalarValue::TimestampNanosecond(None, _) => { Ok(ast::Expr::Value(ast::Value::Null)) @@ -1079,6 +1064,7 @@ mod tests { use arrow::datatypes::{Field, Schema}; use arrow_schema::DataType::Int8; + use datafusion_common::TableReference; use datafusion_expr::{ case, col, cube, exists, grouping_set, interval_datetime_lit, @@ -1246,6 +1232,17 @@ mod tests { Expr::Literal(ScalarValue::Date32(Some(-1))), r#"CAST('1969-12-31' AS DATE)"#, ), + ( + Expr::Literal(ScalarValue::TimestampSecond(Some(10001), None)), + r#"CAST('1970-01-01 02:46:41' AS TIMESTAMP)"#, + ), + ( + Expr::Literal(ScalarValue::TimestampSecond( + Some(10001), + Some("+08:00".into()), + )), + r#"CAST('1970-01-01 10:46:41 +08:00' AS TIMESTAMP)"#, + ), ( Expr::Literal(ScalarValue::TimestampMillisecond(Some(10001), None)), r#"CAST('1970-01-01 00:00:10.001' AS TIMESTAMP)"#, @@ -1257,6 +1254,17 @@ mod tests { )), r#"CAST('1970-01-01 08:00:10.001 +08:00' AS TIMESTAMP)"#, ), + ( + Expr::Literal(ScalarValue::TimestampMicrosecond(Some(10001), None)), + r#"CAST('1970-01-01 00:00:00.010001' AS TIMESTAMP)"#, + ), + ( + Expr::Literal(ScalarValue::TimestampMicrosecond( + Some(10001), + Some("+08:00".into()), + )), + r#"CAST('1970-01-01 08:00:00.010001 +08:00' AS TIMESTAMP)"#, + ), ( Expr::Literal(ScalarValue::TimestampNanosecond(Some(10001), None)), r#"CAST('1970-01-01 00:00:00.000010001' AS TIMESTAMP)"#,