From f35591e5dc14402d81da02587badd0e3c8b3409d Mon Sep 17 00:00:00 2001 From: Jia-Xuan Liu Date: Tue, 18 Jun 2024 23:17:33 +0800 Subject: [PATCH 1/3] support unparse TimestampNanosecond --- datafusion/sql/src/unparser/expr.rs | 60 +++++++++++++++++++++++++++-- 1 file changed, 56 insertions(+), 4 deletions(-) diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 66334b0b410d..3f612e53328b 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -19,11 +19,11 @@ use arrow::util::display::array_value_to_string; use core::fmt; use std::{fmt::Display, vec}; -use arrow_array::{Date32Array, Date64Array}; +use arrow_array::{Date32Array, Date64Array, TimestampNanosecondArray}; use arrow_schema::DataType; use sqlparser::ast::Value::SingleQuotedString; use sqlparser::ast::{ - self, Expr as AstExpr, Function, FunctionArg, Ident, Interval, UnaryOperator, + self, Expr as AstExpr, Function, FunctionArg, Ident, Interval, TimezoneInfo, UnaryOperator, }; use datafusion_common::{ @@ -819,8 +819,49 @@ impl Unparser<'_> { ScalarValue::TimestampMicrosecond(None, _) => { Ok(ast::Expr::Value(ast::Value::Null)) } - ScalarValue::TimestampNanosecond(Some(_ts), _) => { - not_impl_err!("Unsupported scalar: {v:?}") + ScalarValue::TimestampNanosecond(Some(_ts), tz) => { + if let Some(tz) = tz { + let result = 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" + ))?; + Ok(ast::Expr::Cast { + kind: ast::CastKind::Cast, + expr: Box::new(ast::Expr::Value(ast::Value::SingleQuotedString( + result.to_string(), + ))), + data_type: ast::DataType::Timestamp(None, TimezoneInfo::None), + format: None, + }) + } else { + let result = 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" + ))?; + + Ok(ast::Expr::Cast { + kind: ast::CastKind::Cast, + expr: Box::new(ast::Expr::Value(ast::Value::SingleQuotedString( + result.to_string(), + ))), + data_type: ast::DataType::Timestamp(None, TimezoneInfo::None), + format: None, + }) + } } ScalarValue::TimestampNanosecond(None, _) => { Ok(ast::Expr::Value(ast::Value::Null)) @@ -1151,6 +1192,17 @@ mod tests { Expr::Literal(ScalarValue::Date32(Some(-1))), r#"CAST('1969-12-31' AS DATE)"#, ), + ( + Expr::Literal(ScalarValue::TimestampNanosecond(Some(10001), None)), + r#"CAST('1970-01-01 00:00:00.000010001' AS TIMESTAMP)"#, + ), + ( + Expr::Literal(ScalarValue::TimestampNanosecond( + Some(10001), + Some("+08:00".into()), + )), + r#"CAST('1970-01-01 08:00:00.000010001 +08:00' AS TIMESTAMP)"#, + ), (sum(col("a")), r#"sum(a)"#), ( count_udaf() From 740296f270c6ee2cd6c06f6633458d6b902efe8e Mon Sep 17 00:00:00 2001 From: Jia-Xuan Liu Date: Tue, 18 Jun 2024 23:48:23 +0800 Subject: [PATCH 2/3] cargo fmt --- datafusion/sql/src/unparser/expr.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 3f612e53328b..28ef3ec2c893 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -23,7 +23,8 @@ use arrow_array::{Date32Array, Date64Array, TimestampNanosecondArray}; use arrow_schema::DataType; use sqlparser::ast::Value::SingleQuotedString; use sqlparser::ast::{ - self, Expr as AstExpr, Function, FunctionArg, Ident, Interval, TimezoneInfo, UnaryOperator, + self, Expr as AstExpr, Function, FunctionArg, Ident, Interval, TimezoneInfo, + UnaryOperator, }; use datafusion_common::{ From f23878269c7fd46157986197efb7eb14f3eee25b Mon Sep 17 00:00:00 2001 From: Jia-Xuan Liu Date: Wed, 19 Jun 2024 02:25:30 +0800 Subject: [PATCH 3/3] extract the duplicate code --- datafusion/sql/src/unparser/expr.rs | 37 ++++++++++------------------- 1 file changed, 12 insertions(+), 25 deletions(-) diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 28ef3ec2c893..65481aed64f9 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -821,9 +821,8 @@ impl Unparser<'_> { Ok(ast::Expr::Value(ast::Value::Null)) } ScalarValue::TimestampNanosecond(Some(_ts), tz) => { - if let Some(tz) = tz { - let result = v - .to_array()? + let result = if let Some(tz) = tz { + v.to_array()? .as_any() .downcast_ref::() .ok_or(internal_datafusion_err!( @@ -832,18 +831,9 @@ impl Unparser<'_> { .value_as_datetime_with_tz(0, tz.parse()?) .ok_or(internal_datafusion_err!( "Unable to convert TimestampNanosecond to DateTime" - ))?; - Ok(ast::Expr::Cast { - kind: ast::CastKind::Cast, - expr: Box::new(ast::Expr::Value(ast::Value::SingleQuotedString( - result.to_string(), - ))), - data_type: ast::DataType::Timestamp(None, TimezoneInfo::None), - format: None, - }) + ))?.to_string() } else { - let result = v - .to_array()? + v.to_array()? .as_any() .downcast_ref::() .ok_or(internal_datafusion_err!( @@ -852,17 +842,14 @@ impl Unparser<'_> { .value_as_datetime(0) .ok_or(internal_datafusion_err!( "Unable to convert TimestampNanosecond to NaiveDateTime" - ))?; - - Ok(ast::Expr::Cast { - kind: ast::CastKind::Cast, - expr: Box::new(ast::Expr::Value(ast::Value::SingleQuotedString( - result.to_string(), - ))), - data_type: ast::DataType::Timestamp(None, TimezoneInfo::None), - format: None, - }) - } + ))?.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, + }) } ScalarValue::TimestampNanosecond(None, _) => { Ok(ast::Expr::Value(ast::Value::Null))