fix: support joined update and delete

src/binder/delete.rs

@@ -12,16 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::binder::{lower_case_name, Binder, Source};
+use crate::binder::{lower_case_name, Binder};
 use crate::errors::DatabaseError;
 use crate::planner::operator::delete::DeleteOperator;
-use crate::planner::operator::table_scan::TableScanOperator;
 use crate::planner::operator::Operator;
 use crate::planner::{Childrens, LogicalPlan};
 use crate::storage::Transaction;
 use crate::types::value::DataValue;
 use itertools::Itertools;
-use sqlparser::ast::{Expr, TableAlias, TableFactor, TableWithJoins};
+use sqlparser::ast::{Expr, TableFactor, TableWithJoins};
 use std::sync::Arc;
 
 impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A> {
@@ -30,33 +29,19 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
         from: &TableWithJoins,
         selection: &Option<Expr>,
     ) -> Result<LogicalPlan, DatabaseError> {
-        if let TableFactor::Table { name, alias, .. } = &from.relation {
+        if let TableFactor::Table { name, .. } = &from.relation {
             let table_name: Arc<str> = lower_case_name(name)?.into();
-            let mut table_alias = None;
-            let mut alias_idents = None;
-
-            if let Some(TableAlias { name, columns, .. }) = alias {
-                table_alias = Some(name.value.to_lowercase().into());
-                alias_idents = Some(columns);
-            }
-            let Source::Table(table) = self
+            let table = self
                 .context
-                .source_and_bind(table_name.clone(), table_alias.as_ref(), None, true)?
-                .ok_or(DatabaseError::TableNotFound)?
-            else {
-                unreachable!()
-            };
+                .table(table_name.clone())?
+                .ok_or(DatabaseError::TableNotFound)?;
             let primary_keys = table
                 .primary_keys()
                 .iter()
                 .map(|(_, column)| column.clone())
                 .collect_vec();
-            let mut plan = TableScanOperator::build(table_name.clone(), table, true)?;
-
-            if let Some(alias_idents) = alias_idents {
-                plan =
-                    self.bind_alias(plan, alias_idents, table_alias.unwrap(), table_name.clone())?;
-            }
+            self.with_pk(table_name.clone());
+            let mut plan = self.bind_table_ref(from)?;
 
             if let Some(predicate) = selection {
                 plan = self.bind_where(plan, predicate)?;

src/binder/mod.rs

@@ -573,23 +573,15 @@ impl<'a, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'a, '
             }
             Statement::Update(update) => {
                 let table = &update.table;
-                if !table.joins.is_empty() {
-                    unimplemented!()
-                } else {
-                    self.bind_update(table, &update.selection, &update.assignments)?
-                }
+                self.bind_update(table, &update.selection, &update.assignments)?
             }
             Statement::Delete(delete) => {
                 let from = match &delete.from {
                     FromTable::WithFromKeyword(from) | FromTable::WithoutKeyword(from) => from,
                 };
                 let table = &from[0];
 
-                if !table.joins.is_empty() {
-                    unimplemented!()
-                } else {
-                    self.bind_delete(table, &delete.selection)?
-                }
+                self.bind_delete(table, &delete.selection)?
             }
             Statement::Analyze(analyze) => {
                 let table_name = analyze.table_name.as_ref().ok_or_else(|| {

src/binder/update.rs

@@ -16,7 +16,9 @@ use crate::binder::{
     attach_span_from_sqlparser_span_if_absent, attach_span_if_absent, lower_case_name, Binder,
 };
 use crate::errors::DatabaseError;
+use crate::expression::visitor_mut::VisitorMut;
 use crate::expression::ScalarExpression;
+use crate::planner::operator::project::ProjectOperator;
 use crate::planner::operator::update::UpdateOperator;
 use crate::planner::operator::Operator;
 use crate::planner::{Childrens, LogicalPlan};
@@ -29,6 +31,30 @@ use std::borrow::Cow;
 use std::slice;
 use std::sync::Arc;
 
+struct UpdateExprTargetRemapper<'a> {
+    target_schema: &'a [crate::catalog::ColumnRef],
+}
+
+impl VisitorMut<'_> for UpdateExprTargetRemapper<'_> {
+    fn visit_column_ref(
+        &mut self,
+        column: &mut crate::catalog::ColumnRef,
+        position: &mut usize,
+    ) -> Result<(), DatabaseError> {
+        let Some(target_position) = self
+            .target_schema
+            .iter()
+            .position(|target_column| target_column.same_column(column))
+        else {
+            return Err(DatabaseError::UnsupportedStmt(
+                "joined UPDATE SET expressions can only reference target table columns".to_string(),
+            ));
+        };
+        *position = target_position;
+        Ok(())
+    }
+}
+
 impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A> {
     fn single_ident_from_object_name(name: &ObjectName) -> Result<&Ident, DatabaseError> {
         if name.0.len() != 1 {
@@ -51,10 +77,16 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
         // FIXME: Make it better to detect the current BindStep
         self.context.allow_default = true;
         if let TableFactor::Table { name, .. } = &to.relation {
+            let is_joined_update = !to.joins.is_empty();
             let table_name: Arc<str> = lower_case_name(name)?.into();
             self.with_pk(table_name.clone());
 
             let mut plan = self.bind_table_ref(to)?;
+            let (target_schema, target_offset) = Self::resolve_source_columns_in_scope(
+                &self.context,
+                &mut self.table_schema_buf,
+                &table_name,
+            )?;
 
             if let Some(predicate) = selection {
                 plan = self.bind_where(plan, predicate)?;
@@ -96,6 +128,12 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
                                 expr,
                                 Cow::Borrowed(column.datatype()),
                             )?;
+                            if is_joined_update {
+                                UpdateExprTargetRemapper {
+                                    target_schema: &target_schema,
+                                }
+                                .visit(&mut expr)?;
+                            }
                             value_exprs.push((column, expr));
                         }
                         _ => {
@@ -108,6 +146,19 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
                 }
             }
             self.context.allow_default = false;
+            if is_joined_update {
+                let exprs = target_schema
+                    .iter()
+                    .enumerate()
+                    .map(|(index, column)| {
+                        ScalarExpression::column_expr(column.clone(), target_offset + index)
+                    })
+                    .collect();
+                plan = LogicalPlan::new(
+                    Operator::Project(ProjectOperator { exprs }),
+                    Childrens::Only(Box::new(plan)),
+                );
+            }
             Ok(LogicalPlan::new(
                 Operator::Update(UpdateOperator {
                     table_name,

src/execution/dml/update.rs

@@ -110,7 +110,9 @@ impl Update {
                 let mut tuple = arena.result_tuple().clone();
                 let mut is_overwrite = true;
 
-                let old_pk = tuple.pk.clone().ok_or(DatabaseError::PrimaryKeyNotFound)?;
+                let Some(old_pk) = tuple.pk.clone() else {
+                    continue;
+                };
                 for (index_meta, exprs) in index_metas.iter() {
                     let values = Projection::projection(&tuple, exprs)?;
                     let Some(value) = DataValue::values_to_tuple(values) else {

src/execution/dql/join/hash/full_join.rs

@@ -85,7 +85,10 @@ impl JoinProbeState for FullJoinState {
             );
             build_state.is_used = true;
             build_state.has_filted = probe_state.has_filtered;
-            return Ok(Some(Tuple::new(pk.clone(), full_values)));
+            return Ok(Some(Tuple::new(
+                pk.as_ref().or(probe_state.probe_tuple.pk.as_ref()).cloned(),
+                full_values,
+            )));
         }
 
         build_state.is_used = !probe_state.has_filtered;

src/execution/dql/join/hash/inner_join.rs

@@ -55,7 +55,10 @@ impl JoinProbeState for InnerJoinState {
                     .chain(probe_state.probe_tuple.values.iter())
                     .cloned(),
             );
-            return Ok(Some(Tuple::new(pk.clone(), full_values)));
+            return Ok(Some(Tuple::new(
+                pk.as_ref().or(probe_state.probe_tuple.pk.as_ref()).cloned(),
+                full_values,
+            )));
         }
 
         probe_state.finished = true;

src/execution/dql/join/hash/left_join.rs

@@ -65,7 +65,10 @@ impl JoinProbeState for LeftJoinState {
                     .cloned(),
             );
             build_state.is_used = true;
-            return Ok(Some(Tuple::new(pk.clone(), full_values)));
+            return Ok(Some(Tuple::new(
+                pk.as_ref().or(probe_state.probe_tuple.pk.as_ref()).cloned(),
+                full_values,
+            )));
         }
 
         build_state.is_used = !probe_state.has_filtered;

src/execution/dql/join/hash/right_join.rs

@@ -77,7 +77,10 @@ impl JoinProbeState for RightJoinState {
             probe_state.produced = true;
             build_state.is_used = true;
             build_state.has_filted = probe_state.has_filtered;
-            return Ok(Some(Tuple::new(pk.clone(), full_values)));
+            return Ok(Some(Tuple::new(
+                pk.as_ref().or(probe_state.probe_tuple.pk.as_ref()).cloned(),
+                full_values,
+            )));
         }
 
         build_state.is_used = probe_state.produced;

src/execution/dql/join/hash_join.rs

@@ -70,7 +70,7 @@ enum HashJoinState {
 
 impl From<(JoinOperator, LogicalPlan, LogicalPlan)> for HashJoin {
     fn from(
-        (JoinOperator { on, join_type, .. }, mut left_input, mut right_input): (
+        (JoinOperator { on, join_type }, mut left_input, mut right_input): (
             JoinOperator,
             LogicalPlan,
             LogicalPlan,

src/execution/dql/join/nested_loop_join.rs

@@ -120,7 +120,7 @@ struct ActiveLeftState {
 
 impl From<(JoinOperator, LogicalPlan, LogicalPlan)> for NestedLoopJoin {
     fn from(
-        (JoinOperator { on, join_type, .. }, left_input, right_input): (
+        (JoinOperator { on, join_type }, left_input, right_input): (
             JoinOperator,
             LogicalPlan,
             LogicalPlan,

src/optimizer/rule/normalization/pushdown_predicates.rs

@@ -73,6 +73,23 @@ fn plan_output_columns(plan: &LogicalPlan) -> Vec<ColumnRef> {
     }
 }
 
+fn localize_right_filters(
+    filters: &mut [ScalarExpression],
+    left_len: usize,
+) -> Result<(), DatabaseError> {
+    if filters.is_empty() {
+        return Ok(());
+    }
+
+    let mut localizer = PositionShift {
+        delta: -(left_len as isize),
+    };
+    for expr in filters {
+        localizer.visit(expr)?;
+    }
+    Ok(())
+}
+
 /// Comments copied from Spark Catalyst PushPredicateThroughJoin
 ///
 /// Pushes down `Filter` operators where the `condition` can be
@@ -138,6 +155,8 @@ impl NormalizationRule for PushPredicateThroughJoin {
                         new_ops.0 = Some(Operator::Filter(left_filter_op));
                     }
 
+                    let mut right_filters = right_filters;
+                    localize_right_filters(&mut right_filters, left_columns.len())?;
                     if let Some(right_filter_op) = reduce_filters(right_filters, filter_op.having) {
                         new_ops.1 = Some(Operator::Filter(right_filter_op));
                     }
@@ -155,6 +174,8 @@ impl NormalizationRule for PushPredicateThroughJoin {
                         .collect_vec()
                 }
                 JoinType::RightOuter => {
+                    let mut right_filters = right_filters;
+                    localize_right_filters(&mut right_filters, left_columns.len())?;
                     if let Some(right_filter_op) = reduce_filters(right_filters, filter_op.having) {
                         new_ops.1 = Some(Operator::Filter(right_filter_op));
                     }
@@ -415,14 +436,7 @@ impl NormalizationRule for PushJoinPredicateIntoScan {
         } else {
             (Vec::new(), right_filters)
         };
-        if !right_push.is_empty() {
-            let mut localizer = PositionShift {
-                delta: -(left_columns.len() as isize),
-            };
-            for expr in &mut right_push {
-                localizer.visit(expr)?;
-            }
-        }
+        localize_right_filters(&mut right_push, left_columns.len())?;
         if let Some(filter_op) = reduce_filters(right_push, false) {
             new_ops.1 = Some(Operator::Filter(filter_op));
         } else {