From 0b700f28833a0946c08622b3d066a69036f046cb Mon Sep 17 00:00:00 2001 From: jonahgao Date: Tue, 24 Oct 2023 23:46:22 +0800 Subject: [PATCH 1/5] minor: cast the updated value to the data type of target column --- datafusion/sql/src/statement.rs | 41 ++++++++---------- datafusion/sqllogictest/test_files/update.slt | 43 +++++++++++++++++++ 2 files changed, 61 insertions(+), 23 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/update.slt diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index f8504a487a664..b5196c0866388 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -31,9 +31,9 @@ use arrow_schema::DataType; use datafusion_common::file_options::StatementOptions; use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::{ - not_impl_err, plan_datafusion_err, plan_err, unqualified_field_not_found, Column, - Constraints, DFField, DFSchema, DFSchemaRef, DataFusionError, ExprSchema, - OwnedTableReference, Result, SchemaReference, TableReference, ToDFSchema, + not_impl_err, plan_datafusion_err, plan_err, unqualified_field_not_found, + Constraints, DFField, DFSchema, DFSchemaRef, DataFusionError, OwnedTableReference, + Result, SchemaReference, TableReference, ToDFSchema, }; use datafusion_expr::dml::{CopyOptions, CopyTo}; use datafusion_expr::expr::Placeholder; @@ -969,12 +969,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { table_name.clone(), &arrow_schema, )?); - let values = table_schema.fields().iter().map(|f| { - ( - f.name().clone(), - ast::Expr::Identifier(ast::Ident::from(f.name().as_str())), - ) - }); // Overwrite with assignment expressions let mut planner_context = PlannerContext::new(); @@ -992,11 +986,15 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { }) .collect::>>()?; - let values = values - .into_iter() - .map(|(k, v)| { - let val = assign_map.remove(&k).unwrap_or(v); - (k, val) + let values_and_types = table_schema + .fields() + .iter() + .map(|f| { + let col_name = f.name(); + let val = assign_map.remove(col_name).unwrap_or_else(|| { + ast::Expr::Identifier(ast::Ident::from(col_name.as_str())) + }); + (col_name, val, f.data_type()) }) .collect::>(); @@ -1026,25 +1024,22 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // Projection let mut exprs = vec![]; - for (col_name, expr) in values.into_iter() { + for (col_name, expr, dt) in values_and_types.into_iter() { let expr = self.sql_to_expr(expr, &table_schema, &mut planner_context)?; let expr = match expr { datafusion_expr::Expr::Placeholder(Placeholder { ref id, ref data_type, }) => match data_type { - None => { - let dt = table_schema.data_type(&Column::from_name(&col_name))?; - datafusion_expr::Expr::Placeholder(Placeholder::new( - id.clone(), - Some(dt.clone()), - )) - } + None => datafusion_expr::Expr::Placeholder(Placeholder::new( + id.clone(), + Some(dt.clone()), + )), Some(_) => expr, }, _ => expr, }; - let expr = expr.alias(col_name); + let expr = expr.cast_to(dt, source.schema())?.alias(col_name); exprs.push(expr); } let source = project(source, exprs)?; diff --git a/datafusion/sqllogictest/test_files/update.slt b/datafusion/sqllogictest/test_files/update.slt new file mode 100644 index 0000000000000..ece589407c82c --- /dev/null +++ b/datafusion/sqllogictest/test_files/update.slt @@ -0,0 +1,43 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +########## +## Update Tests +########## + +statement ok +create table t1(a int, b varchar, c double); + +# Turn off the optimizer to make the logical plan closer to the initial one +statement ok +set datafusion.optimizer.max_passes = 0; + +query TT +explain update t1 set a=1, b=2, c=3.0; +---- +logical_plan +Dml: op=[Update] table=[t1] +--Projection: CAST(Int64(1) AS Int32) AS a, CAST(Int64(2) AS Utf8) AS b, Float64(3) AS c +----TableScan: t1 + +query TT +explain update t1 set a=c+1, b=a, c=c+1.0; +---- +logical_plan +Dml: op=[Update] table=[t1] +--Projection: CAST(t1.c + CAST(Int64(1) AS Float64) AS Int32) AS a, CAST(t1.a AS Utf8) AS b, t1.c + Float64(1) AS c +----TableScan: t1 From 54034fc94d003f80b591768915b8fa84cfbecb07 Mon Sep 17 00:00:00 2001 From: Jonah Gao Date: Tue, 24 Oct 2023 21:02:19 -0500 Subject: [PATCH 2/5] Update datafusion/sqllogictest/test_files/update.slt Co-authored-by: Alex Huang --- datafusion/sqllogictest/test_files/update.slt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/test_files/update.slt b/datafusion/sqllogictest/test_files/update.slt index ece589407c82c..7afa1f21d9312 100644 --- a/datafusion/sqllogictest/test_files/update.slt +++ b/datafusion/sqllogictest/test_files/update.slt @@ -20,7 +20,7 @@ ########## statement ok -create table t1(a int, b varchar, c double); +create table t1(a int, b varchar, c double, d int); # Turn off the optimizer to make the logical plan closer to the initial one statement ok From def54f22a258b43e5e0cb70a39a2d7d7498cf0fd Mon Sep 17 00:00:00 2001 From: Jonah Gao Date: Tue, 24 Oct 2023 21:02:25 -0500 Subject: [PATCH 3/5] Update datafusion/sqllogictest/test_files/update.slt Co-authored-by: Alex Huang --- datafusion/sqllogictest/test_files/update.slt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/test_files/update.slt b/datafusion/sqllogictest/test_files/update.slt index 7afa1f21d9312..d49d8f8c90730 100644 --- a/datafusion/sqllogictest/test_files/update.slt +++ b/datafusion/sqllogictest/test_files/update.slt @@ -27,7 +27,7 @@ statement ok set datafusion.optimizer.max_passes = 0; query TT -explain update t1 set a=1, b=2, c=3.0; +explain update t1 set a=1, b=2, c=3.0, d=NULL; ---- logical_plan Dml: op=[Update] table=[t1] From 92a642c3a4d13dba73f598243a55f7df4eb2188b Mon Sep 17 00:00:00 2001 From: Jonah Gao Date: Tue, 24 Oct 2023 21:02:31 -0500 Subject: [PATCH 4/5] Update datafusion/sqllogictest/test_files/update.slt Co-authored-by: Alex Huang --- datafusion/sqllogictest/test_files/update.slt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/test_files/update.slt b/datafusion/sqllogictest/test_files/update.slt index d49d8f8c90730..0eafbf28d18d9 100644 --- a/datafusion/sqllogictest/test_files/update.slt +++ b/datafusion/sqllogictest/test_files/update.slt @@ -35,7 +35,7 @@ Dml: op=[Update] table=[t1] ----TableScan: t1 query TT -explain update t1 set a=c+1, b=a, c=c+1.0; +explain update t1 set a=c+1, b=a, c=c+1.0, d=b; ---- logical_plan Dml: op=[Update] table=[t1] From e8c50553f33caee0a29a882b37309b41d1a5e64f Mon Sep 17 00:00:00 2001 From: jonahgao Date: Wed, 25 Oct 2023 10:06:01 +0800 Subject: [PATCH 5/5] fix tests --- datafusion/sqllogictest/test_files/update.slt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/sqllogictest/test_files/update.slt b/datafusion/sqllogictest/test_files/update.slt index 0eafbf28d18d9..4542a262390cd 100644 --- a/datafusion/sqllogictest/test_files/update.slt +++ b/datafusion/sqllogictest/test_files/update.slt @@ -31,7 +31,7 @@ explain update t1 set a=1, b=2, c=3.0, d=NULL; ---- logical_plan Dml: op=[Update] table=[t1] ---Projection: CAST(Int64(1) AS Int32) AS a, CAST(Int64(2) AS Utf8) AS b, Float64(3) AS c +--Projection: CAST(Int64(1) AS Int32) AS a, CAST(Int64(2) AS Utf8) AS b, Float64(3) AS c, CAST(NULL AS Int32) AS d ----TableScan: t1 query TT @@ -39,5 +39,5 @@ explain update t1 set a=c+1, b=a, c=c+1.0, d=b; ---- logical_plan Dml: op=[Update] table=[t1] ---Projection: CAST(t1.c + CAST(Int64(1) AS Float64) AS Int32) AS a, CAST(t1.a AS Utf8) AS b, t1.c + Float64(1) AS c +--Projection: CAST(t1.c + CAST(Int64(1) AS Float64) AS Int32) AS a, CAST(t1.a AS Utf8) AS b, t1.c + Float64(1) AS c, CAST(t1.b AS Int32) AS d ----TableScan: t1