feat: support postgres types transparently in queries

apache · Jul 31, 2023 · 27841c7 · 27841c7
1 parent 01d7dba
commit 27841c7
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 11 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -38,7 +38,9 @@ arrow-buffer = { version = "43.0.0", default-features = false }
 arrow-flight = { version = "43.0.0", features = ["flight-sql-experimental"] }
 arrow-schema = { version = "43.0.0", default-features = false }
 parquet = { version = "43.0.0", features = ["arrow", "async", "object_store"] }
-sqlparser = { version = "0.36.1", features = ["visitor"] }
+# todo: use released version of sqlparser-rs once it is released
+sqlparser = { git = "https://github.com/sqlparser-rs/sqlparser-rs", branch = "main", features = ["visitor"] }
+
 
 [profile.release]
 codegen-units = 1

diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs
@@ -305,20 +305,20 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
 
     fn convert_simple_data_type(&self, sql_type: &SQLDataType) -> Result<DataType> {
         match sql_type {
-            SQLDataType::Boolean => Ok(DataType::Boolean),
+            SQLDataType::Boolean | SQLDataType::Bool=> Ok(DataType::Boolean),
             SQLDataType::TinyInt(_) => Ok(DataType::Int8),
-            SQLDataType::SmallInt(_) => Ok(DataType::Int16),
-            SQLDataType::Int(_) | SQLDataType::Integer(_) => Ok(DataType::Int32),
-            SQLDataType::BigInt(_) => Ok(DataType::Int64),
+            SQLDataType::SmallInt(_) | SQLDataType::Int2(_) => Ok(DataType::Int16),
+            SQLDataType::Int(_) | SQLDataType::Integer(_) | SQLDataType::Int4(_)  => Ok(DataType::Int32),
+            SQLDataType::BigInt(_) | SQLDataType::Int8(_) => Ok(DataType::Int64),
             SQLDataType::UnsignedTinyInt(_) => Ok(DataType::UInt8),
-            SQLDataType::UnsignedSmallInt(_) => Ok(DataType::UInt16),
-            SQLDataType::UnsignedInt(_) | SQLDataType::UnsignedInteger(_) => {
+            SQLDataType::UnsignedSmallInt(_) | SQLDataType::UnsignedInt2(_) => Ok(DataType::UInt16),
+            SQLDataType::UnsignedInt(_) | SQLDataType::UnsignedInteger(_)| SQLDataType::UnsignedInt4(_) => {
                 Ok(DataType::UInt32)
             }
-            SQLDataType::UnsignedBigInt(_) => Ok(DataType::UInt64),
+            SQLDataType::UnsignedBigInt(_) | SQLDataType::UnsignedInt8(_) => Ok(DataType::UInt64),
             SQLDataType::Float(_) => Ok(DataType::Float32),
-            SQLDataType::Real => Ok(DataType::Float32),
-            SQLDataType::Double | SQLDataType::DoublePrecision => Ok(DataType::Float64),
+            SQLDataType::Real | SQLDataType::FLOAT4=> Ok(DataType::Float32),
+            SQLDataType::Double | SQLDataType::DoublePrecision | SQLDataType::FLOAT8=> Ok(DataType::Float64),
             SQLDataType::Char(_)
             | SQLDataType::Varchar(_)
             | SQLDataType::Text

diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs
@@ -436,7 +436,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 self.delete_to_plan(table_name, selection)
             }
 
-            Statement::StartTransaction { modes } => {
+            Statement::StartTransaction { modes, .. } => {
                 let isolation_level: ast::TransactionIsolationLevel = modes
                     .iter()
                     .filter_map(|m: &ast::TransactionMode| match m {

diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs
@@ -4121,6 +4121,25 @@ fn test_multi_grouping_sets() {
     quick_test(sql, expected);
 }
 
+#[test]
+fn test_create_table_with_postgres_type_alias() {
+    let sql = "create table foo_table(int8_col int8, int4_col int4, int2_col int2,\
+     float4_col float4, float8_col float8, bool_col bool);";
+    let expected = "CreateMemoryTable: Bare { table: \"foo_table\" }\n  EmptyRelation";
+    quick_test(sql, expected);
+}
+
+#[test]
+fn test_select_table_with_postgres_type_alias() {
+    let sql = "select age::int8 as int8_col, age::int4 as int4_col, age::int2 as int2_col,\
+    TRY_CAST(age AS FLOAT)::float4 as float4_col, TRY_CAST(age AS FLOAT)::float8 as float8_col, (age = 18) as bool \
+    from person;";
+    let expected = "Projection: CAST(person.age AS Int64) AS int8_col, CAST(person.age AS Int32) AS int4_col, \
+    CAST(person.age AS Int16) AS int2_col, CAST(TRY_CAST(person.age AS Float32) AS Float32) AS float4_col, \
+    CAST(TRY_CAST(person.age AS Float32) AS Float64) AS float8_col, person.age = Int64(18) AS bool\n  TableScan: person";
+    quick_test(sql, expected);
+}
+
 fn assert_field_not_found(err: DataFusionError, name: &str) {
     match err {
         DataFusionError::SchemaError { .. } => {