This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22988-fc846dd36852c8628718175d88ff7bb7c003167e in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 634f0b3004af274b0a3d8230710090447c53f352 Author: Xuanyi Li <[email protected]> AuthorDate: Fri Aug 7 13:15:58 2026 -0700 feat: Add SQL planner, physical planner, and TableProvider hook for MERGE INTO (#22988) ## Which issue does this PR close? Follow-up to #20763 (merged) which added `MergeIntoOp`, `MergeIntoClause`, and proto types. ## Rationale for this change `MERGE INTO` (SQL:2003) is a widely-used DML statement for upsert/conditional update workloads. This PR wires the types introduced in #20763 through the SQL planner, physical planner, and `TableProvider` trait so that table implementations can actually execute merge operations. ## What changes are included in this PR? **`datafusion/session` — `TableProvider` trait extension** - Add `merge_into(source, on, clauses)` async method with a default `not_impl_err` impl so existing providers are unaffected. **`datafusion/sql` — SQL → LogicalPlan** - `statement.rs`: parse `Statement::Merge` into `LogicalPlan::Dml` with `WriteOp::MergeInto`. - Resolve the target table and plan the `USING` source into a `LogicalPlan`. - Build a combined target+source schema to resolve `ON` and `WHEN` expressions. - Convert `ON` condition and `WHEN MATCHED / NOT MATCHED` clauses to DataFusion `Expr`. **`datafusion/expr` — expression plumbing** - `MergeIntoOp::exprs()`: stable iteration over all expressions (ON, then per-clause predicate + action values). - `MergeIntoOp::with_new_exprs()`: rebuild op from a transformed expr vector. - Branch `LogicalPlan::apply_expressions`, `map_expressions`, and `with_new_exprs` on `WriteOp::MergeInto` so optimizers can rewrite merge expressions. Other `WriteOp` variants are unchanged. **`datafusion/core` — physical planner dispatch** - Dispatch `WriteOp::MergeInto` in the physical planner. - Recover the `TableProvider` via `source_as_provider()`, extract the source `ExecutionPlan`, build the target+source merge schema, and call `TableProvider::merge_into`. ## Unsupported MERGE SQL: implementation limitations These are valid or reasonable MERGE forms that this PR rejects to avoid silently producing an incorrect logical plan. They can be revisited once the planner representation / qualifier handling supports them safely. **Target-correlated subqueries when the target is aliased** ```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; ``` Currently rejected with: ```text MERGE subqueries correlated to target alias 't' are not supported ``` The limitation is alias canonicalization: top-level target columns can be rewritten from `t.id` to `target.id`, but target references inside subquery plans appear as outer-reference columns and are not rewritten yet. **Source qualifier collides with the real target table name** ```sql MERGE INTO target AS t USING source AS target ON t.id = target.id WHEN MATCHED THEN DELETE; ``` Currently rejected with: ```text MERGE source may not use the target table name 'target' as a qualifier while the target is aliased as 't'; use a different source alias ``` The limitation is again qualifier canonicalization: rewriting `t.id` to `target.id` would collide with the source alias `target` and could change expression meaning. **Execution still requires a provider implementation** This PR adds the planner and `TableProvider::merge_into` dispatch point. If the target provider does not override the default implementation, execution fails even for supported MERGE syntax: ```sql MERGE INTO target AS t USING source AS s ON t.id = s.id WHEN MATCHED THEN UPDATE SET val = s.val WHEN NOT MATCHED THEN INSERT (id, val) VALUES (s.id, s.val); ``` Default error: ```text MERGE INTO not supported for Base table ``` ## Unsupported MERGE SQL: intentionally unsupported syntax These forms are either dialect-specific, target-table modifier syntax, or outside the initial core/common MERGE surface. This PR rejects them explicitly rather than accepting SQL and dropping semantics. **SQL Server-style `OUTPUT` clause** ```sql MERGE INTO target AS t USING source AS s ON t.id = s.id WHEN MATCHED THEN DELETE OUTPUT deleted.id; ``` Rejected with: ```text MERGE OUTPUT clause is not supported ``` **Target table modifiers / hints / partitions / samples** ```sql MERGE INTO target PARTITION (p0) USING source AS s ON target.id = s.id WHEN MATCHED THEN DELETE; ``` ```sql MERGE INTO target() USING source AS s ON true WHEN MATCHED THEN DELETE; ``` Rejected with: ```text MERGE target table modifiers are not supported ``` **Target alias column lists** ```sql MERGE INTO target AS t(a, b) USING source AS s ON t.a = s.a WHEN MATCHED THEN DELETE; ``` Rejected with: ```text MERGE target alias column lists are not supported ``` **Non-table targets** ```sql MERGE INTO (SELECT * FROM target) AS t USING source AS s ON t.id = s.id WHEN MATCHED THEN DELETE; ``` Rejected with: ```text Cannot MERGE INTO non-table relation! ``` **Oracle action-level predicates** The common conditional form is supported via `WHEN ... AND <predicate>`: ```sql MERGE INTO target AS t USING source AS s ON t.id = s.id WHEN MATCHED AND s.is_active THEN UPDATE SET val = s.val; ``` However, Oracle-style action-level predicates are not represented in `MergeIntoAction` and are explicitly rejected: ```sql MERGE INTO target AS t USING source AS s ON t.id = s.id WHEN MATCHED THEN UPDATE SET val = s.val WHERE s.is_active; ``` ```text MERGE UPDATE WHERE predicates are not supported ``` ```sql MERGE INTO target AS t USING source AS s ON t.id = s.id WHEN MATCHED THEN UPDATE SET val = s.val DELETE WHERE t.val IS NULL; ``` ```text MERGE UPDATE DELETE WHERE predicates are not supported ``` ```sql MERGE INTO target AS t USING source AS s ON t.id = s.id WHEN NOT MATCHED THEN INSERT (id, val) VALUES (s.id, s.val) WHERE s.is_active; ``` ```text MERGE INSERT WHERE predicates are not supported ``` **`INSERT ROW` shorthand** ```sql MERGE INTO target AS t USING source AS s ON t.id = s.id WHEN NOT MATCHED THEN INSERT ROW; ``` Rejected with: ```text MERGE INSERT ROW is not supported ``` **Multiple rows in one MERGE INSERT action** ```sql MERGE INTO target AS t USING source AS s ON t.id = s.id WHEN NOT MATCHED THEN INSERT (id, val) VALUES (s.id, s.val), (s.id + 1, s.val); ``` Rejected with: ```text MERGE INSERT must have exactly one row of values ``` **Tuple assignment in UPDATE** ```sql MERGE INTO target AS t USING source AS s ON t.id = s.id WHEN MATCHED THEN UPDATE SET (id, val) = (s.id, s.val); ``` Rejected with: ```text Tuples are not supported ``` **Invalid clause/action combinations** These are rejected by `sqlparser` before DataFusion planning. For example, this is invalid because `WHEN NOT MATCHED` / `WHEN NOT MATCHED BY TARGET` describes source rows without a target row, so there is no target row to delete: ```sql MERGE INTO target AS t USING source AS s ON t.id = s.id WHEN NOT MATCHED THEN DELETE; ``` Use `NOT MATCHED BY SOURCE` to delete target rows that have no source row: ```sql MERGE INTO target AS t USING source AS s ON t.id = s.id WHEN NOT MATCHED BY SOURCE THEN DELETE; ``` ## Are these changes tested? - SQL planner tests cover supported MERGE planning, target/source qualifier handling, quoted action columns, duplicate action columns, invalid action targets, missing `WHEN`, target modifiers, and unsupported action-level predicates. - `sql_api` tests cover target/source qualifier collisions, target-correlated subquery rejection, supported source-correlated and uncorrelated subqueries through optimization, and boolean validation for `ON` / `WHEN` predicates. - `physical_planner` tests verify that the provider receives the combined target+source logical schema and that physical expression planning can resolve both target and source columns. - Unit tests for `MergeIntoOp::exprs` / `with_new_exprs` are included in `dml.rs`. - Proto round-trip tests for `MergeInto` remain covered in `datafusion/proto/tests/cases/roundtrip_logical_plan.rs`. ## Are there any user-facing changes? - **`TableProvider`** gains a new `merge_into` method. The default implementation returns `not_impl_err`, so existing implementations compile without changes. - `MERGE INTO <table> USING <source> ON <cond> WHEN ...` SQL syntax is now accepted by the DataFusion SQL parser and planner for the supported core forms described above. --------- Co-authored-by: Andrew Lamb <[email protected]> --- datafusion/core/src/physical_planner.rs | 109 ++++++ datafusion/core/tests/sql/sql_api.rs | 96 +++++ datafusion/expr/src/logical_plan/dml.rs | 151 +++++++- datafusion/expr/src/logical_plan/invariants.rs | 8 +- datafusion/expr/src/logical_plan/plan.rs | 14 +- datafusion/expr/src/logical_plan/tree_node.rs | 27 +- .../optimizer/src/analyzer/function_rewrite.rs | 19 +- datafusion/optimizer/src/analyzer/type_coercion.rs | 162 ++++++++- datafusion/optimizer/src/rewrite_set_comparison.rs | 16 +- .../src/simplify_expressions/simplify_exprs.rs | 17 +- datafusion/proto-common/src/generated/mod.rs | 1 + datafusion/proto/src/bytes/mod.rs | 4 + datafusion/proto/src/logical_plan/mod.rs | 7 +- datafusion/session/src/table.rs | 26 +- datafusion/sql/src/statement.rs | 404 ++++++++++++++++++++- datafusion/sql/tests/sql_integration.rs | 98 +++++ 16 files changed, 1136 insertions(+), 23 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index da8e0f2f57..58f00551ac 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -834,6 +834,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(), @@ -3280,6 +3309,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::registry::ExtensionTypeRegistryRef; @@ -3459,6 +3489,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 diff --git a/datafusion/core/tests/sql/sql_api.rs b/datafusion/core/tests/sql/sql_api.rs index e3180210ca..ca18406a8e 100644 --- a/datafusion/core/tests/sql/sql_api.rs +++ b/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(); diff --git a/datafusion/expr/src/logical_plan/dml.rs b/datafusion/expr/src/logical_plan/dml.rs index 5b6403e6e2..7717dfaff7 100644 --- a/datafusion/expr/src/logical_plan/dml.rs +++ b/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}; @@ -307,6 +307,106 @@ pub struct MergeIntoOp { pub clauses: Vec<MergeIntoClause>, } +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 }) + } +} + /// A single WHEN clause within a MERGE INTO statement. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] pub struct MergeIntoClause { @@ -445,4 +545,53 @@ 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, + }, + ], + }; + 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![], + }; + let err = op.with_new_exprs(vec![]).unwrap_err(); + assert!( + err.to_string().contains("expected 1 expressions, got 0"), + "unexpected error: {err}" + ); + } } diff --git a/datafusion/expr/src/logical_plan/invariants.rs b/datafusion/expr/src/logical_plan/invariants.rs index 0889afd08f..d6867d1ceb 100644 --- a/datafusion/expr/src/logical_plan/invariants.rs +++ b/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, \ diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 9ac27b46a7..1a141ea52a 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/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::{ check_aggregate_and_window_nesting, enumerate_grouping_sets, exprlist_to_fields, find_out_reference_exprs, grouping_set_expr_count, grouping_set_to_exprlist, @@ -811,12 +811,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), ))) } diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index c10ac92eef..c4c1d743b5 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/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(_) diff --git a/datafusion/optimizer/src/analyzer/function_rewrite.rs b/datafusion/optimizer/src/analyzer/function_rewrite.rs index 9faa60d939..a66e3ccc0c 100644 --- a/datafusion/optimizer/src/analyzer/function_rewrite.rs +++ b/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,23 @@ impl ApplyFunctionRewrites { schema.merge(&source_schema); } + // MERGE expressions reference the target table, which is not one of + // `plan.inputs()`. Rebuild the target schema from the DML's + // `table_name` and `target` so those columns resolve. + if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + let target_schema = DFSchema::try_from_qualified_schema( + table_name.clone(), + &target.schema(), + )?; + schema.merge(&target_schema); + } + let name_preserver = NamePreserver::new(&plan); plan.map_expressions(|expr| { diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index afd4e980b5..d11c3e7435 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -57,9 +57,10 @@ use datafusion_expr::type_coercion::{ }; use datafusion_expr::utils::merge_schema; use datafusion_expr::{ - Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, Projection, Union, - ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, is_false, - is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, lit, not, + Cast, DmlStatement, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, + Projection, Union, ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, + WriteOp, is_false, is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, + lit, not, }; /// Performs type coercion by determining the schema @@ -128,6 +129,21 @@ fn analyze_internal( schema.merge(&source_schema); } + // MERGE expressions (ON / WHEN clauses) reference the target table, which + // is not one of `plan.inputs()`. Rebuild the target schema from the DML's + // `table_name` and `target` so those columns resolve during coercion. + if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + let target_schema = + DFSchema::try_from_qualified_schema(table_name.clone(), &target.schema())?; + schema.merge(&target_schema); + } + // merge the outer schema for correlated subqueries // like case: // select t2.c2 from t1 where t1.c1 in (select t2.c1 from t2 where t2.c2=t1.c3) @@ -177,10 +193,60 @@ impl<'a> TypeCoercionRewriter<'a> { LogicalPlan::Join(join) => self.coerce_join(join), LogicalPlan::Union(union) => Self::coerce_union(union), LogicalPlan::Limit(limit) => Self::coerce_limit(limit), + LogicalPlan::Dml(dml) => self.coerce_dml(dml), _ => Ok(plan), } } + fn coerce_dml(&self, mut dml: DmlStatement) -> Result<LogicalPlan> { + let WriteOp::MergeInto(merge_op) = &dml.op else { + return Ok(LogicalPlan::Dml(dml)); + }; + + let target_schema = DFSchema::try_from_qualified_schema( + dml.table_name.clone(), + &dml.target.schema(), + )?; + let mut merge_op = (**merge_op).clone(); + merge_op.on = self.coerce_predicate(merge_op.on, "MERGE ON condition")?; + for clause in &mut merge_op.clauses { + clause.predicate = clause + .predicate + .take() + .map(|expr| self.coerce_predicate(expr, "MERGE WHEN condition")) + .transpose()?; + + match &mut clause.action { + datafusion_expr::dml::MergeIntoAction::Update(assignments) => { + for (column, value) in assignments { + let field = target_schema.field_with_unqualified_name(column)?; + *value = value.clone().cast_to(field.data_type(), self.schema)?; + } + } + datafusion_expr::dml::MergeIntoAction::Insert { columns, values } => { + if columns.is_empty() { + for (value, field) in + values.iter_mut().zip(target_schema.fields()) + { + *value = + value.clone().cast_to(field.data_type(), self.schema)?; + } + } else { + for (column, value) in columns.iter().zip(values) { + let field = + target_schema.field_with_unqualified_name(column)?; + *value = + value.clone().cast_to(field.data_type(), self.schema)?; + } + } + } + datafusion_expr::dml::MergeIntoAction::Delete => {} + } + } + dml.op = WriteOp::MergeInto(Box::new(merge_op)); + Ok(LogicalPlan::Dml(dml)) + } + /// Coerce join equality expressions and join filter /// /// Joins must be treated specially as their equality expressions are stored @@ -212,7 +278,7 @@ impl<'a> TypeCoercionRewriter<'a> { // Join filter must be boolean join.filter = join .filter - .map(|expr| self.coerce_join_filter(expr)) + .map(|expr| self.coerce_predicate(expr, "Join condition")) .transpose()?; Ok(LogicalPlan::Join(join)) @@ -280,12 +346,14 @@ impl<'a> TypeCoercionRewriter<'a> { })) } - fn coerce_join_filter(&self, expr: Expr) -> Result<Expr> { + fn coerce_predicate(&self, expr: Expr, description: &str) -> Result<Expr> { let expr_type = expr.get_type(self.schema)?; match expr_type { DataType::Boolean => Ok(expr), DataType::Null => expr.cast_to(&DataType::Boolean, self.schema), - other => plan_err!("Join condition must be boolean type, but got {other:?}"), + other => { + plan_err!("{description} must be boolean type, but got {other:?}") + } } } @@ -1535,6 +1603,88 @@ mod test { ) } + #[test] + fn merge_into_resolves_and_coerces_target_and_source_columns() -> Result<()> { + use datafusion_expr::dml::{ + MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, + }; + use datafusion_expr::logical_plan::table_scan; + use datafusion_expr::{DmlStatement, WriteOp}; + + // Target table `target(id: UInt32)`. + let target_table_name = TableReference::bare("target"); + let target_arrow_schema = + Schema::new(vec![Field::new("id", DataType::UInt32, false)]); + let target_plan = + table_scan(Some(target_table_name.clone()), &target_arrow_schema, None)? + .build()?; + let target_source = match &target_plan { + LogicalPlan::TableScan(ts) => Arc::clone(&ts.source), + _ => unreachable!("table_scan() always builds a TableScan"), + }; + + // Source plan `source(id: Int64)` — deliberately a different numeric + // type than `target.id` so the `ON` comparison needs a CAST. + let source_arrow_schema = + Schema::new(vec![Field::new("id", DataType::Int64, false)]); + let source_plan = + table_scan(Some("source"), &source_arrow_schema, None)?.build()?; + + // `ON target.id = source.id`. Resolving `target.id` requires the + // target schema to be visible to the analyzer, which only sees + // `plan.inputs()` (the source plan) by default. + let on = col("target.id").eq(col("source.id")); + let merge_op = MergeIntoOp { + on, + clauses: vec![ + MergeIntoClause { + kind: MergeIntoClauseKind::Matched, + predicate: None, + action: MergeIntoAction::Update(vec![( + "id".to_string(), + col("source.id"), + )]), + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatched, + predicate: None, + action: MergeIntoAction::Insert { + columns: vec!["id".to_string()], + values: vec![col("source.id")], + }, + }, + ], + }; + let plan = LogicalPlan::Dml(DmlStatement::new( + target_table_name, + target_source, + WriteOp::MergeInto(Box::new(merge_op)), + Arc::new(source_plan), + )); + + let analyzed = Analyzer::with_rules(vec![Arc::new(TypeCoercion::new())]) + .execute_and_check(plan, &ConfigOptions::default(), |_, _| {})?; + let LogicalPlan::Dml(dml) = analyzed else { + panic!("expected Dml"); + }; + let WriteOp::MergeInto(merge_op) = dml.op else { + panic!("expected MergeInto"); + }; + assert_eq!( + merge_op.on.to_string(), + "CAST(target.id AS Int64) = source.id" + ); + let MergeIntoAction::Update(assignments) = &merge_op.clauses[0].action else { + panic!("expected UPDATE"); + }; + assert_eq!(assignments[0].1.to_string(), "CAST(source.id AS UInt32)"); + let MergeIntoAction::Insert { values, .. } = &merge_op.clauses[1].action else { + panic!("expected INSERT"); + }; + assert_eq!(values[0].to_string(), "CAST(source.id AS UInt32)"); + Ok(()) + } + #[test] fn coerce_utf8view_output() -> Result<()> { // Plan A diff --git a/datafusion/optimizer/src/rewrite_set_comparison.rs b/datafusion/optimizer/src/rewrite_set_comparison.rs index c8c35b5187..18712c5335 100644 --- a/datafusion/optimizer/src/rewrite_set_comparison.rs +++ b/datafusion/optimizer/src/rewrite_set_comparison.rs @@ -25,7 +25,7 @@ use datafusion_common::{Column, DFSchema, ExprSchema, Result, ScalarValue, plan_ use datafusion_expr::expr::{self, Exists, SetComparison, SetQuantifier}; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::logical_plan::builder::LogicalPlanBuilder; -use datafusion_expr::{Expr, LogicalPlan, lit}; +use datafusion_expr::{DmlStatement, Expr, LogicalPlan, WriteOp, lit}; use std::sync::Arc; use datafusion_expr::utils::merge_schema; @@ -44,7 +44,19 @@ impl RewriteSetComparison { } fn rewrite_plan(&self, plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> { - let schema = merge_schema(&plan.inputs()); + let mut schema = merge_schema(&plan.inputs()); + if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + schema.merge(&DFSchema::try_from_qualified_schema( + table_name.clone(), + &target.schema(), + )?); + } plan.map_expressions(|expr| { expr.transform_up(|expr| rewrite_set_comparison(expr, &schema)) }) diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs index 3e495f5355..0e72a17abc 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs @@ -21,12 +21,12 @@ use std::sync::Arc; use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{Column, DFSchema, DFSchemaRef, DataFusionError, Result}; -use datafusion_expr::Expr; use datafusion_expr::logical_plan::{Aggregate, LogicalPlan, Projection}; use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{ columnize_expr, find_aggregate_exprs, grouping_set_to_exprlist, merge_schema, }; +use datafusion_expr::{DmlStatement, Expr, WriteOp}; use super::ExprSimplifier; use crate::optimizer::ApplyOrder; @@ -77,7 +77,20 @@ impl SimplifyExpressions { plan: LogicalPlan, config: &dyn OptimizerConfig, ) -> Result<Transformed<LogicalPlan>> { - let schema = if !plan.inputs().is_empty() { + let schema = if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + let mut schema = merge_schema(&plan.inputs()); + schema.merge(&DFSchema::try_from_qualified_schema( + table_name.clone(), + &target.schema(), + )?); + DFSchemaRef::new(schema) + } else if !plan.inputs().is_empty() { DFSchemaRef::new(merge_schema(&plan.inputs())) } else if let LogicalPlan::TableScan(scan) = &plan { // When predicates are pushed into a table scan, there is no input diff --git a/datafusion/proto-common/src/generated/mod.rs b/datafusion/proto-common/src/generated/mod.rs index e5b384c9c5..49d09bf3f4 100644 --- a/datafusion/proto-common/src/generated/mod.rs +++ b/datafusion/proto-common/src/generated/mod.rs @@ -18,6 +18,7 @@ // This code is generated so we don't want to fix any lint violations manually #[allow(clippy::allow_attributes)] #[allow(clippy::all)] +#[allow(clippy::uninlined_format_args)] #[allow(unused_qualifications)] #[rustfmt::skip] pub mod datafusion_proto_common { diff --git a/datafusion/proto/src/bytes/mod.rs b/datafusion/proto/src/bytes/mod.rs index ab013f8dd5..388e373c3f 100644 --- a/datafusion/proto/src/bytes/mod.rs +++ b/datafusion/proto/src/bytes/mod.rs @@ -192,6 +192,10 @@ pub fn physical_plan_to_bytes(plan: Arc<dyn ExecutionPlan>) -> Result<Bytes> { /// Serialize a PhysicalPlan as JSON #[cfg(feature = "json")] +#[expect( + clippy::needless_pass_by_value, + reason = "Preserve the existing public API" +)] pub fn physical_plan_to_json(plan: Arc<dyn ExecutionPlan>) -> Result<String> { let extension_codec = DefaultPhysicalExtensionCodec {}; let proto_converter = DefaultPhysicalProtoConverter {}; diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 653ae9ab05..f273b33431 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -1270,11 +1270,14 @@ impl AsLogicalPlan for LogicalPlanNode { .build() } LogicalPlanType::Dml(dml_node) => { + let table_name = + from_table_reference(dml_node.table_name.as_ref(), "DML ")?; + let target = to_table_source(&dml_node.target, ctx, extension_codec)?; let write_op = from_proto::parse_write_op(dml_node, ctx, extension_codec)?; Ok(LogicalPlan::Dml(DmlStatement::new( - from_table_reference(dml_node.table_name.as_ref(), "DML ")?, - to_table_source(&dml_node.target, ctx, extension_codec)?, + table_name, + target, write_op, Arc::new(into_logical_plan!(dml_node.input, ctx, extension_codec)?), ))) diff --git a/datafusion/session/src/table.rs b/datafusion/session/src/table.rs index 8d9cd92d4c..69e7e731b3 100644 --- a/datafusion/session/src/table.rs +++ b/datafusion/session/src/table.rs @@ -24,11 +24,11 @@ use crate::session::Session; use arrow_schema::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; -use datafusion_expr::dml::InsertOp; +use datafusion_expr::dml::{InsertOp, MergeIntoClause}; use datafusion_expr::{ CreateExternalTable, LogicalPlan, TableProviderFilterPushDown, TableType, }; @@ -379,6 +379,28 @@ 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 `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. + /// + /// Returns an [`ExecutionPlan`] producing a single row with `count` (UInt64). + async fn merge_into( + &self, + _state: &dyn Session, + _source: Arc<dyn ExecutionPlan>, + _merge_schema: DFSchemaRef, + _on: Expr, + _clauses: Vec<MergeIntoClause>, + ) -> Result<Arc<dyn ExecutionPlan>> { + not_impl_err!("MERGE INTO not supported for {} table", self.table_type()) + } } impl dyn TableProvider { diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 3cfbb45688..fd5c34ff5d 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -33,13 +33,16 @@ use arrow::datatypes::{Field, FieldRef, Fields}; use datafusion_common::error::_plan_err; use datafusion_common::format::ExplainStatementOptions; use datafusion_common::parsers::CompressionTypeVariant; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_common::{ Column, Constraint, Constraints, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, SchemaError, SchemaReference, TableReference, ToDFSchema, exec_err, internal_err, not_impl_err, plan_datafusion_err, plan_err, schema_err, unqualified_field_not_found, }; -use datafusion_expr::dml::{CopyTo, InsertOp}; +use datafusion_expr::dml::{ + CopyTo, InsertOp, MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, +}; use datafusion_expr::expr_rewriter::normalize_col_with_schemas_and_ambiguity_check; use datafusion_expr::logical_plan::DdlStatement; use datafusion_expr::logical_plan::builder::project; @@ -1216,6 +1219,8 @@ impl<S: ContextProvider> SqlToRel<'_, S> { self.delete_to_plan(&table_name, selection, limit) } + Statement::Merge(merge) => self.merge_to_plan(merge), + Statement::StartTransaction { modes, begin: false, @@ -2414,6 +2419,403 @@ impl<S: ContextProvider> SqlToRel<'_, S> { Ok(plan) } + fn merge_to_plan(&self, merge: ast::Merge) -> Result<LogicalPlan> { + let ast::Merge { + table, + source, + on, + clauses, + into: _, + merge_token: _, + optimizer_hints, + output, + } = merge; + + if !optimizer_hints.is_empty() { + plan_err!("Optimizer hints not supported")?; + } + + if output.is_some() { + return not_impl_err!("MERGE OUTPUT clause is not supported"); + } + + if clauses.is_empty() { + return plan_err!("MERGE INTO requires at least one WHEN clause"); + } + + // 1. Resolve target table + let (target_table_name, target_alias) = match table { + TableFactor::Table { + name, + alias, + args, + with_hints, + version, + with_ordinality, + partitions, + json_path, + sample, + index_hints, + } => { + if alias + .as_ref() + .is_some_and(|alias| !alias.columns.is_empty()) + { + return not_impl_err!( + "MERGE target alias column lists are not supported" + ); + } + if args.is_some() + || !with_hints.is_empty() + || version.is_some() + || with_ordinality + || !partitions.is_empty() + || json_path.is_some() + || sample.is_some() + || !index_hints.is_empty() + { + return not_impl_err!( + "MERGE target table modifiers are not supported" + ); + } + (name, alias) + } + _ => plan_err!("Cannot MERGE INTO non-table relation!")?, + }; + let target_table_ref = self.object_name_to_table_reference(target_table_name)?; + let target_table_source = self + .context_provider + .get_table_source(target_table_ref.clone())?; + // Use alias as schema qualifier so `t.col` resolves when user writes + // `MERGE INTO target AS t`. Fall back to the table reference itself. + let target_qualifier = target_alias + .as_ref() + .map(|a| { + TableReference::bare(self.ident_normalizer.normalize(a.name.clone())) + }) + .unwrap_or_else(|| target_table_ref.clone()); + let target_schema = Arc::new(DFSchema::try_from_qualified_schema( + target_qualifier.clone(), + &target_table_source.schema(), + )?); + + // 2. Plan the source (USING clause) as a LogicalPlan + let mut planner_context = PlannerContext::new(); + let source_table_with_joins = TableWithJoins { + relation: source, + joins: vec![], + }; + let source_plan = + self.plan_from_tables(vec![source_table_with_joins], &mut planner_context)?; + + // 3. Build a combined schema for resolving expressions in ON and WHEN clauses + let combined_schema = + Arc::new(target_schema.as_ref().join(source_plan.schema())?); + + // 4. Convert the ON condition from sqlparser Expr to datafusion Expr + let on_expr = self.sql_to_expr(*on, &combined_schema, &mut planner_context)?; + + // 5. Convert each WHEN clause + let df_clauses = clauses + .into_iter() + .map(|clause| { + self.merge_clause_to_plan( + clause, + &combined_schema, + &target_schema, + &target_qualifier, + &mut planner_context, + ) + }) + .collect::<Result<Vec<_>>>()?; + + // 6. Build the MERGE operation. Column references to the target may be + // qualified with the SQL alias (`MERGE INTO target AS t ... t.col`). + // Canonicalize those to the real target table qualifier so the stored + // plan is independent of the alias: this lets the analyzer passes and + // proto deserialization rebuild the target schema from `table_name` + // alone, without carrying the alias as extra state. + let mut merge_op = MergeIntoOp { + on: on_expr, + clauses: df_clauses, + }; + if target_qualifier != target_table_ref { + // Target references in correlated subqueries are represented as + // `OuterReferenceColumn`s inside the embedded logical plan. The + // alias canonicalization below only rewrites top-level expression + // columns, so accepting such a subquery would leave the target + // alias in the public MERGE representation. Reject this case until + // the alias can be rewritten scope-safely inside subquery plans. + for expr in merge_op.exprs() { + if Self::has_outer_reference_to_qualifier(expr, &target_qualifier)? { + return not_impl_err!( + "MERGE subqueries correlated to target alias \ + '{target_qualifier}' are not supported" + ); + } + } + + // Canonicalizing target columns to `target_table_ref` is only safe + // when the source does not already use that qualifier. If it does + // (e.g. `MERGE INTO target AS t USING source AS target`), the two + // namespaces would collapse and later resolution could silently + // pick the source column for a target reference. Reject that + // collision rather than change the meaning of the condition. + if source_plan.schema().iter().any(|(qualifier, _)| { + qualifier.is_some_and(|q| q.resolved_eq(&target_table_ref)) + }) { + return plan_err!( + "MERGE source may not use the target table name '{target_table_ref}' \ + as a qualifier while the target is aliased as '{target_qualifier}'; \ + use a different source alias" + ); + } + let canonical = merge_op + .exprs() + .into_iter() + .cloned() + .map(|expr| { + Self::canonicalize_target_qualifier( + expr, + &target_qualifier, + &target_table_ref, + ) + }) + .collect::<Result<Vec<_>>>()?; + merge_op = merge_op.with_new_exprs(canonical)?; + } + + Ok(LogicalPlan::Dml(DmlStatement::new( + target_table_ref, + target_table_source, + WriteOp::MergeInto(Box::new(merge_op)), + Arc::new(source_plan), + ))) + } + + /// Rewrite every [`Expr::Column`] qualified with `from` to instead use + /// `to`, leaving all other columns untouched. Used to canonicalize MERGE + /// target-alias references to the real target table qualifier. + fn canonicalize_target_qualifier( + expr: Expr, + from: &TableReference, + to: &TableReference, + ) -> Result<Expr> { + expr.transform(|expr| match expr { + Expr::Column(col) if col.relation.as_ref() == Some(from) => Ok( + Transformed::yes(Expr::Column(Column::new(Some(to.clone()), col.name))), + ), + other => Ok(Transformed::no(other)), + }) + .map(|transformed| transformed.data) + } + + /// Return true if an expression contains a subquery whose embedded plan + /// has an outer reference qualified by `qualifier`. + fn has_outer_reference_to_qualifier( + expr: &Expr, + qualifier: &TableReference, + ) -> Result<bool> { + let mut found = false; + expr.apply(|expr| { + let subquery = match expr { + Expr::Exists(exists) => Some(&exists.subquery), + Expr::InSubquery(in_subquery) => Some(&in_subquery.subquery), + Expr::SetComparison(set_comparison) => Some(&set_comparison.subquery), + Expr::ScalarSubquery(subquery) => Some(subquery), + _ => None, + }; + + if let Some(subquery) = subquery { + subquery.subquery.apply_with_subqueries(|plan| { + plan.apply_expressions(|expr| { + expr.apply(|expr| { + if let Expr::OuterReferenceColumn(_, column) = expr + && column.relation.as_ref() == Some(qualifier) + { + found = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + })?; + Ok(if found { + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }) + })?; + } + + Ok(if found { + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }) + })?; + Ok(found) + } + + fn merge_target_column_name( + &self, + name: &ObjectName, + target_qualifier: &TableReference, + ) -> Result<String> { + let part = name + .0 + .iter() + .last() + .ok_or_else(|| plan_datafusion_err!("Empty column name"))?; + let ident = part + .as_ident() + .cloned() + .ok_or_else(|| plan_datafusion_err!("Expected simple identifier"))?; + + if name.0.len() > 1 { + let qualifier = self.object_name_to_table_reference(ObjectName( + name.0[..name.0.len() - 1].to_vec(), + ))?; + if !qualifier.resolved_eq(target_qualifier) { + return plan_err!( + "MERGE assignment target '{name}' must reference target table \ + '{target_qualifier}'" + ); + } + } + + Ok(self.ident_normalizer.normalize(ident)) + } + + fn merge_clause_to_plan( + &self, + clause: ast::MergeClause, + combined_schema: &DFSchema, + target_schema: &DFSchema, + target_qualifier: &TableReference, + planner_context: &mut PlannerContext, + ) -> Result<MergeIntoClause> { + let kind = match clause.clause_kind { + ast::MergeClauseKind::Matched => MergeIntoClauseKind::Matched, + ast::MergeClauseKind::NotMatched => MergeIntoClauseKind::NotMatched, + ast::MergeClauseKind::NotMatchedByTarget => { + MergeIntoClauseKind::NotMatchedByTarget + } + ast::MergeClauseKind::NotMatchedBySource => { + MergeIntoClauseKind::NotMatchedBySource + } + }; + + let predicate = clause + .predicate + .map(|p| self.sql_to_expr(p, combined_schema, planner_context)) + .transpose()?; + + let action = match clause.action { + ast::MergeAction::Update(update_expr) => { + if update_expr.update_predicate.is_some() { + return not_impl_err!( + "MERGE UPDATE WHERE predicates are not supported" + ); + } + if update_expr.delete_predicate.is_some() { + return not_impl_err!( + "MERGE UPDATE DELETE WHERE predicates are not supported" + ); + } + let assignments = update_expr + .assignments + .into_iter() + .map(|assign| { + let col_name = match &assign.target { + AssignmentTarget::ColumnName(cols) => { + self.merge_target_column_name(cols, target_qualifier)? + } + _ => plan_err!("Tuples are not supported")?, + }; + // Validate column exists in target + target_schema.field_with_unqualified_name(&col_name)?; + let value = self.sql_to_expr( + assign.value, + combined_schema, + planner_context, + )?; + Ok((col_name, value)) + }) + .collect::<Result<Vec<_>>>()?; + let mut seen = HashSet::new(); + for (column, _) in &assignments { + if !seen.insert(column.as_str()) { + return plan_err!("Duplicate column '{column}' in MERGE UPDATE"); + } + } + MergeIntoAction::Update(assignments) + } + ast::MergeAction::Insert(insert_expr) => { + if insert_expr.insert_predicate.is_some() { + return not_impl_err!( + "MERGE INSERT WHERE predicates are not supported" + ); + } + let columns: Vec<String> = insert_expr + .columns + .iter() + .map(|c| self.merge_target_column_name(c, target_qualifier)) + .collect::<Result<Vec<_>>>()?; + + // Validate: no duplicates, all columns exist in target schema + let mut seen = HashSet::new(); + for col in &columns { + if !seen.insert(col.as_str()) { + return plan_err!("Duplicate column '{col}' in MERGE INSERT"); + } + target_schema.field_with_unqualified_name(col)?; + } + + let num_target_cols = target_schema.fields().len(); + + let values = match insert_expr.kind { + ast::MergeInsertKind::Values(values) => { + if values.rows.len() != 1 { + return plan_err!( + "MERGE INSERT must have exactly one row of values" + ); + } + let row = values.rows.into_iter().next().unwrap().content; + let expected = if columns.is_empty() { + num_target_cols + } else { + columns.len() + }; + if row.len() != expected { + return plan_err!( + "MERGE INSERT has {expected} column(s) but {} value(s)", + row.len() + ); + } + row.into_iter() + .map(|v| { + self.sql_to_expr(v, combined_schema, planner_context) + }) + .collect::<Result<Vec<_>>>()? + } + ast::MergeInsertKind::Row => { + return not_impl_err!("MERGE INSERT ROW is not supported"); + } + }; + + MergeIntoAction::Insert { columns, values } + } + ast::MergeAction::Delete { .. } => MergeIntoAction::Delete, + }; + + Ok(MergeIntoClause { + kind, + predicate, + action, + }) + } + fn insert_to_plan( &self, table_name: ObjectName, diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index a4bf0db910..4f282f5e06 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -3600,6 +3600,104 @@ fn select_groupby_orderby_aggregate_on_non_selected_column_original_issue() { ); } +#[test] +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); + }; + + 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] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN MATCHED THEN UPDATE SET j1_string = j2.j2_string WHERE false", + "MERGE UPDATE WHERE predicates are not supported" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN MATCHED THEN UPDATE SET j1_string = j2.j2_string DELETE WHERE false", + "MERGE UPDATE DELETE WHERE predicates are not supported" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN NOT MATCHED THEN INSERT (j1_id, j1_string) \ + VALUES (j2.j2_id, j2.j2_string) WHERE false", + "MERGE INSERT WHERE predicates are not supported" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN MATCHED THEN UPDATE SET j1_string = 'a', j1_string = 'b'", + "Duplicate column 'j1_string' in MERGE UPDATE" +)] +#[case( + "MERGE INTO j1 AS t USING j2 AS s ON t.j1_id = s.j2_id \ + WHEN MATCHED THEN UPDATE SET s.j1_string = s.j2_string", + "MERGE assignment target 's.j1_string' must reference target table 't'" +)] +#[case( + "MERGE INTO j1 AS t USING j2 AS s ON t.j1_id = s.j2_id \ + WHEN NOT MATCHED THEN INSERT (s.j1_id) VALUES (s.j2_id)", + "MERGE assignment target 's.j1_id' must reference target table 't'" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id", + "MERGE INTO requires at least one WHEN clause" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN NOT MATCHED THEN INSERT (j1_id, J1_ID) VALUES (1, 2)", + "Duplicate column 'j1_id' in MERGE INSERT" +)] +#[case( + "MERGE INTO j1() USING j2 ON true WHEN MATCHED THEN DELETE", + "MERGE target table modifiers are not supported" +)] +#[case( + "MERGE INTO j1 PARTITION (p0) USING j2 ON true WHEN MATCHED THEN DELETE", + "MERGE target table modifiers are not supported" +)] +#[case( + "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_rejects_invalid_actions_and_structure( + #[case] sql: &str, + #[case] expected: &str, +) { + let err = logical_plan(sql).unwrap_err(); + assert!( + err.strip_backtrace().contains(expected), + "unexpected error: {err}" + ); +} + fn logical_plan(sql: &str) -> Result<LogicalPlan> { logical_plan_with_options(sql, ParserOptions::default()) } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
