From a6bda4b0d1e701a38adb0a8219fe662e8265faf0 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Fri, 6 Mar 2026 16:25:37 +0000 Subject: [PATCH 1/2] Add merge_into hook to TableProvider trait Add merge_into async method to TableProvider trait for MERGE INTO DML support. The method accepts: - source: ExecutionPlan representing the USING clause - on: Expr representing the ON join condition - clauses: Vec for WHEN MATCHED/NOT MATCHED actions Default implementation returns not_impl_err for tables that don't support MERGE INTO operations. --- datafusion/catalog/src/table.rs | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/datafusion/catalog/src/table.rs b/datafusion/catalog/src/table.rs index c6468fd5ad131..43eebcb36adca 100644 --- a/datafusion/catalog/src/table.rs +++ b/datafusion/catalog/src/table.rs @@ -28,7 +28,7 @@ use datafusion_common::{Result, internal_err}; use datafusion_expr::Expr; use datafusion_expr::statistics::StatisticsRequest; -use datafusion_expr::dml::InsertOp; +use datafusion_expr::dml::{InsertOp, MergeIntoClause}; use datafusion_expr::{ CreateExternalTable, LogicalPlan, TableProviderFilterPushDown, TableType, }; @@ -379,6 +379,23 @@ pub trait TableProvider: Any + Debug + Sync + Send { async fn truncate(&self, _state: &dyn Session) -> Result> { not_impl_err!("TRUNCATE not supported for {} table", self.table_type()) } + + /// Merge rows from a source into this table. + /// + /// The `source` is an [`ExecutionPlan`] representing the USING clause. + /// The `on` condition is the join predicate from the ON clause. + /// The `clauses` describe the WHEN MATCHED / WHEN NOT MATCHED actions. + /// + /// Returns an [`ExecutionPlan`] producing a single row with `count` (UInt64). + async fn merge_into( + &self, + _state: &dyn Session, + _source: Arc, + _on: Expr, + _clauses: Vec, + ) -> Result> { + not_impl_err!("MERGE INTO not supported for {} table", self.table_type()) + } } impl dyn TableProvider { From e3b27a5d6c20603d75a77c7a89b69ceba37eccdb Mon Sep 17 00:00:00 2001 From: xuanyili Date: Fri, 6 Mar 2026 16:31:08 +0000 Subject: [PATCH 2/2] Add SQL and physical planner support for MERGE INTO Implement merge_to_plan and merge_clause_to_plan in SQL planner: - Parse Statement::Merge into LogicalPlan::Dml with WriteOp::MergeInto - Resolve target table and plan source (USING clause) as LogicalPlan - Build combined schema for target + source to resolve ON and WHEN expressions - Convert ON condition and WHEN clauses to DataFusion Expr - Handle UPDATE, INSERT, and DELETE actions in WHEN clauses Add physical planner dispatch for WriteOp::MergeInto: - Use source_as_provider() to recover the TableProvider from the TableSource - Extract source ExecutionPlan from children - Call TableProvider::merge_into with source plan, ON condition, and clauses - Wrap errors with MERGE INTO operation context Wire MergeInto's expressions through LogicalPlan tree-traversal so optimizers can rewrite them: add MergeIntoOp::exprs() (stable iteration order: on, then per-clause predicate + action value Exprs) and MergeIntoOp::with_new_exprs() to rebuild the op from a transformed expr vector. Branch LogicalPlan::apply_expressions, map_expressions, and with_new_exprs on WriteOp::MergeInto to use these helpers; other WriteOp variants continue to expose no expressions as before. --- datafusion/core/src/physical_planner.rs | 20 ++ datafusion/core/tests/sql/sql_api.rs | 31 ++ datafusion/expr/src/logical_plan/dml.rs | 151 +++++++++- datafusion/expr/src/logical_plan/plan.rs | 14 +- datafusion/expr/src/logical_plan/tree_node.rs | 27 +- .../src/analyzer/function_rewrite.rs | 19 +- .../optimizer/src/analyzer/type_coercion.rs | 77 ++++- datafusion/proto/src/logical_plan/mod.rs | 7 +- datafusion/sql/src/statement.rs | 270 +++++++++++++++++- datafusion/sql/tests/sql_integration.rs | 48 ++++ 10 files changed, 652 insertions(+), 12 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index b6d28e7b21c79..1f335b9b701c7 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -927,6 +927,26 @@ impl DefaultPhysicalPlanner { ); } } + LogicalPlan::Dml(DmlStatement { + table_name, + target, + op: WriteOp::MergeInto(merge_op), + .. + }) => { + let provider = source_as_provider(target)?; + let input_exec = children.one()?; + provider + .merge_into( + session_state, + input_exec, + merge_op.on.clone(), + merge_op.clauses.clone(), + ) + .await + .map_err(|e| { + e.context(format!("MERGE INTO operation on table '{table_name}'")) + })? + } LogicalPlan::Window(Window { window_expr, .. }) => { assert_or_internal_err!( !window_expr.is_empty(), diff --git a/datafusion/core/tests/sql/sql_api.rs b/datafusion/core/tests/sql/sql_api.rs index e3180210ca46b..4f98cd0d34c80 100644 --- a/datafusion/core/tests/sql/sql_api.rs +++ b/datafusion/core/tests/sql/sql_api.rs @@ -208,6 +208,37 @@ async fn ddl_can_not_be_planned_by_session_state() { ); } +#[tokio::test] +async fn merge_into_rejects_source_alias_colliding_with_target_name() { + // Regression test: `MERGE INTO target AS t USING source AS target ...` + // aliases the source to the target table's real name. Both `id` columns + // are identically named. Before the fix, canonicalizing the aliased target + // reference `t.id` to `target.id` collapsed it onto the source's + // `target.id`, so `t.id = target.id` silently became `target.id = target.id` + // (a comparison of the source column with itself). The planner must reject + // this collision instead of changing the meaning of the condition. + let ctx = SessionContext::new(); + ctx.sql("CREATE TABLE target (id INT, val INT)") + .await + .unwrap(); + ctx.sql("CREATE TABLE source (id INT, val INT)") + .await + .unwrap(); + + let err = ctx + .sql( + "MERGE INTO target AS t USING source AS target ON t.id = target.id \ + WHEN MATCHED THEN DELETE", + ) + .await + .unwrap_err(); + + assert_contains!( + err.strip_backtrace(), + "MERGE source may not use the target table name 'target' as a qualifier" + ); +} + #[tokio::test] async fn invalid_wrapped_negation_fails_during_planning() { let ctx = SessionContext::new(); diff --git a/datafusion/expr/src/logical_plan/dml.rs b/datafusion/expr/src/logical_plan/dml.rs index 5b6403e6e2f08..7717dfaff7a33 100644 --- a/datafusion/expr/src/logical_plan/dml.rs +++ b/datafusion/expr/src/logical_plan/dml.rs @@ -23,7 +23,7 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::file_options::file_type::FileType; -use datafusion_common::{DFSchemaRef, TableReference}; +use datafusion_common::{DFSchemaRef, Result, TableReference, internal_err}; use crate::{Expr, LogicalPlan, TableSource}; @@ -307,6 +307,106 @@ pub struct MergeIntoOp { pub clauses: Vec, } +impl MergeIntoOp { + /// Count of top-level [`Expr`]s owned by this operation (no allocation). + /// + /// Matches the length of [`Self::exprs`] and the `exprs` vec consumed by + /// [`Self::with_new_exprs`]. + fn expr_count(&self) -> usize { + 1 + self + .clauses + .iter() + .map(|c| { + c.predicate.is_some() as usize + + match &c.action { + MergeIntoAction::Update(a) => a.len(), + MergeIntoAction::Insert { values, .. } => values.len(), + MergeIntoAction::Delete => 0, + } + }) + .sum::() + } + + /// Top-level [`Expr`]s in stable order: `on`, then per-clause predicate + /// (if any) and action value expressions. + pub fn exprs(&self) -> Vec<&Expr> { + let mut out = Vec::with_capacity(self.expr_count()); + out.push(&self.on); + for clause in &self.clauses { + if let Some(predicate) = &clause.predicate { + out.push(predicate); + } + match &clause.action { + MergeIntoAction::Update(assignments) => { + out.extend(assignments.iter().map(|(_, value)| value)); + } + MergeIntoAction::Insert { values, .. } => { + out.extend(values.iter()); + } + MergeIntoAction::Delete => {} + } + } + out + } + + /// Rebuild this `MergeIntoOp` from a flat vector of new expressions, in + /// the same order produced by [`Self::exprs`]. The clause kinds, action + /// kinds, column lists, and presence/absence of each predicate are + /// preserved from `self`. + pub fn with_new_exprs(&self, exprs: Vec) -> Result { + let expected = self.expr_count(); + if exprs.len() != expected { + return internal_err!( + "MergeIntoOp::with_new_exprs expected {expected} expressions, got {}", + exprs.len() + ); + } + let mut iter = exprs.into_iter(); + let on = iter.next().expect("non-empty by length check"); + let clauses = self + .clauses + .iter() + .map(|clause| { + let predicate = clause + .predicate + .is_some() + .then(|| iter.next().expect("non-empty by length check")); + let action = match &clause.action { + MergeIntoAction::Update(assignments) => { + let assignments = assignments + .iter() + .map(|(name, _)| { + ( + name.clone(), + iter.next().expect("non-empty by length check"), + ) + }) + .collect(); + MergeIntoAction::Update(assignments) + } + MergeIntoAction::Insert { columns, values } => { + let values = values + .iter() + .map(|_| iter.next().expect("non-empty by length check")) + .collect(); + MergeIntoAction::Insert { + columns: columns.clone(), + values, + } + } + MergeIntoAction::Delete => MergeIntoAction::Delete, + }; + MergeIntoClause { + kind: clause.kind, + predicate, + action, + } + }) + .collect(); + Ok(Self { on, clauses }) + } +} + /// A single WHEN clause within a MERGE INTO statement. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] pub struct MergeIntoClause { @@ -445,4 +545,53 @@ mod tests { MergeIntoClauseKind::NotMatchedBySource ); } + + #[test] + fn merge_into_op_exprs_round_trip() { + let op = MergeIntoOp { + on: col("id").eq(col("source_id")), + clauses: vec![ + MergeIntoClause { + kind: MergeIntoClauseKind::Matched, + predicate: Some(col("qty").gt(lit(0_i64))), + action: MergeIntoAction::Update(vec![ + ("qty".to_string(), col("source_qty")), + ("price".to_string(), col("source_price")), + ]), + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatched, + predicate: None, + action: MergeIntoAction::Insert { + columns: vec!["id".to_string(), "qty".to_string()], + values: vec![col("source_id"), col("source_qty")], + }, + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatchedBySource, + predicate: Some(col("active").eq(lit(true))), + action: MergeIntoAction::Delete, + }, + ], + }; + let exprs = op.exprs(); + assert_eq!(exprs.len(), 7); + + let owned: Vec = exprs.into_iter().cloned().collect(); + let rebuilt = op.with_new_exprs(owned).unwrap(); + assert_eq!(op, rebuilt); + } + + #[test] + fn merge_into_op_with_new_exprs_length_mismatch() { + let op = MergeIntoOp { + on: col("id").eq(col("source_id")), + clauses: vec![], + }; + let err = op.with_new_exprs(vec![]).unwrap_err(); + assert!( + err.to_string().contains("expected 1 expressions, got 0"), + "unexpected error: {err}" + ); + } } diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index b6e6cc7683664..40e5241cd887d 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -39,7 +39,7 @@ use crate::expr_rewriter::{ }; use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor}; use crate::logical_plan::extension::UserDefinedLogicalNode; -use crate::logical_plan::{DmlStatement, Statement}; +use crate::logical_plan::{DmlStatement, Statement, WriteOp}; use crate::utils::{ enumerate_grouping_sets, exprlist_to_fields, find_out_reference_exprs, grouping_set_expr_count, grouping_set_to_exprlist, merge_schema, split_conjunction, @@ -810,12 +810,20 @@ impl LogicalPlan { op, .. }) => { - self.assert_no_expressions(expr)?; let input = self.only_input(inputs)?; + let op = match op { + WriteOp::MergeInto(merge_op) => { + WriteOp::MergeInto(Box::new(merge_op.with_new_exprs(expr)?)) + } + other => { + self.assert_no_expressions(expr)?; + other.clone() + } + }; Ok(LogicalPlan::Dml(DmlStatement::new( table_name.clone(), Arc::clone(target), - op.clone(), + op, Arc::new(input), ))) } diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index c10ac92eef4f5..c4c1d743b58b6 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -45,7 +45,7 @@ use crate::{ DistinctOn, DmlStatement, Execute, Explain, Expr, Extension, Filter, Join, Limit, LogicalPlan, Partitioning, Prepare, Projection, RecursiveQuery, Repartition, Sort, Statement, Subquery, SubqueryAlias, TableScan, Union, Unnest, UserDefinedLogicalNode, - Values, Window, builder::unnest_with_options, dml::CopyTo, + Values, Window, WriteOp, builder::unnest_with_options, dml::CopyTo, }; use datafusion_common::tree_node::TreeNodeRefContainer; @@ -480,6 +480,10 @@ impl LogicalPlan { } _ => Ok(TreeNodeRecursion::Continue), }, + LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(merge_op), + .. + }) => merge_op.exprs().apply_ref_elements(f), // plans without expressions LogicalPlan::EmptyRelation(_) | LogicalPlan::RecursiveQuery(_) @@ -719,6 +723,27 @@ impl LogicalPlan { ) })? } + LogicalPlan::Dml(DmlStatement { + table_name, + target, + op: WriteOp::MergeInto(merge_op), + input, + output_schema, + }) => { + let owned_exprs: Vec = + merge_op.exprs().into_iter().cloned().collect(); + owned_exprs.map_elements(f)?.transform_data(|new_exprs| { + Ok(Transformed::no(LogicalPlan::Dml(DmlStatement { + table_name, + target, + op: WriteOp::MergeInto(Box::new( + merge_op.with_new_exprs(new_exprs)?, + )), + input, + output_schema, + }))) + })? + } // plans without expressions LogicalPlan::EmptyRelation(_) | LogicalPlan::RecursiveQuery(_) diff --git a/datafusion/optimizer/src/analyzer/function_rewrite.rs b/datafusion/optimizer/src/analyzer/function_rewrite.rs index 9faa60d939fe3..a66e3ccc0cf8a 100644 --- a/datafusion/optimizer/src/analyzer/function_rewrite.rs +++ b/datafusion/optimizer/src/analyzer/function_rewrite.rs @@ -23,9 +23,9 @@ use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{DFSchema, Result}; use crate::utils::NamePreserver; -use datafusion_expr::LogicalPlan; use datafusion_expr::expr_rewriter::FunctionRewrite; use datafusion_expr::utils::merge_schema; +use datafusion_expr::{DmlStatement, LogicalPlan, WriteOp}; use std::sync::Arc; /// Analyzer rule that invokes [`FunctionRewrite`]s on expressions @@ -58,6 +58,23 @@ impl ApplyFunctionRewrites { schema.merge(&source_schema); } + // MERGE expressions reference the target table, which is not one of + // `plan.inputs()`. Rebuild the target schema from the DML's + // `table_name` and `target` so those columns resolve. + if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + let target_schema = DFSchema::try_from_qualified_schema( + table_name.clone(), + &target.schema(), + )?; + schema.merge(&target_schema); + } + let name_preserver = NamePreserver::new(&plan); plan.map_expressions(|expr| { diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 2503fc807207f..ff334415d2e96 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -57,9 +57,10 @@ use datafusion_expr::type_coercion::{ }; use datafusion_expr::utils::merge_schema; use datafusion_expr::{ - Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, Projection, Union, - ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, is_false, - is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, lit, not, + Cast, DmlStatement, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, + Projection, Union, ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, + WriteOp, is_false, is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, + lit, not, }; /// Performs type coercion by determining the schema @@ -128,6 +129,21 @@ fn analyze_internal( schema.merge(&source_schema); } + // MERGE expressions (ON / WHEN clauses) reference the target table, which + // is not one of `plan.inputs()`. Rebuild the target schema from the DML's + // `table_name` and `target` so those columns resolve during coercion. + if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + let target_schema = + DFSchema::try_from_qualified_schema(table_name.clone(), &target.schema())?; + schema.merge(&target_schema); + } + // merge the outer schema for correlated subqueries // like case: // select t2.c2 from t1 where t1.c1 in (select t2.c1 from t2 where t2.c2=t1.c3) @@ -1487,6 +1503,61 @@ mod test { ) } + #[test] + fn merge_into_resolves_and_coerces_target_and_source_columns() -> Result<()> { + use datafusion_expr::dml::{ + MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, + }; + use datafusion_expr::logical_plan::table_scan; + use datafusion_expr::{DmlStatement, WriteOp}; + + // Target table `target(id: UInt32)`. + let target_table_name = TableReference::bare("target"); + let target_arrow_schema = + Schema::new(vec![Field::new("id", DataType::UInt32, false)]); + let target_plan = + table_scan(Some(target_table_name.clone()), &target_arrow_schema, None)? + .build()?; + let target_source = match &target_plan { + LogicalPlan::TableScan(ts) => Arc::clone(&ts.source), + _ => unreachable!("table_scan() always builds a TableScan"), + }; + + // Source plan `source(id: Int64)` — deliberately a different numeric + // type than `target.id` so the `ON` comparison needs a CAST. + let source_arrow_schema = + Schema::new(vec![Field::new("id", DataType::Int64, false)]); + let source_plan = + table_scan(Some("source"), &source_arrow_schema, None)?.build()?; + + // `ON target.id = source.id`. Resolving `target.id` requires the + // target schema to be visible to the analyzer, which only sees + // `plan.inputs()` (the source plan) by default. + let on = col("target.id").eq(col("source.id")); + let merge_op = MergeIntoOp { + on, + clauses: vec![MergeIntoClause { + kind: MergeIntoClauseKind::Matched, + predicate: None, + action: MergeIntoAction::Delete, + }], + }; + let plan = LogicalPlan::Dml(DmlStatement::new( + target_table_name, + target_source, + WriteOp::MergeInto(Box::new(merge_op)), + Arc::new(source_plan), + )); + + assert_analyzed_plan_eq!( + plan, + @r" + Dml: op=[MergeInto] table=[target] + TableScan: source + " + ) + } + #[test] fn coerce_utf8view_output() -> Result<()> { // Plan A diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 3195b050b3056..465da6bcf23e1 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -1258,11 +1258,14 @@ impl AsLogicalPlan for LogicalPlanNode { .build() } LogicalPlanType::Dml(dml_node) => { + let table_name = + from_table_reference(dml_node.table_name.as_ref(), "DML ")?; + let target = to_table_source(&dml_node.target, ctx, extension_codec)?; let write_op = from_proto::parse_write_op(dml_node, ctx, extension_codec)?; Ok(LogicalPlan::Dml(datafusion_expr::DmlStatement::new( - from_table_reference(dml_node.table_name.as_ref(), "DML ")?, - to_table_source(&dml_node.target, ctx, extension_codec)?, + table_name, + target, write_op, Arc::new(into_logical_plan!(dml_node.input, ctx, extension_codec)?), ))) diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 838228a3e0381..5a7997c0c4eb8 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -33,13 +33,16 @@ use arrow::datatypes::{Field, FieldRef, Fields}; use datafusion_common::error::_plan_err; use datafusion_common::format::ExplainStatementOptions; use datafusion_common::parsers::CompressionTypeVariant; +use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{ Column, Constraint, Constraints, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, SchemaError, SchemaReference, TableReference, ToDFSchema, exec_err, internal_err, not_impl_err, plan_datafusion_err, plan_err, schema_err, unqualified_field_not_found, }; -use datafusion_expr::dml::{CopyTo, InsertOp}; +use datafusion_expr::dml::{ + CopyTo, InsertOp, MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, +}; use datafusion_expr::expr_rewriter::normalize_col_with_schemas_and_ambiguity_check; use datafusion_expr::logical_plan::DdlStatement; use datafusion_expr::logical_plan::builder::project; @@ -1216,6 +1219,8 @@ impl SqlToRel<'_, S> { self.delete_to_plan(&table_name, selection, limit) } + Statement::Merge(merge) => self.merge_to_plan(merge), + Statement::StartTransaction { modes, begin: false, @@ -2406,6 +2411,269 @@ impl SqlToRel<'_, S> { Ok(plan) } + fn merge_to_plan(&self, merge: ast::Merge) -> Result { + let ast::Merge { + table, + source, + on, + clauses, + into: _, + merge_token: _, + optimizer_hints, + output, + } = merge; + + if !optimizer_hints.is_empty() { + plan_err!("Optimizer hints not supported")?; + } + + if output.is_some() { + return not_impl_err!("MERGE OUTPUT clause is not supported"); + } + + // 1. Resolve target table + let (target_table_name, target_alias) = match &table { + TableFactor::Table { name, alias, .. } => (name.clone(), alias.clone()), + _ => plan_err!("Cannot MERGE INTO non-table relation!")?, + }; + let target_table_ref = self.object_name_to_table_reference(target_table_name)?; + let target_table_source = self + .context_provider + .get_table_source(target_table_ref.clone())?; + // Use alias as schema qualifier so `t.col` resolves when user writes + // `MERGE INTO target AS t`. Fall back to the table reference itself. + let target_qualifier = target_alias + .as_ref() + .map(|a| { + TableReference::bare(self.ident_normalizer.normalize(a.name.clone())) + }) + .unwrap_or_else(|| target_table_ref.clone()); + let target_schema = Arc::new(DFSchema::try_from_qualified_schema( + target_qualifier.clone(), + &target_table_source.schema(), + )?); + + // 2. Plan the source (USING clause) as a LogicalPlan + let mut planner_context = PlannerContext::new(); + let source_table_with_joins = TableWithJoins { + relation: source, + joins: vec![], + }; + let source_plan = + self.plan_from_tables(vec![source_table_with_joins], &mut planner_context)?; + + // 3. Build a combined schema for resolving expressions in ON and WHEN clauses + let combined_schema = + Arc::new(target_schema.as_ref().join(source_plan.schema())?); + + // 4. Convert the ON condition from sqlparser Expr to datafusion Expr + let on_expr = self.sql_to_expr(*on, &combined_schema, &mut planner_context)?; + + // 5. Convert each WHEN clause + let df_clauses = clauses + .into_iter() + .map(|clause| { + self.merge_clause_to_plan( + clause, + &combined_schema, + &target_schema, + &target_alias, + &mut planner_context, + ) + }) + .collect::>>()?; + + // 6. Build the MERGE operation. Column references to the target may be + // qualified with the SQL alias (`MERGE INTO target AS t ... t.col`). + // Canonicalize those to the real target table qualifier so the stored + // plan is independent of the alias: this lets the analyzer passes and + // proto deserialization rebuild the target schema from `table_name` + // alone, without carrying the alias as extra state. + let mut merge_op = MergeIntoOp { + on: on_expr, + clauses: df_clauses, + }; + if target_qualifier != target_table_ref { + // Canonicalizing target columns to `target_table_ref` is only safe + // when the source does not already use that qualifier. If it does + // (e.g. `MERGE INTO target AS t USING source AS target`), the two + // namespaces would collapse and later resolution could silently + // pick the source column for a target reference. Reject that + // collision rather than change the meaning of the condition. + if source_plan + .schema() + .iter() + .any(|(qualifier, _)| qualifier == Some(&target_table_ref)) + { + return plan_err!( + "MERGE source may not use the target table name '{target_table_ref}' \ + as a qualifier while the target is aliased as '{target_qualifier}'; \ + use a different source alias" + ); + } + let canonical = merge_op + .exprs() + .into_iter() + .cloned() + .map(|expr| { + Self::canonicalize_target_qualifier( + expr, + &target_qualifier, + &target_table_ref, + ) + }) + .collect::>>()?; + merge_op = merge_op.with_new_exprs(canonical)?; + } + + Ok(LogicalPlan::Dml(DmlStatement::new( + target_table_ref, + target_table_source, + WriteOp::MergeInto(Box::new(merge_op)), + Arc::new(source_plan), + ))) + } + + /// Rewrite every [`Expr::Column`] qualified with `from` to instead use + /// `to`, leaving all other columns untouched. Used to canonicalize MERGE + /// target-alias references to the real target table qualifier. + fn canonicalize_target_qualifier( + expr: Expr, + from: &TableReference, + to: &TableReference, + ) -> Result { + expr.transform(|expr| match expr { + Expr::Column(col) if col.relation.as_ref() == Some(from) => Ok( + Transformed::yes(Expr::Column(Column::new(Some(to.clone()), col.name))), + ), + other => Ok(Transformed::no(other)), + }) + .map(|transformed| transformed.data) + } + + fn ident_from_object_name_last(name: &ObjectName) -> Result { + let part = name + .0 + .iter() + .last() + .ok_or_else(|| plan_datafusion_err!("Empty column name"))?; + part.as_ident() + .cloned() + .ok_or_else(|| plan_datafusion_err!("Expected simple identifier")) + } + + fn merge_clause_to_plan( + &self, + clause: ast::MergeClause, + combined_schema: &DFSchema, + target_schema: &DFSchema, + _target_alias: &Option, + planner_context: &mut PlannerContext, + ) -> Result { + let kind = match clause.clause_kind { + ast::MergeClauseKind::Matched => MergeIntoClauseKind::Matched, + ast::MergeClauseKind::NotMatched => MergeIntoClauseKind::NotMatched, + ast::MergeClauseKind::NotMatchedByTarget => { + MergeIntoClauseKind::NotMatchedByTarget + } + ast::MergeClauseKind::NotMatchedBySource => { + MergeIntoClauseKind::NotMatchedBySource + } + }; + + let predicate = clause + .predicate + .map(|p| self.sql_to_expr(p, combined_schema, planner_context)) + .transpose()?; + + let action = match clause.action { + ast::MergeAction::Update(update_expr) => { + let assignments = update_expr + .assignments + .into_iter() + .map(|assign| { + let col_name = match &assign.target { + AssignmentTarget::ColumnName(cols) => { + let ident = Self::ident_from_object_name_last(cols)?; + self.ident_normalizer.normalize(ident) + } + _ => plan_err!("Tuples are not supported")?, + }; + // Validate column exists in target + target_schema.field_with_unqualified_name(&col_name)?; + let value = self.sql_to_expr( + assign.value, + combined_schema, + planner_context, + )?; + Ok((col_name, value)) + }) + .collect::>>()?; + MergeIntoAction::Update(assignments) + } + ast::MergeAction::Insert(insert_expr) => { + let columns: Vec = insert_expr + .columns + .iter() + .map(|c| { + let ident = Self::ident_from_object_name_last(c)?; + Ok(self.ident_normalizer.normalize(ident)) + }) + .collect::>>()?; + + // Validate: no duplicates, all columns exist in target schema + let mut seen = HashSet::new(); + for col in &columns { + if !seen.insert(col.as_str()) { + return plan_err!("Duplicate column '{col}' in MERGE INSERT"); + } + target_schema.field_with_unqualified_name(col)?; + } + + let num_target_cols = target_schema.fields().len(); + + let values = match insert_expr.kind { + ast::MergeInsertKind::Values(values) => { + if values.rows.len() != 1 { + return plan_err!( + "MERGE INSERT must have exactly one row of values" + ); + } + let row = values.rows.into_iter().next().unwrap().content; + let expected = if columns.is_empty() { + num_target_cols + } else { + columns.len() + }; + if row.len() != expected { + return plan_err!( + "MERGE INSERT has {expected} column(s) but {} value(s)", + row.len() + ); + } + row.into_iter() + .map(|v| { + self.sql_to_expr(v, combined_schema, planner_context) + }) + .collect::>>()? + } + ast::MergeInsertKind::Row => { + return not_impl_err!("MERGE INSERT ROW is not supported"); + } + }; + + MergeIntoAction::Insert { columns, values } + } + ast::MergeAction::Delete { .. } => MergeIntoAction::Delete, + }; + + Ok(MergeIntoClause { + kind, + predicate, + action, + }) + } + fn insert_to_plan( &self, table_name: ObjectName, diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 88b7b43eb73f6..7ff5080e4308a 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -3502,6 +3502,54 @@ fn select_groupby_orderby_aggregate_on_non_selected_column_original_issue() { ); } +#[test] +fn plan_merge_into_canonicalizes_target_alias() { + // The target alias `t` must be rewritten to the real table name `j1` in the + // stored plan, while the source alias `s` is preserved. This keeps the plan + // independent of the SQL alias so analyzer passes and proto deserialization + // can rebuild the target schema from the table name alone. + let sql = "MERGE INTO j1 AS t USING j2 AS s ON t.j1_id = s.j2_id \ + WHEN MATCHED THEN UPDATE SET j1_string = s.j2_string \ + WHEN NOT MATCHED THEN INSERT (j1_id, j1_string) VALUES (s.j2_id, s.j2_string)"; + let plan = logical_plan(sql).unwrap(); + let LogicalPlan::Dml(dml) = &plan else { + panic!("expected Dml, got {plan:?}"); + }; + let datafusion_expr::WriteOp::MergeInto(merge_op) = &dml.op else { + panic!("expected MergeInto, got {:?}", dml.op); + }; + // `t.j1_id` -> `j1.j1_id`, `s.j2_id` left untouched. + assert_eq!(merge_op.on.to_string(), "j1.j1_id = s.j2_id"); + // Source-side column references remain qualified with the source alias. + let exprs: Vec = merge_op.exprs().iter().map(|e| e.to_string()).collect(); + assert_eq!( + exprs, + vec![ + "j1.j1_id = s.j2_id".to_string(), + "s.j2_string".to_string(), + "s.j2_id".to_string(), + "s.j2_string".to_string(), + ] + ); +} + +#[test] +fn plan_merge_into_rejects_source_alias_colliding_with_target_name() { + // Aliasing the source to the target table's real name (`USING j2 AS j1`, + // where `j1` is the target) would make canonicalizing the target alias + // `t` to `j1` collide with the source qualifier, silently collapsing the + // two namespaces. The planner must reject this rather than misresolve. + let sql = "MERGE INTO j1 AS t USING j2 AS j1 ON t.j1_id = j1.j2_id \ + WHEN MATCHED THEN DELETE"; + let err = logical_plan(sql).unwrap_err(); + assert!( + err.strip_backtrace().contains( + "MERGE source may not use the target table name 'j1' as a qualifier" + ), + "unexpected error: {err}" + ); +} + fn logical_plan(sql: &str) -> Result { logical_plan_with_options(sql, ParserOptions::default()) }