feat: Support casting to array of primity type (#2366)

apache · Apr 28, 2022 · cda997e · cda997e
1 parent af6be6e
commit cda997e
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
diff --git a/datafusion/core/src/sql/planner.rs b/datafusion/core/src/sql/planner.rs
@@ -2400,8 +2400,8 @@ fn extract_possible_join_keys(
     }
 }
 
-/// Convert SQL data type to relational representation of data type
-pub fn convert_data_type(sql_type: &SQLDataType) -> Result<DataType> {
+/// Convert SQL simple data type to relational representation of data type
+pub fn convert_simple_data_type(sql_type: &SQLDataType) -> Result<DataType> {
     match sql_type {
         SQLDataType::Boolean => Ok(DataType::Boolean),
         SQLDataType::SmallInt(_) => Ok(DataType::Int16),
@@ -2410,7 +2410,10 @@ pub fn convert_data_type(sql_type: &SQLDataType) -> Result<DataType> {
         SQLDataType::Float(_) => Ok(DataType::Float32),
         SQLDataType::Real => Ok(DataType::Float32),
         SQLDataType::Double => Ok(DataType::Float64),
-        SQLDataType::Char(_) | SQLDataType::Varchar(_) => Ok(DataType::Utf8),
+        SQLDataType::Char(_)
+        | SQLDataType::Varchar(_)
+        | SQLDataType::Text
+        | SQLDataType::String => Ok(DataType::Utf8),
         SQLDataType::Timestamp => Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)),
         SQLDataType::Date => Ok(DataType::Date32),
         SQLDataType::Decimal(precision, scale) => make_decimal_type(*precision, *scale),
@@ -2421,6 +2424,20 @@ pub fn convert_data_type(sql_type: &SQLDataType) -> Result<DataType> {
     }
 }
 
+/// Convert SQL data type to relational representation of data type
+pub fn convert_data_type(sql_type: &SQLDataType) -> Result<DataType> {
+    match sql_type {
+        SQLDataType::Array(inner_sql_type) => {
+            let data_type = convert_simple_data_type(inner_sql_type)?;
+
+            Ok(DataType::List(Box::new(Field::new(
+                "field", data_type, true,
+            ))))
+        }
+        other => convert_simple_data_type(other),
+    }
+}
+
 // Parse number in sql string, convert to Expr::Literal
 fn parse_sql_number(n: &str) -> Result<Expr> {
     match n.parse::<i64>() {

diff --git a/datafusion/core/tests/sql/expr.rs b/datafusion/core/tests/sql/expr.rs
@@ -796,6 +796,11 @@ async fn test_regex_expressions() -> Result<()> {
 
 #[tokio::test]
 async fn test_cast_expressions() -> Result<()> {
+    test_expression!("CAST([1,2,3,4] AS INT[])", "[1, 2, 3, 4]");
+    test_expression!(
+        "CAST([1,2,3,4] AS NUMERIC(10,4)[])",
+        "[1.0000, 2.0000, 3.0000, 4.0000]"
+    );
     test_expression!("CAST('0' AS INT)", "0");
     test_expression!("CAST(NULL AS INT)", "NULL");
     test_expression!("TRY_CAST('0' AS INT)", "0");