This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 634f0b3004 feat: Add SQL planner, physical planner, and TableProvider
hook for MERGE INTO (#22988)
634f0b3004 is described below
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]
