Skip to content

Commit 799ed48

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 2c6eada commit 799ed48

5 files changed

Lines changed: 419 additions & 6 deletions

File tree

datafusion/core/src/physical_planner.rs

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

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(_)

0 commit comments

Comments
 (0)