From 33309578961f9b5ac8f536dd6324a2aedf8c9d4e Mon Sep 17 00:00:00 2001 From: localhost Date: Tue, 16 May 2023 16:14:11 +0800 Subject: [PATCH] chore: add fmt for statement query (#1588) * chore: add fmt for statement query * chore: add test for query display --- src/datatypes/src/data_type.rs | 104 ++++++++++++++++++++++++++++++++ src/sql/src/statements/query.rs | 55 +++++++++++++++++ 2 files changed, 159 insertions(+) diff --git a/src/datatypes/src/data_type.rs b/src/datatypes/src/data_type.rs index 7723e0ee6f..c2e7a0ed6b 100644 --- a/src/datatypes/src/data_type.rs +++ b/src/datatypes/src/data_type.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::fmt; use std::sync::Arc; use arrow::datatypes::{DataType as ArrowDataType, TimeUnit as ArrowTimeUnit}; @@ -62,6 +63,32 @@ pub enum ConcreteDataType { Dictionary(DictionaryType), } +impl fmt::Display for ConcreteDataType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ConcreteDataType::Null(_) => write!(f, "Null"), + ConcreteDataType::Boolean(_) => write!(f, "Boolean"), + ConcreteDataType::Int8(_) => write!(f, "Int8"), + ConcreteDataType::Int16(_) => write!(f, "Int16"), + ConcreteDataType::Int32(_) => write!(f, "Int32"), + ConcreteDataType::Int64(_) => write!(f, "Int64"), + ConcreteDataType::UInt8(_) => write!(f, "UInt8"), + ConcreteDataType::UInt16(_) => write!(f, "UInt16"), + ConcreteDataType::UInt32(_) => write!(f, "UInt32"), + ConcreteDataType::UInt64(_) => write!(f, "UInt64"), + ConcreteDataType::Float32(_) => write!(f, "Float32"), + ConcreteDataType::Float64(_) => write!(f, "Float64"), + ConcreteDataType::Binary(_) => write!(f, "Binary"), + ConcreteDataType::String(_) => write!(f, "String"), + ConcreteDataType::Date(_) => write!(f, "Date"), + ConcreteDataType::DateTime(_) => write!(f, "DateTime"), + ConcreteDataType::Timestamp(_) => write!(f, "Timestamp"), + ConcreteDataType::List(_) => write!(f, "List"), + ConcreteDataType::Dictionary(_) => write!(f, "Dictionary"), + } + } +} + // TODO(yingwen): Refactor these `is_xxx()` methods, such as adding a `properties()` method // returning all these properties to the `DataType` trait impl ConcreteDataType { @@ -514,4 +541,81 @@ mod tests { ); assert!(ConcreteDataType::int32_datatype().as_list().is_none()); } + + #[test] + fn test_display_concrete_data_type() { + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::Null).to_string(), + "Null" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::Boolean).to_string(), + "Boolean" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::Binary).to_string(), + "Binary" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::LargeBinary).to_string(), + "Binary" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::Int8).to_string(), + "Int8" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::Int16).to_string(), + "Int16" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::Int32).to_string(), + "Int32" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::Int64).to_string(), + "Int64" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::UInt8).to_string(), + "UInt8" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::UInt16).to_string(), + "UInt16" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::UInt32).to_string(), + "UInt32" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::UInt64).to_string(), + "UInt64" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::Float32).to_string(), + "Float32" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::Float64).to_string(), + "Float64" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::Utf8).to_string(), + "String" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::List(Arc::new(Field::new( + "item", + ArrowDataType::Int32, + true, + )))) + .to_string(), + "List" + ); + assert_eq!( + ConcreteDataType::from_arrow_type(&ArrowDataType::Date32).to_string(), + "Date" + ); + } } diff --git a/src/sql/src/statements/query.rs b/src/sql/src/statements/query.rs index 6e00962dca..6b0b9bbde0 100644 --- a/src/sql/src/statements/query.rs +++ b/src/sql/src/statements/query.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::fmt; + use datatypes::prelude::ConcreteDataType; use sqlparser::ast::Query as SpQuery; @@ -53,3 +55,56 @@ impl Query { &mut self.param_types } } + +impl fmt::Display for Query { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{} ", self.inner)?; + write!(f, "[")?; + for i in 0..self.param_types.len() { + write!(f, "{}", self.param_types[i])?; + if i != self.param_types.len() - 1 { + write!(f, ",")?; + } + } + write!(f, "]")?; + Ok(()) + } +} + +#[cfg(test)] +mod test { + + use sqlparser::dialect::GenericDialect; + + use super::Query; + use crate::parser::ParserContext; + use crate::statements::statement::Statement; + + fn create_query(sql: &str) -> Option> { + match ParserContext::create_with_dialect(sql, &GenericDialect {}) + .unwrap() + .remove(0) + { + Statement::Query(query) => Some(query), + _ => None, + } + } + + #[test] + fn test_query_display() { + assert_eq!( + create_query("select * from abc where x = 1 and y = 7") + .unwrap() + .to_string(), + "SELECT * FROM abc WHERE x = 1 AND y = 7 []" + ); + assert_eq!( + create_query( + "select * from abc left join bcd where abc.a = 1 and bcd.d = 7 and abc.id = bcd.id" + ) + .unwrap() + .to_string(), + "SELECT * FROM abc LEFT JOIN bcd WHERE abc.a = 1 AND bcd.d = 7 AND abc.id = bcd.id []" + ); + } +}