kosiew commented on code in PR #22988:
URL: https://github.com/apache/datafusion/pull/22988#discussion_r3637075437
##########
datafusion/sql/src/statement.rs:
##########
@@ -2406,6 +2411,269 @@ 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 {
+ // 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 == 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"
+ );
+ }
+ 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)
+ }
+
+ fn ident_from_object_name_last(name: &ObjectName) -> Result<Ident> {
+ let part = name
+ .0
+ .iter()
+ .last()
+ .ok_or_else(|| plan_datafusion_err!("Empty column name"))?;
+ part.as_ident()
+ .cloned()
+ .ok_or_else(|| plan_datafusion_err!("Expected simple identifier"))
+ }
+
+ fn merge_clause_to_plan(
+ &self,
+ clause: ast::MergeClause,
+ combined_schema: &DFSchema,
+ target_schema: &DFSchema,
+ _target_alias: &Option<ast::TableAlias>,
+ 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) => {
Review Comment:
It looks like `sqlparser` preserves action-level predicates on
`MergeUpdateExpr` (`update_predicate` and `delete_predicate`) and
`MergeInsertExpr` (`insert_predicate`), but this conversion only reads the
assignments, columns, and values.
As a result, statements like `WHEN MATCHED THEN UPDATE SET ... WHERE false`
are planned as unconditional updates, and INSERT `WHERE` predicates are dropped
as well. That means accepted SQL can modify a different set of rows than the
user requested.
Could we either preserve these predicates in `MergeIntoClause` /
`MergeIntoAction` and pass them through to providers, or reject any non-`None`
action-level predicate during planning until the logical representation
supports them?
It would also be good to add regression tests covering UPDATE `WHERE`,
UPDATE `... DELETE WHERE`, and INSERT `WHERE`.
##########
datafusion/sql/src/statement.rs:
##########
@@ -2406,6 +2411,269 @@ 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 {
+ // 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 == 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"
+ );
+ }
+ 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 {
Review Comment:
This rewrite only handles `Expr::Column`. Target references inside MERGE
subqueries become `Expr::OuterReferenceColumn` in the embedded logical plan, so
they keep the SQL alias instead of being canonicalized.
For example:
```sql
MERGE INTO target AS t
USING source AS s
ON EXISTS (SELECT 1 FROM source x WHERE x.id = t.id)
WHEN MATCHED THEN DELETE
```
After `SessionContext::sql()`, the nested subquery still contains
`outer_ref(t.id)`, while ordinary target references have already been rewritten
to `target.id`. Since `MergeIntoOp` does not retain the target alias or an
alias-to-target mapping, providers cannot reliably resolve that qualifier from
the public logical representation.
Could we either canonicalize target-qualified `OuterReferenceColumn`s inside
embedded subquery plans using a scope-aware traversal, or reject
target-correlated MERGE subqueries until the logical representation retains
enough information to resolve them?
A regression test that inspects the analyzed `MergeIntoOp` and the nested
subquery after `SessionContext::sql()` would help catch this. Checking only
that planning succeeds would not detect this issue.
--
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]
