Skip to content

Commit 98dd55f

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 98dd55f

10 files changed

Lines changed: 564 additions & 17 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: 188 additions & 2 deletions
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

@@ -299,12 +299,129 @@ impl Display for InsertOp {
299299
}
300300

301301
/// Describes a MERGE INTO operation's parameters.
302-
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
302+
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
303303
pub struct MergeIntoOp {
304304
/// The join condition from `ON <expr>`.
305305
pub on: Expr,
306306
/// The WHEN clauses, in the order they appeared in the SQL.
307307
pub clauses: Vec<MergeIntoClause>,
308+
/// Schema of the target table (qualified with the alias or table name used
309+
/// in the SQL). Stored here so the analyzer can build a combined
310+
/// target+source schema for type coercion and function rewrites.
311+
pub target_schema: DFSchemaRef,
312+
}
313+
314+
impl PartialOrd for MergeIntoOp {
315+
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
316+
match self.on.partial_cmp(&other.on) {
317+
Some(Ordering::Equal) => self.clauses.partial_cmp(&other.clauses),
318+
cmp => cmp,
319+
}
320+
}
321+
}
322+
323+
impl MergeIntoOp {
324+
/// Count of top-level [`Expr`]s owned by this operation (no allocation).
325+
///
326+
/// Matches the length of [`Self::exprs`] and the `exprs` vec consumed by
327+
/// [`Self::with_new_exprs`].
328+
fn expr_count(&self) -> usize {
329+
1 + self
330+
.clauses
331+
.iter()
332+
.map(|c| {
333+
c.predicate.is_some() as usize
334+
+ match &c.action {
335+
MergeIntoAction::Update(a) => a.len(),
336+
MergeIntoAction::Insert { values, .. } => values.len(),
337+
MergeIntoAction::Delete => 0,
338+
}
339+
})
340+
.sum::<usize>()
341+
}
342+
343+
/// Top-level [`Expr`]s in stable order: `on`, then per-clause predicate
344+
/// (if any) and action value expressions.
345+
pub fn exprs(&self) -> Vec<&Expr> {
346+
let mut out = Vec::with_capacity(self.expr_count());
347+
out.push(&self.on);
348+
for clause in &self.clauses {
349+
if let Some(predicate) = &clause.predicate {
350+
out.push(predicate);
351+
}
352+
match &clause.action {
353+
MergeIntoAction::Update(assignments) => {
354+
out.extend(assignments.iter().map(|(_, value)| value));
355+
}
356+
MergeIntoAction::Insert { values, .. } => {
357+
out.extend(values.iter());
358+
}
359+
MergeIntoAction::Delete => {}
360+
}
361+
}
362+
out
363+
}
364+
365+
/// Rebuild this `MergeIntoOp` from a flat vector of new expressions, in
366+
/// the same order produced by [`Self::exprs`]. The clause kinds, action
367+
/// kinds, column lists, and presence/absence of each predicate are
368+
/// preserved from `self`.
369+
pub fn with_new_exprs(&self, exprs: Vec<Expr>) -> Result<Self> {
370+
let expected = self.expr_count();
371+
if exprs.len() != expected {
372+
return internal_err!(
373+
"MergeIntoOp::with_new_exprs expected {expected} expressions, got {}",
374+
exprs.len()
375+
);
376+
}
377+
let mut iter = exprs.into_iter();
378+
let on = iter.next().expect("non-empty by length check");
379+
let clauses = self
380+
.clauses
381+
.iter()
382+
.map(|clause| {
383+
let predicate = clause
384+
.predicate
385+
.is_some()
386+
.then(|| iter.next().expect("non-empty by length check"));
387+
let action = match &clause.action {
388+
MergeIntoAction::Update(assignments) => {
389+
let assignments = assignments
390+
.iter()
391+
.map(|(name, _)| {
392+
(
393+
name.clone(),
394+
iter.next().expect("non-empty by length check"),
395+
)
396+
})
397+
.collect();
398+
MergeIntoAction::Update(assignments)
399+
}
400+
MergeIntoAction::Insert { columns, values } => {
401+
let values = values
402+
.iter()
403+
.map(|_| iter.next().expect("non-empty by length check"))
404+
.collect();
405+
MergeIntoAction::Insert {
406+
columns: columns.clone(),
407+
values,
408+
}
409+
}
410+
MergeIntoAction::Delete => MergeIntoAction::Delete,
411+
};
412+
MergeIntoClause {
413+
kind: clause.kind,
414+
predicate,
415+
action,
416+
}
417+
})
418+
.collect();
419+
Ok(Self {
420+
on,
421+
clauses,
422+
target_schema: Arc::clone(&self.target_schema),
423+
})
424+
}
308425
}
309426

310427
/// A single WHEN clause within a MERGE INTO statement.
@@ -400,6 +517,23 @@ fn make_count_schema() -> DFSchemaRef {
400517
mod tests {
401518
use super::*;
402519
use crate::{col, lit};
520+
use arrow::datatypes::DataType;
521+
use datafusion_common::DFSchema;
522+
523+
fn dummy_target_schema() -> DFSchemaRef {
524+
Arc::new(
525+
DFSchema::try_from_qualified_schema(
526+
"target",
527+
&Schema::new(vec![
528+
Field::new("id", DataType::Int64, false),
529+
Field::new("qty", DataType::Int64, true),
530+
Field::new("price", DataType::Float64, true),
531+
Field::new("active", DataType::Boolean, true),
532+
]),
533+
)
534+
.unwrap(),
535+
)
536+
}
403537

404538
#[test]
405539
fn write_op_merge_into_name_and_display() {
@@ -413,6 +547,7 @@ mod tests {
413547
col("source_qty"),
414548
)]),
415549
}],
550+
target_schema: dummy_target_schema(),
416551
}));
417552
assert_eq!(op.name(), "MergeInto");
418553
assert_eq!(format!("{op}"), "MergeInto");
@@ -445,4 +580,55 @@ mod tests {
445580
MergeIntoClauseKind::NotMatchedBySource
446581
);
447582
}
583+
584+
#[test]
585+
fn merge_into_op_exprs_round_trip() {
586+
let op = MergeIntoOp {
587+
on: col("id").eq(col("source_id")),
588+
clauses: vec![
589+
MergeIntoClause {
590+
kind: MergeIntoClauseKind::Matched,
591+
predicate: Some(col("qty").gt(lit(0_i64))),
592+
action: MergeIntoAction::Update(vec![
593+
("qty".to_string(), col("source_qty")),
594+
("price".to_string(), col("source_price")),
595+
]),
596+
},
597+
MergeIntoClause {
598+
kind: MergeIntoClauseKind::NotMatched,
599+
predicate: None,
600+
action: MergeIntoAction::Insert {
601+
columns: vec!["id".to_string(), "qty".to_string()],
602+
values: vec![col("source_id"), col("source_qty")],
603+
},
604+
},
605+
MergeIntoClause {
606+
kind: MergeIntoClauseKind::NotMatchedBySource,
607+
predicate: Some(col("active").eq(lit(true))),
608+
action: MergeIntoAction::Delete,
609+
},
610+
],
611+
target_schema: dummy_target_schema(),
612+
};
613+
let exprs = op.exprs();
614+
assert_eq!(exprs.len(), 7);
615+
616+
let owned: Vec<Expr> = exprs.into_iter().cloned().collect();
617+
let rebuilt = op.with_new_exprs(owned).unwrap();
618+
assert_eq!(op, rebuilt);
619+
}
620+
621+
#[test]
622+
fn merge_into_op_with_new_exprs_length_mismatch() {
623+
let op = MergeIntoOp {
624+
on: col("id").eq(col("source_id")),
625+
clauses: vec![],
626+
target_schema: dummy_target_schema(),
627+
};
628+
let err = op.with_new_exprs(vec![]).unwrap_err();
629+
assert!(
630+
err.to_string().contains("expected 1 expressions, got 0"),
631+
"unexpected error: {err}"
632+
);
633+
}
448634
}

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: 9 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,14 @@ impl ApplyFunctionRewrites {
5858
schema.merge(&source_schema);
5959
}
6060

61+
if let LogicalPlan::Dml(DmlStatement {
62+
op: WriteOp::MergeInto(merge_op),
63+
..
64+
}) = &plan
65+
{
66+
schema.merge(&merge_op.target_schema);
67+
}
68+
6169
let name_preserver = NamePreserver::new(&plan);
6270

6371
plan.map_expressions(|expr| {

0 commit comments

Comments
 (0)