wirybeaver commented on code in PR #22988:
URL: https://github.com/apache/datafusion/pull/22988#discussion_r3623797024
##########
datafusion/sql/src/statement.rs:
##########
@@ -2406,6 +2411,252 @@ 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");
+ }
+
+ // 1. Resolve target table
+ let (target_table_name, target_alias) = match &table {
+ TableFactor::Table { name, alias, .. } => (name.clone(),
alias.clone()),
+ _ => 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_alias,
+ &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 {
+ 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(
Review Comment:
Thanks — you're right, and I confirmed it before fixing. I reproduced the
exact scenario via `SessionContext::sql()`: for `MERGE INTO target AS t USING
source AS target ON t.id = target.id`, the stored condition collapsed to
`target.id = target.id` and the plan analyzed silently, so `t.id = target.id`
did become a source-vs-source comparison. Good catch.
Fix: canonicalizing the aliased target reference to `target_table_ref` is
inherently lossy when the source already uses that qualifier, so rather than
silently change the condition's meaning I now reject that collision during
planning:
```rust
if target_qualifier != target_table_ref {
// Canonicalizing target columns to `target_table_ref` is only safe when
// the source does not already use that qualifier ...
if source_plan
.schema()
.iter()
.any(|(qualifier, _)| qualifier == Some(&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"
);
}
// ... canonicalize
}
```
Regression tests (both assert explicit rejection, per your note that
"checking only that planning succeeds would not catch silent misresolution"):
- `merge_into_rejects_source_alias_colliding_with_target_name` — a
`SessionContext::sql()` test in `datafusion/core/tests/sql/sql_api.rs`. Target
and source both have an identically named `id` column and the source is aliased
to the target table's real name (`USING source AS target`).
- `plan_merge_into_rejects_source_alias_colliding_with_target_name` — the
same collision at the SQL-planner level for fast coverage.
The non-colliding aliased case still canonicalizes as before, covered by
`plan_merge_into_canonicalizes_target_alias`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]