Consolidate MERGE INTO regression tests

Author: wirybeaver

datafusion/core/src/physical_planner.rs

@@ -3475,10 +3475,8 @@ mod tests {
 
     #[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 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),
@@ -3490,8 +3488,7 @@ mod tests {
 
         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)",
+             WHEN MATCHED AND t.id > s.id THEN DELETE",
         )
         .await?
         .create_physical_plan()
@@ -3500,17 +3497,17 @@ mod tests {
         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!(*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"))?,
-            2
+            1
         );
         assert_contains!(physical_on, "index: 0");
-        assert_contains!(physical_on, "index: 2");
+        assert_contains!(physical_on, "index: 1");
         Ok(())
     }
 

datafusion/core/tests/sql/sql_api.rs

@@ -208,67 +208,64 @@ async fn ddl_can_not_be_planned_by_session_state() {
     );
 }
 
-#[tokio::test]
-async fn merge_into_rejects_source_alias_colliding_with_target_name() {
-    // Regression test: `MERGE INTO target AS t USING source AS target ...`
-    // aliases the source to the target table's real name. Both `id` columns
-    // are identically named. Before the fix, canonicalizing the aliased target
-    // reference `t.id` to `target.id` collapsed it onto the source's
-    // `target.id`, so `t.id = target.id` silently became `target.id = target.id`
-    // (a comparison of the source column with itself). The planner must reject
-    // this collision instead of changing the meaning of the condition.
+async fn merge_into_context() -> SessionContext {
     let ctx = SessionContext::new();
-    ctx.sql("CREATE TABLE target (id INT, val INT)")
+    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();
-    ctx.sql("CREATE TABLE source (id INT, val INT)")
+        .unwrap()
+        .create_physical_plan()
         .await
-        .unwrap();
+        .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"] {
-        let err = ctx
-            .sql(&format!(
+        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"
-            ))
-            .await
-            .unwrap_err();
-
-        assert_contains!(
-            err.strip_backtrace(),
+            ),
             &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 = 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 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"
-    );
+    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 correlation and uncorrelated set-comparison subqueries do not
-    // require target-alias canonicalization and remain supported through
+    // Source-correlated and uncorrelated subqueries remain supported through
     // logical optimization.
     for sql in [
         "MERGE INTO target AS t USING source AS s \
@@ -278,29 +275,14 @@ async fn merge_into_rejects_subqueries_correlated_to_target_alias() {
          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"
-        );
+        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 = 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 ctx = merge_into_context().await;
 
     for (sql, expected) in [
         (
@@ -312,31 +294,14 @@ async fn merge_into_requires_boolean_conditions() {
              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 INTO not supported for Base table"
-    );
+            "MERGE INTO not supported for Base table",
+        ),
+    ] {
+        assert_merge_physical_error(&ctx, sql, expected).await;
+    }
 }
 
 #[tokio::test]

datafusion/sql/tests/sql_integration.rs

@@ -3601,51 +3601,37 @@ fn select_groupby_orderby_aggregate_on_non_selected_column_original_issue() {
 }
 
 #[test]
-fn plan_merge_into_canonicalizes_target_alias() {
-    // The target alias `t` must be rewritten to the real table name `j1` in the
-    // stored plan, while the source alias `s` is preserved. This keeps the plan
-    // independent of the SQL alias so analyzer passes and proto deserialization
-    // can rebuild the target schema from the table name alone.
-    let sql = "MERGE INTO j1 AS t USING j2 AS s ON t.j1_id = s.j2_id \
-               WHEN MATCHED THEN UPDATE SET j1_string = s.j2_string \
-               WHEN NOT MATCHED THEN INSERT (j1_id, j1_string) VALUES (s.j2_id, s.j2_string)";
-    let plan = logical_plan(sql).unwrap();
+fn plan_merge_into_canonicalizes_qualifiers_and_preserves_quoted_columns() {
+    let plan = logical_plan(
+        "MERGE INTO person_quoted_cols AS t USING j2 AS s ON t.id = s.j2_id \
+         WHEN MATCHED THEN UPDATE SET \"First Name\" = s.j2_string \
+         WHEN NOT MATCHED THEN INSERT (id, \"Age\") VALUES (s.j2_id, 42)",
+    )
+    .unwrap();
     let LogicalPlan::Dml(dml) = &plan else {
         panic!("expected Dml, got {plan:?}");
     };
     let datafusion_expr::WriteOp::MergeInto(merge_op) = &dml.op else {
         panic!("expected MergeInto, got {:?}", dml.op);
     };
-    // `t.j1_id` -> `j1.j1_id`, `s.j2_id` left untouched.
-    assert_eq!(merge_op.on.to_string(), "j1.j1_id = s.j2_id");
-    // Source-side column references remain qualified with the source alias.
-    let exprs: Vec<String> = merge_op.exprs().iter().map(|e| e.to_string()).collect();
-    assert_eq!(
-        exprs,
-        vec![
-            "j1.j1_id = s.j2_id".to_string(),
-            "s.j2_string".to_string(),
-            "s.j2_id".to_string(),
-            "s.j2_string".to_string(),
-        ]
-    );
-}
 
-#[test]
-fn plan_merge_into_rejects_source_alias_colliding_with_target_name() {
-    // Aliasing the source to the target table's real name (`USING j2 AS j1`,
-    // where `j1` is the target) would make canonicalizing the target alias
-    // `t` to `j1` collide with the source qualifier, silently collapsing the
-    // two namespaces. The planner must reject this rather than misresolve.
-    let sql = "MERGE INTO j1 AS t USING j2 AS j1 ON t.j1_id = j1.j2_id \
-               WHEN MATCHED THEN DELETE";
-    let err = logical_plan(sql).unwrap_err();
-    assert!(
-        err.strip_backtrace().contains(
-            "MERGE source may not use the target table name 'j1' as a qualifier"
-        ),
-        "unexpected error: {err}"
-    );
+    assert_eq!(merge_op.on.to_string(), "person_quoted_cols.id = s.j2_id");
+
+    let datafusion_expr::dml::MergeIntoAction::Update(assignments) =
+        &merge_op.clauses[0].action
+    else {
+        panic!("expected UPDATE");
+    };
+    assert_eq!(assignments[0].0, "First Name");
+    assert_eq!(assignments[0].1.to_string(), "s.j2_string");
+
+    let datafusion_expr::dml::MergeIntoAction::Insert { columns, values } =
+        &merge_op.clauses[1].action
+    else {
+        panic!("expected INSERT");
+    };
+    assert_eq!(columns, &["id".to_string(), "Age".to_string()]);
+    assert_eq!(values[0].to_string(), "s.j2_id");
 }
 
 #[rstest]
@@ -3665,18 +3651,6 @@ fn plan_merge_into_rejects_source_alias_colliding_with_target_name() {
      VALUES (j2.j2_id, j2.j2_string) WHERE false",
     "MERGE INSERT WHERE predicates are not supported"
 )]
-fn plan_merge_into_rejects_unsupported_action_predicates(
-    #[case] sql: &str,
-    #[case] expected: &str,
-) {
-    let err = logical_plan(sql).unwrap_err();
-    assert!(
-        err.strip_backtrace().contains(expected),
-        "unexpected error: {err}"
-    );
-}
-
-#[rstest]
 #[case(
     "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \
      WHEN MATCHED THEN UPDATE SET j1_string = 'a', j1_string = 'b'",
@@ -3713,7 +3687,7 @@ fn plan_merge_into_rejects_unsupported_action_predicates(
     "MERGE INTO j1 AS t(a) USING j2 ON true WHEN MATCHED THEN DELETE",
     "MERGE target alias column lists are not supported"
 )]
-fn plan_merge_into_validates_action_targets_and_structure(
+fn plan_merge_into_rejects_invalid_actions_and_structure(
     #[case] sql: &str,
     #[case] expected: &str,
 ) {
@@ -3724,36 +3698,6 @@ fn plan_merge_into_validates_action_targets_and_structure(
     );
 }
 
-#[test]
-fn plan_merge_into_preserves_quoted_action_columns() {
-    let plan = logical_plan(
-        "MERGE INTO person_quoted_cols AS t USING j2 AS s ON t.id = s.j2_id \
-         WHEN MATCHED THEN UPDATE SET \"First Name\" = s.j2_string \
-         WHEN NOT MATCHED THEN INSERT (id, \"Age\") VALUES (s.j2_id, 42)",
-    )
-    .unwrap();
-    let LogicalPlan::Dml(dml) = plan else {
-        panic!("expected Dml");
-    };
-    let datafusion_expr::WriteOp::MergeInto(merge_op) = dml.op else {
-        panic!("expected MergeInto");
-    };
-
-    let datafusion_expr::dml::MergeIntoAction::Update(assignments) =
-        &merge_op.clauses[0].action
-    else {
-        panic!("expected UPDATE");
-    };
-    assert_eq!(assignments[0].0, "First Name");
-
-    let datafusion_expr::dml::MergeIntoAction::Insert { columns, .. } =
-        &merge_op.clauses[1].action
-    else {
-        panic!("expected INSERT");
-    };
-    assert_eq!(columns, &["id".to_string(), "Age".to_string()]);
-}
-
 fn logical_plan(sql: &str) -> Result<LogicalPlan> {
     logical_plan_with_options(sql, ParserOptions::default())
 }