diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index 5d07133799ffc..2b0ce021132d1 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -17,7 +17,7 @@ //! [`MemTable`] for querying `Vec` by DataFusion. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::fmt::Debug; use std::sync::Arc; @@ -25,19 +25,24 @@ use crate::TableProvider; use arrow::array::{ Array, ArrayRef, BooleanArray, RecordBatch as ArrowRecordBatch, UInt64Array, + new_empty_array, }; use arrow::compute::kernels::zip::zip; use arrow::compute::{and, filter_record_batch}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use arrow::record_batch::RecordBatch; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; use datafusion_common::error::Result; -use datafusion_common::{Constraints, DFSchema, SchemaExt, not_impl_err, plan_err}; +use datafusion_common::{ + Constraints, DFSchema, DFSchemaRef, ScalarValue, SchemaExt, not_impl_err, plan_err, +}; use datafusion_datasource::memory::{MemSink, MemorySourceConfig}; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::source::DataSourceExec; -use datafusion_expr::dml::InsertOp; +use datafusion_expr::dml::{ + InsertOp, MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, +}; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; -use datafusion_expr::{Expr, SortExpr, TableType}; +use datafusion_expr::{Expr, ExprSchemable, SortExpr, TableType}; use datafusion_physical_expr::{ LexOrdering, PhysicalExpr, create_physical_expr, create_physical_sort_exprs, }; @@ -461,6 +466,553 @@ impl TableProvider for MemTable { Ok(Arc::new(DmlResultExec::new(total_updated))) } + + async fn merge_into( + &self, + state: &dyn Session, + source: Arc, + merge_schema: DFSchemaRef, + on: Expr, + clauses: Vec, + ) -> Result> { + if self.batches.is_empty() { + return plan_err!("No partitions provided, expected at least one partition"); + } + + let source_schema = source.schema(); + let source_partitions = collect_partitioned(source, state.task_ctx()).await?; + let source_rows = partitioned_batches_to_rows(&source_partitions)?; + + let mut target_batches = vec![]; + for partition_data in &self.batches { + let partition = partition_data.read().await; + target_batches.extend(partition.iter().cloned()); + } + let target_rows = batches_to_rows(&target_batches)?; + + let target_width = self.schema.fields().len(); + let source_width = source_schema.fields().len(); + if merge_schema.fields().len() != target_width + source_width { + return plan_err!( + "MERGE INTO schema mismatch: expected {} target/source columns, got {}", + target_width + source_width, + merge_schema.fields().len() + ); + } + + let merge_arrow_schema = Arc::new(merge_schema.as_arrow().clone()); + let on = on.cast_to(&DataType::Boolean, merge_schema.as_ref())?; + let on = state.create_physical_expr(on, merge_schema.as_ref())?; + let clauses = compile_merge_clauses(self, state, merge_schema.as_ref(), clauses)?; + + let null_target = null_row_for_schema(&self.schema)?; + let null_source = null_row_for_schema(&source_schema)?; + + let mut target_matches: Vec> = vec![None; target_rows.len()]; + let mut source_matched = vec![false; source_rows.len()]; + + for (target_idx, target_row) in target_rows.iter().enumerate() { + for (source_idx, source_row) in source_rows.iter().enumerate() { + let combined = combined_row_batch( + Arc::clone(&merge_arrow_schema), + target_row, + source_row, + )?; + if evaluate_merge_predicate(&on, &combined)? { + if let Some(first_source_idx) = target_matches[target_idx] { + return plan_err!( + "MERGE INTO matched target row {target_idx} with more than one source row ({first_source_idx} and {source_idx})" + ); + } + target_matches[target_idx] = Some(source_idx); + source_matched[source_idx] = true; + } + } + } + + let default_batch = one_row_empty_batch()?; + let mut merged_rows = + Vec::with_capacity(target_rows.len().saturating_add(source_rows.len())); + let mut rows_affected = 0_u64; + + for (target_idx, target_row) in target_rows.iter().enumerate() { + let (source_row, clause_kind) = + if let Some(source_idx) = target_matches[target_idx] { + (&source_rows[source_idx], MergeIntoClauseKind::Matched) + } else { + (&null_source, MergeIntoClauseKind::NotMatchedBySource) + }; + + let combined = combined_row_batch( + Arc::clone(&merge_arrow_schema), + target_row, + source_row, + )?; + let application = apply_first_merge_clause( + &clauses, + clause_kind, + &combined, + &default_batch, + Some(target_row), + )?; + if application.affected { + rows_affected += 1; + } + if let Some(row) = application.row { + merged_rows.push(row); + } + } + + for (source_idx, source_row) in source_rows.iter().enumerate() { + if source_matched[source_idx] { + continue; + } + + let combined = combined_row_batch( + Arc::clone(&merge_arrow_schema), + &null_target, + source_row, + )?; + let application = apply_first_merge_clause( + &clauses, + MergeIntoClauseKind::NotMatchedByTarget, + &combined, + &default_batch, + None, + )?; + if application.affected { + rows_affected += 1; + } + if let Some(row) = application.row { + merged_rows.push(row); + } + } + + let merged_batch = rows_to_batch(Arc::clone(&self.schema), &merged_rows)?; + + *self.sort_order.lock() = vec![]; + let mut wrote_first_partition = false; + for partition_data in &self.batches { + let mut partition = partition_data.write().await; + if !wrote_first_partition { + if merged_batch.num_rows() == 0 { + partition.clear(); + } else { + *partition = vec![merged_batch.clone()]; + } + wrote_first_partition = true; + } else { + partition.clear(); + } + } + + Ok(Arc::new(DmlResultExec::new(rows_affected))) + } +} + +struct CompiledMergeClause { + kind: MergeIntoClauseKind, + predicate: Option>, + action: CompiledMergeAction, +} + +enum CompiledMergeAction { + Update(Vec), + Insert(Vec), + Delete, +} + +struct CompiledMergeAssignment { + target_index: usize, + data_type: DataType, + expr: Arc, +} + +enum CompiledInsertValue { + MergeExpr { + data_type: DataType, + expr: Arc, + }, + DefaultExpr { + data_type: DataType, + expr: Arc, + }, + Null(ScalarValue), +} + +struct MergeApplication { + row: Option>, + affected: bool, +} + +fn compile_merge_clauses( + table: &MemTable, + state: &dyn Session, + merge_schema: &DFSchema, + clauses: Vec, +) -> Result> { + let empty_schema = DFSchema::empty(); + clauses + .into_iter() + .map(|clause| { + let predicate = clause + .predicate + .map(|predicate| { + let predicate = + predicate.cast_to(&DataType::Boolean, merge_schema)?; + state.create_physical_expr(predicate, merge_schema) + }) + .transpose()?; + + let action = match (clause.kind.canonical(), clause.action) { + (MergeIntoClauseKind::Matched, MergeIntoAction::Update(assignments)) + | ( + MergeIntoClauseKind::NotMatchedBySource, + MergeIntoAction::Update(assignments), + ) => CompiledMergeAction::Update(compile_merge_assignments( + table, + state, + merge_schema, + assignments, + )?), + (MergeIntoClauseKind::Matched, MergeIntoAction::Delete) + | (MergeIntoClauseKind::NotMatchedBySource, MergeIntoAction::Delete) => { + CompiledMergeAction::Delete + } + ( + MergeIntoClauseKind::NotMatchedByTarget, + MergeIntoAction::Insert { columns, values }, + ) => CompiledMergeAction::Insert(compile_merge_insert_values( + table, + state, + merge_schema, + &empty_schema, + columns, + values, + )?), + (MergeIntoClauseKind::Matched, MergeIntoAction::Insert { .. }) => { + return plan_err!("MERGE MATCHED INSERT is not supported"); + } + (MergeIntoClauseKind::NotMatchedByTarget, MergeIntoAction::Update(_)) => { + return plan_err!("MERGE NOT MATCHED UPDATE is not supported"); + } + (MergeIntoClauseKind::NotMatchedByTarget, MergeIntoAction::Delete) => { + return plan_err!("MERGE NOT MATCHED DELETE is not supported"); + } + ( + MergeIntoClauseKind::NotMatchedBySource, + MergeIntoAction::Insert { .. }, + ) => { + return plan_err!( + "MERGE NOT MATCHED BY SOURCE INSERT is not supported" + ); + } + (MergeIntoClauseKind::NotMatched, _) => { + unreachable!("canonical() never returns NotMatched") + } + }; + + Ok(CompiledMergeClause { + kind: clause.kind, + predicate, + action, + }) + }) + .collect() +} + +fn compile_merge_assignments( + table: &MemTable, + state: &dyn Session, + merge_schema: &DFSchema, + assignments: Vec<(String, Expr)>, +) -> Result> { + let available_columns = table.available_column_names(); + let mut seen = HashSet::new(); + assignments + .into_iter() + .map(|(column, value)| { + if !seen.insert(column.clone()) { + return plan_err!("Duplicate column '{column}' in MERGE UPDATE"); + } + let (target_index, field) = + table.target_field(&column).ok_or_else(|| { + datafusion_common::DataFusionError::Plan(format!( + "MERGE UPDATE failed: column '{column}' does not exist. Available columns: {}", + available_columns.join(", ") + )) + })?; + let value = value.cast_to(field.data_type(), merge_schema)?; + let expr = state.create_physical_expr(value, merge_schema)?; + Ok(CompiledMergeAssignment { + target_index, + data_type: field.data_type().clone(), + expr, + }) + }) + .collect() +} + +fn compile_merge_insert_values( + table: &MemTable, + state: &dyn Session, + merge_schema: &DFSchema, + empty_schema: &DFSchema, + columns: Vec, + values: Vec, +) -> Result> { + let target_width = table.schema.fields().len(); + if columns.is_empty() { + if values.len() != target_width { + return plan_err!( + "MERGE INSERT has {target_width} column(s) but {} value(s)", + values.len() + ); + } + return values + .into_iter() + .zip(table.schema.fields()) + .map(|(value, field)| { + let value = value.cast_to(field.data_type(), merge_schema)?; + let expr = state.create_physical_expr(value, merge_schema)?; + Ok(CompiledInsertValue::MergeExpr { + data_type: field.data_type().clone(), + expr, + }) + }) + .collect(); + } + + if columns.len() != values.len() { + return plan_err!( + "MERGE INSERT has {} column(s) but {} value(s)", + columns.len(), + values.len() + ); + } + + let mut insert_values = table + .schema + .fields() + .iter() + .map(|field| { + if let Some(default) = table.column_defaults.get(field.name()) { + let default = default.clone().cast_to(field.data_type(), empty_schema)?; + let expr = state.create_physical_expr(default, empty_schema)?; + Ok(CompiledInsertValue::DefaultExpr { + data_type: field.data_type().clone(), + expr, + }) + } else { + Ok(CompiledInsertValue::Null(ScalarValue::try_new_null( + field.data_type(), + )?)) + } + }) + .collect::>>()?; + + let available_columns = table.available_column_names(); + let mut seen = HashSet::new(); + for (column, value) in columns.into_iter().zip(values) { + if !seen.insert(column.clone()) { + return plan_err!("Duplicate column '{column}' in MERGE INSERT"); + } + let (target_index, field) = table.target_field(&column).ok_or_else(|| { + datafusion_common::DataFusionError::Plan(format!( + "MERGE INSERT failed: column '{column}' does not exist. Available columns: {}", + available_columns.join(", ") + )) + })?; + let value = value.cast_to(field.data_type(), merge_schema)?; + let expr = state.create_physical_expr(value, merge_schema)?; + insert_values[target_index] = CompiledInsertValue::MergeExpr { + data_type: field.data_type().clone(), + expr, + }; + } + + Ok(insert_values) +} + +impl MemTable { + fn target_field(&self, column: &str) -> Option<(usize, &Field)> { + self.schema + .fields() + .iter() + .enumerate() + .find_map(|(idx, field)| { + (field.name() == column).then_some((idx, field.as_ref())) + }) + } + + fn available_column_names(&self) -> Vec<&str> { + self.schema + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect() + } +} + +fn apply_first_merge_clause( + clauses: &[CompiledMergeClause], + clause_kind: MergeIntoClauseKind, + combined: &RecordBatch, + default_batch: &RecordBatch, + base_target_row: Option<&[ScalarValue]>, +) -> Result { + for clause in clauses { + if clause.kind.canonical() != clause_kind.canonical() { + continue; + } + if let Some(predicate) = &clause.predicate + && !evaluate_merge_predicate(predicate, combined)? + { + continue; + } + + return match &clause.action { + CompiledMergeAction::Update(assignments) => { + let Some(base_target_row) = base_target_row else { + return plan_err!("MERGE UPDATE requires a target row"); + }; + let mut row = base_target_row.to_vec(); + for assignment in assignments { + row[assignment.target_index] = evaluate_merge_value( + &assignment.expr, + combined, + &assignment.data_type, + )?; + } + Ok(MergeApplication { + row: Some(row), + affected: true, + }) + } + CompiledMergeAction::Delete => Ok(MergeApplication { + row: None, + affected: true, + }), + CompiledMergeAction::Insert(values) => { + let row = values + .iter() + .map(|value| match value { + CompiledInsertValue::MergeExpr { data_type, expr } => { + evaluate_merge_value(expr, combined, data_type) + } + CompiledInsertValue::DefaultExpr { data_type, expr } => { + evaluate_merge_value(expr, default_batch, data_type) + } + CompiledInsertValue::Null(value) => Ok(value.clone()), + }) + .collect::>>()?; + Ok(MergeApplication { + row: Some(row), + affected: true, + }) + } + }; + } + + Ok(MergeApplication { + row: base_target_row.map(|row| row.to_vec()), + affected: false, + }) +} + +fn evaluate_merge_predicate( + predicate: &Arc, + batch: &RecordBatch, +) -> Result { + let array = predicate.evaluate(batch)?.into_array(batch.num_rows())?; + let bool_array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + datafusion_common::DataFusionError::Internal( + "MERGE predicate did not evaluate to boolean".to_string(), + ) + })?; + Ok(!bool_array.is_null(0) && bool_array.value(0)) +} + +fn evaluate_merge_value( + expr: &Arc, + batch: &RecordBatch, + data_type: &DataType, +) -> Result { + let array = expr.evaluate(batch)?.into_array(batch.num_rows())?; + ScalarValue::try_from_array(array.as_ref(), 0)?.cast_to(data_type) +} + +fn combined_row_batch( + schema: SchemaRef, + target_row: &[ScalarValue], + source_row: &[ScalarValue], +) -> Result { + let columns = target_row + .iter() + .chain(source_row.iter()) + .map(ScalarValue::to_array) + .collect::>>()?; + Ok(ArrowRecordBatch::try_new(schema, columns)?) +} + +fn null_row_for_schema(schema: &SchemaRef) -> Result> { + schema + .fields() + .iter() + .map(|field| ScalarValue::try_new_null(field.data_type())) + .collect() +} + +fn partitioned_batches_to_rows( + partitions: &[Vec], +) -> Result>> { + let mut rows = vec![]; + for partition in partitions { + rows.extend(batches_to_rows(partition)?); + } + Ok(rows) +} + +fn batches_to_rows(batches: &[RecordBatch]) -> Result>> { + let mut rows = vec![]; + for batch in batches { + for row_idx in 0..batch.num_rows() { + let row = batch + .columns() + .iter() + .map(|column| ScalarValue::try_from_array(column.as_ref(), row_idx)) + .collect::>>()?; + rows.push(row); + } + } + Ok(rows) +} + +fn rows_to_batch(schema: SchemaRef, rows: &[Vec]) -> Result { + let columns = schema + .fields() + .iter() + .enumerate() + .map(|(column_idx, field)| { + if rows.is_empty() { + return Ok(new_empty_array(field.data_type())); + } + + ScalarValue::iter_to_array(rows.iter().map(|row| row[column_idx].clone())) + }) + .collect::>>()?; + Ok(ArrowRecordBatch::try_new(schema, columns)?) +} + +fn one_row_empty_batch() -> Result { + Ok(ArrowRecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &RecordBatchOptions::new().with_row_count(Some(1)), + )?) } /// Evaluate filter expressions against a batch and return a combined boolean mask. diff --git a/datafusion/core/tests/sql/sql_api.rs b/datafusion/core/tests/sql/sql_api.rs index ca18406a8e40d..a09700a7f3eac 100644 --- a/datafusion/core/tests/sql/sql_api.rs +++ b/datafusion/core/tests/sql/sql_api.rs @@ -266,7 +266,8 @@ async fn merge_into_rejects_subqueries_correlated_to_target_alias() { .await; // Source-correlated and uncorrelated subqueries remain supported through - // logical optimization. + // logical optimization, but basic MemTable MERGE execution does not yet + // evaluate subquery predicates. for sql in [ "MERGE INTO target AS t USING source AS s \ ON EXISTS (SELECT 1 FROM source AS x WHERE x.id = s.id) \ @@ -275,8 +276,12 @@ async fn merge_into_rejects_subqueries_correlated_to_target_alias() { ON t.id = ANY (SELECT id FROM source) \ WHEN MATCHED THEN DELETE", ] { - assert_merge_physical_error(&ctx, sql, "MERGE INTO not supported for Base table") - .await; + assert_merge_physical_error( + &ctx, + sql, + "Physical plan does not support logical expression", + ) + .await; } } @@ -294,14 +299,19 @@ async fn merge_into_requires_boolean_conditions() { WHEN MATCHED AND 1 THEN DELETE", "MERGE WHEN condition must be boolean type, but got Int64", ), - ( - "MERGE INTO target USING source ON NULL \ - WHEN MATCHED AND NULL THEN DELETE", - "MERGE INTO not supported for Base table", - ), ] { assert_merge_physical_error(&ctx, sql, expected).await; } + + ctx.sql( + "MERGE INTO target USING source ON NULL \ + WHEN MATCHED AND NULL THEN DELETE", + ) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); } #[tokio::test] diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 09f41c94f64fa..a5a31ad47c11a 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -24,7 +24,7 @@ use crate::{ Aggregate, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr, Filter, Join, Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery, Repartition, Sort, Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan, Unnest, Values, - Window, expr_vec_fmt, + Window, WriteOp, expr_vec_fmt, }; use crate::dml::CopyTo; @@ -404,11 +404,22 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { }) } LogicalPlan::Dml(DmlStatement { table_name, op, .. }) => { - json!({ - "Node Type": "Projection", + let mut object = json!({ + "Node Type": "Dml", "Operation": op.name(), "Table Name": table_name.table() - }) + }); + if let WriteOp::MergeInto(merge_op) = op { + object["On"] = serde_json::Value::String(merge_op.on.to_string()); + object["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 7717dfaff7a33..e1a9640374761 100644 --- a/datafusion/expr/src/logical_plan/dml.rs +++ b/datafusion/expr/src/logical_plan/dml.rs @@ -407,6 +407,23 @@ impl MergeIntoOp { } } +impl Display for MergeIntoOp { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "on=[{}]", self.on)?; + if !self.clauses.is_empty() { + write!(f, " clauses=[")?; + for (i, clause) in self.clauses.iter().enumerate() { + if i > 0 { + write!(f, ", ")?; + } + write!(f, "{clause}")?; + } + write!(f, "]")?; + } + Ok(()) + } +} + /// A single WHEN clause within a MERGE INTO statement. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] pub struct MergeIntoClause { @@ -418,6 +435,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 @@ -444,6 +471,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 @@ -488,6 +526,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)]) @@ -516,6 +586,13 @@ 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]" + ); } #[test] diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 1a141ea52a13a..da6329eb2c493 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -2027,7 +2027,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 4f282f5e067fe..1d6427922f97f 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -3632,6 +3632,20 @@ fn plan_merge_into_canonicalizes_qualifiers_and_preserves_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=[person_quoted_cols.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": "person_quoted_cols.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 f868bcbdc4862..dd11d2940989e 100644 --- a/datafusion/sqllogictest/test_files/merge_into.slt +++ b/datafusion/sqllogictest/test_files/merge_into.slt @@ -18,8 +18,7 @@ ########## ## MERGE INTO Tests ## -## Note that MERGE INTO planning is supported, but the built-in MemTable does not -## (yet) support execution. These tests verify planning +## MERGE INTO planning and basic MemTable execution are supported. ########## statement ok @@ -55,12 +54,11 @@ 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' -02)caused by -03)This feature is not implemented: MERGE INTO not supported for Base table +physical_plan +01)CooperativeExec +02)--DmlResultExec: rows_affected=2 # Simple MATCHED DELETE query TT @@ -68,12 +66,11 @@ 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' -02)caused by -03)This feature is not implemented: MERGE INTO not supported for Base table +physical_plan +01)CooperativeExec +02)--DmlResultExec: rows_affected=2 # Aliased target and source: alias is canonicalized to the table name query TT @@ -82,13 +79,12 @@ 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=[target.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 -01)MERGE INTO operation on table 'target' -02)caused by -03)This feature is not implemented: MERGE INTO not supported for Base table +physical_plan +01)CooperativeExec +02)--DmlResultExec: rows_affected=2 # WHEN NOT MATCHED THEN DELETE is rejected by the parser (no target row exists); query error DELETE is not allowed in a NOT MATCHED merge clause at Line: 2, Column: 23 @@ -101,12 +97,11 @@ 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' -02)caused by -03)This feature is not implemented: MERGE INTO not supported for Base table +physical_plan +01)CooperativeExec +02)--DmlResultExec: rows_affected=0 # Subquery as the USING source query TT @@ -115,15 +110,14 @@ 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)]] 05)--------TableScan: source projection=[id, val] -physical_plan_error -01)MERGE INTO operation on table 'target' -02)caused by -03)This feature is not implemented: MERGE INTO not supported for Base table +physical_plan +01)CooperativeExec +02)--DmlResultExec: rows_affected=0 # INSERT without an explicit column list requires values for all target columns query TT @@ -131,21 +125,156 @@ 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' -02)caused by -03)This feature is not implemented: MERGE INTO not supported for Base table +physical_plan +01)CooperativeExec +02)--DmlResultExec: rows_affected=2 -# Execution fails: the default TableProvider does not implement merge_into -statement error -merge into target using source on target.id = source.id +########## +# MemTable execution +########## + +statement ok +create table exec_target(id int, val varchar, qty int); + +statement ok +insert into exec_target values (1, 'old', 10), (2, 'keep', 20); + +statement ok +create table exec_source(id int, val varchar, qty int, flag boolean); + +statement ok +insert into exec_source values (1, 'new', 100, true), (3, 'add', 30, true); + +query I +merge into exec_target using exec_source on exec_target.id = exec_source.id +when matched and exec_source.flag then update set val = exec_source.val, qty = exec_source.qty +when matched then update set val = 'second' +when not matched then insert (id, val, qty) values (exec_source.id, exec_source.val, exec_source.qty); +---- +2 + +query ITI rowsort +select * from exec_target; +---- +1 new 100 +2 keep 20 +3 add 30 + +statement ok +create table delete_target(id int, val varchar); + +statement ok +insert into delete_target values (1, 'one'), (2, 'two'), (3, 'three'); + +statement ok +create table delete_source(id int); + +statement ok +insert into delete_source values (2); + +query I +merge into delete_target using delete_source on delete_target.id = delete_source.id when matched then delete; ---- -DataFusion error: MERGE INTO operation on table 'target' -caused by -This feature is not implemented: MERGE INTO not supported for Base table +1 + +query IT rowsort +select * from delete_target; +---- +1 one +3 three + +statement ok +create table source_miss_target(id int, val varchar); + +statement ok +insert into source_miss_target values (1, 'one'), (2, 'two'), (3, 'three'); + +statement ok +create table source_miss_source(id int); + +statement ok +insert into source_miss_source values (2); + +query I +merge into source_miss_target using source_miss_source on source_miss_target.id = source_miss_source.id +when not matched by source and source_miss_target.id = 1 then update set val = 'stale' +when not matched by source and source_miss_target.id = 3 then delete; +---- +2 + +query IT rowsort +select * from source_miss_target; +---- +1 stale +2 two + +statement ok +create table null_pred_target(id int, val varchar); + +statement ok +insert into null_pred_target values (1, 'old'); + +statement ok +create table null_pred_source(id int, val varchar); + +statement ok +insert into null_pred_source values (1, 'new'), (2, 'add'); + +query I +merge into null_pred_target using null_pred_source on null +when matched then update set val = null_pred_source.val +when not matched and null then insert (id, val) values (null_pred_source.id, null_pred_source.val); +---- +0 + +query IT rowsort +select * from null_pred_target; +---- +1 old + +statement ok +create table default_target(id int, val varchar default 'missing', qty int); + +statement ok +create table default_source(id int); + +statement ok +insert into default_source values (10); + +query I +merge into default_target using default_source on false +when not matched then insert (id) values (default_source.id); +---- +1 + +query ITI +select * from default_target; +---- +10 missing NULL + +statement ok +create table dup_target(id int, val varchar); + +statement ok +insert into dup_target values (1, 'old'); + +statement ok +create table dup_source(id int, val varchar); + +statement ok +insert into dup_source values (1, 'new'), (1, 'newer'); + +statement error MERGE INTO matched target row 0 with more than one source row +merge into dup_target using dup_source on dup_target.id = dup_source.id +when matched then update set val = dup_source.val; + +query IT +select * from dup_target; +---- +1 old ########## @@ -241,6 +370,42 @@ statement error DataFusion error: This feature is not implemented: MERGE INSERT merge into target using source on target.id = source.id when not matched then insert row; +statement ok +drop table exec_target; + +statement ok +drop table exec_source; + +statement ok +drop table delete_target; + +statement ok +drop table delete_source; + +statement ok +drop table source_miss_target; + +statement ok +drop table source_miss_source; + +statement ok +drop table null_pred_target; + +statement ok +drop table null_pred_source; + +statement ok +drop table default_target; + +statement ok +drop table default_source; + +statement ok +drop table dup_target; + +statement ok +drop table dup_source; + statement ok drop table target; diff --git a/docs/source/user-guide/sql/dml.md b/docs/source/user-guide/sql/dml.md index 4934bc2674375..13c9d6ee8f4ab 100644 --- a/docs/source/user-guide/sql/dml.md +++ b/docs/source/user-guide/sql/dml.md @@ -136,3 +136,59 @@ INSERT INTO table_name { VALUES ( expression [, ...] | 2 | +-------+ ``` + +## MERGE INTO + +Merges rows from a source relation into a target table. Each target row may +match at most one source row. DataFusion applies only the first matching +`WHEN` clause for each row and returns the number of rows inserted, updated, or +deleted. + +
+MERGE INTO target_table [ AS target_alias ]
+USING { source_table | ( query ) } [ AS source_alias ]
+ON condition
+merge_clause [ ... ]
+
+ +`merge_clause` can be: + +```sql +WHEN MATCHED [ AND condition ] THEN UPDATE SET column = expression [, ...] +WHEN MATCHED [ AND condition ] THEN DELETE +WHEN NOT MATCHED [ BY TARGET ] [ AND condition ] THEN INSERT [(column [, ...])] VALUES (expression [, ...]) +WHEN NOT MATCHED BY SOURCE [ AND condition ] THEN UPDATE SET column = expression [, ...] +WHEN NOT MATCHED BY SOURCE [ AND condition ] THEN DELETE +``` + +### Examples + +```sql +> MERGE INTO inventory AS t + USING updates AS s + ON t.id = s.id + WHEN MATCHED AND s.deleted THEN DELETE + WHEN MATCHED THEN UPDATE SET qty = s.qty + WHEN NOT MATCHED THEN INSERT (id, qty) VALUES (s.id, s.qty); ++-------+ +| count | ++-------+ +| 3 | ++-------+ +``` + +### Provider support + +`MERGE INTO` execution depends on the target table provider. DataFusion's +in-memory `MemTable` supports basic `MERGE INTO` execution. Other providers +must implement `TableProvider::merge_into`; otherwise planning the statement +returns an unsupported-operation error. + +The SQL planner currently rejects these `MERGE INTO` forms: + +- target table modifiers, such as `MERGE INTO target PARTITION (...)` +- target alias column lists, such as `MERGE INTO target AS t(a, b)` +- `UPDATE ... WHERE` and `UPDATE ... DELETE WHERE` predicates inside a merge + action +- `INSERT ... WHERE` predicates inside a merge action +- `INSERT ROW`