Skip to content

Commit a9d66e2

Browse files
jonahgaoWeijun-H
andauthored
minor: add a datatype casting for the updated value (#7922)
* minor: cast the updated value to the data type of target column * Update datafusion/sqllogictest/test_files/update.slt Co-authored-by: Alex Huang <huangweijun1001@gmail.com> * Update datafusion/sqllogictest/test_files/update.slt Co-authored-by: Alex Huang <huangweijun1001@gmail.com> * Update datafusion/sqllogictest/test_files/update.slt Co-authored-by: Alex Huang <huangweijun1001@gmail.com> * fix tests --------- Co-authored-by: Alex Huang <huangweijun1001@gmail.com>
1 parent 30e5f42 commit a9d66e2

File tree

2 files changed

+61
-23
lines changed

2 files changed

+61
-23
lines changed

datafusion/sql/src/statement.rs

Lines changed: 18 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -31,9 +31,9 @@ use arrow_schema::DataType;
3131
use datafusion_common::file_options::StatementOptions;
3232
use datafusion_common::parsers::CompressionTypeVariant;
3333
use datafusion_common::{
34-
not_impl_err, plan_datafusion_err, plan_err, unqualified_field_not_found, Column,
35-
Constraints, DFField, DFSchema, DFSchemaRef, DataFusionError, ExprSchema,
36-
OwnedTableReference, Result, SchemaReference, TableReference, ToDFSchema,
34+
not_impl_err, plan_datafusion_err, plan_err, unqualified_field_not_found,
35+
Constraints, DFField, DFSchema, DFSchemaRef, DataFusionError, OwnedTableReference,
36+
Result, SchemaReference, TableReference, ToDFSchema,
3737
};
3838
use datafusion_expr::dml::{CopyOptions, CopyTo};
3939
use datafusion_expr::expr::Placeholder;
@@ -969,12 +969,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
969969
table_name.clone(),
970970
&arrow_schema,
971971
)?);
972-
let values = table_schema.fields().iter().map(|f| {
973-
(
974-
f.name().clone(),
975-
ast::Expr::Identifier(ast::Ident::from(f.name().as_str())),
976-
)
977-
});
978972

979973
// Overwrite with assignment expressions
980974
let mut planner_context = PlannerContext::new();
@@ -992,11 +986,15 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
992986
})
993987
.collect::<Result<HashMap<String, Expr>>>()?;
994988

995-
let values = values
996-
.into_iter()
997-
.map(|(k, v)| {
998-
let val = assign_map.remove(&k).unwrap_or(v);
999-
(k, val)
989+
let values_and_types = table_schema
990+
.fields()
991+
.iter()
992+
.map(|f| {
993+
let col_name = f.name();
994+
let val = assign_map.remove(col_name).unwrap_or_else(|| {
995+
ast::Expr::Identifier(ast::Ident::from(col_name.as_str()))
996+
});
997+
(col_name, val, f.data_type())
1000998
})
1001999
.collect::<Vec<_>>();
10021000

@@ -1026,25 +1024,22 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
10261024

10271025
// Projection
10281026
let mut exprs = vec![];
1029-
for (col_name, expr) in values.into_iter() {
1027+
for (col_name, expr, dt) in values_and_types.into_iter() {
10301028
let expr = self.sql_to_expr(expr, &table_schema, &mut planner_context)?;
10311029
let expr = match expr {
10321030
datafusion_expr::Expr::Placeholder(Placeholder {
10331031
ref id,
10341032
ref data_type,
10351033
}) => match data_type {
1036-
None => {
1037-
let dt = table_schema.data_type(&Column::from_name(&col_name))?;
1038-
datafusion_expr::Expr::Placeholder(Placeholder::new(
1039-
id.clone(),
1040-
Some(dt.clone()),
1041-
))
1042-
}
1034+
None => datafusion_expr::Expr::Placeholder(Placeholder::new(
1035+
id.clone(),
1036+
Some(dt.clone()),
1037+
)),
10431038
Some(_) => expr,
10441039
},
10451040
_ => expr,
10461041
};
1047-
let expr = expr.alias(col_name);
1042+
let expr = expr.cast_to(dt, source.schema())?.alias(col_name);
10481043
exprs.push(expr);
10491044
}
10501045
let source = project(source, exprs)?;
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
##########
19+
## Update Tests
20+
##########
21+
22+
statement ok
23+
create table t1(a int, b varchar, c double, d int);
24+
25+
# Turn off the optimizer to make the logical plan closer to the initial one
26+
statement ok
27+
set datafusion.optimizer.max_passes = 0;
28+
29+
query TT
30+
explain update t1 set a=1, b=2, c=3.0, d=NULL;
31+
----
32+
logical_plan
33+
Dml: op=[Update] table=[t1]
34+
--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
35+
----TableScan: t1
36+
37+
query TT
38+
explain update t1 set a=c+1, b=a, c=c+1.0, d=b;
39+
----
40+
logical_plan
41+
Dml: op=[Update] table=[t1]
42+
--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
43+
----TableScan: t1

0 commit comments

Comments
 (0)