[arrow-datafusion] branch master updated: Ensure insert projections are of correct type (#5049)

[av...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

avantgardner pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new c313601b5 Ensure insert projections are of correct type (#5049)
c313601b5 is described below

commit c313601b5ae47dc2fd7618d6a6d89947fdbca99b
Author: Brent Gardner <bg...@squarelabs.net>
AuthorDate: Tue Jan 24 13:15:36 2023 -0700

    Ensure insert projections are of correct type (#5049)
    
    * Ensure insert projections are of correct type
---
 datafusion/sql/src/statement.rs          | 21 +++++++++++++--------
 datafusion/sql/tests/integration_test.rs |  2 +-
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs
index fb50ef7bc..2d7791dd5 100644
--- a/datafusion/sql/src/statement.rs
+++ b/datafusion/sql/src/statement.rs
@@ -35,8 +35,8 @@ use datafusion_expr::utils::expr_to_columns;
 use datafusion_expr::{
     cast, col, CreateCatalog, CreateCatalogSchema,
     CreateExternalTable as PlanCreateExternalTable, CreateMemoryTable, CreateView,
-    DescribeTable, DmlStatement, DropTable, DropView, Explain, Filter, LogicalPlan,
-    LogicalPlanBuilder, PlanType, SetVariable, ToStringifiedPlan, WriteOp,
+    DescribeTable, DmlStatement, DropTable, DropView, Explain, ExprSchemable, Filter,
+    LogicalPlan, LogicalPlanBuilder, PlanType, SetVariable, ToStringifiedPlan, WriteOp,
 };
 use sqlparser::ast;
 use sqlparser::ast::{
@@ -753,7 +753,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             .schema_provider
             .get_table_provider((&table_name).into())?;
         let arrow_schema = (*provider.schema()).clone();
-        let table_schema = Arc::new(DFSchema::try_from(arrow_schema)?);
+        let table_schema = DFSchema::try_from(arrow_schema)?;
 
         // infer types for Values clause... other types should be resolvable the regular way
         let mut prepare_param_data_types = BTreeMap::new();
@@ -792,19 +792,24 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 "Column count doesn't match insert query!".to_owned(),
             ))?;
         }
-        let exprs: Vec<_> = columns
+        let values_schema = source.schema();
+        let exprs = columns
             .iter()
             .zip(source.schema().fields().iter())
             .map(|(c, f)| {
-                datafusion_expr::Expr::Column(Column::from(f.name().clone()))
-                    .alias(c.value.clone())
+                let col_name = c.value.clone();
+                let col = table_schema.field_with_name(None, col_name.as_str())?;
+                let expr = datafusion_expr::Expr::Column(Column::from(f.name().clone()))
+                    .alias(col_name)
+                    .cast_to(col.data_type(), values_schema)?;
+                Ok(expr)
             })
-            .collect();
+            .collect::<Result<Vec<datafusion_expr::Expr>>>()?;
         let source = project(source, exprs)?;
 
         let plan = LogicalPlan::Dml(DmlStatement {
             table_name,
-            table_schema,
+            table_schema: Arc::new(table_schema),
             op: WriteOp::Insert,
             input: Arc::new(source),
         });
diff --git a/datafusion/sql/tests/integration_test.rs b/datafusion/sql/tests/integration_test.rs
index c771a3ec5..c3725ec05 100644
--- a/datafusion/sql/tests/integration_test.rs
+++ b/datafusion/sql/tests/integration_test.rs
@@ -163,7 +163,7 @@ fn plan_insert() {
         "insert into person (id, first_name, last_name) values (1, 'Alan', 'Turing')";
     let plan = r#"
 Dml: op=[Insert] table=[person]
-  Projection: column1 AS id, column2 AS first_name, column3 AS last_name
+  Projection: CAST(column1 AS id AS UInt32), column2 AS first_name, column3 AS last_name
     Values: (Int64(1), Utf8("Alan"), Utf8("Turing"))
     "#
     .trim();