Skip to content

Commit e3b27a5

Browse files
committed
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.
1 parent a6bda4b commit e3b27a5

10 files changed

Lines changed: 652 additions & 12 deletions

File tree

datafusion/core/src/physical_planner.rs

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -927,6 +927,26 @@ impl DefaultPhysicalPlanner {
927927
);
928928
}
929929
}
930+
LogicalPlan::Dml(DmlStatement {
931+
table_name,
932+
target,
933+
op: WriteOp::MergeInto(merge_op),
934+
..
935+
}) => {
936+
let provider = source_as_provider(target)?;
937+
let input_exec = children.one()?;
938+
provider
939+
.merge_into(
940+
session_state,
941+
input_exec,
942+
merge_op.on.clone(),
943+
merge_op.clauses.clone(),
944+
)
945+
.await
946+
.map_err(|e| {
947+
e.context(format!("MERGE INTO operation on table '{table_name}'"))
948+
})?
949+
}
930950
LogicalPlan::Window(Window { window_expr, .. }) => {
931951
assert_or_internal_err!(
932952
!window_expr.is_empty(),

datafusion/core/tests/sql/sql_api.rs

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,37 @@ async fn ddl_can_not_be_planned_by_session_state() {
208208
);
209209
}
210210

211+
#[tokio::test]
212+
async fn merge_into_rejects_source_alias_colliding_with_target_name() {
213+
// Regression test: `MERGE INTO target AS t USING source AS target ...`
214+
// aliases the source to the target table's real name. Both `id` columns
215+
// are identically named. Before the fix, canonicalizing the aliased target
216+
// reference `t.id` to `target.id` collapsed it onto the source's
217+
// `target.id`, so `t.id = target.id` silently became `target.id = target.id`
218+
// (a comparison of the source column with itself). The planner must reject
219+
// this collision instead of changing the meaning of the condition.
220+
let ctx = SessionContext::new();
221+
ctx.sql("CREATE TABLE target (id INT, val INT)")
222+
.await
223+
.unwrap();
224+
ctx.sql("CREATE TABLE source (id INT, val INT)")
225+
.await
226+
.unwrap();
227+
228+
let err = ctx
229+
.sql(
230+
"MERGE INTO target AS t USING source AS target ON t.id = target.id \
231+
WHEN MATCHED THEN DELETE",
232+
)
233+
.await
234+
.unwrap_err();
235+
236+
assert_contains!(
237+
err.strip_backtrace(),
238+
"MERGE source may not use the target table name 'target' as a qualifier"
239+
);
240+
}
241+
211242
#[tokio::test]
212243
async fn invalid_wrapped_negation_fails_during_planning() {
213244
let ctx = SessionContext::new();

datafusion/expr/src/logical_plan/dml.rs

Lines changed: 150 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ use std::sync::Arc;
2323

2424
use arrow::datatypes::{DataType, Field, Schema};
2525
use datafusion_common::file_options::file_type::FileType;
26-
use datafusion_common::{DFSchemaRef, TableReference};
26+
use datafusion_common::{DFSchemaRef, Result, TableReference, internal_err};
2727

2828
use crate::{Expr, LogicalPlan, TableSource};
2929

@@ -307,6 +307,106 @@ pub struct MergeIntoOp {
307307
pub clauses: Vec<MergeIntoClause>,
308308
}
309309

310+
impl MergeIntoOp {
311+
/// Count of top-level [`Expr`]s owned by this operation (no allocation).
312+
///
313+
/// Matches the length of [`Self::exprs`] and the `exprs` vec consumed by
314+
/// [`Self::with_new_exprs`].
315+
fn expr_count(&self) -> usize {
316+
1 + self
317+
.clauses
318+
.iter()
319+
.map(|c| {
320+
c.predicate.is_some() as usize
321+
+ match &c.action {
322+
MergeIntoAction::Update(a) => a.len(),
323+
MergeIntoAction::Insert { values, .. } => values.len(),
324+
MergeIntoAction::Delete => 0,
325+
}
326+
})
327+
.sum::<usize>()
328+
}
329+
330+
/// Top-level [`Expr`]s in stable order: `on`, then per-clause predicate
331+
/// (if any) and action value expressions.
332+
pub fn exprs(&self) -> Vec<&Expr> {
333+
let mut out = Vec::with_capacity(self.expr_count());
334+
out.push(&self.on);
335+
for clause in &self.clauses {
336+
if let Some(predicate) = &clause.predicate {
337+
out.push(predicate);
338+
}
339+
match &clause.action {
340+
MergeIntoAction::Update(assignments) => {
341+
out.extend(assignments.iter().map(|(_, value)| value));
342+
}
343+
MergeIntoAction::Insert { values, .. } => {
344+
out.extend(values.iter());
345+
}
346+
MergeIntoAction::Delete => {}
347+
}
348+
}
349+
out
350+
}
351+
352+
/// Rebuild this `MergeIntoOp` from a flat vector of new expressions, in
353+
/// the same order produced by [`Self::exprs`]. The clause kinds, action
354+
/// kinds, column lists, and presence/absence of each predicate are
355+
/// preserved from `self`.
356+
pub fn with_new_exprs(&self, exprs: Vec<Expr>) -> Result<Self> {
357+
let expected = self.expr_count();
358+
if exprs.len() != expected {
359+
return internal_err!(
360+
"MergeIntoOp::with_new_exprs expected {expected} expressions, got {}",
361+
exprs.len()
362+
);
363+
}
364+
let mut iter = exprs.into_iter();
365+
let on = iter.next().expect("non-empty by length check");
366+
let clauses = self
367+
.clauses
368+
.iter()
369+
.map(|clause| {
370+
let predicate = clause
371+
.predicate
372+
.is_some()
373+
.then(|| iter.next().expect("non-empty by length check"));
374+
let action = match &clause.action {
375+
MergeIntoAction::Update(assignments) => {
376+
let assignments = assignments
377+
.iter()
378+
.map(|(name, _)| {
379+
(
380+
name.clone(),
381+
iter.next().expect("non-empty by length check"),
382+
)
383+
})
384+
.collect();
385+
MergeIntoAction::Update(assignments)
386+
}
387+
MergeIntoAction::Insert { columns, values } => {
388+
let values = values
389+
.iter()
390+
.map(|_| iter.next().expect("non-empty by length check"))
391+
.collect();
392+
MergeIntoAction::Insert {
393+
columns: columns.clone(),
394+
values,
395+
}
396+
}
397+
MergeIntoAction::Delete => MergeIntoAction::Delete,
398+
};
399+
MergeIntoClause {
400+
kind: clause.kind,
401+
predicate,
402+
action,
403+
}
404+
})
405+
.collect();
406+
Ok(Self { on, clauses })
407+
}
408+
}
409+
310410
/// A single WHEN clause within a MERGE INTO statement.
311411
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
312412
pub struct MergeIntoClause {
@@ -445,4 +545,53 @@ mod tests {
445545
MergeIntoClauseKind::NotMatchedBySource
446546
);
447547
}
548+
549+
#[test]
550+
fn merge_into_op_exprs_round_trip() {
551+
let op = MergeIntoOp {
552+
on: col("id").eq(col("source_id")),
553+
clauses: vec![
554+
MergeIntoClause {
555+
kind: MergeIntoClauseKind::Matched,
556+
predicate: Some(col("qty").gt(lit(0_i64))),
557+
action: MergeIntoAction::Update(vec![
558+
("qty".to_string(), col("source_qty")),
559+
("price".to_string(), col("source_price")),
560+
]),
561+
},
562+
MergeIntoClause {
563+
kind: MergeIntoClauseKind::NotMatched,
564+
predicate: None,
565+
action: MergeIntoAction::Insert {
566+
columns: vec!["id".to_string(), "qty".to_string()],
567+
values: vec![col("source_id"), col("source_qty")],
568+
},
569+
},
570+
MergeIntoClause {
571+
kind: MergeIntoClauseKind::NotMatchedBySource,
572+
predicate: Some(col("active").eq(lit(true))),
573+
action: MergeIntoAction::Delete,
574+
},
575+
],
576+
};
577+
let exprs = op.exprs();
578+
assert_eq!(exprs.len(), 7);
579+
580+
let owned: Vec<Expr> = exprs.into_iter().cloned().collect();
581+
let rebuilt = op.with_new_exprs(owned).unwrap();
582+
assert_eq!(op, rebuilt);
583+
}
584+
585+
#[test]
586+
fn merge_into_op_with_new_exprs_length_mismatch() {
587+
let op = MergeIntoOp {
588+
on: col("id").eq(col("source_id")),
589+
clauses: vec![],
590+
};
591+
let err = op.with_new_exprs(vec![]).unwrap_err();
592+
assert!(
593+
err.to_string().contains("expected 1 expressions, got 0"),
594+
"unexpected error: {err}"
595+
);
596+
}
448597
}

datafusion/expr/src/logical_plan/plan.rs

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ use crate::expr_rewriter::{
3939
};
4040
use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor};
4141
use crate::logical_plan::extension::UserDefinedLogicalNode;
42-
use crate::logical_plan::{DmlStatement, Statement};
42+
use crate::logical_plan::{DmlStatement, Statement, WriteOp};
4343
use crate::utils::{
4444
enumerate_grouping_sets, exprlist_to_fields, find_out_reference_exprs,
4545
grouping_set_expr_count, grouping_set_to_exprlist, merge_schema, split_conjunction,
@@ -810,12 +810,20 @@ impl LogicalPlan {
810810
op,
811811
..
812812
}) => {
813-
self.assert_no_expressions(expr)?;
814813
let input = self.only_input(inputs)?;
814+
let op = match op {
815+
WriteOp::MergeInto(merge_op) => {
816+
WriteOp::MergeInto(Box::new(merge_op.with_new_exprs(expr)?))
817+
}
818+
other => {
819+
self.assert_no_expressions(expr)?;
820+
other.clone()
821+
}
822+
};
815823
Ok(LogicalPlan::Dml(DmlStatement::new(
816824
table_name.clone(),
817825
Arc::clone(target),
818-
op.clone(),
826+
op,
819827
Arc::new(input),
820828
)))
821829
}

datafusion/expr/src/logical_plan/tree_node.rs

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ use crate::{
4545
DistinctOn, DmlStatement, Execute, Explain, Expr, Extension, Filter, Join, Limit,
4646
LogicalPlan, Partitioning, Prepare, Projection, RecursiveQuery, Repartition, Sort,
4747
Statement, Subquery, SubqueryAlias, TableScan, Union, Unnest, UserDefinedLogicalNode,
48-
Values, Window, builder::unnest_with_options, dml::CopyTo,
48+
Values, Window, WriteOp, builder::unnest_with_options, dml::CopyTo,
4949
};
5050
use datafusion_common::tree_node::TreeNodeRefContainer;
5151

@@ -480,6 +480,10 @@ impl LogicalPlan {
480480
}
481481
_ => Ok(TreeNodeRecursion::Continue),
482482
},
483+
LogicalPlan::Dml(DmlStatement {
484+
op: WriteOp::MergeInto(merge_op),
485+
..
486+
}) => merge_op.exprs().apply_ref_elements(f),
483487
// plans without expressions
484488
LogicalPlan::EmptyRelation(_)
485489
| LogicalPlan::RecursiveQuery(_)
@@ -719,6 +723,27 @@ impl LogicalPlan {
719723
)
720724
})?
721725
}
726+
LogicalPlan::Dml(DmlStatement {
727+
table_name,
728+
target,
729+
op: WriteOp::MergeInto(merge_op),
730+
input,
731+
output_schema,
732+
}) => {
733+
let owned_exprs: Vec<Expr> =
734+
merge_op.exprs().into_iter().cloned().collect();
735+
owned_exprs.map_elements(f)?.transform_data(|new_exprs| {
736+
Ok(Transformed::no(LogicalPlan::Dml(DmlStatement {
737+
table_name,
738+
target,
739+
op: WriteOp::MergeInto(Box::new(
740+
merge_op.with_new_exprs(new_exprs)?,
741+
)),
742+
input,
743+
output_schema,
744+
})))
745+
})?
746+
}
722747
// plans without expressions
723748
LogicalPlan::EmptyRelation(_)
724749
| LogicalPlan::RecursiveQuery(_)

datafusion/optimizer/src/analyzer/function_rewrite.rs

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,9 @@ use datafusion_common::tree_node::{Transformed, TreeNode};
2323
use datafusion_common::{DFSchema, Result};
2424

2525
use crate::utils::NamePreserver;
26-
use datafusion_expr::LogicalPlan;
2726
use datafusion_expr::expr_rewriter::FunctionRewrite;
2827
use datafusion_expr::utils::merge_schema;
28+
use datafusion_expr::{DmlStatement, LogicalPlan, WriteOp};
2929
use std::sync::Arc;
3030

3131
/// Analyzer rule that invokes [`FunctionRewrite`]s on expressions
@@ -58,6 +58,23 @@ impl ApplyFunctionRewrites {
5858
schema.merge(&source_schema);
5959
}
6060

61+
// MERGE expressions reference the target table, which is not one of
62+
// `plan.inputs()`. Rebuild the target schema from the DML's
63+
// `table_name` and `target` so those columns resolve.
64+
if let LogicalPlan::Dml(DmlStatement {
65+
op: WriteOp::MergeInto(_),
66+
table_name,
67+
target,
68+
..
69+
}) = &plan
70+
{
71+
let target_schema = DFSchema::try_from_qualified_schema(
72+
table_name.clone(),
73+
&target.schema(),
74+
)?;
75+
schema.merge(&target_schema);
76+
}
77+
6178
let name_preserver = NamePreserver::new(&plan);
6279

6380
plan.map_expressions(|expr| {

0 commit comments

Comments
 (0)