diff --git a/Cargo.lock b/Cargo.lock index da096ef3a34ae..eda64d8b9ef9a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2714,6 +2714,7 @@ dependencies = [ "recursive", "regex", "rstest", + "serde_json", "sqlparser", "stacker", ] 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..cdddb17a698fa 100644 --- a/datafusion/expr/src/logical_plan/dml.rs +++ b/datafusion/expr/src/logical_plan/dml.rs @@ -23,6 +23,7 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::file_options::file_type::FileType; +use datafusion_common::utils::quote_identifier; use datafusion_common::{ DFSchema, DFSchemaRef, Result, TableReference, internal_err, plan_err, }; @@ -479,6 +480,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 +504,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 +540,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 +595,46 @@ 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, "{} = {value}", quote_identifier(column))?; + } + Ok(()) + } + Self::Insert { columns, values } => { + write!(f, "INSERT")?; + if !columns.is_empty() { + write!( + f, + " ({})", + columns + .iter() + .map(|column| quote_identifier(column)) + .collect::>() + .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 +664,27 @@ 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=[]"); + + let update = MergeIntoAction::Update(vec![("a, b".to_string(), lit(1_i64))]); + assert_eq!(update.to_string(), r#"UPDATE SET "a, b" = Int64(1)"#); + let insert = MergeIntoAction::Insert { + columns: vec!["id".to_string(), "a\"b".to_string()], + values: vec![lit(1_i64), lit(2_i64)], + }; + assert_eq!( + insert.to_string(), + r#"INSERT (id, "a""b") VALUES (Int64(1), Int64(2))"# + ); } #[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/Cargo.toml b/datafusion/sql/Cargo.toml index 318f8934639aa..305a6d03d556e 100644 --- a/datafusion/sql/Cargo.toml +++ b/datafusion/sql/Cargo.toml @@ -75,3 +75,4 @@ env_logger = { workspace = true } insta = { workspace = true } itertools = { workspace = true } rstest = { workspace = true } +serde_json = { workspace = true } diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 3e683a54daf53..669b62afc0c08 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -3979,6 +3979,40 @@ 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, + r#"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: serde_json::Value = + serde_json::from_str(&plan.display_pg_json().to_string()).unwrap(); + let root = &json[0]["Plan"]; + assert_eq!(root["Node Type"], "Dml"); + assert_eq!(root["Operation"], "MergeInto"); + assert_eq!(root["Table Name"], "person_quoted_cols"); + assert_eq!(root["On"], "t.id = s.j2_id"); + assert_eq!( + root["Clauses"], + serde_json::json!([ + r#"WHEN MATCHED THEN UPDATE SET "First Name" = s.j2_string"#, + r#"WHEN NOT MATCHED THEN INSERT (id, "Age") VALUES (s.j2_id, Int64(42))"#, + ]) + ); +} + +#[test] +fn plan_update_json_has_dml_node_without_merge_fields() { + let plan = logical_plan("UPDATE person SET age = 42").unwrap(); + let json: serde_json::Value = + serde_json::from_str(&plan.display_pg_json().to_string()).unwrap(); + let root = json[0]["Plan"].as_object().unwrap(); + assert_eq!(root["Node Type"], "Dml"); + assert_eq!(root["Operation"], "Update"); + assert_eq!(root["Table Name"], "person"); + assert!(!root.contains_key("On")); + assert!(!root.contains_key("Clauses")); } #[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'