From d79d1e8f2d6a459f43ed13904f5efeb0ea494291 Mon Sep 17 00:00:00 2001
From: jonahgao <jonahgaox@gmail.com>
Date: Tue, 7 Nov 2023 23:44:38 +0800
Subject: [PATCH] feat: support target table alias in update statement

---
 datafusion/sql/src/statement.rs               | 18 ++++++++++++++----
 datafusion/sqllogictest/test_files/update.slt | 15 ++++++++++++++-
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs
index 9d9c55361a5e9..116624c6f7b97 100644
--- a/datafusion/sql/src/statement.rs
+++ b/datafusion/sql/src/statement.rs
@@ -31,7 +31,7 @@ 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,
+    not_impl_err, plan_datafusion_err, plan_err, unqualified_field_not_found, Column,
     Constraints, DFField, DFSchema, DFSchemaRef, DataFusionError, OwnedTableReference,
     Result, SchemaReference, TableReference, ToDFSchema,
 };
@@ -970,8 +970,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         from: Option<TableWithJoins>,
         predicate_expr: Option<Expr>,
     ) -> Result<LogicalPlan> {
-        let table_name = match &table.relation {
-            TableFactor::Table { name, .. } => name.clone(),
+        let (table_name, table_alias) = match &table.relation {
+            TableFactor::Table { name, alias, .. } => (name.clone(), alias.clone()),
             _ => plan_err!("Cannot update non-table relation!")?,
         };
 
@@ -1047,7 +1047,17 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                         // Cast to target column type, if necessary
                         expr.cast_to(field.data_type(), source.schema())?
                     }
-                    None => datafusion_expr::Expr::Column(field.qualified_column()),
+                    None => {
+                        // If the target table has an alias, use it to qualify the column name
+                        if let Some(alias) = &table_alias {
+                            datafusion_expr::Expr::Column(Column::new(
+                                Some(self.normalizer.normalize(alias.name.clone())),
+                                field.name(),
+                            ))
+                        } else {
+                            datafusion_expr::Expr::Column(field.qualified_column())
+                        }
+                    }
                 };
                 Ok(expr.alias(field.name()))
             })
diff --git a/datafusion/sqllogictest/test_files/update.slt b/datafusion/sqllogictest/test_files/update.slt
index cb8c6a4fac28a..c88082fc7272e 100644
--- a/datafusion/sqllogictest/test_files/update.slt
+++ b/datafusion/sqllogictest/test_files/update.slt
@@ -76,4 +76,17 @@ create table t3(a int, b varchar, c double, d int);
 
 # set from mutiple tables, sqlparser only supports from one table
 query error DataFusion error: SQL error: ParserError\("Expected end of statement, found: ,"\)
-explain update t1 set b = t2.b, c = t3.a, d = 1 from t2, t3 where t1.a = t2.a and t1.a = t3.a;
\ No newline at end of file
+explain update t1 set b = t2.b, c = t3.a, d = 1 from t2, t3 where t1.a = t2.a and t1.a = t3.a;
+
+# test table alias
+query TT
+explain update t1 as T set b = t2.b, c = t.a, d = 1 from t2 where t.a = t2.a and t.b > 'foo' and t2.c > 1.0;
+----
+logical_plan
+Dml: op=[Update] table=[t1]
+--Projection: t.a AS a, t2.b AS b, CAST(t.a AS Float64) AS c, CAST(Int64(1) AS Int32) AS d
+----Filter: t.a = t2.a AND t.b > Utf8("foo") AND t2.c > Float64(1)
+------CrossJoin:
+--------SubqueryAlias: t
+----------TableScan: t1
+--------TableScan: t2
\ No newline at end of file