Fix MERGE INTO planner edge cases

Author: wirybeaver

datafusion/catalog/src/table.rs

@@ -24,7 +24,7 @@ use crate::session::Session;
 use arrow::datatypes::SchemaRef;
 use async_trait::async_trait;
 use datafusion_common::{Constraints, Statistics, not_impl_err};
-use datafusion_common::{Result, internal_err};
+use datafusion_common::{DFSchemaRef, Result, internal_err};
 use datafusion_expr::Expr;
 use datafusion_expr::statistics::StatisticsRequest;
 
@@ -383,6 +383,10 @@ pub trait TableProvider: Any + Debug + Sync + Send {
     /// Merge rows from a source into this table.
     ///
     /// The `source` is an [`ExecutionPlan`] representing the USING clause.
+    /// The `merge_schema` contains the target columns followed by the source
+    /// columns, preserving their logical qualifiers. Providers can use this
+    /// schema to resolve the logical expressions against the combined rows
+    /// they construct while executing the merge.
     /// The `on` condition is the join predicate from the ON clause.
     /// The `clauses` describe the WHEN MATCHED / WHEN NOT MATCHED actions.
     ///
@@ -391,6 +395,7 @@ pub trait TableProvider: Any + Debug + Sync + Send {
         &self,
         _state: &dyn Session,
         _source: Arc<dyn ExecutionPlan>,
+        _merge_schema: DFSchemaRef,
         _on: Expr,
         _clauses: Vec<MergeIntoClause>,
     ) -> Result<Arc<dyn ExecutionPlan>> {

datafusion/core/src/physical_planner.rs

@@ -931,14 +931,23 @@ impl DefaultPhysicalPlanner {
                 table_name,
                 target,
                 op: WriteOp::MergeInto(merge_op),
+                input,
                 ..
             }) => {
-                let provider = source_as_provider(target)?;
+                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(),
                     )
@@ -3262,6 +3271,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::{
@@ -3285,6 +3295,88 @@ 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),
+            Field::new("val", DataType::Int32, true),
+        ]));
+        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.val IS NULL THEN UPDATE SET val = s.val \
+             WHEN NOT MATCHED THEN INSERT (id, val) VALUES (s.id, s.val)",
+        )
+        .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, 2);
+        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"))?,
+            2
+        );
+        assert_contains!(physical_on, "index: 0");
+        assert_contains!(physical_on, "index: 2");
+        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

@@ -225,17 +225,117 @@ async fn merge_into_rejects_source_alias_colliding_with_target_name() {
         .await
         .unwrap();
 
+    for target_ref in ["target", "public.target", "datafusion.public.target"] {
+        let err = ctx
+            .sql(&format!(
+                "MERGE INTO {target_ref} AS t USING source AS target \
+                 ON t.id = target.id WHEN MATCHED THEN DELETE"
+            ))
+            .await
+            .unwrap_err();
+
+        assert_contains!(
+            err.strip_backtrace(),
+            &format!(
+                "MERGE source may not use the target table name '{target_ref}' \
+                 as a qualifier"
+            )
+        );
+    }
+}
+
+#[tokio::test]
+async fn merge_into_rejects_subqueries_correlated_to_target_alias() {
+    let ctx = SessionContext::new();
+    ctx.sql("CREATE TABLE target (id INT, val INT)")
+        .await
+        .unwrap();
+    ctx.sql("CREATE TABLE source (id INT, val INT)")
+        .await
+        .unwrap();
+
     let err = ctx
         .sql(
-            "MERGE INTO target AS t USING source AS target ON t.id = target.id \
+            "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",
         )
         .await
         .unwrap_err();
+    assert_contains!(
+        err.strip_backtrace(),
+        "MERGE subqueries correlated to target alias 't' are not supported"
+    );
+
+    // Source correlation and uncorrelated set-comparison subqueries do not
+    // require target-alias canonicalization and 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",
+    ] {
+        let err = ctx
+            .sql(sql)
+            .await
+            .unwrap()
+            .create_physical_plan()
+            .await
+            .unwrap_err();
+        assert_contains!(
+            err.strip_backtrace(),
+            "MERGE INTO not supported for Base table"
+        );
+    }
+}
 
+#[tokio::test]
+async fn merge_into_requires_boolean_conditions() {
+    let ctx = SessionContext::new();
+    ctx.sql("CREATE TABLE target (id INT, val INT)")
+        .await
+        .unwrap();
+    ctx.sql("CREATE TABLE source (id INT, val INT)")
+        .await
+        .unwrap();
+
+    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",
+        ),
+    ] {
+        let err = ctx
+            .sql(sql)
+            .await
+            .unwrap()
+            .create_physical_plan()
+            .await
+            .unwrap_err();
+        assert_contains!(err.strip_backtrace(), expected);
+    }
+
+    let err = ctx
+        .sql(
+            "MERGE INTO target USING source ON NULL \
+             WHEN MATCHED AND NULL THEN DELETE",
+        )
+        .await
+        .unwrap()
+        .create_physical_plan()
+        .await
+        .unwrap_err();
     assert_contains!(
         err.strip_backtrace(),
-        "MERGE source may not use the target table name 'target' as a qualifier"
+        "MERGE INTO not supported for Base table"
     );
 }
 

datafusion/expr/src/logical_plan/invariants.rs

@@ -21,7 +21,7 @@ use datafusion_common::{
 };
 
 use crate::{
-    Aggregate, Expr, Filter, Join, JoinType, LogicalPlan, Window,
+    Aggregate, DmlStatement, Expr, Filter, Join, JoinType, LogicalPlan, Window, WriteOp,
     expr::{Exists, InSubquery, SetComparison},
     expr_rewriter::strip_outer_reference,
     utils::{collect_subquery_cols, split_conjunction},
@@ -253,7 +253,11 @@ pub fn check_subquery_expr(
             | LogicalPlan::TableScan(_)
             | LogicalPlan::Window(_)
             | LogicalPlan::Aggregate(_)
-            | LogicalPlan::Join(_) => Ok(()),
+            | LogicalPlan::Join(_)
+            | LogicalPlan::Dml(DmlStatement {
+                op: WriteOp::MergeInto(_),
+                ..
+            }) => Ok(()),
             _ => plan_err!(
                 "In/Exist/SetComparison subquery can only be used in \
                 Projection, Filter, TableScan, Window functions, Aggregate and Join plan nodes, \