From 44e50bdd0ee09fa935571143f1db7fa5c7d49d9e Mon Sep 17 00:00:00 2001 From: xuanyili Date: Fri, 6 Mar 2026 16:25:37 +0000 Subject: [PATCH 1/5] 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/session/src/table.rs | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/datafusion/session/src/table.rs b/datafusion/session/src/table.rs index 8d9cd92d4c664..c0b95b4a46c00 100644 --- a/datafusion/session/src/table.rs +++ b/datafusion/session/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 c0c26cb0103652f3bf11c39b8728850e50158129 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Fri, 6 Mar 2026 16:31:08 +0000 Subject: [PATCH 2/5] 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 4e914556b4cc0..9f4d7d60d6b6c 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -967,6 +967,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 9ac27b46a78e6..1a141ea52a13a 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::{ check_aggregate_and_window_nesting, enumerate_grouping_sets, exprlist_to_fields, find_out_reference_exprs, grouping_set_expr_count, grouping_set_to_exprlist, @@ -811,12 +811,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 afd4e980b5424..566522e9273fd 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) @@ -1535,6 +1551,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 732676a3c0a0f..e58e24d281051 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -1270,11 +1270,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 ae7579c8c4dfb..3c71e9b5c713e 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, @@ -2414,6 +2419,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 a4bf0db910774..8334755094a3e 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -3600,6 +3600,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()) } From c4c0544ee09548e6feafae0cfd207c9ad8dc305a Mon Sep 17 00:00:00 2001 From: xuanyili Date: Fri, 24 Jul 2026 07:14:02 +0000 Subject: [PATCH 3/5] Fix MERGE INTO planner edge cases --- datafusion/core/src/physical_planner.rs | 94 ++++++- datafusion/core/tests/sql/sql_api.rs | 104 +++++++- .../expr/src/logical_plan/invariants.rs | 8 +- datafusion/ffi/src/table_provider.rs | 244 +++++++++++++++++- datafusion/ffi/src/tests/mod.rs | 124 ++++++++- datafusion/ffi/tests/ffi_integration.rs | 58 ++++- .../optimizer/src/analyzer/type_coercion.rs | 109 ++++++-- .../optimizer/src/rewrite_set_comparison.rs | 16 +- .../simplify_expressions/simplify_exprs.rs | 17 +- .../proto/src/logical_plan/from_proto.rs | 2 +- datafusion/session/src/table.rs | 7 +- datafusion/sql/src/statement.rs | 172 ++++++++++-- datafusion/sql/tests/sql_integration.rs | 106 ++++++++ 13 files changed, 1005 insertions(+), 56 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 9f4d7d60d6b6c..7b09a61b85d46 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -971,14 +971,23 @@ impl DefaultPhysicalPlanner { table_name, target, op: WriteOp::MergeInto(merge_op), + input, .. }) => { - let provider = source_as_provider(target)?; + let provider = source_as_provider(target).map_err(|e| { + e.context(format!("MERGE INTO operation on table '{table_name}'")) + })?; let input_exec = children.one()?; + let target_schema = DFSchema::try_from_qualified_schema( + table_name.clone(), + &target.schema(), + )?; + let merge_schema = Arc::new(target_schema.join(input.schema())?); provider .merge_into( session_state, input_exec, + merge_schema, merge_op.on.clone(), merge_op.clauses.clone(), ) @@ -3398,6 +3407,7 @@ mod tests { use datafusion_execution::TaskContext; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_expr::builder::subquery_alias; + use datafusion_expr::dml::MergeIntoClause; use datafusion_expr::expr::AggregateFunctionParams; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::{ @@ -3422,6 +3432,88 @@ mod tests { .build() } + #[derive(Debug)] + struct CaptureMergeProvider { + schema: SchemaRef, + captured: Mutex>, + } + + #[async_trait] + impl TableProvider for CaptureMergeProvider { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema)))) + } + + async fn merge_into( + &self, + state: &dyn Session, + source: Arc, + merge_schema: DFSchemaRef, + on: Expr, + clauses: Vec, + ) -> Result> { + let physical_on = state.create_physical_expr(on, &merge_schema)?; + *self.captured.lock().await = + Some((merge_schema, format!("{physical_on:?}"), clauses.len())); + Ok(source) + } + } + + #[tokio::test] + async fn merge_into_provider_receives_combined_logical_schema() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("val", DataType::Int32, true), + ])); + let target = Arc::new(CaptureMergeProvider { + schema: Arc::clone(&schema), + captured: Mutex::new(None), + }); + let source = Arc::new(MemTable::try_new(Arc::clone(&schema), vec![vec![]])?); + let ctx = SessionContext::new(); + ctx.register_table("target", target.clone())?; + ctx.register_table("source", source)?; + + ctx.sql( + "MERGE INTO target AS t USING source AS s ON t.id = s.id \ + WHEN MATCHED AND t.val IS NULL THEN UPDATE SET val = s.val \ + WHEN NOT MATCHED THEN INSERT (id, val) VALUES (s.id, s.val)", + ) + .await? + .create_physical_plan() + .await?; + + let captured = target.captured.lock().await; + let (merge_schema, physical_on, clause_count) = + captured.as_ref().expect("merge_into should be called"); + assert_eq!(*clause_count, 2); + assert_eq!( + merge_schema.index_of_column(&Column::new(Some("target"), "id"))?, + 0 + ); + assert_eq!( + merge_schema.index_of_column(&Column::new(Some("s"), "id"))?, + 2 + ); + assert_contains!(physical_on, "index: 0"); + assert_contains!(physical_on, "index: 2"); + Ok(()) + } + async fn plan(logical_plan: &LogicalPlan) -> Result> { let session_state = make_session_state(); // optimize the logical plan diff --git a/datafusion/core/tests/sql/sql_api.rs b/datafusion/core/tests/sql/sql_api.rs index 4f98cd0d34c80..f578dae12c43b 100644 --- a/datafusion/core/tests/sql/sql_api.rs +++ b/datafusion/core/tests/sql/sql_api.rs @@ -225,17 +225,117 @@ async fn merge_into_rejects_source_alias_colliding_with_target_name() { .await .unwrap(); + for target_ref in ["target", "public.target", "datafusion.public.target"] { + let err = ctx + .sql(&format!( + "MERGE INTO {target_ref} AS t USING source AS target \ + ON t.id = target.id WHEN MATCHED THEN DELETE" + )) + .await + .unwrap_err(); + + assert_contains!( + err.strip_backtrace(), + &format!( + "MERGE source may not use the target table name '{target_ref}' \ + as a qualifier" + ) + ); + } +} + +#[tokio::test] +async fn merge_into_rejects_subqueries_correlated_to_target_alias() { + 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 \ + "MERGE INTO target AS t USING source AS s \ + ON EXISTS (SELECT 1 FROM source AS x WHERE x.id = t.id) \ WHEN MATCHED THEN DELETE", ) .await .unwrap_err(); + assert_contains!( + err.strip_backtrace(), + "MERGE subqueries correlated to target alias 't' are not supported" + ); + + // Source correlation and uncorrelated set-comparison subqueries do not + // require target-alias canonicalization and remain supported through + // logical optimization. + for sql in [ + "MERGE INTO target AS t USING source AS s \ + ON EXISTS (SELECT 1 FROM source AS x WHERE x.id = s.id) \ + WHEN MATCHED THEN DELETE", + "MERGE INTO target AS t USING source AS s \ + ON t.id = ANY (SELECT id FROM source) \ + WHEN MATCHED THEN DELETE", + ] { + let err = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap_err(); + assert_contains!( + err.strip_backtrace(), + "MERGE INTO not supported for Base table" + ); + } +} +#[tokio::test] +async fn merge_into_requires_boolean_conditions() { + 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(); + + for (sql, expected) in [ + ( + "MERGE INTO target USING source ON 1 WHEN MATCHED THEN DELETE", + "MERGE ON condition must be boolean type, but got Int64", + ), + ( + "MERGE INTO target USING source ON true \ + WHEN MATCHED AND 1 THEN DELETE", + "MERGE WHEN condition must be boolean type, but got Int64", + ), + ] { + let err = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap_err(); + assert_contains!(err.strip_backtrace(), expected); + } + + let err = ctx + .sql( + "MERGE INTO target USING source ON NULL \ + WHEN MATCHED AND NULL THEN DELETE", + ) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap_err(); assert_contains!( err.strip_backtrace(), - "MERGE source may not use the target table name 'target' as a qualifier" + "MERGE INTO not supported for Base table" ); } diff --git a/datafusion/expr/src/logical_plan/invariants.rs b/datafusion/expr/src/logical_plan/invariants.rs index 0889afd08fee4..d6867d1ceb112 100644 --- a/datafusion/expr/src/logical_plan/invariants.rs +++ b/datafusion/expr/src/logical_plan/invariants.rs @@ -21,7 +21,7 @@ use datafusion_common::{ }; use crate::{ - Aggregate, Expr, Filter, Join, JoinType, LogicalPlan, Window, + Aggregate, DmlStatement, Expr, Filter, Join, JoinType, LogicalPlan, Window, WriteOp, expr::{Exists, InSubquery, SetComparison}, expr_rewriter::strip_outer_reference, utils::{collect_subquery_cols, split_conjunction}, @@ -253,7 +253,11 @@ pub fn check_subquery_expr( | LogicalPlan::TableScan(_) | LogicalPlan::Window(_) | LogicalPlan::Aggregate(_) - | LogicalPlan::Join(_) => Ok(()), + | LogicalPlan::Join(_) + | LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + .. + }) => Ok(()), _ => plan_err!( "In/Exist/SetComparison subquery can only be used in \ Projection, Filter, TableScan, Window functions, Aggregate and Join plan nodes, \ diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index 5a4b2fa27256f..35c0889c44a24 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -22,18 +22,20 @@ use arrow::datatypes::SchemaRef; use async_ffi::{FfiFuture, FutureExt}; use async_trait::async_trait; use datafusion_catalog::{Session, TableProvider}; -use datafusion_common::Statistics; use datafusion_common::error::{DataFusionError, Result}; +use datafusion_common::{DFSchemaRef, Statistics}; use datafusion_execution::TaskContext; -use datafusion_expr::dml::InsertOp; +use datafusion_expr::dml::{InsertOp, MergeIntoClause, MergeIntoOp}; use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; use datafusion_physical_plan::ExecutionPlan; -use datafusion_proto::logical_plan::from_proto::parse_exprs; -use datafusion_proto::logical_plan::to_proto::serialize_exprs; +use datafusion_proto::logical_plan::from_proto::{parse_exprs, parse_merge_into_op}; +use datafusion_proto::logical_plan::to_proto::{ + serialize_exprs, serialize_merge_into_op, +}; use datafusion_proto::logical_plan::{ DefaultLogicalExtensionCodec, LogicalExtensionCodec, }; -use datafusion_proto::protobuf::LogicalExprList; +use datafusion_proto::protobuf::{DfSchema, LogicalExprList, MergeIntoOpNode}; use prost::Message; use stabby::vec::Vec as SVec; @@ -135,6 +137,17 @@ pub struct FFI_TableProvider { insert_op: FFI_InsertOp, ) -> FfiFuture>, + /// Merge source rows into the table. The schema and MERGE operation are + /// prost-encoded to preserve logical qualifiers and expressions across the + /// FFI boundary. + merge_into: unsafe extern "C" fn( + provider: &Self, + session: FFI_SessionRef, + source: &FFI_ExecutionPlan, + merge_schema_serialized: SVec, + merge_op_serialized: SVec, + ) -> FfiFuture>, + /// Snapshot the provider's table-level statistics. [`FFI_Option::None`] /// corresponds to [`TableProvider::statistics`] returning `None`; /// `Some(bytes)` is a prost-encoded `datafusion_proto_common::Statistics`. @@ -336,6 +349,60 @@ unsafe extern "C" fn insert_into_fn_wrapper( .into_ffi() } +unsafe extern "C" fn merge_into_fn_wrapper( + provider: &FFI_TableProvider, + session: FFI_SessionRef, + source: &FFI_ExecutionPlan, + merge_schema_serialized: SVec, + merge_op_serialized: SVec, +) -> FfiFuture> { + let task_ctx: Result, DataFusionError> = + (&provider.logical_codec.task_ctx_provider).try_into(); + let runtime = provider.runtime().clone(); + let logical_codec: Arc = (&provider.logical_codec).into(); + let internal_provider = Arc::clone(provider.inner()); + let source = source.clone(); + + async move { + let mut foreign_session = None; + let session = sresult_return!( + session + .as_local() + .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>) + .unwrap_or_else(|| { + foreign_session = Some(ForeignSession::try_from(&session)?); + Ok(foreign_session.as_ref().unwrap()) + }) + ); + + let source = sresult_return!(>::try_from(&source)); + let task_ctx = sresult_return!(task_ctx); + + let merge_schema = + sresult_return!(DfSchema::decode(merge_schema_serialized.as_ref())); + let merge_schema: DFSchemaRef = sresult_return!( + DFSchemaRef::try_from(merge_schema) + .map_err(|e| DataFusionError::Plan(e.to_string())) + ); + + let merge_op = + sresult_return!(MergeIntoOpNode::decode(merge_op_serialized.as_ref())); + let MergeIntoOp { on, clauses } = sresult_return!( + parse_merge_into_op(&merge_op, task_ctx.as_ref(), logical_codec.as_ref()) + .map_err(|e| DataFusionError::Plan(e.to_string())) + ); + + let plan = sresult_return!( + internal_provider + .merge_into(session, source, merge_schema, on, clauses) + .await + ); + + FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime.clone())) + } + .into_ffi() +} + unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_TableProvider) { unsafe { debug_assert!(!provider.private_data.is_null()); @@ -361,6 +428,7 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_TableProvider) -> FFI_Table table_type: table_type_fn_wrapper, supports_filters_pushdown: provider.supports_filters_pushdown, insert_into: provider.insert_into, + merge_into: provider.merge_into, statistics: statistics_fn_wrapper, logical_codec: provider.logical_codec.clone(), clone: clone_fn_wrapper, @@ -422,6 +490,7 @@ impl FFI_TableProvider { false => None, }, insert_into: insert_into_fn_wrapper, + merge_into: merge_into_fn_wrapper, statistics: statistics_fn_wrapper, logical_codec, clone: clone_fn_wrapper, @@ -577,16 +646,63 @@ impl TableProvider for ForeignTableProvider { Ok(plan) } + + async fn merge_into( + &self, + session: &dyn Session, + source: Arc, + merge_schema: DFSchemaRef, + on: Expr, + clauses: Vec, + ) -> Result> { + let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); + + let rc = Handle::try_current().ok(); + let source = FFI_ExecutionPlan::new(source, rc); + + let merge_schema = DfSchema::try_from(&merge_schema) + .map_err(|e| DataFusionError::Plan(e.to_string()))?; + let merge_schema_serialized: SVec = + merge_schema.encode_to_vec().into_iter().collect(); + + let codec: Arc = (&self.0.logical_codec).into(); + let merge_op = + serialize_merge_into_op(&MergeIntoOp { on, clauses }, codec.as_ref()) + .map_err(|e| DataFusionError::Plan(e.to_string()))?; + let merge_op_serialized = merge_op.encode_to_vec().into_iter().collect(); + + let plan = unsafe { + let maybe_plan = (self.0.merge_into)( + &self.0, + session, + &source, + merge_schema_serialized, + merge_op_serialized, + ) + .await; + + >::try_from(&df_result!(maybe_plan)?)? + }; + + Ok(plan) + } } #[cfg(test)] mod tests { - use arrow::datatypes::Schema; + use std::sync::Mutex; + + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::{Column, DFSchema}; + use datafusion::logical_expr::dml::{MergeIntoAction, MergeIntoClauseKind}; + use datafusion::physical_plan::empty::EmptyExec; use datafusion::prelude::{SessionContext, col, lit}; use datafusion_execution::TaskContextProvider; use super::*; + type CapturedMerge = (usize, usize, String, Vec); + fn create_test_table_provider() -> Result> { use arrow::datatypes::Field; use datafusion::arrow::array::Float32Array; @@ -613,6 +729,122 @@ mod tests { )?)) } + #[derive(Debug)] + struct CaptureMergeProvider { + schema: SchemaRef, + captured: Mutex>, + } + + #[async_trait] + impl TableProvider for CaptureMergeProvider { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema)))) + } + + async fn merge_into( + &self, + state: &dyn Session, + source: Arc, + merge_schema: DFSchemaRef, + on: Expr, + clauses: Vec, + ) -> Result> { + let target_id = + merge_schema.index_of_column(&Column::new(Some("target"), "id"))?; + let source_id = + merge_schema.index_of_column(&Column::new(Some("source"), "id"))?; + let physical_on = state.create_physical_expr(on, &merge_schema)?; + let clauses = clauses.iter().map(|clause| format!("{clause:?}")).collect(); + *self.captured.lock().unwrap() = + Some((target_id, source_id, format!("{physical_on:?}"), clauses)); + Ok(source) + } + } + + #[tokio::test] + async fn test_round_trip_ffi_table_provider_merge_into() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("val", DataType::Int64, true), + ])); + let provider = Arc::new(CaptureMergeProvider { + schema: Arc::clone(&schema), + captured: Mutex::new(None), + }); + let ctx = Arc::new(SessionContext::new()); + let task_ctx_provider = Arc::clone(&ctx) as Arc; + let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + + let mut ffi_provider = FFI_TableProvider::new( + Arc::clone(&provider) as Arc, + true, + None, + task_ctx_provider, + None, + ); + ffi_provider.library_marker_id = crate::mock_foreign_marker_id; + let foreign: Arc = (&ffi_provider).into(); + + let target_schema = DFSchema::try_from_qualified_schema("target", &schema)?; + let source_schema = DFSchema::try_from_qualified_schema("source", &schema)?; + let merge_schema = Arc::new(target_schema.join(&source_schema)?); + let source: Arc = + Arc::new(EmptyExec::new(Arc::clone(&schema))); + let clauses = vec![ + MergeIntoClause { + kind: MergeIntoClauseKind::Matched, + predicate: Some(col("target.val").is_null()), + action: MergeIntoAction::Update(vec![( + "val".to_string(), + col("source.val"), + )]), + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatched, + predicate: None, + action: MergeIntoAction::Insert { + columns: vec!["id".to_string(), "val".to_string()], + values: vec![col("source.id"), col("source.val")], + }, + }, + ]; + + foreign + .merge_into( + &ctx.state(), + source, + merge_schema, + col("target.id").eq(col("source.id")), + clauses, + ) + .await?; + + let captured = provider.captured.lock().unwrap(); + let (target_id, source_id, physical_on, clauses) = + captured.as_ref().expect("merge_into should be called"); + assert_eq!((*target_id, *source_id), (0, 2)); + assert!(physical_on.contains("index: 0")); + assert!(physical_on.contains("index: 2")); + assert_eq!(clauses.len(), 2); + assert!(clauses[0].contains("source")); + assert!(clauses[1].contains("Insert")); + Ok(()) + } + #[tokio::test] async fn test_round_trip_ffi_table_provider_scan() -> Result<()> { let provider = create_test_table_provider()?; diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index d372dcf9177e6..73b543c083b74 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -25,8 +25,9 @@ use catalog::create_catalog_provider; use datafusion_catalog::MemTable; use datafusion_catalog::{Session, TableProvider}; use datafusion_common::stats::Precision; +use datafusion_common::{Column, DFSchemaRef, Result, ScalarValue}; use datafusion_common::{ColumnStatistics, Statistics}; -use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::dml::{MergeIntoAction, MergeIntoClause, MergeIntoClauseKind}; use datafusion_expr::{Expr, TableType}; use datafusion_physical_plan::ExecutionPlan; use sync_provider::create_sync_table_provider; @@ -112,6 +113,9 @@ pub struct ForeignLibraryModule { pub create_table_with_statistics: extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_TableProvider, + pub create_merge_table: + extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_TableProvider, + pub create_physical_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule, pub create_context_aware_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule, @@ -241,6 +245,123 @@ pub(crate) extern "C" fn create_table_with_statistics( FFI_TableProvider::new_with_ffi_codec(provider, true, None, codec) } +#[derive(Debug)] +struct MergeTableProvider { + schema: Arc, +} + +#[async_trait] +impl TableProvider for MergeTableProvider { + fn schema(&self) -> Arc { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema)))) + } + + async fn merge_into( + &self, + _state: &dyn Session, + source: Arc, + merge_schema: DFSchemaRef, + on: Expr, + clauses: Vec, + ) -> Result> { + let target_id = + merge_schema.index_of_column(&Column::new(Some("target"), "id"))?; + let source_id = + merge_schema.index_of_column(&Column::new(Some("source"), "id"))?; + if (target_id, source_id) != (0, 2) { + return datafusion_common::plan_err!( + "unexpected MERGE schema indices: target={target_id}, source={source_id}" + ); + } + + if on.to_string() != "target.id = source.id" { + return datafusion_common::plan_err!( + "unexpected logical MERGE condition: {on}" + ); + } + + let [matched, not_matched] = clauses.as_slice() else { + return datafusion_common::plan_err!( + "expected two MERGE clauses, got {}", + clauses.len() + ); + }; + if matched.kind != MergeIntoClauseKind::Matched + || matched + .predicate + .as_ref() + .map(ToString::to_string) + .as_deref() + != Some("target.val IS NULL") + { + return datafusion_common::plan_err!( + "unexpected matched MERGE clause: {matched:?}" + ); + } + let MergeIntoAction::Update(assignments) = &matched.action else { + return datafusion_common::plan_err!( + "expected MERGE update action, got {:?}", + matched.action + ); + }; + if assignments.len() != 1 + || assignments[0].0 != "val" + || assignments[0].1.to_string() != "source.val" + { + return datafusion_common::plan_err!( + "unexpected MERGE update assignments: {assignments:?}" + ); + } + + if not_matched.kind != MergeIntoClauseKind::NotMatched { + return datafusion_common::plan_err!( + "unexpected not-matched MERGE clause: {not_matched:?}" + ); + } + let MergeIntoAction::Insert { columns, values } = ¬_matched.action else { + return datafusion_common::plan_err!( + "expected MERGE insert action, got {:?}", + not_matched.action + ); + }; + if columns != &["id".to_string(), "val".to_string()] + || values.iter().map(ToString::to_string).collect::>() + != ["source.id".to_string(), "source.val".to_string()] + { + return datafusion_common::plan_err!( + "unexpected MERGE insert action: {not_matched:?}" + ); + } + + Ok(source) + } +} + +pub(crate) extern "C" fn create_merge_table( + codec: FFI_LogicalExtensionCodec, +) -> FFI_TableProvider { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("val", DataType::Int64, true), + ])); + let provider = Arc::new(MergeTableProvider { schema }); + FFI_TableProvider::new_with_ffi_codec(provider, true, None, codec) +} + /// This defines the entry point for using the module. #[unsafe(no_mangle)] pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { @@ -261,6 +382,7 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { create_empty_exec, create_exec_with_statistics, create_table_with_statistics, + create_merge_table, create_physical_optimizer_rule: physical_optimizer::create_physical_optimizer_rule, create_context_aware_optimizer_rule: diff --git a/datafusion/ffi/tests/ffi_integration.rs b/datafusion/ffi/tests/ffi_integration.rs index 86f953e262ead..899fa528712b0 100644 --- a/datafusion/ffi/tests/ffi_integration.rs +++ b/datafusion/ffi/tests/ffi_integration.rs @@ -24,12 +24,13 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; - use arrow::datatypes::Schema; + use arrow::datatypes::{DataType, Field, Schema}; use datafusion::catalog::{TableProvider, TableProviderFactory}; use datafusion::error::Result; - use datafusion_common::TableReference; - use datafusion_common::ToDFSchema; - use datafusion_expr::CreateExternalTable; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion_common::{DFSchema, TableReference, ToDFSchema}; + use datafusion_expr::dml::{MergeIntoAction, MergeIntoClause, MergeIntoClauseKind}; + use datafusion_expr::{CreateExternalTable, col}; use datafusion_ffi::tests::create_record_batch; use datafusion_ffi::tests::utils::get_module; @@ -95,6 +96,55 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_ffi_table_provider_merge_into_cross_library() -> Result<()> { + let module = get_module()?; + let (ctx, codec) = super::utils::ctx_and_codec(); + let ffi_provider = (module.create_merge_table)(codec); + let foreign: Arc = (&ffi_provider).into(); + + let target_schema = + DFSchema::try_from_qualified_schema("target", &foreign.schema())?; + let source_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("val", DataType::Int64, true), + ])); + let source_df_schema = + DFSchema::try_from_qualified_schema("source", &source_schema)?; + let merge_schema = Arc::new(target_schema.join(&source_df_schema)?); + let source = Arc::new(EmptyExec::new(Arc::clone(&source_schema))); + let clauses = vec![ + MergeIntoClause { + kind: MergeIntoClauseKind::Matched, + predicate: Some(col("target.val").is_null()), + action: MergeIntoAction::Update(vec![( + "val".to_string(), + col("source.val"), + )]), + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatched, + predicate: None, + action: MergeIntoAction::Insert { + columns: vec!["id".to_string(), "val".to_string()], + values: vec![col("source.id"), col("source.val")], + }, + }, + ]; + + let result = foreign + .merge_into( + &ctx.state(), + source, + merge_schema, + col("target.id").eq(col("source.id")), + clauses, + ) + .await?; + assert_eq!(result.schema(), source_schema); + Ok(()) + } + #[tokio::test] async fn test_table_provider_factory() -> Result<()> { let table_provider_module = get_module()?; diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 566522e9273fd..d11c3e7435fde 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -193,10 +193,60 @@ impl<'a> TypeCoercionRewriter<'a> { LogicalPlan::Join(join) => self.coerce_join(join), LogicalPlan::Union(union) => Self::coerce_union(union), LogicalPlan::Limit(limit) => Self::coerce_limit(limit), + LogicalPlan::Dml(dml) => self.coerce_dml(dml), _ => Ok(plan), } } + fn coerce_dml(&self, mut dml: DmlStatement) -> Result { + let WriteOp::MergeInto(merge_op) = &dml.op else { + return Ok(LogicalPlan::Dml(dml)); + }; + + let target_schema = DFSchema::try_from_qualified_schema( + dml.table_name.clone(), + &dml.target.schema(), + )?; + let mut merge_op = (**merge_op).clone(); + merge_op.on = self.coerce_predicate(merge_op.on, "MERGE ON condition")?; + for clause in &mut merge_op.clauses { + clause.predicate = clause + .predicate + .take() + .map(|expr| self.coerce_predicate(expr, "MERGE WHEN condition")) + .transpose()?; + + match &mut clause.action { + datafusion_expr::dml::MergeIntoAction::Update(assignments) => { + for (column, value) in assignments { + let field = target_schema.field_with_unqualified_name(column)?; + *value = value.clone().cast_to(field.data_type(), self.schema)?; + } + } + datafusion_expr::dml::MergeIntoAction::Insert { columns, values } => { + if columns.is_empty() { + for (value, field) in + values.iter_mut().zip(target_schema.fields()) + { + *value = + value.clone().cast_to(field.data_type(), self.schema)?; + } + } else { + for (column, value) in columns.iter().zip(values) { + let field = + target_schema.field_with_unqualified_name(column)?; + *value = + value.clone().cast_to(field.data_type(), self.schema)?; + } + } + } + datafusion_expr::dml::MergeIntoAction::Delete => {} + } + } + dml.op = WriteOp::MergeInto(Box::new(merge_op)); + Ok(LogicalPlan::Dml(dml)) + } + /// Coerce join equality expressions and join filter /// /// Joins must be treated specially as their equality expressions are stored @@ -228,7 +278,7 @@ impl<'a> TypeCoercionRewriter<'a> { // Join filter must be boolean join.filter = join .filter - .map(|expr| self.coerce_join_filter(expr)) + .map(|expr| self.coerce_predicate(expr, "Join condition")) .transpose()?; Ok(LogicalPlan::Join(join)) @@ -296,12 +346,14 @@ impl<'a> TypeCoercionRewriter<'a> { })) } - fn coerce_join_filter(&self, expr: Expr) -> Result { + fn coerce_predicate(&self, expr: Expr, description: &str) -> Result { let expr_type = expr.get_type(self.schema)?; match expr_type { DataType::Boolean => Ok(expr), DataType::Null => expr.cast_to(&DataType::Boolean, self.schema), - other => plan_err!("Join condition must be boolean type, but got {other:?}"), + other => { + plan_err!("{description} must be boolean type, but got {other:?}") + } } } @@ -1584,11 +1636,24 @@ mod test { 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, - }], + clauses: vec![ + MergeIntoClause { + kind: MergeIntoClauseKind::Matched, + predicate: None, + action: MergeIntoAction::Update(vec![( + "id".to_string(), + col("source.id"), + )]), + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatched, + predicate: None, + action: MergeIntoAction::Insert { + columns: vec!["id".to_string()], + values: vec![col("source.id")], + }, + }, + ], }; let plan = LogicalPlan::Dml(DmlStatement::new( target_table_name, @@ -1597,13 +1662,27 @@ mod test { Arc::new(source_plan), )); - assert_analyzed_plan_eq!( - plan, - @r" - Dml: op=[MergeInto] table=[target] - TableScan: source - " - ) + let analyzed = Analyzer::with_rules(vec![Arc::new(TypeCoercion::new())]) + .execute_and_check(plan, &ConfigOptions::default(), |_, _| {})?; + let LogicalPlan::Dml(dml) = analyzed else { + panic!("expected Dml"); + }; + let WriteOp::MergeInto(merge_op) = dml.op else { + panic!("expected MergeInto"); + }; + assert_eq!( + merge_op.on.to_string(), + "CAST(target.id AS Int64) = source.id" + ); + let MergeIntoAction::Update(assignments) = &merge_op.clauses[0].action else { + panic!("expected UPDATE"); + }; + assert_eq!(assignments[0].1.to_string(), "CAST(source.id AS UInt32)"); + let MergeIntoAction::Insert { values, .. } = &merge_op.clauses[1].action else { + panic!("expected INSERT"); + }; + assert_eq!(values[0].to_string(), "CAST(source.id AS UInt32)"); + Ok(()) } #[test] diff --git a/datafusion/optimizer/src/rewrite_set_comparison.rs b/datafusion/optimizer/src/rewrite_set_comparison.rs index c8c35b518743a..18712c5335205 100644 --- a/datafusion/optimizer/src/rewrite_set_comparison.rs +++ b/datafusion/optimizer/src/rewrite_set_comparison.rs @@ -25,7 +25,7 @@ use datafusion_common::{Column, DFSchema, ExprSchema, Result, ScalarValue, plan_ use datafusion_expr::expr::{self, Exists, SetComparison, SetQuantifier}; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::logical_plan::builder::LogicalPlanBuilder; -use datafusion_expr::{Expr, LogicalPlan, lit}; +use datafusion_expr::{DmlStatement, Expr, LogicalPlan, WriteOp, lit}; use std::sync::Arc; use datafusion_expr::utils::merge_schema; @@ -44,7 +44,19 @@ impl RewriteSetComparison { } fn rewrite_plan(&self, plan: LogicalPlan) -> Result> { - let schema = merge_schema(&plan.inputs()); + let mut schema = merge_schema(&plan.inputs()); + if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + schema.merge(&DFSchema::try_from_qualified_schema( + table_name.clone(), + &target.schema(), + )?); + } plan.map_expressions(|expr| { expr.transform_up(|expr| rewrite_set_comparison(expr, &schema)) }) diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs index 3e495f5355103..0e72a17abc9f7 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs @@ -21,12 +21,12 @@ use std::sync::Arc; use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{Column, DFSchema, DFSchemaRef, DataFusionError, Result}; -use datafusion_expr::Expr; use datafusion_expr::logical_plan::{Aggregate, LogicalPlan, Projection}; use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{ columnize_expr, find_aggregate_exprs, grouping_set_to_exprlist, merge_schema, }; +use datafusion_expr::{DmlStatement, Expr, WriteOp}; use super::ExprSimplifier; use crate::optimizer::ApplyOrder; @@ -77,7 +77,20 @@ impl SimplifyExpressions { plan: LogicalPlan, config: &dyn OptimizerConfig, ) -> Result> { - let schema = if !plan.inputs().is_empty() { + let schema = if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + let mut schema = merge_schema(&plan.inputs()); + schema.merge(&DFSchema::try_from_qualified_schema( + table_name.clone(), + &target.schema(), + )?); + DFSchemaRef::new(schema) + } else if !plan.inputs().is_empty() { DFSchemaRef::new(merge_schema(&plan.inputs())) } else if let LogicalPlan::TableScan(scan) = &plan { // When predicates are pushed into a table scan, there is no input diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 6d9a73e06ff45..4669240036073 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -290,7 +290,7 @@ pub fn parse_write_op( }) } -fn parse_merge_into_op( +pub fn parse_merge_into_op( op: &protobuf::MergeIntoOpNode, ctx: &TaskContext, codec: &dyn LogicalExtensionCodec, diff --git a/datafusion/session/src/table.rs b/datafusion/session/src/table.rs index c0b95b4a46c00..69e7e731b32c7 100644 --- a/datafusion/session/src/table.rs +++ b/datafusion/session/src/table.rs @@ -24,7 +24,7 @@ use crate::session::Session; use arrow_schema::SchemaRef; use async_trait::async_trait; use datafusion_common::{Constraints, Statistics, not_impl_err}; -use datafusion_common::{Result, internal_err}; +use datafusion_common::{DFSchemaRef, Result, internal_err}; use datafusion_expr::Expr; use datafusion_expr::statistics::StatisticsRequest; @@ -383,6 +383,10 @@ pub trait TableProvider: Any + Debug + Sync + Send { /// Merge rows from a source into this table. /// /// The `source` is an [`ExecutionPlan`] representing the USING clause. + /// The `merge_schema` contains the target columns followed by the source + /// columns, preserving their logical qualifiers. Providers can use this + /// schema to resolve the logical expressions against the combined rows + /// they construct while executing the merge. /// The `on` condition is the join predicate from the ON clause. /// The `clauses` describe the WHEN MATCHED / WHEN NOT MATCHED actions. /// @@ -391,6 +395,7 @@ pub trait TableProvider: Any + Debug + Sync + Send { &self, _state: &dyn Session, _source: Arc, + _merge_schema: DFSchemaRef, _on: Expr, _clauses: Vec, ) -> Result> { diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 3c71e9b5c713e..961a710ffb20a 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -33,7 +33,7 @@ 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::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_common::{ Column, Constraint, Constraints, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, SchemaError, SchemaReference, TableReference, ToDFSchema, exec_err, @@ -2439,9 +2439,47 @@ impl SqlToRel<'_, S> { return not_impl_err!("MERGE OUTPUT clause is not supported"); } + if clauses.is_empty() { + return plan_err!("MERGE INTO requires at least one WHEN clause"); + } + // 1. Resolve target table - let (target_table_name, target_alias) = match &table { - TableFactor::Table { name, alias, .. } => (name.clone(), alias.clone()), + let (target_table_name, target_alias) = match table { + TableFactor::Table { + name, + alias, + args, + with_hints, + version, + with_ordinality, + partitions, + json_path, + sample, + index_hints, + } => { + if alias + .as_ref() + .is_some_and(|alias| !alias.columns.is_empty()) + { + return not_impl_err!( + "MERGE target alias column lists are not supported" + ); + } + if args.is_some() + || !with_hints.is_empty() + || version.is_some() + || with_ordinality + || !partitions.is_empty() + || json_path.is_some() + || sample.is_some() + || !index_hints.is_empty() + { + return not_impl_err!( + "MERGE target table modifiers are not supported" + ); + } + (name, alias) + } _ => plan_err!("Cannot MERGE INTO non-table relation!")?, }; let target_table_ref = self.object_name_to_table_reference(target_table_name)?; @@ -2485,7 +2523,7 @@ impl SqlToRel<'_, S> { clause, &combined_schema, &target_schema, - &target_alias, + &target_qualifier, &mut planner_context, ) }) @@ -2502,17 +2540,30 @@ impl SqlToRel<'_, S> { clauses: df_clauses, }; if target_qualifier != target_table_ref { + // Target references in correlated subqueries are represented as + // `OuterReferenceColumn`s inside the embedded logical plan. The + // alias canonicalization below only rewrites top-level expression + // columns, so accepting such a subquery would leave the target + // alias in the public MERGE representation. Reject this case until + // the alias can be rewritten scope-safely inside subquery plans. + for expr in merge_op.exprs() { + if Self::has_outer_reference_to_qualifier(expr, &target_qualifier)? { + return not_impl_err!( + "MERGE subqueries correlated to target alias \ + '{target_qualifier}' are not supported" + ); + } + } + // 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)) - { + if source_plan.schema().iter().any(|(qualifier, _)| { + qualifier.is_some_and(|q| q.resolved_eq(&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}'; \ @@ -2559,15 +2610,81 @@ impl SqlToRel<'_, S> { .map(|transformed| transformed.data) } - fn ident_from_object_name_last(name: &ObjectName) -> Result { + /// Return true if an expression contains a subquery whose embedded plan + /// has an outer reference qualified by `qualifier`. + fn has_outer_reference_to_qualifier( + expr: &Expr, + qualifier: &TableReference, + ) -> Result { + let mut found = false; + expr.apply(|expr| { + let subquery = match expr { + Expr::Exists(exists) => Some(&exists.subquery), + Expr::InSubquery(in_subquery) => Some(&in_subquery.subquery), + Expr::SetComparison(set_comparison) => Some(&set_comparison.subquery), + Expr::ScalarSubquery(subquery) => Some(subquery), + _ => None, + }; + + if let Some(subquery) = subquery { + subquery.subquery.apply_with_subqueries(|plan| { + plan.apply_expressions(|expr| { + expr.apply(|expr| { + if let Expr::OuterReferenceColumn(_, column) = expr + && column.relation.as_ref() == Some(qualifier) + { + found = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + })?; + Ok(if found { + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }) + })?; + } + + Ok(if found { + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }) + })?; + Ok(found) + } + + fn merge_target_column_name( + &self, + name: &ObjectName, + target_qualifier: &TableReference, + ) -> Result { let part = name .0 .iter() .last() .ok_or_else(|| plan_datafusion_err!("Empty column name"))?; - part.as_ident() + let ident = part + .as_ident() .cloned() - .ok_or_else(|| plan_datafusion_err!("Expected simple identifier")) + .ok_or_else(|| plan_datafusion_err!("Expected simple identifier"))?; + + if name.0.len() > 1 { + let qualifier = self.object_name_to_table_reference(ObjectName( + name.0[..name.0.len() - 1].to_vec(), + ))?; + if !qualifier.resolved_eq(target_qualifier) { + return plan_err!( + "MERGE assignment target '{name}' must reference target table \ + '{target_qualifier}'" + ); + } + } + + Ok(self.ident_normalizer.normalize(ident)) } fn merge_clause_to_plan( @@ -2575,7 +2692,7 @@ impl SqlToRel<'_, S> { clause: ast::MergeClause, combined_schema: &DFSchema, target_schema: &DFSchema, - _target_alias: &Option, + target_qualifier: &TableReference, planner_context: &mut PlannerContext, ) -> Result { let kind = match clause.clause_kind { @@ -2596,14 +2713,23 @@ impl SqlToRel<'_, S> { let action = match clause.action { ast::MergeAction::Update(update_expr) => { + if update_expr.update_predicate.is_some() { + return not_impl_err!( + "MERGE UPDATE WHERE predicates are not supported" + ); + } + if update_expr.delete_predicate.is_some() { + return not_impl_err!( + "MERGE UPDATE DELETE WHERE predicates are not supported" + ); + } 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) + self.merge_target_column_name(cols, target_qualifier)? } _ => plan_err!("Tuples are not supported")?, }; @@ -2617,16 +2743,24 @@ impl SqlToRel<'_, S> { Ok((col_name, value)) }) .collect::>>()?; + let mut seen = HashSet::new(); + for (column, _) in &assignments { + if !seen.insert(column.as_str()) { + return plan_err!("Duplicate column '{column}' in MERGE UPDATE"); + } + } MergeIntoAction::Update(assignments) } ast::MergeAction::Insert(insert_expr) => { + if insert_expr.insert_predicate.is_some() { + return not_impl_err!( + "MERGE INSERT WHERE predicates are not supported" + ); + } let columns: Vec = insert_expr .columns .iter() - .map(|c| { - let ident = Self::ident_from_object_name_last(c)?; - Ok(self.ident_normalizer.normalize(ident)) - }) + .map(|c| self.merge_target_column_name(c, target_qualifier)) .collect::>>()?; // Validate: no duplicates, all columns exist in target schema diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 8334755094a3e..ffd84183cf256 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -3648,6 +3648,112 @@ fn plan_merge_into_rejects_source_alias_colliding_with_target_name() { ); } +#[rstest] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN MATCHED THEN UPDATE SET j1_string = j2.j2_string WHERE false", + "MERGE UPDATE WHERE predicates are not supported" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN MATCHED THEN UPDATE SET j1_string = j2.j2_string DELETE WHERE false", + "MERGE UPDATE DELETE WHERE predicates are not supported" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN NOT MATCHED THEN INSERT (j1_id, j1_string) \ + VALUES (j2.j2_id, j2.j2_string) WHERE false", + "MERGE INSERT WHERE predicates are not supported" +)] +fn plan_merge_into_rejects_unsupported_action_predicates( + #[case] sql: &str, + #[case] expected: &str, +) { + let err = logical_plan(sql).unwrap_err(); + assert!( + err.strip_backtrace().contains(expected), + "unexpected error: {err}" + ); +} + +#[rstest] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN MATCHED THEN UPDATE SET j1_string = 'a', j1_string = 'b'", + "Duplicate column 'j1_string' in MERGE UPDATE" +)] +#[case( + "MERGE INTO j1 AS t USING j2 AS s ON t.j1_id = s.j2_id \ + WHEN MATCHED THEN UPDATE SET s.j1_string = s.j2_string", + "MERGE assignment target 's.j1_string' must reference target table 't'" +)] +#[case( + "MERGE INTO j1 AS t USING j2 AS s ON t.j1_id = s.j2_id \ + WHEN NOT MATCHED THEN INSERT (s.j1_id) VALUES (s.j2_id)", + "MERGE assignment target 's.j1_id' must reference target table 't'" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id", + "MERGE INTO requires at least one WHEN clause" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN NOT MATCHED THEN INSERT (j1_id, J1_ID) VALUES (1, 2)", + "Duplicate column 'j1_id' in MERGE INSERT" +)] +#[case( + "MERGE INTO j1() USING j2 ON true WHEN MATCHED THEN DELETE", + "MERGE target table modifiers are not supported" +)] +#[case( + "MERGE INTO j1 PARTITION (p0) USING j2 ON true WHEN MATCHED THEN DELETE", + "MERGE target table modifiers are not supported" +)] +#[case( + "MERGE INTO j1 AS t(a) USING j2 ON true WHEN MATCHED THEN DELETE", + "MERGE target alias column lists are not supported" +)] +fn plan_merge_into_validates_action_targets_and_structure( + #[case] sql: &str, + #[case] expected: &str, +) { + let err = logical_plan(sql).unwrap_err(); + assert!( + err.strip_backtrace().contains(expected), + "unexpected error: {err}" + ); +} + +#[test] +fn plan_merge_into_preserves_quoted_action_columns() { + let plan = logical_plan( + "MERGE INTO person_quoted_cols AS t USING j2 AS s ON t.id = s.j2_id \ + WHEN MATCHED THEN UPDATE SET \"First Name\" = s.j2_string \ + WHEN NOT MATCHED THEN INSERT (id, \"Age\") VALUES (s.j2_id, 42)", + ) + .unwrap(); + let LogicalPlan::Dml(dml) = plan else { + panic!("expected Dml"); + }; + let datafusion_expr::WriteOp::MergeInto(merge_op) = dml.op else { + panic!("expected MergeInto"); + }; + + let datafusion_expr::dml::MergeIntoAction::Update(assignments) = + &merge_op.clauses[0].action + else { + panic!("expected UPDATE"); + }; + assert_eq!(assignments[0].0, "First Name"); + + let datafusion_expr::dml::MergeIntoAction::Insert { columns, .. } = + &merge_op.clauses[1].action + else { + panic!("expected INSERT"); + }; + assert_eq!(columns, &["id".to_string(), "Age".to_string()]); +} + fn logical_plan(sql: &str) -> Result { logical_plan_with_options(sql, ParserOptions::default()) } From 3b9ebdb8fc587b90c801adc7eb072bde80b36475 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Fri, 24 Jul 2026 14:38:26 +0000 Subject: [PATCH 4/5] Move MERGE INTO FFI support to follow-up --- datafusion/ffi/src/table_provider.rs | 244 +----------------- datafusion/ffi/src/tests/mod.rs | 124 +-------- datafusion/ffi/tests/ffi_integration.rs | 58 +---- .../proto/src/logical_plan/from_proto.rs | 2 +- 4 files changed, 12 insertions(+), 416 deletions(-) diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index 35c0889c44a24..5a4b2fa27256f 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -22,20 +22,18 @@ use arrow::datatypes::SchemaRef; use async_ffi::{FfiFuture, FutureExt}; use async_trait::async_trait; use datafusion_catalog::{Session, TableProvider}; +use datafusion_common::Statistics; use datafusion_common::error::{DataFusionError, Result}; -use datafusion_common::{DFSchemaRef, Statistics}; use datafusion_execution::TaskContext; -use datafusion_expr::dml::{InsertOp, MergeIntoClause, MergeIntoOp}; +use datafusion_expr::dml::InsertOp; use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; use datafusion_physical_plan::ExecutionPlan; -use datafusion_proto::logical_plan::from_proto::{parse_exprs, parse_merge_into_op}; -use datafusion_proto::logical_plan::to_proto::{ - serialize_exprs, serialize_merge_into_op, -}; +use datafusion_proto::logical_plan::from_proto::parse_exprs; +use datafusion_proto::logical_plan::to_proto::serialize_exprs; use datafusion_proto::logical_plan::{ DefaultLogicalExtensionCodec, LogicalExtensionCodec, }; -use datafusion_proto::protobuf::{DfSchema, LogicalExprList, MergeIntoOpNode}; +use datafusion_proto::protobuf::LogicalExprList; use prost::Message; use stabby::vec::Vec as SVec; @@ -137,17 +135,6 @@ pub struct FFI_TableProvider { insert_op: FFI_InsertOp, ) -> FfiFuture>, - /// Merge source rows into the table. The schema and MERGE operation are - /// prost-encoded to preserve logical qualifiers and expressions across the - /// FFI boundary. - merge_into: unsafe extern "C" fn( - provider: &Self, - session: FFI_SessionRef, - source: &FFI_ExecutionPlan, - merge_schema_serialized: SVec, - merge_op_serialized: SVec, - ) -> FfiFuture>, - /// Snapshot the provider's table-level statistics. [`FFI_Option::None`] /// corresponds to [`TableProvider::statistics`] returning `None`; /// `Some(bytes)` is a prost-encoded `datafusion_proto_common::Statistics`. @@ -349,60 +336,6 @@ unsafe extern "C" fn insert_into_fn_wrapper( .into_ffi() } -unsafe extern "C" fn merge_into_fn_wrapper( - provider: &FFI_TableProvider, - session: FFI_SessionRef, - source: &FFI_ExecutionPlan, - merge_schema_serialized: SVec, - merge_op_serialized: SVec, -) -> FfiFuture> { - let task_ctx: Result, DataFusionError> = - (&provider.logical_codec.task_ctx_provider).try_into(); - let runtime = provider.runtime().clone(); - let logical_codec: Arc = (&provider.logical_codec).into(); - let internal_provider = Arc::clone(provider.inner()); - let source = source.clone(); - - async move { - let mut foreign_session = None; - let session = sresult_return!( - session - .as_local() - .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>) - .unwrap_or_else(|| { - foreign_session = Some(ForeignSession::try_from(&session)?); - Ok(foreign_session.as_ref().unwrap()) - }) - ); - - let source = sresult_return!(>::try_from(&source)); - let task_ctx = sresult_return!(task_ctx); - - let merge_schema = - sresult_return!(DfSchema::decode(merge_schema_serialized.as_ref())); - let merge_schema: DFSchemaRef = sresult_return!( - DFSchemaRef::try_from(merge_schema) - .map_err(|e| DataFusionError::Plan(e.to_string())) - ); - - let merge_op = - sresult_return!(MergeIntoOpNode::decode(merge_op_serialized.as_ref())); - let MergeIntoOp { on, clauses } = sresult_return!( - parse_merge_into_op(&merge_op, task_ctx.as_ref(), logical_codec.as_ref()) - .map_err(|e| DataFusionError::Plan(e.to_string())) - ); - - let plan = sresult_return!( - internal_provider - .merge_into(session, source, merge_schema, on, clauses) - .await - ); - - FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime.clone())) - } - .into_ffi() -} - unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_TableProvider) { unsafe { debug_assert!(!provider.private_data.is_null()); @@ -428,7 +361,6 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_TableProvider) -> FFI_Table table_type: table_type_fn_wrapper, supports_filters_pushdown: provider.supports_filters_pushdown, insert_into: provider.insert_into, - merge_into: provider.merge_into, statistics: statistics_fn_wrapper, logical_codec: provider.logical_codec.clone(), clone: clone_fn_wrapper, @@ -490,7 +422,6 @@ impl FFI_TableProvider { false => None, }, insert_into: insert_into_fn_wrapper, - merge_into: merge_into_fn_wrapper, statistics: statistics_fn_wrapper, logical_codec, clone: clone_fn_wrapper, @@ -646,63 +577,16 @@ impl TableProvider for ForeignTableProvider { Ok(plan) } - - async fn merge_into( - &self, - session: &dyn Session, - source: Arc, - merge_schema: DFSchemaRef, - on: Expr, - clauses: Vec, - ) -> Result> { - let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); - - let rc = Handle::try_current().ok(); - let source = FFI_ExecutionPlan::new(source, rc); - - let merge_schema = DfSchema::try_from(&merge_schema) - .map_err(|e| DataFusionError::Plan(e.to_string()))?; - let merge_schema_serialized: SVec = - merge_schema.encode_to_vec().into_iter().collect(); - - let codec: Arc = (&self.0.logical_codec).into(); - let merge_op = - serialize_merge_into_op(&MergeIntoOp { on, clauses }, codec.as_ref()) - .map_err(|e| DataFusionError::Plan(e.to_string()))?; - let merge_op_serialized = merge_op.encode_to_vec().into_iter().collect(); - - let plan = unsafe { - let maybe_plan = (self.0.merge_into)( - &self.0, - session, - &source, - merge_schema_serialized, - merge_op_serialized, - ) - .await; - - >::try_from(&df_result!(maybe_plan)?)? - }; - - Ok(plan) - } } #[cfg(test)] mod tests { - use std::sync::Mutex; - - use arrow::datatypes::{DataType, Field, Schema}; - use datafusion::common::{Column, DFSchema}; - use datafusion::logical_expr::dml::{MergeIntoAction, MergeIntoClauseKind}; - use datafusion::physical_plan::empty::EmptyExec; + use arrow::datatypes::Schema; use datafusion::prelude::{SessionContext, col, lit}; use datafusion_execution::TaskContextProvider; use super::*; - type CapturedMerge = (usize, usize, String, Vec); - fn create_test_table_provider() -> Result> { use arrow::datatypes::Field; use datafusion::arrow::array::Float32Array; @@ -729,122 +613,6 @@ mod tests { )?)) } - #[derive(Debug)] - struct CaptureMergeProvider { - schema: SchemaRef, - captured: Mutex>, - } - - #[async_trait] - impl TableProvider for CaptureMergeProvider { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } - - fn table_type(&self) -> TableType { - TableType::Base - } - - async fn scan( - &self, - _state: &dyn Session, - _projection: Option<&Vec>, - _filters: &[Expr], - _limit: Option, - ) -> Result> { - Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema)))) - } - - async fn merge_into( - &self, - state: &dyn Session, - source: Arc, - merge_schema: DFSchemaRef, - on: Expr, - clauses: Vec, - ) -> Result> { - let target_id = - merge_schema.index_of_column(&Column::new(Some("target"), "id"))?; - let source_id = - merge_schema.index_of_column(&Column::new(Some("source"), "id"))?; - let physical_on = state.create_physical_expr(on, &merge_schema)?; - let clauses = clauses.iter().map(|clause| format!("{clause:?}")).collect(); - *self.captured.lock().unwrap() = - Some((target_id, source_id, format!("{physical_on:?}"), clauses)); - Ok(source) - } - } - - #[tokio::test] - async fn test_round_trip_ffi_table_provider_merge_into() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("val", DataType::Int64, true), - ])); - let provider = Arc::new(CaptureMergeProvider { - schema: Arc::clone(&schema), - captured: Mutex::new(None), - }); - let ctx = Arc::new(SessionContext::new()); - let task_ctx_provider = Arc::clone(&ctx) as Arc; - let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); - - let mut ffi_provider = FFI_TableProvider::new( - Arc::clone(&provider) as Arc, - true, - None, - task_ctx_provider, - None, - ); - ffi_provider.library_marker_id = crate::mock_foreign_marker_id; - let foreign: Arc = (&ffi_provider).into(); - - let target_schema = DFSchema::try_from_qualified_schema("target", &schema)?; - let source_schema = DFSchema::try_from_qualified_schema("source", &schema)?; - let merge_schema = Arc::new(target_schema.join(&source_schema)?); - let source: Arc = - Arc::new(EmptyExec::new(Arc::clone(&schema))); - let clauses = vec![ - MergeIntoClause { - kind: MergeIntoClauseKind::Matched, - predicate: Some(col("target.val").is_null()), - action: MergeIntoAction::Update(vec![( - "val".to_string(), - col("source.val"), - )]), - }, - MergeIntoClause { - kind: MergeIntoClauseKind::NotMatched, - predicate: None, - action: MergeIntoAction::Insert { - columns: vec!["id".to_string(), "val".to_string()], - values: vec![col("source.id"), col("source.val")], - }, - }, - ]; - - foreign - .merge_into( - &ctx.state(), - source, - merge_schema, - col("target.id").eq(col("source.id")), - clauses, - ) - .await?; - - let captured = provider.captured.lock().unwrap(); - let (target_id, source_id, physical_on, clauses) = - captured.as_ref().expect("merge_into should be called"); - assert_eq!((*target_id, *source_id), (0, 2)); - assert!(physical_on.contains("index: 0")); - assert!(physical_on.contains("index: 2")); - assert_eq!(clauses.len(), 2); - assert!(clauses[0].contains("source")); - assert!(clauses[1].contains("Insert")); - Ok(()) - } - #[tokio::test] async fn test_round_trip_ffi_table_provider_scan() -> Result<()> { let provider = create_test_table_provider()?; diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 73b543c083b74..d372dcf9177e6 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -25,9 +25,8 @@ use catalog::create_catalog_provider; use datafusion_catalog::MemTable; use datafusion_catalog::{Session, TableProvider}; use datafusion_common::stats::Precision; -use datafusion_common::{Column, DFSchemaRef, Result, ScalarValue}; use datafusion_common::{ColumnStatistics, Statistics}; -use datafusion_expr::dml::{MergeIntoAction, MergeIntoClause, MergeIntoClauseKind}; +use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{Expr, TableType}; use datafusion_physical_plan::ExecutionPlan; use sync_provider::create_sync_table_provider; @@ -113,9 +112,6 @@ pub struct ForeignLibraryModule { pub create_table_with_statistics: extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_TableProvider, - pub create_merge_table: - extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_TableProvider, - pub create_physical_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule, pub create_context_aware_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule, @@ -245,123 +241,6 @@ pub(crate) extern "C" fn create_table_with_statistics( FFI_TableProvider::new_with_ffi_codec(provider, true, None, codec) } -#[derive(Debug)] -struct MergeTableProvider { - schema: Arc, -} - -#[async_trait] -impl TableProvider for MergeTableProvider { - fn schema(&self) -> Arc { - Arc::clone(&self.schema) - } - - fn table_type(&self) -> TableType { - TableType::Base - } - - async fn scan( - &self, - _state: &dyn Session, - _projection: Option<&Vec>, - _filters: &[Expr], - _limit: Option, - ) -> Result> { - Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema)))) - } - - async fn merge_into( - &self, - _state: &dyn Session, - source: Arc, - merge_schema: DFSchemaRef, - on: Expr, - clauses: Vec, - ) -> Result> { - let target_id = - merge_schema.index_of_column(&Column::new(Some("target"), "id"))?; - let source_id = - merge_schema.index_of_column(&Column::new(Some("source"), "id"))?; - if (target_id, source_id) != (0, 2) { - return datafusion_common::plan_err!( - "unexpected MERGE schema indices: target={target_id}, source={source_id}" - ); - } - - if on.to_string() != "target.id = source.id" { - return datafusion_common::plan_err!( - "unexpected logical MERGE condition: {on}" - ); - } - - let [matched, not_matched] = clauses.as_slice() else { - return datafusion_common::plan_err!( - "expected two MERGE clauses, got {}", - clauses.len() - ); - }; - if matched.kind != MergeIntoClauseKind::Matched - || matched - .predicate - .as_ref() - .map(ToString::to_string) - .as_deref() - != Some("target.val IS NULL") - { - return datafusion_common::plan_err!( - "unexpected matched MERGE clause: {matched:?}" - ); - } - let MergeIntoAction::Update(assignments) = &matched.action else { - return datafusion_common::plan_err!( - "expected MERGE update action, got {:?}", - matched.action - ); - }; - if assignments.len() != 1 - || assignments[0].0 != "val" - || assignments[0].1.to_string() != "source.val" - { - return datafusion_common::plan_err!( - "unexpected MERGE update assignments: {assignments:?}" - ); - } - - if not_matched.kind != MergeIntoClauseKind::NotMatched { - return datafusion_common::plan_err!( - "unexpected not-matched MERGE clause: {not_matched:?}" - ); - } - let MergeIntoAction::Insert { columns, values } = ¬_matched.action else { - return datafusion_common::plan_err!( - "expected MERGE insert action, got {:?}", - not_matched.action - ); - }; - if columns != &["id".to_string(), "val".to_string()] - || values.iter().map(ToString::to_string).collect::>() - != ["source.id".to_string(), "source.val".to_string()] - { - return datafusion_common::plan_err!( - "unexpected MERGE insert action: {not_matched:?}" - ); - } - - Ok(source) - } -} - -pub(crate) extern "C" fn create_merge_table( - codec: FFI_LogicalExtensionCodec, -) -> FFI_TableProvider { - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("val", DataType::Int64, true), - ])); - let provider = Arc::new(MergeTableProvider { schema }); - FFI_TableProvider::new_with_ffi_codec(provider, true, None, codec) -} - /// This defines the entry point for using the module. #[unsafe(no_mangle)] pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { @@ -382,7 +261,6 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { create_empty_exec, create_exec_with_statistics, create_table_with_statistics, - create_merge_table, create_physical_optimizer_rule: physical_optimizer::create_physical_optimizer_rule, create_context_aware_optimizer_rule: diff --git a/datafusion/ffi/tests/ffi_integration.rs b/datafusion/ffi/tests/ffi_integration.rs index 899fa528712b0..86f953e262ead 100644 --- a/datafusion/ffi/tests/ffi_integration.rs +++ b/datafusion/ffi/tests/ffi_integration.rs @@ -24,13 +24,12 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::datatypes::Schema; use datafusion::catalog::{TableProvider, TableProviderFactory}; use datafusion::error::Result; - use datafusion::physical_plan::empty::EmptyExec; - use datafusion_common::{DFSchema, TableReference, ToDFSchema}; - use datafusion_expr::dml::{MergeIntoAction, MergeIntoClause, MergeIntoClauseKind}; - use datafusion_expr::{CreateExternalTable, col}; + use datafusion_common::TableReference; + use datafusion_common::ToDFSchema; + use datafusion_expr::CreateExternalTable; use datafusion_ffi::tests::create_record_batch; use datafusion_ffi::tests::utils::get_module; @@ -96,55 +95,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn test_ffi_table_provider_merge_into_cross_library() -> Result<()> { - let module = get_module()?; - let (ctx, codec) = super::utils::ctx_and_codec(); - let ffi_provider = (module.create_merge_table)(codec); - let foreign: Arc = (&ffi_provider).into(); - - let target_schema = - DFSchema::try_from_qualified_schema("target", &foreign.schema())?; - let source_schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("val", DataType::Int64, true), - ])); - let source_df_schema = - DFSchema::try_from_qualified_schema("source", &source_schema)?; - let merge_schema = Arc::new(target_schema.join(&source_df_schema)?); - let source = Arc::new(EmptyExec::new(Arc::clone(&source_schema))); - let clauses = vec![ - MergeIntoClause { - kind: MergeIntoClauseKind::Matched, - predicate: Some(col("target.val").is_null()), - action: MergeIntoAction::Update(vec![( - "val".to_string(), - col("source.val"), - )]), - }, - MergeIntoClause { - kind: MergeIntoClauseKind::NotMatched, - predicate: None, - action: MergeIntoAction::Insert { - columns: vec!["id".to_string(), "val".to_string()], - values: vec![col("source.id"), col("source.val")], - }, - }, - ]; - - let result = foreign - .merge_into( - &ctx.state(), - source, - merge_schema, - col("target.id").eq(col("source.id")), - clauses, - ) - .await?; - assert_eq!(result.schema(), source_schema); - Ok(()) - } - #[tokio::test] async fn test_table_provider_factory() -> Result<()> { let table_provider_module = get_module()?; diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 4669240036073..6d9a73e06ff45 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -290,7 +290,7 @@ pub fn parse_write_op( }) } -pub fn parse_merge_into_op( +fn parse_merge_into_op( op: &protobuf::MergeIntoOpNode, ctx: &TaskContext, codec: &dyn LogicalExtensionCodec, From 3a76cd38a2d8965892e2cb49fd14aac4bee66917 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Fri, 24 Jul 2026 16:59:38 +0000 Subject: [PATCH 5/5] Consolidate MERGE INTO regression tests --- datafusion/core/src/physical_planner.rs | 15 ++- datafusion/core/tests/sql/sql_api.rs | 137 +++++++++--------------- datafusion/sql/tests/sql_integration.rs | 106 +++++------------- 3 files changed, 82 insertions(+), 176 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 7b09a61b85d46..dc336b9ffbe15 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -3475,10 +3475,8 @@ mod tests { #[tokio::test] async fn merge_into_provider_receives_combined_logical_schema() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("val", DataType::Int32, true), - ])); + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); let target = Arc::new(CaptureMergeProvider { schema: Arc::clone(&schema), captured: Mutex::new(None), @@ -3490,8 +3488,7 @@ mod tests { ctx.sql( "MERGE INTO target AS t USING source AS s ON t.id = s.id \ - WHEN MATCHED AND t.val IS NULL THEN UPDATE SET val = s.val \ - WHEN NOT MATCHED THEN INSERT (id, val) VALUES (s.id, s.val)", + WHEN MATCHED AND t.id > s.id THEN DELETE", ) .await? .create_physical_plan() @@ -3500,17 +3497,17 @@ mod tests { let captured = target.captured.lock().await; let (merge_schema, physical_on, clause_count) = captured.as_ref().expect("merge_into should be called"); - assert_eq!(*clause_count, 2); + assert_eq!(*clause_count, 1); assert_eq!( merge_schema.index_of_column(&Column::new(Some("target"), "id"))?, 0 ); assert_eq!( merge_schema.index_of_column(&Column::new(Some("s"), "id"))?, - 2 + 1 ); assert_contains!(physical_on, "index: 0"); - assert_contains!(physical_on, "index: 2"); + assert_contains!(physical_on, "index: 1"); Ok(()) } diff --git a/datafusion/core/tests/sql/sql_api.rs b/datafusion/core/tests/sql/sql_api.rs index f578dae12c43b..ca18406a8e40d 100644 --- a/datafusion/core/tests/sql/sql_api.rs +++ b/datafusion/core/tests/sql/sql_api.rs @@ -208,67 +208,64 @@ 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. +async fn merge_into_context() -> SessionContext { let ctx = SessionContext::new(); - ctx.sql("CREATE TABLE target (id INT, val INT)") + ctx.sql("CREATE TABLE target (id INT)").await.unwrap(); + ctx.sql("CREATE TABLE source (id INT)").await.unwrap(); + ctx +} + +async fn assert_merge_sql_error(ctx: &SessionContext, sql: &str, expected: &str) { + let err = ctx.sql(sql).await.unwrap_err(); + assert_contains!(err.strip_backtrace(), expected); +} + +async fn assert_merge_physical_error(ctx: &SessionContext, sql: &str, expected: &str) { + let err = ctx + .sql(sql) .await - .unwrap(); - ctx.sql("CREATE TABLE source (id INT, val INT)") + .unwrap() + .create_physical_plan() .await - .unwrap(); + .unwrap_err(); + assert_contains!(err.strip_backtrace(), expected); +} + +#[tokio::test] +async fn merge_into_rejects_source_alias_colliding_with_target_name() { + // Canonicalizing `t.id` to `target.id` must not collapse it onto a source + // that also uses `target` as its qualifier. + let ctx = merge_into_context().await; for target_ref in ["target", "public.target", "datafusion.public.target"] { - let err = ctx - .sql(&format!( + assert_merge_sql_error( + &ctx, + &format!( "MERGE INTO {target_ref} AS t USING source AS target \ ON t.id = target.id WHEN MATCHED THEN DELETE" - )) - .await - .unwrap_err(); - - assert_contains!( - err.strip_backtrace(), + ), &format!( "MERGE source may not use the target table name '{target_ref}' \ as a qualifier" - ) - ); + ), + ) + .await; } } #[tokio::test] async fn merge_into_rejects_subqueries_correlated_to_target_alias() { - 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 s \ - ON EXISTS (SELECT 1 FROM source AS x WHERE x.id = t.id) \ - WHEN MATCHED THEN DELETE", - ) - .await - .unwrap_err(); - assert_contains!( - err.strip_backtrace(), - "MERGE subqueries correlated to target alias 't' are not supported" - ); + let ctx = merge_into_context().await; + assert_merge_sql_error( + &ctx, + "MERGE INTO target AS t USING source AS s \ + ON EXISTS (SELECT 1 FROM source AS x WHERE x.id = t.id) \ + WHEN MATCHED THEN DELETE", + "MERGE subqueries correlated to target alias 't' are not supported", + ) + .await; - // Source correlation and uncorrelated set-comparison subqueries do not - // require target-alias canonicalization and remain supported through + // Source-correlated and uncorrelated subqueries remain supported through // logical optimization. for sql in [ "MERGE INTO target AS t USING source AS s \ @@ -278,29 +275,14 @@ async fn merge_into_rejects_subqueries_correlated_to_target_alias() { ON t.id = ANY (SELECT id FROM source) \ WHEN MATCHED THEN DELETE", ] { - let err = ctx - .sql(sql) - .await - .unwrap() - .create_physical_plan() - .await - .unwrap_err(); - assert_contains!( - err.strip_backtrace(), - "MERGE INTO not supported for Base table" - ); + assert_merge_physical_error(&ctx, sql, "MERGE INTO not supported for Base table") + .await; } } #[tokio::test] async fn merge_into_requires_boolean_conditions() { - 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 ctx = merge_into_context().await; for (sql, expected) in [ ( @@ -312,31 +294,14 @@ async fn merge_into_requires_boolean_conditions() { WHEN MATCHED AND 1 THEN DELETE", "MERGE WHEN condition must be boolean type, but got Int64", ), - ] { - let err = ctx - .sql(sql) - .await - .unwrap() - .create_physical_plan() - .await - .unwrap_err(); - assert_contains!(err.strip_backtrace(), expected); - } - - let err = ctx - .sql( + ( "MERGE INTO target USING source ON NULL \ WHEN MATCHED AND NULL THEN DELETE", - ) - .await - .unwrap() - .create_physical_plan() - .await - .unwrap_err(); - assert_contains!( - err.strip_backtrace(), - "MERGE INTO not supported for Base table" - ); + "MERGE INTO not supported for Base table", + ), + ] { + assert_merge_physical_error(&ctx, sql, expected).await; + } } #[tokio::test] diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index ffd84183cf256..4f282f5e067fe 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -3601,51 +3601,37 @@ 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(); +fn plan_merge_into_canonicalizes_qualifiers_and_preserves_quoted_columns() { + let plan = logical_plan( + "MERGE INTO person_quoted_cols AS t USING j2 AS s ON t.id = s.j2_id \ + WHEN MATCHED THEN UPDATE SET \"First Name\" = s.j2_string \ + WHEN NOT MATCHED THEN INSERT (id, \"Age\") VALUES (s.j2_id, 42)", + ) + .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}" - ); + assert_eq!(merge_op.on.to_string(), "person_quoted_cols.id = s.j2_id"); + + let datafusion_expr::dml::MergeIntoAction::Update(assignments) = + &merge_op.clauses[0].action + else { + panic!("expected UPDATE"); + }; + assert_eq!(assignments[0].0, "First Name"); + assert_eq!(assignments[0].1.to_string(), "s.j2_string"); + + let datafusion_expr::dml::MergeIntoAction::Insert { columns, values } = + &merge_op.clauses[1].action + else { + panic!("expected INSERT"); + }; + assert_eq!(columns, &["id".to_string(), "Age".to_string()]); + assert_eq!(values[0].to_string(), "s.j2_id"); } #[rstest] @@ -3665,18 +3651,6 @@ fn plan_merge_into_rejects_source_alias_colliding_with_target_name() { VALUES (j2.j2_id, j2.j2_string) WHERE false", "MERGE INSERT WHERE predicates are not supported" )] -fn plan_merge_into_rejects_unsupported_action_predicates( - #[case] sql: &str, - #[case] expected: &str, -) { - let err = logical_plan(sql).unwrap_err(); - assert!( - err.strip_backtrace().contains(expected), - "unexpected error: {err}" - ); -} - -#[rstest] #[case( "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ WHEN MATCHED THEN UPDATE SET j1_string = 'a', j1_string = 'b'", @@ -3713,7 +3687,7 @@ fn plan_merge_into_rejects_unsupported_action_predicates( "MERGE INTO j1 AS t(a) USING j2 ON true WHEN MATCHED THEN DELETE", "MERGE target alias column lists are not supported" )] -fn plan_merge_into_validates_action_targets_and_structure( +fn plan_merge_into_rejects_invalid_actions_and_structure( #[case] sql: &str, #[case] expected: &str, ) { @@ -3724,36 +3698,6 @@ fn plan_merge_into_validates_action_targets_and_structure( ); } -#[test] -fn plan_merge_into_preserves_quoted_action_columns() { - let plan = logical_plan( - "MERGE INTO person_quoted_cols AS t USING j2 AS s ON t.id = s.j2_id \ - WHEN MATCHED THEN UPDATE SET \"First Name\" = s.j2_string \ - WHEN NOT MATCHED THEN INSERT (id, \"Age\") VALUES (s.j2_id, 42)", - ) - .unwrap(); - let LogicalPlan::Dml(dml) = plan else { - panic!("expected Dml"); - }; - let datafusion_expr::WriteOp::MergeInto(merge_op) = dml.op else { - panic!("expected MergeInto"); - }; - - let datafusion_expr::dml::MergeIntoAction::Update(assignments) = - &merge_op.clauses[0].action - else { - panic!("expected UPDATE"); - }; - assert_eq!(assignments[0].0, "First Name"); - - let datafusion_expr::dml::MergeIntoAction::Insert { columns, .. } = - &merge_op.clauses[1].action - else { - panic!("expected INSERT"); - }; - assert_eq!(columns, &["id".to_string(), "Age".to_string()]); -} - fn logical_plan(sql: &str) -> Result { logical_plan_with_options(sql, ParserOptions::default()) }