kosiew commented on code in PR #22988:
URL: https://github.com/apache/datafusion/pull/22988#discussion_r3650019540
##########
datafusion/sql/src/statement.rs:
##########
@@ -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)
Review Comment:
`has_outer_reference_to_qualifier` currently decides whether a subquery is
correlated to the MERGE target by comparing the `OuterReferenceColumn`
qualifier with the target alias. The problem is that qualifiers are scope
local, so the same alias can legally refer to a relation introduced inside a
nested subquery.
For example, this valid query is currently rejected:
```sql
MERGE INTO target AS t
USING source AS s
ON EXISTS (
SELECT 1
FROM source AS t
WHERE EXISTS (
SELECT 1 FROM source AS x WHERE x.id = t.id
)
)
WHEN MATCHED THEN DELETE
```
The deepest `outer_ref(t.id)` refers to the immediately enclosing `source AS
t`, not the MERGE target, but it is still treated as target correlation and
produces:
```text
This feature is not implemented: MERGE subqueries correlated to target alias
't' are not supported
```
This is a pretty narrow false positive, and there is an easy workaround by
using a different inner alias such as `source AS tt`. With that change,
planning succeeds and reaches `TableProvider::merge_into`, where the expected
Base provider `NotImplemented` error is returned.
I think the guard should become scope aware by accounting for aliases
introduced by each embedded plan. It would also be good to add this nested
alias shadowing case to the `SessionContext::sql()` regression tests so we do
not accidentally regress it later.
--
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]
