From 6a21de74c666f9c0abb8302190636cc78af467d1 Mon Sep 17 00:00:00 2001 From: Nachiket Roy Date: Sun, 6 Sep 2026 19:52:54 +0530 Subject: [PATCH 1/3] feat: implement support for INSERT ... ON CONFLICT (DO NOTHING / DO UPDATE) clauses --- datafusion/catalog/src/memory/table.rs | 615 +++++++++++++++++- datafusion/sql/src/statement.rs | 280 +++++++- datafusion/sql/tests/common/mod.rs | 30 +- datafusion/sql/tests/sql_integration.rs | 70 ++ .../test_files/insert_on_conflict.slt | 166 +++++ .../sqllogictest/test_files/merge_into.slt | 52 +- 6 files changed, 1167 insertions(+), 46 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/insert_on_conflict.slt diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index d817aa8b7788a..4763657622dc9 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::future::ready; use std::sync::Arc; @@ -27,17 +27,23 @@ use crate::TableProvider; use arrow::array::{ Array, ArrayRef, BooleanArray, RecordBatch as ArrowRecordBatch, UInt64Array, }; +use arrow::compute::concat_batches; use arrow::compute::kernels::zip::zip; -use arrow::compute::{and, filter_record_batch}; +use arrow::compute::{and, cast, filter_record_batch}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::error::Result; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{Constraints, DFSchema, SchemaExt, not_impl_err, plan_err}; +use datafusion_common::{ + Column, Constraints, DFSchema, DFSchemaRef, DataFusionError, ScalarValue, SchemaExt, + exec_err, 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_physical_expr::{ @@ -47,7 +53,7 @@ use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, - PhysicalExpr, PlanProperties, ReplaceChildrenOptions, collect_partitioned, + PhysicalExpr, PlanProperties, ReplaceChildrenOptions, collect, collect_partitioned, }; use datafusion_session::Session; @@ -269,6 +275,22 @@ impl TableProvider for MemTable { { self.update_boxed(state, assignments, filters) } + + fn merge_into<'life0, 'life1, 'async_trait>( + &'life0 self, + state: &'life1 dyn Session, + source: Arc, + merge_schema: DFSchemaRef, + on: Expr, + clauses: Vec, + ) -> BoxFuture<'async_trait, Result>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + self.merge_into_boxed(state, source, merge_schema, on, clauses) + } } impl MemTable { @@ -534,7 +556,7 @@ impl MemTable { let column_name = field.name(); let original_column = batch.column_by_name(column_name).ok_or_else(|| { - datafusion_common::DataFusionError::Internal(format!( + DataFusionError::Internal(format!( "Column '{column_name}' not found in batch" )) })?; @@ -571,6 +593,585 @@ impl MemTable { Ok(Arc::new(DmlResultExec::new(total_updated))) } + + fn merge_into_boxed<'a>( + &'a self, + state: &'a dyn Session, + source: Arc, + merge_schema: DFSchemaRef, + on: Expr, + clauses: Vec, + ) -> BoxFuture<'a, Result>> { + Box::pin(self.merge_into_inner(state, source, merge_schema, on, clauses)) + } + + async fn merge_into_inner( + &self, + state: &dyn Session, + source: Arc, + merge_schema: DFSchemaRef, + on: Expr, + clauses: Vec, + ) -> Result> { + // Collect source batches + let source_batches = collect(source, state.task_ctx()).await?; + let total_source_rows: usize = source_batches.iter().map(|b| b.num_rows()).sum(); + if total_source_rows == 0 { + return Ok(Arc::new(DmlResultExec::new(0))); + } + + // Lock all partitions in global ascending order to avoid deadlock + let mut partitions = Vec::with_capacity(self.batches.len()); + for p in &self.batches { + partitions.push(p.write().await); + } + + *self.sort_order.lock() = vec![]; + + let target_num_cols = self.schema.fields().len(); + let target_schema_ref = Arc::clone(&self.schema); + + // Helper to extract equi-join column indices from `on` + let equi_keys = extract_equi_join_keys(&on, &merge_schema, target_num_cols)?; + let target_key_indices: Vec = equi_keys.iter().map(|(t, _)| *t).collect(); + let source_key_indices: Vec = equi_keys.iter().map(|(_, s)| *s).collect(); + + // Build target hash index across all partitions: + // Key -> (partition_idx, batch_idx, row_idx) + let mut target_key_map: HashMap, (usize, usize, usize)> = + HashMap::new(); + for (p_idx, partition) in partitions.iter().enumerate() { + for (b_idx, batch) in partition.iter().enumerate() { + for r_idx in 0..batch.num_rows() { + let key = target_key_indices + .iter() + .map(|&col_idx| { + ScalarValue::try_from_array(batch.column(col_idx), r_idx) + }) + .collect::>>()?; + + // NULL in any conflict key column never conflicts + if key.iter().any(|v| v.is_null()) { + continue; + } + + if target_key_map + .insert(key.clone(), (p_idx, b_idx, r_idx)) + .is_some() + { + return exec_err!( + "Table contains duplicate rows for conflict key '{key:?}', but ON CONFLICT requires unique keys" + ); + } + } + } + } + + // Find WHEN MATCHED, WHEN NOT MATCHED, and WHEN NOT MATCHED BY SOURCE clauses + let matched_clause = clauses + .iter() + .find(|c| c.kind == MergeIntoClauseKind::Matched); + let not_matched_clause = clauses.iter().find(|c| { + c.kind == MergeIntoClauseKind::NotMatched + || c.kind == MergeIntoClauseKind::NotMatchedByTarget + }); + let not_matched_by_source_clause = clauses + .iter() + .find(|c| c.kind == MergeIntoClauseKind::NotMatchedBySource); + + let has_update = matched_clause + .is_some_and(|c| matches!(c.action, MergeIntoAction::Update(_))); + let mut seen_incoming_keys = HashSet::new(); + + let mut matched_target_rows = HashSet::new(); + let mut row_updates: HashMap<(usize, usize, usize), Vec<(String, ScalarValue)>> = + HashMap::new(); + let mut row_deletions: HashSet<(usize, usize, usize)> = HashSet::new(); + let mut rows_to_insert: Vec = Vec::new(); + let mut affected_count: u64 = 0; + + for source_batch in &source_batches { + if source_batch.num_rows() == 0 { + continue; + } + + for s_r_idx in 0..source_batch.num_rows() { + let source_key = source_key_indices + .iter() + .map(|&col_idx| { + ScalarValue::try_from_array(source_batch.column(col_idx), s_r_idx) + }) + .collect::>>()?; + let has_null_key = source_key.iter().any(|v| v.is_null()); + + if !has_null_key { + if has_update { + if !seen_incoming_keys.insert(source_key.clone()) { + return exec_err!( + "ON CONFLICT DO UPDATE command cannot affect row a second time" + ); + } + } else { + // DO NOTHING: coalesce intra-batch duplicates that don't match target + if !target_key_map.contains_key(&source_key) + && !seen_incoming_keys.insert(source_key.clone()) + { + continue; + } + } + } + + let target_match = if has_null_key { + None + } else { + target_key_map.get(&source_key).copied() + }; + + if let Some(target_loc) = target_match { + if !matched_target_rows.insert(target_loc) { + return exec_err!( + "ON CONFLICT DO UPDATE command cannot affect row a second time" + ); + } + + if let Some(clause) = matched_clause { + let (p_idx, b_idx, r_idx) = target_loc; + let target_batch = &partitions[p_idx][b_idx]; + + let combined_batch = create_combined_row_batch( + &merge_schema, + target_batch, + r_idx, + source_batch, + s_r_idx, + )?; + + let predicate_passed = match &clause.predicate { + Some(pred) => { + let phys_pred = create_physical_expr( + pred, + &merge_schema, + state.execution_props(), + &PhysicalPlanningContext::default(), + )?; + let result = phys_pred.evaluate(&combined_batch)?; + let arr = result.into_array(1)?; + let bool_arr = arr + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal( + "Predicate did not evaluate to boolean" + .to_string(), + ) + })?; + bool_arr.value(0) && !bool_arr.is_null(0) + } + None => true, + }; + + if predicate_passed { + match &clause.action { + MergeIntoAction::Update(assignments) => { + let mut new_vals = + Vec::with_capacity(assignments.len()); + for (col_name, expr) in assignments { + let phys_expr = create_physical_expr( + expr, + &merge_schema, + state.execution_props(), + &PhysicalPlanningContext::default(), + )?; + let res = phys_expr.evaluate(&combined_batch)?; + let val = res.into_array(1)?; + let sv = ScalarValue::try_from_array(&val, 0)?; + new_vals.push((col_name.clone(), sv)); + } + row_updates.insert(target_loc, new_vals); + affected_count += 1; + } + MergeIntoAction::Delete => { + row_deletions.insert(target_loc); + affected_count += 1; + } + MergeIntoAction::Insert { .. } => {} + } + } + } + } else if let Some(clause) = not_matched_clause + && let MergeIntoAction::Insert { columns, values } = &clause.action + { + let not_matched_batch = create_not_matched_row_batch( + &merge_schema, + &target_schema_ref, + source_batch, + s_r_idx, + )?; + + let predicate_passed = match &clause.predicate { + Some(pred) => { + let phys_pred = create_physical_expr( + pred, + &merge_schema, + state.execution_props(), + &PhysicalPlanningContext::default(), + )?; + let result = phys_pred.evaluate(¬_matched_batch)?; + let arr = result.into_array(1)?; + let bool_arr = arr + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal( + "Predicate did not evaluate to boolean" + .to_string(), + ) + })?; + bool_arr.value(0) && !bool_arr.is_null(0) + } + None => true, + }; + + if predicate_passed { + let insert_col_names: Vec = if columns.is_empty() { + target_schema_ref + .fields() + .iter() + .map(|f| f.name().clone()) + .collect() + } else { + columns.clone() + }; + + let mut evaluated_cols = HashMap::with_capacity(values.len()); + for (col_name, expr) in insert_col_names.iter().zip(values.iter()) + { + let phys_expr = create_physical_expr( + expr, + &merge_schema, + state.execution_props(), + &PhysicalPlanningContext::default(), + )?; + let res = phys_expr.evaluate(¬_matched_batch)?; + let arr = res.into_array(1)?; + evaluated_cols.insert(col_name.clone(), arr); + } + + let mut row_cols = Vec::with_capacity(target_num_cols); + for field in target_schema_ref.fields() { + if let Some(arr) = evaluated_cols.remove(field.name()) { + let target_type = field.data_type(); + let casted_arr = if arr.data_type() == target_type { + arr + } else { + cast(&arr, target_type)? + }; + row_cols.push(casted_arr); + } else { + row_cols.push(arrow::array::new_null_array( + field.data_type(), + 1, + )); + } + } + + let projected_row = RecordBatch::try_new( + Arc::clone(&target_schema_ref), + row_cols, + )?; + rows_to_insert.push(projected_row); + affected_count += 1; + } + } + } + } + + // Handle WHEN NOT MATCHED BY SOURCE clauses + if let Some(clause) = not_matched_by_source_clause { + for (p_idx, partition) in partitions.iter().enumerate() { + for (b_idx, batch) in partition.iter().enumerate() { + for r_idx in 0..batch.num_rows() { + let target_loc = (p_idx, b_idx, r_idx); + if !matched_target_rows.contains(&target_loc) { + match &clause.action { + MergeIntoAction::Delete => { + row_deletions.insert(target_loc); + affected_count += 1; + } + MergeIntoAction::Update(assignments) => { + let combined_batch = + create_combined_row_batch_with_null_source( + &merge_schema, + batch, + r_idx, + )?; + let predicate_passed = match &clause.predicate { + Some(pred) => { + let phys_pred = create_physical_expr( + pred, + &merge_schema, + state.execution_props(), + &PhysicalPlanningContext::default(), + )?; + let result = + phys_pred.evaluate(&combined_batch)?; + let arr = result.into_array(1)?; + let bool_arr = arr + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal( + "Predicate did not evaluate to boolean" + .to_string(), + ) + })?; + bool_arr.value(0) && !bool_arr.is_null(0) + } + None => true, + }; + if predicate_passed { + let mut new_vals = + Vec::with_capacity(assignments.len()); + for (col_name, expr) in assignments { + let phys_expr = create_physical_expr( + expr, + &merge_schema, + state.execution_props(), + &PhysicalPlanningContext::default(), + )?; + let res = + phys_expr.evaluate(&combined_batch)?; + let val = res.into_array(1)?; + let sv = + ScalarValue::try_from_array(&val, 0)?; + new_vals.push((col_name.clone(), sv)); + } + row_updates.insert(target_loc, new_vals); + affected_count += 1; + } + } + _ => {} + } + } + } + } + } + } + + // Apply updates and deletions to partitions + for (p_idx, partition) in partitions.iter_mut().enumerate() { + let mut new_partition = Vec::with_capacity(partition.len()); + + for (b_idx, batch) in partition.iter().enumerate() { + let updates_for_batch: HashMap> = + row_updates + .iter() + .filter_map(|(&(p, b, r), vals)| { + if p == p_idx && b == b_idx { + Some((r, vals)) + } else { + None + } + }) + .collect(); + + let deletions_for_batch: HashSet = row_deletions + .iter() + .filter_map(|&(p, b, r)| { + if p == p_idx && b == b_idx { + Some(r) + } else { + None + } + }) + .collect(); + + if updates_for_batch.is_empty() && deletions_for_batch.is_empty() { + new_partition.push(batch.clone()); + continue; + } + + let remaining_rows: Vec = (0..batch.num_rows()) + .filter(|r| !deletions_for_batch.contains(r)) + .collect(); + + if remaining_rows.is_empty() { + continue; + } + + let mut new_columns = Vec::with_capacity(batch.num_columns()); + for (col_idx, field) in target_schema_ref.fields().iter().enumerate() { + let col_name = field.name(); + let orig_col = batch.column(col_idx); + + let mut new_scalars = Vec::with_capacity(remaining_rows.len()); + for &r_idx in &remaining_rows { + if let Some(vals) = updates_for_batch.get(&r_idx) { + if let Some((_, new_sv)) = + vals.iter().find(|(name, _)| name == col_name) + { + new_scalars.push(new_sv.clone()); + } else { + new_scalars + .push(ScalarValue::try_from_array(orig_col, r_idx)?); + } + } else { + new_scalars + .push(ScalarValue::try_from_array(orig_col, r_idx)?); + } + } + let new_col = ScalarValue::iter_to_array(new_scalars)?; + new_columns.push(new_col); + } + new_partition.push(ArrowRecordBatch::try_new( + Arc::clone(&target_schema_ref), + new_columns, + )?); + } + + **partition = new_partition; + } + + // Append rows_to_insert to partition 0 + if !rows_to_insert.is_empty() { + let combined_inserts = concat_batches(&target_schema_ref, &rows_to_insert)?; + if partitions.is_empty() { + return not_impl_err!("MemTable has no partitions"); + } + partitions[0].push(combined_inserts); + } + + Ok(Arc::new(DmlResultExec::new(affected_count))) + } +} + +fn create_combined_row_batch( + merge_schema: &DFSchema, + target_batch: &RecordBatch, + target_row_idx: usize, + source_batch: &RecordBatch, + source_row_idx: usize, +) -> Result { + let mut columns = Vec::with_capacity(merge_schema.fields().len()); + for col in target_batch.columns() { + columns.push(col.slice(target_row_idx, 1)); + } + for col in source_batch.columns() { + columns.push(col.slice(source_row_idx, 1)); + } + let arrow_schema = Arc::new(merge_schema.as_arrow().clone()); + Ok(RecordBatch::try_new(arrow_schema, columns)?) +} + +fn create_not_matched_row_batch( + merge_schema: &DFSchema, + target_schema: &SchemaRef, + source_batch: &RecordBatch, + source_row_idx: usize, +) -> Result { + let mut columns = Vec::with_capacity(merge_schema.fields().len()); + for field in target_schema.fields() { + columns.push(arrow::array::new_null_array(field.data_type(), 1)); + } + for col in source_batch.columns() { + columns.push(col.slice(source_row_idx, 1)); + } + let arrow_schema = Arc::new(merge_schema.as_arrow().clone()); + Ok(RecordBatch::try_new(arrow_schema, columns)?) +} + +fn create_combined_row_batch_with_null_source( + merge_schema: &DFSchema, + target_batch: &RecordBatch, + target_row_idx: usize, +) -> Result { + let mut columns = Vec::with_capacity(merge_schema.fields().len()); + for col in target_batch.columns() { + columns.push(col.slice(target_row_idx, 1)); + } + let target_num_cols = target_batch.num_columns(); + for field in &merge_schema.fields()[target_num_cols..] { + columns.push(arrow::array::new_null_array(field.data_type(), 1)); + } + let arrow_schema = Arc::new(merge_schema.as_arrow().clone()); + Ok(RecordBatch::try_new(arrow_schema, columns)?) +} + +fn extract_equi_join_keys( + on: &Expr, + merge_schema: &DFSchema, + target_num_cols: usize, +) -> Result> { + let mut pairs = Vec::new(); + collect_equi_keys(on, merge_schema, target_num_cols, &mut pairs)?; + if pairs.is_empty() { + return plan_err!( + "MemTable MERGE INTO requires at least one equi-join condition in ON clause" + ); + } + Ok(pairs) +} + +fn extract_column(expr: &Expr) -> Option<&Column> { + match expr { + Expr::Column(c) => Some(c), + Expr::Alias(datafusion_expr::expr::Alias { expr: inner, .. }) => { + extract_column(inner.as_ref()) + } + Expr::Cast(datafusion_expr::Cast { expr: inner, .. }) => { + extract_column(inner.as_ref()) + } + _ => None, + } +} + +fn collect_equi_keys( + expr: &Expr, + merge_schema: &DFSchema, + target_num_cols: usize, + pairs: &mut Vec<(usize, usize)>, +) -> Result<()> { + match expr { + Expr::Alias(datafusion_expr::expr::Alias { expr: inner, .. }) => { + collect_equi_keys(inner.as_ref(), merge_schema, target_num_cols, pairs) + } + Expr::BinaryExpr(datafusion_expr::BinaryExpr { left, op, right }) => match op { + datafusion_expr::Operator::And => { + collect_equi_keys(left.as_ref(), merge_schema, target_num_cols, pairs)?; + collect_equi_keys(right.as_ref(), merge_schema, target_num_cols, pairs)?; + Ok(()) + } + datafusion_expr::Operator::Eq => { + let left_col = extract_column(left.as_ref()); + let right_col = extract_column(right.as_ref()); + if let (Some(c1), Some(c2)) = (left_col, right_col) { + let idx1 = merge_schema.index_of_column(c1)?; + let idx2 = merge_schema.index_of_column(c2)?; + if idx1 < target_num_cols && idx2 >= target_num_cols { + pairs.push((idx1, idx2 - target_num_cols)); + Ok(()) + } else if idx2 < target_num_cols && idx1 >= target_num_cols { + pairs.push((idx2, idx1 - target_num_cols)); + Ok(()) + } else { + plan_err!( + "ON equality condition must compare a target column with a source column: {expr}" + ) + } + } else { + plan_err!( + "MemTable MERGE INTO requires column equality conditions in ON clause, found: {expr}" + ) + } + } + _ => plan_err!( + "MemTable MERGE INTO only supports AND and EQ in ON condition, found: {expr}" + ), + }, + _ => plan_err!( + "MemTable MERGE INTO requires binary equality expressions in ON condition, found: {expr}" + ), + } } /// Evaluate filter expressions against a batch and return a combined boolean mask. @@ -602,7 +1203,7 @@ fn evaluate_filters_to_mask( .as_any() .downcast_ref::() .ok_or_else(|| { - datafusion_common::DataFusionError::Internal( + DataFusionError::Internal( "Filter did not evaluate to boolean".to_string(), ) })? diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 1a9072212f2f3..6eab1b71b2d0a 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -65,11 +65,11 @@ use sqlparser::ast::{ UniqueConstraint, Update, UpdateTableFromKind, ValueWithSpan, }; use sqlparser::ast::{ - Assignment, AssignmentTarget, ColumnDef, CreateIndex, CreateTable, + Assignment, AssignmentTarget, ColumnDef, ConflictTarget, CreateIndex, CreateTable, CreateTableOptions, Delete, DescribeAlias, Expr as SQLExpr, FromTable, Ident, Insert, - ObjectName, ObjectType, Query, SchemaName, SetExpr, ShowCreateObject, - ShowStatementFilter, Statement, TableConstraint, TableFactor, TableWithJoins, - TransactionMode, UnaryOperator, Value, + ObjectName, ObjectType, OnConflictAction, OnInsert, Query, SchemaName, SetExpr, + ShowCreateObject, ShowStatementFilter, Statement, TableConstraint, TableFactor, + TableWithJoins, TransactionMode, UnaryOperator, Value, }; use sqlparser::parser::ParserError::ParserError; @@ -1075,9 +1075,7 @@ impl SqlToRel<'_, S> { if !after_columns.is_empty() { plan_err!("After-columns clause not supported")?; } - if on.is_some() { - plan_err!("Insert-on clause not supported")?; - } + if returning.is_some() { plan_err!("Insert-returning clause not supported")?; } @@ -1125,7 +1123,14 @@ impl SqlToRel<'_, S> { // optional keywords don't change behavior let _ = into; let _ = has_table_keyword; - self.insert_to_plan(table_name, columns, source, overwrite, replace_into) + self.insert_to_plan( + table_name, + columns, + source, + overwrite, + replace_into, + on, + ) } Statement::Update(Update { table, @@ -2821,6 +2826,7 @@ impl SqlToRel<'_, S> { source: Box, overwrite: bool, replace_into: bool, + on: Option, ) -> Result { // Do a table lookup to verify the table exists let table_name = self.object_name_to_table_reference(table_name)?; @@ -2958,6 +2964,264 @@ impl SqlToRel<'_, S> { .collect::>>()?; let source = project(source, exprs)?; + if let Some(on_insert) = on { + if overwrite { + return plan_err!( + "INSERT OVERWRITE cannot be combined with ON CONFLICT clause" + ); + } + if replace_into { + return plan_err!( + "REPLACE INTO cannot be combined with ON CONFLICT clause" + ); + } + + let on_conflict = match on_insert { + OnInsert::OnConflict(on_conflict) => on_conflict, + OnInsert::DuplicateKeyUpdate(_) => { + return not_impl_err!("ON DUPLICATE KEY UPDATE is not supported"); + } + _ => return not_impl_err!("Unsupported ON clause"), + }; + + // Target table cannot be named or aliased as 'excluded' + if table_name.table().eq_ignore_ascii_case("excluded") { + return plan_err!( + "Target table cannot be named or aliased as 'excluded' when using ON CONFLICT" + ); + } + + // Extract and validate conflict target columns + let target_col_names: Vec = match on_conflict.conflict_target { + Some(ConflictTarget::Columns(cols)) => { + if cols.is_empty() { + return plan_err!( + "ON CONFLICT target column list cannot be empty" + ); + } + cols.into_iter() + .map(|c| self.ident_normalizer.normalize(c)) + .collect() + } + Some(ConflictTarget::OnConstraint(_)) => { + return not_impl_err!( + "ON CONFLICT ON CONSTRAINT is not supported because table constraints do not store constraint names" + ); + } + None => { + // In PostgreSQL, ON CONFLICT without target is only allowed with DO NOTHING. + if matches!(on_conflict.action, OnConflictAction::DoUpdate(_)) { + return plan_err!( + "ON CONFLICT DO UPDATE requires a conflict target specification" + ); + } + // For DO NOTHING without target: + // If table has unique/PK constraints, infer from the first constraint; + // otherwise, if no constraints exist, no conflict can occur -> normal insert. + let inferred_target = + table_source.constraints().and_then(|constraints| { + constraints.iter().find_map(|c| match c { + Constraint::PrimaryKey(indices) + | Constraint::Unique(indices) + if !indices.is_empty() => + { + Some( + indices + .iter() + .map(|&i| { + table_schema.field(i).name().clone() + }) + .collect::>(), + ) + } + _ => None, + }) + }); + + match inferred_target { + Some(cols) => cols, + None => { + // No constraints to conflict on: DO NOTHING is equivalent to a regular insert. + return Ok(LogicalPlan::Dml(DmlStatement::new( + table_name, + Arc::clone(&table_source), + WriteOp::Insert(InsertOp::Append), + Arc::new(source), + ))); + } + } + } + }; + + // Validate that target columns exist in the table schema + let mut target_indices = Vec::with_capacity(target_col_names.len()); + for col_name in &target_col_names { + let idx = table_schema + .index_of_column_by_name(None, col_name) + .ok_or_else(|| { + unqualified_field_not_found(col_name, &table_schema) + })?; + target_indices.push(idx); + } + + // If table provider provides non-empty constraints, validate against them + if let Some(constraints) = + table_source.constraints().filter(|c| !c.is_empty()) + { + let mut sorted_target_indices = target_indices.clone(); + sorted_target_indices.sort_unstable(); + + let matches_constraint = constraints.iter().any(|c| { + let mut c_indices = match c { + Constraint::PrimaryKey(indices) | Constraint::Unique(indices) => { + indices.clone() + } + }; + c_indices.sort_unstable(); + c_indices == sorted_target_indices + }); + + if !matches_constraint { + return plan_err!( + "There is no unique or exclusion constraint matching the ON CONFLICT specification" + ); + } + } + + // Wrap source in SubqueryAlias with qualifier "excluded" + let excluded_qualifier = TableReference::bare("excluded"); + let source_aliased = + LogicalPlan::SubqueryAlias(datafusion_expr::SubqueryAlias::try_new( + Arc::new(source), + excluded_qualifier.clone(), + )?); + + // Construct ON condition: target.col = excluded.col AND ... + let mut on_expr: Option = None; + for col_name in &target_col_names { + let target_col = + Expr::Column(Column::new(Some(table_name.clone()), col_name.clone())); + let excluded_col = Expr::Column(Column::new( + Some(excluded_qualifier.clone()), + col_name.clone(), + )); + let eq_expr = target_col.eq(excluded_col); + on_expr = match on_expr { + Some(prev) => Some(prev.and(eq_expr)), + None => Some(eq_expr), + }; + } + let on_expr = on_expr.ok_or_else(|| { + plan_datafusion_err!("Expected at least one conflict column") + })?; + + // Columns for the WHEN NOT MATCHED THEN INSERT clause: all table columns from excluded + let insert_columns: Vec = table_schema + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(); + let insert_values: Vec = table_schema + .fields() + .iter() + .map(|f| { + Expr::Column(Column::new( + Some(excluded_qualifier.clone()), + f.name().clone(), + )) + }) + .collect(); + + let target_schema_qualified = Arc::new(DFSchema::try_from_qualified_schema( + table_name.clone(), + &table_source.schema(), + )?); + let combined_schema = Arc::new( + target_schema_qualified + .as_ref() + .join(source_aliased.schema())?, + ); + + let clauses = match on_conflict.action { + OnConflictAction::DoNothing => { + vec![MergeIntoClause { + kind: MergeIntoClauseKind::NotMatched, + predicate: None, + action: MergeIntoAction::Insert { + columns: insert_columns, + values: insert_values, + }, + }] + } + OnConflictAction::DoUpdate(do_update) => { + let predicate = do_update + .selection + .map(|p| { + self.sql_to_expr(p, &combined_schema, &mut planner_context) + }) + .transpose()?; + + let assignments = do_update + .assignments + .into_iter() + .map(|assign| { + let col_name = match &assign.target { + AssignmentTarget::ColumnName(cols) => { + self.merge_target_column_name(cols, &table_name)? + } + _ => plan_err!("Tuples are not supported")?, + }; + target_schema_qualified + .field_with_unqualified_name(&col_name)?; + let value = self.sql_to_expr( + assign.value, + &combined_schema, + &mut planner_context, + )?; + Ok((col_name, value)) + }) + .collect::>>()?; + + let mut seen = HashSet::new(); + for (col, _) in &assignments { + if !seen.insert(col.as_str()) { + return plan_err!( + "Duplicate column '{col}' in ON CONFLICT DO UPDATE" + ); + } + } + + vec![ + MergeIntoClause { + kind: MergeIntoClauseKind::Matched, + predicate, + action: MergeIntoAction::Update(assignments), + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatched, + predicate: None, + action: MergeIntoAction::Insert { + columns: insert_columns, + values: insert_values, + }, + }, + ] + } + }; + + let merge_op = MergeIntoOp { + on: on_expr, + clauses, + }; + + return Ok(LogicalPlan::Dml(DmlStatement::new( + table_name, + Arc::clone(&table_source), + WriteOp::MergeInto(Box::new(merge_op)), + Arc::new(source_aliased), + ))); + } + let insert_op = match (overwrite, replace_into) { (false, false) => InsertOp::Append, (true, false) => InsertOp::Overwrite, diff --git a/datafusion/sql/tests/common/mod.rs b/datafusion/sql/tests/common/mod.rs index e7c819bbf64a6..1089246cf8bfa 100644 --- a/datafusion/sql/tests/common/mod.rs +++ b/datafusion/sql/tests/common/mod.rs @@ -25,7 +25,9 @@ use arrow::datatypes::*; use datafusion_common::config::ConfigOptions; use datafusion_common::datatype::DataTypeExt; use datafusion_common::file_options::file_type::FileType; -use datafusion_common::{DFSchema, GetExt, Result, TableReference, plan_err}; +use datafusion_common::{ + Constraint, Constraints, DFSchema, GetExt, Result, TableReference, plan_err, +}; use datafusion_expr::planner::{ExprPlanner, PlannerResult, TypePlanner}; use datafusion_expr::{ AggregateUDF, Expr, HigherOrderUDF, ScalarUDF, TableSource, WindowUDF, @@ -278,6 +280,15 @@ impl ContextProvider for MockContextProvider { DataType::UInt32, false, )])), + "table_with_pk" => { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + ])); + let constraints = + Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]); + return Ok(Arc::new(EmptyTable::with_constraints(schema, constraints))); + } _ => plan_err!("No table named: {} found", name.table()), }; @@ -357,11 +368,22 @@ impl ContextProvider for MockContextProvider { struct EmptyTable { table_schema: SchemaRef, + constraints: Option, } impl EmptyTable { fn new(table_schema: SchemaRef) -> Self { - Self { table_schema } + Self { + table_schema, + constraints: None, + } + } + + fn with_constraints(table_schema: SchemaRef, constraints: Constraints) -> Self { + Self { + table_schema, + constraints: Some(constraints), + } } } @@ -369,6 +391,10 @@ impl TableSource for EmptyTable { fn schema(&self) -> SchemaRef { Arc::clone(&self.table_schema) } + + fn constraints(&self) -> Option<&Constraints> { + self.constraints.as_ref() + } } #[derive(Debug)] diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 9f57aaafb0686..fb4ee97c70c91 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -719,6 +719,37 @@ fn plan_insert_no_target_columns() { ); } +#[test] +fn plan_insert_on_conflict_do_nothing() { + let sql = + "INSERT INTO test_decimal (id, price) VALUES (1, 2) ON CONFLICT (id) DO NOTHING"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r" + Dml: op=[MergeInto] table=[test_decimal] + SubqueryAlias: excluded + Projection: column1 AS id, column2 AS price + Values: (CAST(Int64(1) AS Int32), CAST(Int64(2) AS Decimal128(10, 2))) + " + ); +} + +#[test] +fn plan_insert_on_conflict_do_update() { + let sql = "INSERT INTO test_decimal (id, price) VALUES (1, 2) ON CONFLICT (id) DO UPDATE SET price = EXCLUDED.price"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r" + Dml: op=[MergeInto] table=[test_decimal] + SubqueryAlias: excluded + Projection: column1 AS id, column2 AS price + Values: (CAST(Int64(1) AS Int32), CAST(Int64(2) AS Decimal128(10, 2))) + " + ); +} + #[rstest] #[case::duplicate_columns( "INSERT INTO test_decimal (id, price, price) VALUES (1, 2, 3), (4, 5, 6)", @@ -745,12 +776,51 @@ fn plan_insert_no_target_columns() { "INSERT INTO person (id, first_name, last_name) VALUES ($id, $first_name, $last_name)", "Error during planning: Can't parse placeholder: $id" )] +#[case::on_conflict_with_overwrite( + "INSERT OVERWRITE test_decimal (id, price) VALUES (1, 2) ON CONFLICT (id) DO NOTHING", + "Error during planning: INSERT OVERWRITE cannot be combined with ON CONFLICT clause" +)] +#[case::on_conflict_on_constraint( + "INSERT INTO test_decimal (id, price) VALUES (1, 2) ON CONFLICT ON CONSTRAINT some_constraint DO NOTHING", + "This feature is not implemented: ON CONFLICT ON CONSTRAINT is not supported because table constraints do not store constraint names" +)] +#[case::on_conflict_nonexistent_column( + "INSERT INTO test_decimal (id, price) VALUES (1, 2) ON CONFLICT (nonexistent) DO NOTHING", + "Schema error: No field named nonexistent.\nValid fields are id, price." +)] +#[case::on_conflict_do_update_without_target( + "INSERT INTO test_decimal (id, price) VALUES (1, 2) ON CONFLICT DO UPDATE SET price = 10", + "Error during planning: ON CONFLICT DO UPDATE requires a conflict target specification" +)] +#[case::on_conflict_duplicate_assignment( + "INSERT INTO test_decimal (id, price) VALUES (1, 2) ON CONFLICT (id) DO UPDATE SET price = 10, price = 20", + "Error during planning: Duplicate column 'price' in ON CONFLICT DO UPDATE" +)] +#[case::on_conflict_constraint_mismatch( + "INSERT INTO table_with_pk (id, name) VALUES (1, 'a') ON CONFLICT (name) DO NOTHING", + "Error during planning: There is no unique or exclusion constraint matching the ON CONFLICT specification" +)] #[test] fn test_insert_schema_errors(#[case] sql: &str, #[case] error: &str) { let err = logical_plan(sql).unwrap_err(); assert_eq!(err.strip_backtrace(), error) } +#[test] +fn plan_insert_on_conflict_with_matching_pk() { + let sql = "INSERT INTO table_with_pk (id, name) VALUES (1, 'a') ON CONFLICT (id) DO NOTHING"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Dml: op=[MergeInto] table=[table_with_pk] + SubqueryAlias: excluded + Projection: column1 AS id, column2 AS name + Values: (CAST(Int64(1) AS Int32), Utf8("a")) + "# + ); +} + #[test] fn plan_update() { let sql = "update person set last_name='Kay' where id=1"; diff --git a/datafusion/sqllogictest/test_files/insert_on_conflict.slt b/datafusion/sqllogictest/test_files/insert_on_conflict.slt new file mode 100644 index 0000000000000..2dc019f3b83f2 --- /dev/null +++ b/datafusion/sqllogictest/test_files/insert_on_conflict.slt @@ -0,0 +1,166 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +########## +## INSERT ... ON CONFLICT Tests +########## + +statement ok +create table users(id int, name varchar, score int); + +# 1. Basic insert +statement ok +insert into users values (1, 'Alice', 100), (2, 'Bob', 200); + +query ITI rowsort +select * from users; +---- +1 Alice 100 +2 Bob 200 + +# 2. ON CONFLICT DO NOTHING: conflicting row skipped, non-conflicting inserted +statement ok +insert into users values (2, 'Robert', 250), (3, 'Charlie', 300) +on conflict (id) do nothing; + +query ITI rowsort +select * from users; +---- +1 Alice 100 +2 Bob 200 +3 Charlie 300 + +# 3. ON CONFLICT DO UPDATE: conflicting row updated, non-conflicting inserted +statement ok +insert into users values (1, 'Alicia', 150), (4, 'David', 400) +on conflict (id) do update set name = excluded.name, score = excluded.score; + +query ITI rowsort +select * from users; +---- +1 Alicia 150 +2 Bob 200 +3 Charlie 300 +4 David 400 + +# 4. ON CONFLICT DO UPDATE with WHERE clause +# Update only if excluded score is higher than current score +statement ok +insert into users values (1, 'Alicia Low', 50) +on conflict (id) do update set name = excluded.name, score = excluded.score +where excluded.score > users.score; + +query ITI rowsort +select * from users; +---- +1 Alicia 150 +2 Bob 200 +3 Charlie 300 +4 David 400 + +# 5. NULL conflict keys never conflict (NULL != NULL) +statement ok +insert into users values (null, 'Null1', 10), (null, 'Null2', 20) +on conflict (id) do nothing; + +query ITI rowsort +select * from users; +---- +1 Alicia 150 +2 Bob 200 +3 Charlie 300 +4 David 400 +NULL Null1 10 +NULL Null2 20 + +# 6. Intra-batch duplicate keys: DO UPDATE must fail +statement error ON CONFLICT DO UPDATE command cannot affect row a second time +insert into users values (10, 'Dup1', 100), (10, 'Dup2', 200) +on conflict (id) do update set name = excluded.name; + +# 7. Intra-batch duplicate keys: DO NOTHING coalesces/deduplicates +statement ok +insert into users values (10, 'First', 100), (10, 'Second', 200) +on conflict (id) do nothing; + +query ITI rowsort +select * from users where id = 10; +---- +10 First 100 + +# Clean up users table +statement ok +drop table users; + +# 8. Composite conflict keys (multi-column) +statement ok +create table composite_t(a int, b int, val varchar); + +statement ok +insert into composite_t values (1, 10, 'first'), (1, 20, 'second'); + +statement ok +insert into composite_t values (1, 10, 'updated'), (2, 10, 'third') +on conflict (a, b) do update set val = excluded.val; + +query IIT rowsort +select * from composite_t; +---- +1 10 updated +1 20 second +2 10 third + +statement ok +drop table composite_t; + +# 9. Runtime duplicate key detection on unconstrained table +statement ok +create table unconstrained_dups(id int, val varchar); + +statement ok +insert into unconstrained_dups values (1, 'first'), (1, 'duplicate'); + +statement error Table contains duplicate rows for conflict key +insert into unconstrained_dups values (1, 'incoming') +on conflict (id) do nothing; + +statement ok +drop table unconstrained_dups; + +# 10. Planning errors: Mutual exclusivity and validation + +statement ok +create table err_t(id int, val varchar); + +# Cannot combine with INSERT OVERWRITE +statement error INSERT OVERWRITE cannot be combined with ON CONFLICT clause +insert overwrite err_t values (1, 'a') on conflict (id) do nothing; + +# Target table aliased as excluded +statement error DataFusion error: SQL error: ParserError\("Expected: SELECT, VALUES, or a subquery in the query body, found: as at Line: 1, Column: 19"\) +insert into err_t as excluded values (1, 'a') on conflict (id) do nothing; + +# Non-existent conflict column +statement error Schema error: No field named nonexistent +insert into err_t values (1, 'a') on conflict (nonexistent) do nothing; + +# Duplicate column in DO UPDATE +statement error Duplicate column 'val' in ON CONFLICT DO UPDATE +insert into err_t values (1, 'a') on conflict (id) do update set val = 'x', val = 'y'; + +statement ok +drop table err_t; diff --git a/datafusion/sqllogictest/test_files/merge_into.slt b/datafusion/sqllogictest/test_files/merge_into.slt index f868bcbdc4862..eece82c447dd2 100644 --- a/datafusion/sqllogictest/test_files/merge_into.slt +++ b/datafusion/sqllogictest/test_files/merge_into.slt @@ -57,10 +57,9 @@ when not matched then insert (id, val) values (source.id, source.val); logical_plan 01)Dml: op=[MergeInto] table=[target] 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 @@ -70,10 +69,9 @@ when matched then delete; logical_plan 01)Dml: op=[MergeInto] table=[target] 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 @@ -85,10 +83,9 @@ logical_plan 01)Dml: op=[MergeInto] table=[target] 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 @@ -103,10 +100,9 @@ when not matched by source then delete; logical_plan 01)Dml: op=[MergeInto] table=[target] 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 @@ -120,10 +116,9 @@ logical_plan 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 @@ -133,19 +128,18 @@ when not matched then insert values (source.id, source.val, 0); logical_plan 01)Dml: op=[MergeInto] table=[target] 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 +# Execution: MATCHED DELETE +statement ok merge into target using source on target.id = source.id when matched then delete; + +query ITI rowsort +select * from target; ---- -DataFusion error: MERGE INTO operation on table 'target' -caused by -This feature is not implemented: MERGE INTO not supported for Base table ########## From 77f940bc79c458a57bf68427d4ab5a02fcd00c49 Mon Sep 17 00:00:00 2001 From: Nachiket Roy Date: Sun, 6 Sep 2026 20:06:25 +0530 Subject: [PATCH 2/3] fix: allow duplicate matches in DO NOTHING and test multi-constraint inference --- datafusion/catalog/src/memory/table.rs | 8 ++++- datafusion/sql/src/statement.rs | 11 +++++-- datafusion/sql/tests/common/mod.rs | 12 +++++++ datafusion/sql/tests/sql_integration.rs | 32 +++++++++++++++++++ .../test_files/insert_on_conflict.slt | 15 +++++++++ 5 files changed, 75 insertions(+), 3 deletions(-) diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index 4763657622dc9..5044fdcd0ba07 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -728,7 +728,8 @@ impl MemTable { }; if let Some(target_loc) = target_match { - if !matched_target_rows.insert(target_loc) { + let already_matched = !matched_target_rows.insert(target_loc); + if already_matched && matched_clause.is_some() { return exec_err!( "ON CONFLICT DO UPDATE command cannot affect row a second time" ); @@ -1112,6 +1113,11 @@ fn extract_equi_join_keys( Ok(pairs) } +/// Extracts a column reference from an expression, unwrapping aliases and casts. +/// +/// Note: Unwrapping `Expr::Cast` assumes type-coercion casts inserted by the DataFusion planner +/// preserve equi-join uniqueness semantics (e.g. natural type widening). A lossy or non-injective +/// cast in a general MERGE ON condition could map distinct source values to the same key. fn extract_column(expr: &Expr) -> Option<&Column> { match expr { Expr::Column(c) => Some(c), diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 6eab1b71b2d0a..c103f78386557 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -3016,8 +3016,15 @@ impl SqlToRel<'_, S> { ); } // For DO NOTHING without target: - // If table has unique/PK constraints, infer from the first constraint; - // otherwise, if no constraints exist, no conflict can occur -> normal insert. + // Infer conflict target from the first declared Unique or PrimaryKey constraint. + // Note: In PostgreSQL, an unqualified ON CONFLICT DO NOTHING (with no conflict target) + // suppresses violations against *any* unique or exclusion constraint on the table. + // In DataFusion, the operation is desugared into MERGE INTO, which joins on a single + // conjunction of equi-join keys. Inferring the target from the first matching constraint + // is a deliberate simplification for tables with unique/PK constraints; on tables with + // multiple disjoint constraints, conflicts against constraints other than the first will + // not be matched by this equi-join and may be rejected downstream by the table provider. + // If no constraints exist on the table, no conflict can occur -> normal insert. let inferred_target = table_source.constraints().and_then(|constraints| { constraints.iter().find_map(|c| match c { diff --git a/datafusion/sql/tests/common/mod.rs b/datafusion/sql/tests/common/mod.rs index 1089246cf8bfa..c5d3ef8c33f10 100644 --- a/datafusion/sql/tests/common/mod.rs +++ b/datafusion/sql/tests/common/mod.rs @@ -289,6 +289,18 @@ impl ContextProvider for MockContextProvider { Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]); return Ok(Arc::new(EmptyTable::with_constraints(schema, constraints))); } + "table_with_multi_constraints" => { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("email", DataType::Utf8, false), + Field::new("name", DataType::Utf8, false), + ])); + let constraints = Constraints::new_unverified(vec![ + Constraint::PrimaryKey(vec![0]), + Constraint::Unique(vec![1]), + ]); + return Ok(Arc::new(EmptyTable::with_constraints(schema, constraints))); + } _ => plan_err!("No table named: {} found", name.table()), }; diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index fb4ee97c70c91..ca4c11a9ed7f6 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -821,6 +821,38 @@ fn plan_insert_on_conflict_with_matching_pk() { ); } +#[test] +fn plan_insert_on_conflict_unqualified_multi_constraints() { + // Unqualified ON CONFLICT DO NOTHING infers the target from the first constraint (PrimaryKey(id)) + let sql = "INSERT INTO table_with_multi_constraints (id, email, name) VALUES (1, 'a@b.com', 'a') ON CONFLICT DO NOTHING"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Dml: op=[MergeInto] table=[table_with_multi_constraints] + SubqueryAlias: excluded + Projection: column1 AS id, column2 AS email, column3 AS name + Values: (CAST(Int64(1) AS Int32), Utf8("a@b.com"), Utf8("a")) + "# + ); +} + +#[test] +fn plan_insert_on_conflict_explicit_second_constraint() { + // Explicit target on the second constraint (Unique(email)) matches and plans successfully + let sql = "INSERT INTO table_with_multi_constraints (id, email, name) VALUES (1, 'a@b.com', 'a') ON CONFLICT (email) DO NOTHING"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Dml: op=[MergeInto] table=[table_with_multi_constraints] + SubqueryAlias: excluded + Projection: column1 AS id, column2 AS email, column3 AS name + Values: (CAST(Int64(1) AS Int32), Utf8("a@b.com"), Utf8("a")) + "# + ); +} + #[test] fn plan_update() { let sql = "update person set last_name='Kay' where id=1"; diff --git a/datafusion/sqllogictest/test_files/insert_on_conflict.slt b/datafusion/sqllogictest/test_files/insert_on_conflict.slt index 2dc019f3b83f2..0dffdf55cfc40 100644 --- a/datafusion/sqllogictest/test_files/insert_on_conflict.slt +++ b/datafusion/sqllogictest/test_files/insert_on_conflict.slt @@ -102,6 +102,21 @@ select * from users where id = 10; ---- 10 First 100 +# Pre-existing duplicate key hit twice in DO NOTHING: silently skips both +statement ok +insert into users values (10, 'Third', 300), (10, 'Fourth', 400) +on conflict (id) do nothing; + +query ITI rowsort +select * from users where id = 10; +---- +10 First 100 + +# Pre-existing duplicate key hit twice in DO UPDATE: errors as row cannot be affected second time +statement error ON CONFLICT DO UPDATE command cannot affect row a second time +insert into users values (10, 'Update1', 1), (10, 'Update2', 2) +on conflict (id) do update set name = excluded.name; + # Clean up users table statement ok drop table users; From 0994ef45435368d7077c2c96a877dea5da3fbb3a Mon Sep 17 00:00:00 2001 From: Nachiket Roy Date: Sun, 6 Sep 2026 20:36:15 +0530 Subject: [PATCH 3/3] fix(merge_into): validate equi-join condition before empty check and preserve Base table error prefix --- datafusion/catalog/src/memory/table.rs | 36 +++++++++++++------------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index 5044fdcd0ba07..9b45d539c8a9e 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -613,6 +613,14 @@ impl MemTable { on: Expr, clauses: Vec, ) -> Result> { + let target_num_cols = self.schema.fields().len(); + let target_schema_ref = Arc::clone(&self.schema); + + // Helper to extract equi-join column indices from `on` + let equi_keys = extract_equi_join_keys(&on, &merge_schema, target_num_cols)?; + let target_key_indices: Vec = equi_keys.iter().map(|(t, _)| *t).collect(); + let source_key_indices: Vec = equi_keys.iter().map(|(_, s)| *s).collect(); + // Collect source batches let source_batches = collect(source, state.task_ctx()).await?; let total_source_rows: usize = source_batches.iter().map(|b| b.num_rows()).sum(); @@ -628,14 +636,6 @@ impl MemTable { *self.sort_order.lock() = vec![]; - let target_num_cols = self.schema.fields().len(); - let target_schema_ref = Arc::clone(&self.schema); - - // Helper to extract equi-join column indices from `on` - let equi_keys = extract_equi_join_keys(&on, &merge_schema, target_num_cols)?; - let target_key_indices: Vec = equi_keys.iter().map(|(t, _)| *t).collect(); - let source_key_indices: Vec = equi_keys.iter().map(|(_, s)| *s).collect(); - // Build target hash index across all partitions: // Key -> (partition_idx, batch_idx, row_idx) let mut target_key_map: HashMap, (usize, usize, usize)> = @@ -1106,8 +1106,8 @@ fn extract_equi_join_keys( let mut pairs = Vec::new(); collect_equi_keys(on, merge_schema, target_num_cols, &mut pairs)?; if pairs.is_empty() { - return plan_err!( - "MemTable MERGE INTO requires at least one equi-join condition in ON clause" + return not_impl_err!( + "MERGE INTO not supported for Base table: requires at least one equi-join condition in ON clause" ); } Ok(pairs) @@ -1160,22 +1160,22 @@ fn collect_equi_keys( pairs.push((idx2, idx1 - target_num_cols)); Ok(()) } else { - plan_err!( - "ON equality condition must compare a target column with a source column: {expr}" + not_impl_err!( + "MERGE INTO not supported for Base table: ON equality condition must compare a target column with a source column: {expr}" ) } } else { - plan_err!( - "MemTable MERGE INTO requires column equality conditions in ON clause, found: {expr}" + not_impl_err!( + "MERGE INTO not supported for Base table: requires column equality conditions in ON clause, found: {expr}" ) } } - _ => plan_err!( - "MemTable MERGE INTO only supports AND and EQ in ON condition, found: {expr}" + _ => not_impl_err!( + "MERGE INTO not supported for Base table: only supports AND and EQ in ON condition, found: {expr}" ), }, - _ => plan_err!( - "MemTable MERGE INTO requires binary equality expressions in ON condition, found: {expr}" + _ => not_impl_err!( + "MERGE INTO not supported for Base table: requires binary equality expressions in ON condition, found: {expr}" ), } }