feat: Add MERGE INTO support to FFI TableProvider

datafusion/core/src/physical_planner.rs

@@ -967,6 +967,35 @@ impl DefaultPhysicalPlanner {
                     );
                 }
             }
+            LogicalPlan::Dml(DmlStatement {
+                table_name,
+                target,
+                op: WriteOp::MergeInto(merge_op),
+                input,
+                ..
+            }) => {
+                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(),
+                    )
+                    .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(),
@@ -3378,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::{
@@ -3402,6 +3432,85 @@ mod tests {
             .build()
     }
 
+    #[derive(Debug)]
+    struct CaptureMergeProvider {
+        schema: SchemaRef,
+        captured: Mutex<Option<(DFSchemaRef, String, usize)>>,
+    }
+
+    #[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<usize>>,
+            _filters: &[Expr],
+            _limit: Option<usize>,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema))))
+        }
+
+        async fn merge_into(
+            &self,
+            state: &dyn Session,
+            source: Arc<dyn ExecutionPlan>,
+            merge_schema: DFSchemaRef,
+            on: Expr,
+            clauses: Vec<MergeIntoClause>,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            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)]));
+        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.id > s.id THEN DELETE",
+        )
+        .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, 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"))?,
+            1
+        );
+        assert_contains!(physical_on, "index: 0");
+        assert_contains!(physical_on, "index: 1");
+        Ok(())
+    }
+
     async fn plan(logical_plan: &LogicalPlan) -> Result<Arc<dyn ExecutionPlan>> {
         let session_state = make_session_state();
         // optimize the logical plan

datafusion/core/tests/sql/sql_api.rs

@@ -208,6 +208,102 @@ async fn ddl_can_not_be_planned_by_session_state() {
     );
 }
 
+async fn merge_into_context() -> SessionContext {
+    let ctx = SessionContext::new();
+    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()
+        .create_physical_plan()
+        .await
+        .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"] {
+        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"
+            ),
+            &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 = 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-correlated and uncorrelated subqueries 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",
+    ] {
+        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 = merge_into_context().await;
+
+    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",
+        ),
+        (
+            "MERGE INTO target USING source ON NULL \
+             WHEN MATCHED AND NULL THEN DELETE",
+            "MERGE INTO not supported for Base table",
+        ),
+    ] {
+        assert_merge_physical_error(&ctx, sql, expected).await;
+    }
+}
+
 #[tokio::test]
 async fn invalid_wrapped_negation_fails_during_planning() {
     let ctx = SessionContext::new();