diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 2dcaefe9a2af5..843c06c6fd50b 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -24,7 +24,7 @@ use crate::{ Aggregate, AsOfJoin, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter, Join, Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition, Sort, Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest, - Values, Window, expr_vec_fmt, + Values, Window, WriteOp, expr_vec_fmt, }; use crate::dml::CopyTo; @@ -463,11 +463,27 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { ) } LogicalPlan::Dml(DmlStatement { table_name, op, .. }) => { - pg_fields!( - "Node Type": "Projection", + let mut object = pg_fields!( + "Node Type": "Dml", "Operation": op.name(), "Table Name": table_name.table() - ) + ); + if let WriteOp::MergeInto(merge_op) = op { + object.push("On", serde_json::Value::String(merge_op.on.to_string())); + object.push( + "Clauses", + serde_json::Value::Array( + merge_op + .clauses + .iter() + .map(|clause| { + serde_json::Value::String(clause.to_string()) + }) + .collect(), + ), + ); + } + object } LogicalPlan::Copy(CopyTo { input: _, diff --git a/datafusion/expr/src/logical_plan/dml.rs b/datafusion/expr/src/logical_plan/dml.rs index f128159988dca..1fabe8e5c5439 100644 --- a/datafusion/expr/src/logical_plan/dml.rs +++ b/datafusion/expr/src/logical_plan/dml.rs @@ -479,6 +479,19 @@ impl MergeIntoOp { } } +impl Display for MergeIntoOp { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "on=[{}] clauses=[", self.on)?; + for (i, clause) in self.clauses.iter().enumerate() { + if i > 0 { + write!(f, ", ")?; + } + write!(f, "{clause}")?; + } + write!(f, "]") + } +} + /// A single WHEN clause within a MERGE INTO statement. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] pub struct MergeIntoClause { @@ -490,6 +503,16 @@ pub struct MergeIntoClause { pub action: MergeIntoAction, } +impl Display for MergeIntoClause { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "WHEN {}", self.kind)?; + if let Some(predicate) = &self.predicate { + write!(f, " AND {predicate}")?; + } + write!(f, " THEN {}", self.action) + } +} + /// Which rows a MERGE WHEN clause applies to. /// /// Mirrors `sqlparser::ast::MergeClauseKind` so that the SQL spelling is @@ -516,6 +539,17 @@ pub enum MergeIntoClauseKind { NotMatchedBySource, } +impl Display for MergeIntoClauseKind { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::Matched => write!(f, "MATCHED"), + Self::NotMatched => write!(f, "NOT MATCHED"), + Self::NotMatchedByTarget => write!(f, "NOT MATCHED BY TARGET"), + Self::NotMatchedBySource => write!(f, "NOT MATCHED BY SOURCE"), + } + } +} + impl MergeIntoClauseKind { /// True if this clause fires on a source row that has no matching target /// row. Returns `true` for both [`NotMatched`](Self::NotMatched) and @@ -560,6 +594,38 @@ pub enum MergeIntoAction { Delete, } +impl Display for MergeIntoAction { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::Update(assignments) => { + write!(f, "UPDATE SET ")?; + for (i, (column, value)) in assignments.iter().enumerate() { + if i > 0 { + write!(f, ", ")?; + } + write!(f, "{column} = {value}")?; + } + Ok(()) + } + Self::Insert { columns, values } => { + write!(f, "INSERT")?; + if !columns.is_empty() { + write!(f, " ({})", columns.join(", "))?; + } + write!(f, " VALUES (")?; + for (i, value) in values.iter().enumerate() { + if i > 0 { + write!(f, ", ")?; + } + write!(f, "{value}")?; + } + write!(f, ")") + } + Self::Delete => write!(f, "DELETE"), + } + } +} + fn make_count_schema() -> DFSchemaRef { Arc::new( Schema::new(vec![Field::new("count", DataType::UInt64, false)]) @@ -589,6 +655,16 @@ mod tests { ))); assert_eq!(op.name(), "MergeInto"); assert_eq!(format!("{op}"), "MergeInto"); + let WriteOp::MergeInto(merge_op) = &op else { + unreachable!("constructed as MergeInto") + }; + assert_eq!( + merge_op.to_string(), + "on=[id = source_id] clauses=[WHEN MATCHED AND qty > Int64(0) THEN UPDATE SET qty = source_qty]" + ); + + let empty = MergeIntoOp::new("target", lit(true), vec![]); + assert_eq!(empty.to_string(), "on=[Boolean(true)] clauses=[]"); } #[test] diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index ba6b1a6a26ad3..9d10b435341c6 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -2194,7 +2194,11 @@ impl LogicalPlan { Ok(()) } LogicalPlan::Dml(DmlStatement { table_name, op, .. }) => { - write!(f, "Dml: op=[{op}] table=[{table_name}]") + write!(f, "Dml: op=[{op}] table=[{table_name}]")?; + if let WriteOp::MergeInto(merge_op) = op { + write!(f, " {merge_op}")?; + } + Ok(()) } LogicalPlan::Copy(CopyTo { input: _, diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 3e683a54daf53..fe70f7ea99eb0 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -3979,6 +3979,20 @@ fn plan_merge_into_preserves_target_qualifier_and_quoted_columns() { }; assert_eq!(columns, &["id".to_string(), "Age".to_string()]); assert_eq!(values[0].to_string(), "s.j2_id"); + + let display = plan.display_indent().to_string(); + assert_contains!( + &display, + "Dml: op=[MergeInto] table=[person_quoted_cols] on=[t.id = s.j2_id] clauses=[WHEN MATCHED THEN UPDATE SET First Name = s.j2_string, WHEN NOT MATCHED THEN INSERT (id, Age) VALUES (s.j2_id, Int64(42))]" + ); + + let json = plan.display_pg_json().to_string(); + assert_contains!(&json, r#""Node Type": "Dml""#); + assert_contains!(&json, r#""On": "t.id = s.j2_id""#); + assert_contains!( + &json, + r#""WHEN MATCHED THEN UPDATE SET First Name = s.j2_string""# + ); } #[rstest] diff --git a/datafusion/sqllogictest/test_files/merge_into.slt b/datafusion/sqllogictest/test_files/merge_into.slt index e70c90f9bafa5..ede86ea4cd94c 100644 --- a/datafusion/sqllogictest/test_files/merge_into.slt +++ b/datafusion/sqllogictest/test_files/merge_into.slt @@ -64,7 +64,7 @@ when matched then update set val = source.val when not matched then insert (id, val) values (source.id, source.val); ---- logical_plan -01)Dml: op=[MergeInto] table=[target] +01)Dml: op=[MergeInto] table=[target] on=[target.id = source.id] clauses=[WHEN MATCHED THEN UPDATE SET val = source.val, WHEN NOT MATCHED THEN INSERT (id, val) VALUES (source.id, source.val)] 02)--TableScan: source projection=[id, val, is_active] physical_plan_error 01)MERGE INTO operation on table 'target' @@ -77,7 +77,7 @@ explain merge into target using source on target.id = source.id when matched then delete; ---- logical_plan -01)Dml: op=[MergeInto] table=[target] +01)Dml: op=[MergeInto] table=[target] on=[target.id = source.id] clauses=[WHEN MATCHED THEN DELETE] 02)--TableScan: source projection=[id, val, is_active] physical_plan_error 01)MERGE INTO operation on table 'target' @@ -91,7 +91,7 @@ when matched and s.is_active then update set val = s.val when not matched by source then delete; ---- logical_plan -01)Dml: op=[MergeInto] table=[target] +01)Dml: op=[MergeInto] table=[target] on=[t.id = s.id] clauses=[WHEN MATCHED AND s.is_active THEN UPDATE SET val = s.val, WHEN NOT MATCHED BY SOURCE THEN DELETE] 02)--SubqueryAlias: s 03)----TableScan: source projection=[id, val, is_active] physical_plan_error @@ -110,7 +110,7 @@ explain merge into target using source on target.id = source.id when not matched by source then delete; ---- logical_plan -01)Dml: op=[MergeInto] table=[target] +01)Dml: op=[MergeInto] table=[target] on=[target.id = source.id] clauses=[WHEN NOT MATCHED BY SOURCE THEN DELETE] 02)--TableScan: source projection=[id, val, is_active] physical_plan_error 01)MERGE INTO operation on table 'target' @@ -124,7 +124,7 @@ on target.id = s.id when matched then update set val = s.val; ---- logical_plan -01)Dml: op=[MergeInto] table=[target] +01)Dml: op=[MergeInto] table=[target] on=[target.id = s.id] clauses=[WHEN MATCHED THEN UPDATE SET val = s.val] 02)--SubqueryAlias: s 03)----Projection: source.id, max(source.val) AS val 04)------Aggregate: groupBy=[[source.id]], aggr=[[max(source.val)]] @@ -140,7 +140,7 @@ explain merge into target using source on target.id = source.id when not matched then insert values (source.id, source.val, 0); ---- logical_plan -01)Dml: op=[MergeInto] table=[target] +01)Dml: op=[MergeInto] table=[target] on=[target.id = source.id] clauses=[WHEN NOT MATCHED THEN INSERT VALUES (source.id, source.val, Int32(0) AS Int64(0))] 02)--TableScan: source projection=[id, val, is_active] physical_plan_error 01)MERGE INTO operation on table 'target'