feat: Add SQL planner, physical planner, and TableProvider hook for MERGE INTO

datafusion/catalog/src/table.rs

@@ -28,7 +28,7 @@ use datafusion_common::{Result, internal_err};
 use datafusion_expr::Expr;
 use datafusion_expr::statistics::StatisticsRequest;
 
-use datafusion_expr::dml::InsertOp;
+use datafusion_expr::dml::{InsertOp, MergeIntoClause};
 use datafusion_expr::{
     CreateExternalTable, LogicalPlan, TableProviderFilterPushDown, TableType,
 };
@@ -379,6 +379,23 @@ pub trait TableProvider: Any + Debug + Sync + Send {
     async fn truncate(&self, _state: &dyn Session) -> Result<Arc<dyn ExecutionPlan>> {
         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<dyn ExecutionPlan>,
+        _on: Expr,
+        _clauses: Vec<MergeIntoClause>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        not_impl_err!("MERGE INTO not supported for {} table", self.table_type())
+    }
 }
 
 impl dyn TableProvider {

datafusion/core/src/physical_planner.rs

@@ -927,6 +927,26 @@ impl DefaultPhysicalPlanner {
                     );
                 }
             }
+            LogicalPlan::Dml(DmlStatement {
+                table_name,
+                target,
+                op: WriteOp::MergeInto(merge_op),
+                ..
+            }) => {
+                let provider = source_as_provider(target)?;
+                let input_exec = children.one()?;
+                provider
+                    .merge_into(
+                        session_state,
+                        input_exec,
+                        merge_op.on.clone(),
+                        merge_op.clauses.clone(),
+                    )
+                    .await
+                    .map_err(|e| {
+                        e.context(format!("MERGE INTO operation on table '{table_name}'"))
+                    })?
+            }
             LogicalPlan::Window(Window { window_expr, .. }) => {
                 assert_or_internal_err!(
                     !window_expr.is_empty(),

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};
 
@@ -299,12 +299,129 @@ impl Display for InsertOp {
 }
 
 /// Describes a MERGE INTO operation's parameters.
-#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct MergeIntoOp {
     /// The join condition from `ON <expr>`.
     pub on: Expr,
     /// The WHEN clauses, in the order they appeared in the SQL.
     pub clauses: Vec<MergeIntoClause>,
+    /// Schema of the target table (qualified with the alias or table name used
+    /// in the SQL). Stored here so the analyzer can build a combined
+    /// target+source schema for type coercion and function rewrites.
+    pub target_schema: DFSchemaRef,
+}
+
+impl PartialOrd for MergeIntoOp {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        match self.on.partial_cmp(&other.on) {
+            Some(Ordering::Equal) => self.clauses.partial_cmp(&other.clauses),
+            cmp => cmp,
+        }
+    }
+}
+
+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::<usize>()
+    }
+
+    /// 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<Expr>) -> Result<Self> {
+        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,
+            target_schema: Arc::clone(&self.target_schema),
+        })
+    }
 }
 
 /// A single WHEN clause within a MERGE INTO statement.
@@ -400,6 +517,23 @@ fn make_count_schema() -> DFSchemaRef {
 mod tests {
     use super::*;
     use crate::{col, lit};
+    use arrow::datatypes::DataType;
+    use datafusion_common::DFSchema;
+
+    fn dummy_target_schema() -> DFSchemaRef {
+        Arc::new(
+            DFSchema::try_from_qualified_schema(
+                "target",
+                &Schema::new(vec![
+                    Field::new("id", DataType::Int64, false),
+                    Field::new("qty", DataType::Int64, true),
+                    Field::new("price", DataType::Float64, true),
+                    Field::new("active", DataType::Boolean, true),
+                ]),
+            )
+            .unwrap(),
+        )
+    }
 
     #[test]
     fn write_op_merge_into_name_and_display() {
@@ -413,6 +547,7 @@ mod tests {
                     col("source_qty"),
                 )]),
             }],
+            target_schema: dummy_target_schema(),
         }));
         assert_eq!(op.name(), "MergeInto");
         assert_eq!(format!("{op}"), "MergeInto");
@@ -445,4 +580,55 @@ 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,
+                },
+            ],
+            target_schema: dummy_target_schema(),
+        };
+        let exprs = op.exprs();
+        assert_eq!(exprs.len(), 7);
+
+        let owned: Vec<Expr> = 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![],
+            target_schema: dummy_target_schema(),
+        };
+        let err = op.with_new_exprs(vec![]).unwrap_err();
+        assert!(
+            err.to_string().contains("expected 1 expressions, got 0"),
+            "unexpected error: {err}"
+        );
+    }
 }

datafusion/expr/src/logical_plan/plan.rs

@@ -39,7 +39,7 @@ use crate::expr_rewriter::{
 };
 use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor};
 use crate::logical_plan::extension::UserDefinedLogicalNode;
-use crate::logical_plan::{DmlStatement, Statement};
+use crate::logical_plan::{DmlStatement, Statement, WriteOp};
 use crate::utils::{
     enumerate_grouping_sets, exprlist_to_fields, find_out_reference_exprs,
     grouping_set_expr_count, grouping_set_to_exprlist, merge_schema, split_conjunction,
@@ -810,12 +810,20 @@ impl LogicalPlan {
                 op,
                 ..
             }) => {
-                self.assert_no_expressions(expr)?;
                 let input = self.only_input(inputs)?;
+                let op = match op {
+                    WriteOp::MergeInto(merge_op) => {
+                        WriteOp::MergeInto(Box::new(merge_op.with_new_exprs(expr)?))
+                    }
+                    other => {
+                        self.assert_no_expressions(expr)?;
+                        other.clone()
+                    }
+                };
                 Ok(LogicalPlan::Dml(DmlStatement::new(
                     table_name.clone(),
                     Arc::clone(target),
-                    op.clone(),
+                    op,
                     Arc::new(input),
                 )))
             }

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<Expr> =
+                    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(_)

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,14 @@ impl ApplyFunctionRewrites {
             schema.merge(&source_schema);
         }
 
+        if let LogicalPlan::Dml(DmlStatement {
+            op: WriteOp::MergeInto(merge_op),
+            ..
+        }) = &plan
+        {
+            schema.merge(&merge_op.target_schema);
+        }
+
         let name_preserver = NamePreserver::new(&plan);
 
         plan.map_expressions(|expr| {