diff --git a/bindings/python/README.md b/bindings/python/README.md index a00c6c534..515fb4cca 100644 --- a/bindings/python/README.md +++ b/bindings/python/README.md @@ -79,8 +79,8 @@ Operators such as sorting and aggregation may still buffer data during execution `sql_stream_with_options(sql, options)` applies `SQLOptions` as `sql_with_options()` does: they restrict statements delegated to DataFusion, but not statements Paimon -handles itself, such as `DELETE`, `UPDATE`, `MERGE INTO` or `DROP TABLE` on a -Paimon table. +handles itself, such as `DELETE`, `UPDATE` or `DROP TABLE` on a Paimon table. +`MERGE INTO` uses DataFusion's native planner and honors the DML restriction. ### Native Vector Search diff --git a/crates/integrations/datafusion/src/merge_into.rs b/crates/integrations/datafusion/src/merge_into.rs index 348dc40cd..6dc79a301 100644 --- a/crates/integrations/datafusion/src/merge_into.rs +++ b/crates/integrations/datafusion/src/merge_into.rs @@ -22,29 +22,797 @@ //! - **Append-only tables** (no PK, no deletion vectors): copy-on-write file rewriting //! via [`paimon::table::CopyOnWriteMergeWriter`]. -use std::collections::{HashMap, HashSet}; -use std::ops::ControlFlow; +use std::collections::HashSet; +use std::fmt; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; +use async_trait::async_trait; use datafusion::arrow::array::{ - Array, ArrayRef, Int32Array, RecordBatch, UInt32Array, UInt64Array, + Array, ArrayRef, BooleanArray, Int32Array, RecordBatch, UInt32Array, UInt64Array, }; -use datafusion::arrow::compute; -use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; +use datafusion::arrow::compute::{cast, filter_record_batch}; +use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema, SchemaRef}; +use datafusion::arrow::record_batch::RecordBatchOptions; +use datafusion::catalog::Session; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{Column, DFSchema, DFSchemaRef}; +use datafusion::datasource::sink::{DataSink, DataSinkExec}; +use datafusion::datasource::MemTable; use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::execution::SessionStateBuilder; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::dml::{MergeIntoAction, MergeIntoClause, MergeIntoClauseKind}; +use datafusion::logical_expr::execution_props::ExecutionProps; +use datafusion::logical_expr::{col, expr::Case, lit, Expr, JoinType, Operator}; +use datafusion::physical_expr::{expressions::Column as PhysicalColumn, PhysicalExpr}; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::{DisplayAs, ExecutionPlan}; use datafusion::prelude::{DataFrame, SessionContext}; -use datafusion::sql::sqlparser::ast::{ - visit_expressions_mut, AssignmentTarget, BinaryOperator, Expr as SqlExpr, Ident, Merge, - MergeAction, MergeClauseKind, MergeInsertExpr, MergeInsertKind, TableFactor, VisitMut, -}; use futures::TryStreamExt; use paimon::spec::{datums_to_binary_row, extract_datum_from_arrow, CoreOptions, DataField}; -use paimon::table::{CopyOnWriteMergeWriter, DataSplitBuilder, Table, WriteBuilder}; +use paimon::table::{CopyOnWriteMergeWriter, DataEvolutionDeleteWriter, DataSplitBuilder, Table}; use crate::error::to_datafusion_error; -use crate::sql_context::{normalize_schema_identifier, SQLContext}; +use crate::sql_context::SQLContext; +use crate::table::PaimonTableProvider; + +mod normalize; +pub(crate) use normalize::qualify_not_matched_exprs; + +const PRESENT: &str = "__paimon_merge_present"; +const FILE: &str = "__paimon_file_idx"; +const OFFSET: &str = "__paimon_row_offset"; + +fn column_name(index: usize, target_len: usize) -> String { + if index < target_len { + format!("__paimon_merge_t_{index}") + } else { + format!("__paimon_merge_s_{}", index - target_len) + } +} + +/// Bind columns by their position in the native target+source schema. SQL aliases +/// and quoted names have already been resolved by DataFusion's SQL planner. +fn bind_expr( + expr: Expr, + schema: &DFSchema, + target_len: usize, + source_only: bool, +) -> DFResult { + expr.transform(|expr| match expr { + Expr::Column(column) => { + let index = schema.index_of_column(&column)?; + if source_only && index < target_len { + return Err(DataFusionError::Plan( + "MERGE NOT MATCHED expressions may only reference source columns".into(), + )); + } + Ok(Transformed::yes(col(column_name(index, target_len)))) + } + Expr::ScalarSubquery(_) + | Expr::Exists(_) + | Expr::InSubquery(_) + | Expr::OuterReferenceColumn(_, _) => Err(DataFusionError::NotImplemented( + "Paimon MERGE expressions containing subqueries are not supported".into(), + )), + other => Ok(Transformed::no(other)), + }) + .map(|transformed| transformed.data) +} + +pub(crate) fn plan_merge( + state: &dyn Session, + table: Table, + mut source: Arc, + mut merge_schema: DFSchemaRef, + on: Expr, + mut clauses: Vec, +) -> DFResult> { + if table.is_branch_reference() { + return Err(DataFusionError::Plan(format!( + "MERGE INTO on Paimon branch '{}' is not supported", + table.branch() + ))); + } + if table.is_time_traveled() { + return Err(DataFusionError::Plan( + "Cannot execute MERGE INTO on a time-travel table".into(), + )); + } + let de = CoreOptions::new(table.schema().options()).data_evolution_enabled(); + if !table.schema().trimmed_primary_keys().is_empty() { + return Err(DataFusionError::Plan( + "MERGE INTO on primary-key tables without data-evolution is not supported".into(), + )); + } + let target_len = merge_schema.fields().len() - source.schema().fields().len(); + // DataFusion 55 retains every DML input field. Add a physical projection + // before executing the source so scan pushdown can avoid unused payloads. + let mut required = on.column_refs(); + for clause in &clauses { + if let Some(predicate) = &clause.predicate { + required.extend(predicate.column_refs()); + } + match &clause.action { + MergeIntoAction::Update(assignments) => { + for (_, value) in assignments { + required.extend(value.column_refs()); + } + } + MergeIntoAction::Insert { values, .. } => { + for value in values { + required.extend(value.column_refs()); + } + } + MergeIntoAction::Delete => {} + } + } + let mut source_indices = required + .into_iter() + .map(|column| merge_schema.index_of_column(column)) + .collect::>>()?; + source_indices.retain(|index| *index >= target_len); + source_indices.sort_unstable(); + source_indices.dedup(); + if source_indices.len() < source.schema().fields().len() { + let physical_schema = source.schema(); + let expressions = source_indices + .iter() + .map(|index| { + let source_index = index - target_len; + let name = physical_schema.field(source_index).name().clone(); + ( + Arc::new(PhysicalColumn::new(&name, source_index)) as Arc, + name, + ) + }) + .collect::>(); + source = Arc::new(ProjectionExec::try_new(expressions, source)?); + let fields = (0..target_len) + .chain(source_indices) + .map(|index| { + let (qualifier, field) = merge_schema.qualified_field(index); + (qualifier.cloned(), field.clone()) + }) + .collect(); + merge_schema = Arc::new(DFSchema::new_with_metadata( + fields, + merge_schema.metadata().clone(), + )?); + } + let mut update_columns = Vec::new(); + let mut matched_count = 0; + for clause in &mut clauses { + match clause.kind { + MergeIntoClauseKind::NotMatchedBySource => { + return Err(DataFusionError::NotImplemented(format!( + "WHEN NOT MATCHED BY SOURCE is not yet supported for {} MERGE INTO", + if de { "data evolution" } else { "CoW" } + ))); + } + MergeIntoClauseKind::Matched => { + if de && matched_count > 0 { + return Err(DataFusionError::NotImplemented( + "Multiple WHEN MATCHED clauses are not yet supported".into(), + )); + } + if de && clause.predicate.is_some() { + return Err(DataFusionError::NotImplemented( + "WHEN MATCHED AND is not yet supported".into(), + )); + } + matched_count += 1; + } + _ => {} + } + if let Some(predicate) = clause.predicate.take() { + clause.predicate = Some(bind_expr( + predicate, + &merge_schema, + target_len, + clause.kind != MergeIntoClauseKind::Matched, + )?); + } + match &mut clause.action { + MergeIntoAction::Update(assignments) => { + let columns = assignments + .iter() + .map(|(c, _)| c.clone()) + .collect::>(); + validate_update_columns(&columns, table.schema().fields())?; + if de { + // Validate storage constraints without scanning or writing. + let _ = table + .new_write_builder() + .new_data_evolution_writer(columns.clone()) + .map_err(to_datafusion_error)?; + } + for (name, value) in assignments { + if !update_columns.contains(name) { + update_columns.push(name.clone()); + } + *value = bind_expr(value.clone(), &merge_schema, target_len, false)?; + } + } + MergeIntoAction::Insert { columns, values } => { + if columns.is_empty() { + *columns = table + .schema() + .fields() + .iter() + .map(|f| f.name().to_string()) + .collect(); + } + validate_update_columns(columns, table.schema().fields())?; + for value in values { + *value = bind_expr(value.clone(), &merge_schema, target_len, true)?; + } + } + MergeIntoAction::Delete if de => { + let _ = table + .new_write_builder() + .new_delete() + .map_err(to_datafusion_error)?; + } + MergeIntoAction::Delete => {} + } + } + let on = bind_expr(on, &merge_schema, target_len, false)?; + let sink = PaimonMergeSink { + table, + schema: source.schema(), + merge_schema, + target_len, + on, + clauses, + update_columns, + de, + execution_props: state.execution_props().clone(), + }; + Ok(Arc::new(DataSinkExec::new(source, Arc::new(sink), None))) +} + +#[derive(Debug)] +struct PaimonMergeSink { + table: Table, + schema: SchemaRef, + merge_schema: DFSchemaRef, + target_len: usize, + on: Expr, + clauses: Vec, + update_columns: Vec, + de: bool, + execution_props: ExecutionProps, +} + +impl DisplayAs for PaimonMergeSink { + fn fmt_as( + &self, + _: datafusion::physical_plan::DisplayFormatType, + f: &mut fmt::Formatter, + ) -> fmt::Result { + write!(f, "PaimonMerge: table={}", self.table.identifier()) + } +} + +#[async_trait] +impl DataSink for PaimonMergeSink { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + async fn write_all( + &self, + data: SendableRecordBatchStream, + context: &Arc, + ) -> DFResult { + // Capture the source once, including self-merges and volatile expressions. + // Each conflict retry rebuilds the target scan and writers from fresh metadata. + let source: Vec = data.try_collect().await?; + retry_on_conflict( + if self.de { + "MERGE INTO" + } else { + "CoW MERGE INTO" + }, + if self.de { + is_row_id_conflict + } else { + is_delete_conflict + }, + || self.execute_once(&source, context), + ) + .await + } +} + +impl PaimonMergeSink { + async fn execute_once( + &self, + source: &[RecordBatch], + context: &Arc, + ) -> DFResult { + let ctx = SessionContext::new_with_state( + SessionStateBuilder::new() + .with_default_features() + .with_config(context.session_config().clone()) + .with_runtime_env(context.runtime_env()) + .with_execution_props(self.execution_props.clone()) + .build(), + ); + let target_schema = Arc::new(Schema::new( + self.merge_schema.fields()[..self.target_len].to_vec(), + )); + let source_schema = Arc::new(Schema::new( + self.merge_schema.fields()[self.target_len..].to_vec(), + )); + let source_df = memory_relation( + &ctx, + source, + source_schema, + self.target_len, + self.target_len, + )?; + let mut cow_writer = None; + let mut read_snapshot = None; + let (target_df, read_table) = if self.de { + let snapshot = self + .table + .resolve_read_snapshot() + .await + .map_err(to_datafusion_error)?; + read_snapshot = snapshot.as_ref().map(|s| s.id()); + let read_table = self.table.copy_with_pinned_snapshot(snapshot.as_ref()); + // Keep the target as a scan with a projection, so updating a small + // column does not read unrelated BLOB payloads or other wide fields. + let target = ctx + .read_table(Arc::new(PaimonTableProvider::try_new(read_table.clone())?))? + .select(self.target_projection())?; + (target, read_table) + } else { + let partitions = self.source_partitions(source)?; + let writer = + CopyOnWriteMergeWriter::new(&self.table, self.update_columns.clone(), partitions) + .await + .map_err(to_datafusion_error)?; + let target = read_cow_target_batches(&self.table, &writer).await?; + cow_writer = Some(writer); + let mut fields = target_schema.fields().to_vec(); + fields.push(Arc::new(Field::new(FILE, ArrowDataType::Int32, false))); + fields.push(Arc::new(Field::new(OFFSET, ArrowDataType::UInt32, false))); + let target = memory_relation( + &ctx, + &target, + Arc::new(Schema::new(fields)), + 0, + self.target_len, + )?; + (target, self.table.clone()) + }; + let mut target_projection = target_df + .schema() + .fields() + .iter() + .map(|f| col(f.name())) + .collect::>(); + target_projection.push(lit(true).alias(PRESENT)); + let joined = source_df.join_on( + target_df.select(target_projection)?, + JoinType::Left, + [self.on.clone()], + )?; + let schema = joined.schema().clone(); + let state = ctx.state(); + let action = Expr::Case(Case::new( + None, + self.clauses + .iter() + .enumerate() + .map(|(index, clause)| { + let applicable = if clause.kind == MergeIntoClauseKind::Matched { + col(PRESENT).is_not_null() + } else { + col(PRESENT).is_null() + }; + // AND does not guarantee short-circuit evaluation for mixed + // Arrow batches. CASE keeps predicates in their row scope. + let condition = + clause + .predicate + .clone() + .map_or(applicable.clone(), |predicate| { + Expr::Case(Case::new( + None, + vec![(Box::new(applicable), Box::new(predicate))], + Some(Box::new(lit(false))), + )) + }); + (Box::new(condition), Box::new(lit(index as i32))) + }) + .collect(), + Some(Box::new(lit(-1i32))), + )); + let action = state.create_physical_expr(action, &schema)?; + let mut selected = vec![Vec::new(); self.clauses.len()]; + let mut seen_targets = HashSet::new(); + let matched = self + .clauses + .iter() + .filter(|c| c.kind == MergeIntoClauseKind::Matched) + .collect::>(); + let deduplicate_delete = matched.len() == 1 + && matched[0].predicate.is_none() + && matches!(matched[0].action, MergeIntoAction::Delete); + for batch in joined.collect().await? { + let actions = action.evaluate(&batch)?.into_array(batch.num_rows())?; + let actions = actions + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Internal("Invalid MERGE action type".into()))?; + for (index, clause) in self.clauses.iter().enumerate() { + let mask = BooleanArray::from_iter( + (0..batch.num_rows()).map(|row| Some(actions.value(row) == index as i32)), + ); + let mut batch = filter_record_batch(&batch, &mask)?; + if batch.num_rows() == 0 { + continue; + } + if clause.kind == MergeIntoClauseKind::Matched { + let mut keep = Vec::with_capacity(batch.num_rows()); + for row in 0..batch.num_rows() { + let key = if self.de { + let row_id = batch + .column_by_name(&column_name(self.target_len - 1, self.target_len)) + .unwrap(); + ( + 0, + row_id + .as_any() + .downcast_ref::() + .unwrap() + .value(row) as u64, + ) + } else { + let (file, offset) = extract_tracking_columns(&batch)?; + (file.value(row), offset.value(row) as u64) + }; + if !seen_targets.insert(key) { + if deduplicate_delete { + keep.push(false); + continue; + } + return Err(DataFusionError::Execution( + "MERGE INTO found duplicate UPDATE/DELETE matches for a target row" + .into(), + )); + } + keep.push(true); + } + batch = filter_record_batch(&batch, &BooleanArray::from(keep))?; + } + selected[index].push(batch); + } + } + + // Evaluate and validate every action before staging files or publishing. + let table_schema = paimon::arrow::build_target_arrow_schema(self.table.schema().fields()) + .map_err(to_datafusion_error)?; + let mut updates = Vec::new(); + let mut deletes = Vec::new(); + let mut inserts = Vec::new(); + let mut count = 0; + for (clause, batches) in self.clauses.iter().zip(selected) { + for batch in batches { + count += batch.num_rows() as u64; + match &clause.action { + MergeIntoAction::Update(assignments) => { + let mut fields = Vec::new(); + let mut columns = Vec::new(); + if self.de { + fields.push(Field::new("_ROW_ID", ArrowDataType::Int64, false)); + columns.push( + batch + .column_by_name(&column_name( + self.target_len - 1, + self.target_len, + )) + .unwrap() + .clone(), + ); + } + for name in &self.update_columns { + let target_index = table_schema.index_of(name)?; + let value = assignments + .iter() + .find(|(c, _)| c == name) + .map(|(_, e)| e.clone()) + .unwrap_or_else(|| col(column_name(target_index, self.target_len))); + fields.push(table_schema.field(target_index).clone()); + columns.push(evaluate( + &state, + &schema, + &batch, + value, + table_schema.field(target_index).data_type(), + )?); + } + updates.push(( + batch, + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?, + )); + } + MergeIntoAction::Delete => deletes.push(batch), + MergeIntoAction::Insert { columns, values } => { + let mut arrays = Vec::new(); + for field in table_schema.fields() { + let expr = columns + .iter() + .position(|c| c == field.name()) + .map(|i| values[i].clone()) + .unwrap_or_else(|| lit(datafusion::common::ScalarValue::Null)); + arrays.push(evaluate( + &state, + &schema, + &batch, + expr, + field.data_type(), + )?); + } + inserts.push(RecordBatch::try_new(table_schema.clone(), arrays)?); + } + } + } + } + let wb = self.table.new_write_builder(); + let mut messages = Vec::new(); + let commit = wb.try_new_commit().map_err(to_datafusion_error)?; + let staged: DFResult<()> = async { + if let Some(mut writer) = cow_writer { + for (batch_index, (tracking, _)) in updates.iter().enumerate() { + let (files, offsets) = extract_tracking_columns(tracking)?; + for row in 0..tracking.num_rows() { + writer.add_matched_update( + files.value(row) as usize, + offsets.value(row) as usize, + batch_index, + row, + ); + } + } + for batch in deletes { + let (files, offsets) = extract_tracking_columns(&batch)?; + for row in 0..batch.num_rows() { + writer.add_matched_delete( + files.value(row) as usize, + offsets.value(row) as usize, + ); + } + } + writer.set_update_batches(updates.into_iter().map(|(_, values)| values).collect()); + messages.extend(writer.prepare_commit().await.map_err(to_datafusion_error)?); + } else { + if !updates.is_empty() { + let mut writer = wb + .new_data_evolution_writer(self.update_columns.clone()) + .map_err(to_datafusion_error)?; + if let Some(snapshot) = read_snapshot { + writer.pin_read_snapshot(snapshot); + } + for (_, values) in updates { + writer + .add_matched_batch(values) + .map_err(to_datafusion_error)?; + } + messages.extend(writer.prepare_commit().await.map_err(to_datafusion_error)?); + } + if !deletes.is_empty() { + let mut writer = + DataEvolutionDeleteWriter::new(&read_table).map_err(to_datafusion_error)?; + for batch in deletes { + let row_id = batch + .column_by_name(&column_name(self.target_len - 1, self.target_len)) + .unwrap(); + let ids = row_id + .as_any() + .downcast_ref::() + .unwrap(); + writer + .add_row_ids(ids.values().iter().copied()) + .map_err(to_datafusion_error)?; + } + messages.extend(writer.prepare_commit().await.map_err(to_datafusion_error)?); + } + } + if !inserts.is_empty() { + let mut writer = wb.new_write().map_err(to_datafusion_error)?; + let result = async { + for batch in inserts { + writer + .write_arrow_batch(&batch) + .await + .map_err(to_datafusion_error)?; + } + writer.prepare_commit().await.map_err(to_datafusion_error) + } + .await; + writer.close().await; + messages.extend(result?); + } + Ok(()) + } + .await; + if let Err(error) = staged { + // These files have never been submitted. Once commit starts, its + // outcome may be unknown and deleting output would be unsafe. + if let Err(cleanup) = commit.abort(&messages).await { + log::warn!("Failed to abort unsubmitted MERGE files: {cleanup}"); + } + return Err(error); + } + if !messages.is_empty() { + commit.commit(messages).await.map_err(to_datafusion_error)?; + } + Ok(count) + } + + fn target_projection(&self) -> Vec { + let mut used = self.on.column_refs(); + for clause in &self.clauses { + if let Some(predicate) = &clause.predicate { + used.extend(predicate.column_refs()); + } + match &clause.action { + MergeIntoAction::Update(assignments) => { + for (_, expr) in assignments { + used.extend(expr.column_refs()); + } + } + MergeIntoAction::Insert { values, .. } => { + for expr in values { + used.extend(expr.column_refs()); + } + } + MergeIntoAction::Delete => {} + } + } + (0..self.target_len) + .filter(|index| { + *index == self.target_len - 1 + || used + .iter() + .any(|c| c.name == column_name(*index, self.target_len)) + }) + .map(|index| { + let field = self.merge_schema.field(index); + datafusion::logical_expr::expr_fn::cast( + Expr::Column(Column::new_unqualified(field.name())), + field.data_type().clone(), + ) + .alias(column_name(index, self.target_len)) + }) + .collect() + } + + fn source_partitions(&self, source: &[RecordBatch]) -> DFResult>>> { + let keys = self.table.schema().partition_keys(); + if keys.is_empty() { + return Ok(None); + } + let mut indices = Vec::new(); + for key in keys { + let target_index = self + .table + .schema() + .fields() + .iter() + .position(|f| f.name() == key) + .unwrap(); + let target = column_name(target_index, self.target_len); + let Some(index) = partition_source_index( + &self.on, + &target, + self.target_len, + self.merge_schema.fields().len(), + ) else { + return Ok(None); + }; + indices.push(index); + } + let partitions = source + .iter() + .map(|b| b.project(&indices).map_err(Into::into)) + .collect::>>()?; + build_partition_set_from_batches(&self.table, &partitions) + } +} + +fn partition_source_index( + expr: &Expr, + target: &str, + target_len: usize, + combined_len: usize, +) -> Option { + if let Expr::BinaryExpr(binary) = expr { + if binary.op == Operator::And { + return partition_source_index(&binary.left, target, target_len, combined_len).or_else( + || partition_source_index(&binary.right, target, target_len, combined_len), + ); + } + if binary.op == Operator::Eq { + for (left, right) in [(&binary.left, &binary.right), (&binary.right, &binary.left)] { + if matches!(left.as_ref(), Expr::Column(c) if c.name == target) { + if let Expr::Column(column) = right.as_ref() { + return (target_len..combined_len) + .find(|i| column.name == column_name(*i, target_len)) + .map(|i| i - target_len); + } + } + } + } + } + None +} + +fn evaluate( + state: &impl Session, + schema: &DFSchema, + batch: &RecordBatch, + expr: Expr, + target_type: &ArrowDataType, +) -> DFResult { + let value = state + .create_physical_expr(expr, schema)? + .evaluate(batch)? + .into_array(batch.num_rows())?; + cast(value.as_ref(), target_type).map_err(Into::into) +} + +fn memory_relation( + ctx: &SessionContext, + batches: &[RecordBatch], + schema: SchemaRef, + start: usize, + target_len: usize, +) -> DFResult { + let renamed = Arc::new(Schema::new( + schema + .fields() + .iter() + .enumerate() + .map(|(i, f)| { + let name = if start == 0 && i >= target_len { + f.name().clone() + } else { + column_name(start + i, target_len) + }; + f.as_ref().clone().with_name(name) + }) + .collect::>(), + )); + let batches = batches + .iter() + .map(|batch| { + let columns = schema + .fields() + .iter() + .enumerate() + .map(|(i, field)| { + // File tracking columns are appended after user fields. + // Position preserves their identity even when a user column + // has the same name as one of those internal fields. + let array = batch.column(i); + cast(array.as_ref(), field.data_type()).map_err(Into::into) + }) + .collect::>>()?; + let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); + RecordBatch::try_new_with_options(renamed.clone(), columns, &options) + .map_err(Into::into) + }) + .collect::>>()?; + ctx.read_table(Arc::new(MemTable::try_new(renamed, vec![batches])?)) +} /// Maximum number of retries when DML conflicts with concurrent compaction. const DML_MAX_RETRIES: u32 = 5; @@ -92,14 +860,14 @@ impl Drop for TempTableTracker<'_> { } /// Retry a DML operation on conflict, using `is_retryable` to detect retryable errors. -pub(crate) async fn retry_on_conflict( +pub(crate) async fn retry_on_conflict( op_name: &str, is_retryable: fn(&DataFusionError) -> bool, mut action: F, -) -> DFResult +) -> DFResult where F: FnMut() -> Fut, - Fut: std::future::Future>, + Fut: std::future::Future>, { let mut last_err = None; for _ in 0..DML_MAX_RETRIES { @@ -120,35 +888,15 @@ where )))) } -/// Execute a MERGE INTO statement on a Paimon table. -/// -/// Dispatches to the appropriate execution path based on table type: -/// - Data evolution tables → partial-column writes via `DataEvolutionWriter` -/// - Append-only tables (no PK) → copy-on-write file rewriting via `CopyOnWriteMergeWriter` -pub(crate) async fn execute_merge_into( - ctx: &SQLContext, - merge: &Merge, - table: Table, - enable_ident_normalization: bool, -) -> DFResult { - let schema = table.schema(); - let core_options = CoreOptions::new(schema.options()); - - if core_options.data_evolution_enabled() { - execute_data_evolution_merge(ctx, merge, table, enable_ident_normalization).await - } else if schema.trimmed_primary_keys().is_empty() { - execute_cow_merge(ctx, merge, table, enable_ident_normalization).await - } else { - Err(DataFusionError::Plan( - "MERGE INTO on primary-key tables without data-evolution is not supported".to_string(), - )) - } -} - /// Check if a DataFusion error is caused by a row ID conflict during commit. pub(crate) fn is_row_id_conflict(err: &DataFusionError) -> bool { match err { - DataFusionError::External(e) => e.to_string().contains("Row ID conflict"), + DataFusionError::External(e) => { + matches!( + e.downcast_ref::(), + Some(paimon::Error::RowIdExistenceConflict { .. }) + ) || e.to_string().contains("Row ID conflict") + } _ => false, } } @@ -156,1517 +904,292 @@ pub(crate) fn is_row_id_conflict(err: &DataFusionError) -> bool { /// Check if a DataFusion error is caused by a delete conflict during commit. pub(crate) fn is_delete_conflict(err: &DataFusionError) -> bool { match err { - DataFusionError::External(e) => e.to_string().contains("Delete conflict"), + DataFusionError::External(e) => { + matches!( + e.downcast_ref::(), + Some(paimon::Error::RowIdExistenceConflict { .. }) + ) || e.to_string().contains("Delete conflict") + } _ => false, } } -// --------------------------------------------------------------------------- -// Data evolution path (existing) -// --------------------------------------------------------------------------- +/// Extract _ROW_ID + __upd_X columns from matched batches, renaming __upd_X → X. +pub(crate) fn project_update_columns( + batches: &[RecordBatch], + update_columns: &[String], +) -> DFResult> { + let mut result = Vec::new(); + for batch in batches { + let row_id_idx = batch + .schema() + .index_of("_ROW_ID") + .map_err(|e| DataFusionError::Internal(format!("_ROW_ID not found: {e}")))?; -/// Execute MERGE INTO on a data evolution table with retry on row ID conflict. -async fn execute_data_evolution_merge( - ctx: &SQLContext, - merge: &Merge, - table: Table, - enable_ident_normalization: bool, -) -> DFResult { - retry_on_conflict("MERGE INTO", is_row_id_conflict, || { - execute_merge_into_once(ctx, merge, &table, enable_ident_normalization) - }) - .await -} + let mut columns = vec![batch.column(row_id_idx).clone()]; + let mut fields = vec![batch.schema().field(row_id_idx).clone()]; -// --------------------------------------------------------------------------- -// Copy-on-Write path (append-only tables, no PK) -// --------------------------------------------------------------------------- + for col in update_columns { + let prefixed = format!("__upd_{col}"); + let idx = batch.schema().index_of(&prefixed).map_err(|e| { + DataFusionError::Internal(format!("Column {prefixed} not found: {e}")) + })?; + columns.push(batch.column(idx).clone()); + fields.push(Field::new( + col, + batch.schema().field(idx).data_type().clone(), + true, + )); + } -/// Parsed CoW merge clauses — supports UPDATE, DELETE, and INSERT. -struct CowMergeClauses { - /// Ordered list of WHEN MATCHED clauses (preserves SQL ordering for correct semantics). - matched: Vec, - inserts: Vec, + let schema = Arc::new(Schema::new(fields)); + let projected = RecordBatch::try_new(schema, columns)?; + result.push(projected); + } + Ok(result) } -/// A single WHEN MATCHED clause with optional predicate. -struct CowMatchedClause { - action: CowMatchedAction, - predicate: Option, +pub(crate) fn validate_update_columns( + columns: &[String], + table_fields: &[DataField], +) -> DFResult<()> { + validate_target_columns(columns, table_fields, "UPDATE") } -enum CowMatchedAction { - Update(MergeUpdateClause), - Delete, -} +fn validate_target_columns( + columns: &[String], + table_fields: &[DataField], + operation: &str, +) -> DFResult<()> { + let target_columns: HashSet<&str> = table_fields.iter().map(|field| field.name()).collect(); + let mut seen = HashSet::new(); -/// Parse MERGE clauses for the CoW path (supports DELETE unlike the data-evolution parser). -fn extract_cow_merge_clauses( - merge: &Merge, - enable_ident_normalization: bool, -) -> DFResult { - let mut matched: Vec = Vec::new(); - let mut inserts: Vec = Vec::new(); - - for clause in &merge.clauses { - match clause.clause_kind { - MergeClauseKind::Matched => { - let predicate = clause.predicate.as_ref().map(|p| p.to_string()); - match &clause.action { - MergeAction::Update(update_expr) => { - let mut columns = Vec::new(); - let mut exprs = Vec::new(); - for assignment in &update_expr.assignments { - let col_name = match &assignment.target { - AssignmentTarget::ColumnName(name) => name - .0 - .last() - .and_then(|p| p.as_ident()) - .map(|ident| { - normalize_schema_identifier( - ident, - enable_ident_normalization, - ) - }) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Invalid column name in SET: {name}" - )) - })?, - AssignmentTarget::Tuple(_) => { - return Err(DataFusionError::Plan( - "Tuple assignment in MERGE INTO SET is not supported" - .to_string(), - )); - } - }; - columns.push(col_name); - exprs.push(assignment.value.to_string()); - } - matched.push(CowMatchedClause { - action: CowMatchedAction::Update(MergeUpdateClause { columns, exprs }), - predicate, - }); - } - MergeAction::Delete { .. } => { - matched.push(CowMatchedClause { - action: CowMatchedAction::Delete, - predicate, - }); - } - MergeAction::Insert(_) => { - return Err(DataFusionError::Plan( - "WHEN MATCHED THEN INSERT is not valid SQL".to_string(), - )); - } - } - } - MergeClauseKind::NotMatched | MergeClauseKind::NotMatchedByTarget => { - match &clause.action { - MergeAction::Insert(insert_expr) => { - inserts.push(extract_merge_insert_clause( - insert_expr, - clause.predicate.as_ref(), - enable_ident_normalization, - )?); - } - _ => { - return Err(DataFusionError::Plan( - "WHEN NOT MATCHED only supports INSERT".to_string(), - )); - } - } - } - MergeClauseKind::NotMatchedBySource => { - return Err(DataFusionError::Plan( - "WHEN NOT MATCHED BY SOURCE is not yet supported for CoW MERGE INTO" - .to_string(), - )); - } + for column in columns { + if !seen.insert(column.as_str()) { + return Err(DataFusionError::Plan(format!( + "Duplicate column '{column}' in {operation}" + ))); + } + if !target_columns.contains(column.as_str()) { + return Err(DataFusionError::Plan(format!( + "Unknown column '{column}' in {operation}" + ))); } } - if matched.is_empty() && inserts.is_empty() { - return Err(DataFusionError::Plan( - "MERGE INTO requires at least one WHEN MATCHED or WHEN NOT MATCHED clause".to_string(), - )); - } - - Ok(CowMergeClauses { matched, inserts }) + Ok(()) } -/// Execute MERGE INTO on an append-only table with retry on delete conflict. -async fn execute_cow_merge( - ctx: &SQLContext, - merge: &Merge, - table: Table, - enable_ident_normalization: bool, -) -> DFResult { - retry_on_conflict("CoW MERGE INTO", is_delete_conflict, || { - execute_cow_merge_once(ctx, merge, &table, enable_ident_normalization) - }) - .await +pub(crate) fn extract_tracking_columns( + batch: &RecordBatch, +) -> DFResult<(&Int32Array, &UInt32Array)> { + let file_idx_col = batch + .column_by_name("__paimon_file_idx") + .ok_or_else(|| DataFusionError::Internal("__paimon_file_idx not found".to_string()))? + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Internal("__paimon_file_idx is not Int32".to_string()))?; + + let row_offset_col = batch + .column_by_name("__paimon_row_offset") + .ok_or_else(|| DataFusionError::Internal("__paimon_row_offset not found".to_string()))? + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal("__paimon_row_offset is not UInt32".to_string()) + })?; + + Ok((file_idx_col, row_offset_col)) } -/// Execute a single attempt of CoW MERGE INTO. -async fn execute_cow_merge_once( +/// Read all files from a table via the CoW writer's file index, attach `__paimon_file_idx` +/// and `__paimon_row_offset` tracking columns, and register the result as a MemTable. +/// +/// Returns `(has_data, cow_table_name)`. The caller is responsible for deregistering +/// via `TempTableTracker`. +/// +/// Note: all matching partition files are loaded into memory at once. For partitions +/// with many large files this may cause significant memory pressure. A future +/// optimization could stream or batch-process files instead of materializing everything. +pub(crate) async fn register_cow_target_table( ctx: &SQLContext, - merge: &Merge, table: &Table, - enable_ident_normalization: bool, -) -> DFResult { - let (source_ref, source_alias) = extract_source_ref(&merge.source)?; - let (target_ref, target_alias) = extract_table_ref(&merge.table)?; - let s_alias = source_alias.as_deref().unwrap_or(&source_ref); - let t_alias = target_alias.as_deref().unwrap_or("__cow_t"); - - // Rewrite column qualifiers before rendering SQL, preserving literal contents. - let mut aliased_merge = merge.clone(); - rewrite_table_references( - &mut aliased_merge.on, - &target_ref, - t_alias, - &source_ref, - s_alias, - ); - rewrite_table_references( - &mut aliased_merge.clauses, - &target_ref, - t_alias, - &source_ref, - s_alias, - ); - let clauses = extract_cow_merge_clauses(&aliased_merge, enable_ident_normalization)?; - validate_merge_insert_columns(&clauses.inserts, table.schema().fields())?; - for matched in &clauses.matched { - if let CowMatchedAction::Update(update) = &matched.action { - validate_update_columns(&update.columns, table.schema().fields())?; - } - } - - // Collect the union of all update columns across matched clauses (preserving order) - let mut update_columns: Vec = Vec::new(); - for mc in &clauses.matched { - if let CowMatchedAction::Update(upd) = &mc.action { - for col in &upd.columns { - if !update_columns.contains(col) { - update_columns.push(col.clone()); - } - } - } + writer: &CopyOnWriteMergeWriter, + temp_tracker: &mut TempTableTracker<'_>, +) -> DFResult<(bool, String)> { + let all_batches = read_cow_target_batches(table, writer).await?; + let has_data = !all_batches.is_empty(); + let table_name = next_cow_table_name("__cow_target"); + if let Some(batch) = all_batches.first() { + let mem_table = + datafusion::datasource::MemTable::try_new(batch.schema(), vec![all_batches])?; + ctx.register_temp_table(&table_name, Arc::new(mem_table))?; + temp_tracker.register(&table_name); } + Ok((has_data, table_name)) +} - let on_expr = &merge.on; - let on_condition = aliased_merge.on.to_string(); - - // Build partition filter from source data to avoid scanning all partitions - let partition_set = - build_source_partition_set(ctx, table, &source_ref, s_alias, t_alias, on_expr).await?; - - let mut writer = CopyOnWriteMergeWriter::new(table, update_columns.clone(), partition_set) - .await - .map_err(to_datafusion_error)?; - - // Read each target file individually, attach __paimon_file_idx and __paimon_row_offset - let mut temp_tracker = TempTableTracker::new(ctx); - let (has_target_data, cow_table_name) = - register_cow_target_table(ctx, table, &writer, &mut temp_tracker).await?; - let wb = table.new_write_builder(); - - let merge_ctx = CowMergeContext { - source_ref: &source_ref, - s_alias, - t_alias, - on_condition: &on_condition, - has_target_data, - cow_table_name, - update_columns: &update_columns, - }; - - let result = execute_cow_merge_inner( - ctx, - &clauses, - &mut writer, - &wb, - table, - &merge_ctx, - &mut temp_tracker, - ) - .await; - - let (insert_messages, total_count) = result?; - - // CoW rewrite: prepare_commit consumes the writer - let cow_messages = writer.prepare_commit().await.map_err(to_datafusion_error)?; +/// Read a captured CoW file index, retaining physical file/row identity. +async fn read_cow_target_batches( + table: &Table, + writer: &CopyOnWriteMergeWriter, +) -> DFResult> { + let file_index = writer.file_index(); + // Read all files in parallel + let read_futures: Vec<_> = file_index + .iter() + .enumerate() + .map(|(file_idx, file_info)| async move { + let single_split = DataSplitBuilder::new() + .with_snapshot(file_info.snapshot_id) + .with_partition( + paimon::spec::BinaryRow::from_serialized_bytes(&file_info.partition) + .map_err(to_datafusion_error)?, + ) + .with_bucket(file_info.bucket) + .with_bucket_path(file_info.bucket_path.clone()) + .with_total_buckets(file_info.total_buckets) + .with_data_files(vec![file_info.file_meta.clone()]) + .build() + .map_err(to_datafusion_error)?; - let mut all_messages = cow_messages; - all_messages.extend(insert_messages); + let read = table + .new_read_builder() + .new_read() + .map_err(to_datafusion_error)?; + let batches: Vec = read + .to_arrow(&[single_split]) + .map_err(to_datafusion_error)? + .try_collect() + .await + .map_err(to_datafusion_error)?; - if !all_messages.is_empty() { - wb.try_new_commit() - .map_err(to_datafusion_error)? - .commit(all_messages) - .await - .map_err(to_datafusion_error)?; - } + Ok::<_, DataFusionError>((file_idx, batches)) + }) + .collect(); - ok_result(ctx.ctx(), total_count) -} + let file_results = futures::future::try_join_all(read_futures).await?; -/// Context for CoW merge inner execution — groups join-related parameters. -struct CowMergeContext<'a> { - source_ref: &'a str, - s_alias: &'a str, - t_alias: &'a str, - on_condition: &'a str, - has_target_data: bool, - cow_table_name: String, - update_columns: &'a [String], -} + let mut all_batches: Vec = Vec::new(); -/// Inner function that populates the CoW writer with matched operations and handles INSERT. -/// Returns (insert_commit_messages, total_affected_count). -async fn execute_cow_merge_inner( - ctx: &SQLContext, - clauses: &CowMergeClauses, - writer: &mut CopyOnWriteMergeWriter, - wb: &WriteBuilder<'_>, - table: &Table, - merge_ctx: &CowMergeContext<'_>, - temp_tracker: &mut TempTableTracker<'_>, -) -> DFResult<(Vec, u64)> { - let source_ref = merge_ctx.source_ref; - let s_alias = merge_ctx.s_alias; - let t_alias = merge_ctx.t_alias; - let on_condition = merge_ctx.on_condition; - let has_target_data = merge_ctx.has_target_data; - let cow_table_name = &merge_ctx.cow_table_name; - let cow_target_name = cow_table_name.clone(); - let update_columns = merge_ctx.update_columns; - let mut insert_messages = Vec::new(); - let mut total_count: u64 = 0; - - if has_target_data && !clauses.matched.is_empty() { - let mut update_value_batches: Vec = Vec::new(); - let mut update_batch_counter: usize = 0; - // Track consumed predicates for correct multi-clause ordering: - // each clause only applies to rows NOT matched by any previous clause. - let mut consumed_predicates: Vec = Vec::new(); - - for mc in &clauses.matched { - // Build WHERE clause: exclude rows consumed by previous clauses, then apply this predicate - let mut conditions: Vec = Vec::new(); - for prev in &consumed_predicates { - // Only TRUE consumes a row; FALSE and NULL fall through. - conditions.push(format!("({prev}) IS NOT TRUE")); - } - if let Some(ref pred) = mc.predicate { - // Preserve OR grouping when combined with previous-clause exclusions. - conditions.push(format!("({pred})")); - consumed_predicates.push(pred.clone()); - } else { - consumed_predicates.push("TRUE".to_string()); + for (file_idx, batches) in file_results { + let mut row_offset = 0u32; + for batch in batches { + let num_rows = batch.num_rows(); + if num_rows == 0 { + continue; } - let where_clause = if conditions.is_empty() { - String::new() - } else { - format!(" WHERE {}", conditions.join(" AND ")) - }; - - match &mc.action { - CowMatchedAction::Update(upd) => { - let mut select_parts = vec![ - format!("{t_alias}.\"__paimon_file_idx\""), - format!("{t_alias}.\"__paimon_row_offset\""), - ]; - let clause_col_map: HashMap<&str, &str> = upd - .columns - .iter() - .zip(upd.exprs.iter()) - .map(|(c, e)| (c.as_str(), e.as_str())) - .collect(); - for col in update_columns { - let quoted_alias = quote_identifier(&format!("__upd_{col}")); - if let Some(expr) = clause_col_map.get(&col.as_str()) { - select_parts.push(format!("{expr} AS {quoted_alias}")); - } else { - select_parts.push(format!( - "{t_alias}.{} AS {quoted_alias}", - quote_identifier(col) - )); - } - } - let select_clause = select_parts.join(", "); - let join_sql = format!( - "SELECT {select_clause} FROM {source_ref} AS {s_alias} \ - INNER JOIN {cow_target_name} AS {t_alias} ON {on_condition}{where_clause}" - ); - let join_result = ctx.ctx().sql(&join_sql).await?.collect().await?; - - for batch in &join_result { - if batch.num_rows() == 0 { - continue; - } + let file_idx_i32 = i32::try_from(file_idx).map_err(|_| { + DataFusionError::Internal(format!("file_idx {file_idx} exceeds i32 range")) + })?; + let num_rows_u32 = u32::try_from(num_rows).map_err(|_| { + DataFusionError::Internal(format!("batch num_rows {num_rows} exceeds u32 range")) + })?; + let file_idx_col = Arc::new(Int32Array::from(vec![file_idx_i32; num_rows])); + let end_offset = row_offset.checked_add(num_rows_u32).ok_or_else(|| { + DataFusionError::Internal(format!( + "row_offset overflow: {row_offset} + {num_rows_u32}" + )) + })?; + let row_offset_col = Arc::new(UInt32Array::from( + (row_offset..end_offset).collect::>(), + )); - let (file_idx_col, row_offset_col) = extract_tracking_columns(batch)?; + let mut fields: Vec = batch + .schema() + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + fields.push(Field::new("__paimon_file_idx", ArrowDataType::Int32, false)); + fields.push(Field::new( + "__paimon_row_offset", + ArrowDataType::UInt32, + false, + )); + let augmented_schema = Arc::new(Schema::new(fields)); - let mut upd_fields = Vec::new(); - let mut upd_columns: Vec> = Vec::new(); - for col in update_columns { - let prefixed = format!("__upd_{col}"); - let idx = batch.schema().index_of(&prefixed).map_err(|e| { - DataFusionError::Internal(format!( - "Column {prefixed} not found: {e}" - )) - })?; - upd_fields.push(Field::new( - col, - batch.schema().field(idx).data_type().clone(), - true, - )); - upd_columns.push(batch.column(idx).clone()); - } - let upd_schema = Arc::new(Schema::new(upd_fields)); - let upd_batch = RecordBatch::try_new(upd_schema, upd_columns)?; - - let current_batch_idx = update_batch_counter; - update_value_batches.push(upd_batch); - update_batch_counter += 1; - - for row in 0..batch.num_rows() { - let file_idx = file_idx_col.value(row) as usize; - let row_offset = row_offset_col.value(row) as usize; - writer.add_matched_update(file_idx, row_offset, current_batch_idx, row); - total_count += 1; - } - } - } - CowMatchedAction::Delete => { - let select_clause = format!( - "{t_alias}.\"__paimon_file_idx\", {t_alias}.\"__paimon_row_offset\"" - ); - let join_sql = format!( - "SELECT {select_clause} FROM {source_ref} AS {s_alias} \ - INNER JOIN {cow_target_name} AS {t_alias} ON {on_condition}{where_clause}" - ); - - let join_result = ctx.ctx().sql(&join_sql).await?.collect().await?; - - for batch in &join_result { - if batch.num_rows() == 0 { - continue; - } + let mut columns: Vec> = batch.columns().to_vec(); + columns.push(file_idx_col); + columns.push(row_offset_col); - let (file_idx_col, row_offset_col) = extract_tracking_columns(batch)?; + let augmented = RecordBatch::try_new(augmented_schema.clone(), columns) + .map_err(|e| DataFusionError::Internal(format!("Failed to augment batch: {e}")))?; - for row in 0..batch.num_rows() { - let file_idx = file_idx_col.value(row) as usize; - let row_offset = row_offset_col.value(row) as usize; - writer.add_matched_delete(file_idx, row_offset); - total_count += 1; - } - } - } - } + all_batches.push(augmented); + row_offset = end_offset; } + } - if !update_value_batches.is_empty() { - writer.set_update_batches(update_value_batches); - } + Ok(all_batches) +} + +/// Build a partition set from Arrow batches containing partition column values. +/// +/// Converts each row's partition columns into serialized `BinaryRow` bytes. +/// Returns `None` for non-partitioned tables. +pub(crate) fn build_partition_set_from_batches( + table: &Table, + batches: &[RecordBatch], +) -> DFResult>>> { + let partition_keys = table.schema().partition_keys(); + if partition_keys.is_empty() { + return Ok(None); } - // Handle NOT MATCHED → INSERT - if !clauses.inserts.is_empty() { - let insert_sql = if has_target_data { - format!( - "SELECT {s_alias}.* FROM {source_ref} AS {s_alias} \ - LEFT JOIN {cow_target_name} AS {t_alias} ON {on_condition} \ - WHERE {t_alias}.\"__paimon_file_idx\" IS NULL" - ) - } else { - format!("SELECT * FROM {source_ref} AS {s_alias}") - }; + let partition_fields = table.schema().partition_fields(); + let mut partition_set = HashSet::new(); - let not_matched_batches = ctx.ctx().sql(&insert_sql).await?.collect().await?; - - if !not_matched_batches.is_empty() { - let insert_batches = build_insert_batches( - ctx, - ¬_matched_batches, - &clauses.inserts, - s_alias, - &[], - table.schema().fields(), - temp_tracker, - ) - .await?; - - let insert_count: usize = insert_batches.iter().map(|b| b.num_rows()).sum(); - if insert_count > 0 { - let mut table_write = wb.new_write().map_err(to_datafusion_error)?; - for batch in &insert_batches { - table_write - .write_arrow_batch(batch) - .await - .map_err(to_datafusion_error)?; - } - let msgs = table_write - .prepare_commit() - .await - .map_err(to_datafusion_error)?; - insert_messages.extend(msgs); - total_count += insert_count as u64; - } + for batch in batches { + for row in 0..batch.num_rows() { + let datums: Vec<(Option, paimon::spec::DataType)> = + partition_fields + .iter() + .enumerate() + .map(|(col_idx, field)| { + let datum = + extract_datum_from_arrow(batch, row, col_idx, field.data_type()) + .map_err(to_datafusion_error)?; + Ok((datum, field.data_type().clone())) + }) + .collect::>()?; + let refs: Vec<(&Option, &paimon::spec::DataType)> = + datums.iter().map(|(d, t)| (d, t)).collect(); + partition_set.insert(datums_to_binary_row(&refs)); } } - Ok((insert_messages, total_count)) + Ok(Some(partition_set)) } -// --------------------------------------------------------------------------- -// Data evolution path helpers -// --------------------------------------------------------------------------- - -async fn execute_merge_into_once( +/// Query a table for distinct partition values matching an optional WHERE clause. +/// +/// Returns `None` for non-partitioned tables. +pub(crate) async fn build_partition_set_from_where( ctx: &SQLContext, - merge: &Merge, table: &Table, - enable_ident_normalization: bool, -) -> DFResult { - // 1. Parse all MERGE clauses - let parsed = extract_merge_clauses(merge, enable_ident_normalization)?; - validate_merge_insert_columns(&parsed.inserts, table.schema().fields())?; - if let Some(update) = &parsed.update { - validate_update_columns(&update.columns, table.schema().fields())?; + table_ref: &str, + where_clause: Option<&str>, +) -> DFResult>>> { + let partition_keys = table.schema().partition_keys(); + if partition_keys.is_empty() { + return Ok(None); } - // Validate preconditions early and create writer (before executing any SQL) - let wb = table.new_write_builder(); - let update_writer = if let Some(ref upd) = parsed.update { - Some( - wb.new_data_evolution_writer(upd.columns.clone()) - .map_err(to_datafusion_error)?, - ) - } else { - None - }; - let delete_writer = if parsed.delete { - Some(wb.new_delete().map_err(to_datafusion_error)?) - } else { - None + let cols = partition_keys + .iter() + .map(|k| quote_identifier(k)) + .collect::>() + .join(", "); + let where_part = match where_clause { + Some(w) => format!(" WHERE {w}"), + None => String::new(), }; + let sql = format!("SELECT DISTINCT {cols} FROM {table_ref}{where_part}"); + let batches = ctx.ctx().sql(&sql).await?.collect().await?; - let (target_ref, target_alias) = extract_table_ref(&merge.table)?; - let (source_ref, source_alias) = extract_source_ref(&merge.source)?; - let on_condition = merge.on.to_string(); - let t_alias = target_alias.as_deref().unwrap_or(&target_ref); - let s_alias = source_alias.as_deref().unwrap_or(&source_ref); - - // 2. Build a single LEFT JOIN: source LEFT JOIN target - // _ROW_ID IS NOT NULL → matched (UPDATE path) - // _ROW_ID IS NULL → not matched (INSERT path) - let mut select_parts = vec![format!("{t_alias}.\"_ROW_ID\"")]; - - // Add update expressions (prefixed to avoid collisions) - if let Some(ref upd) = parsed.update { - for (col, expr) in upd.columns.iter().zip(upd.exprs.iter()) { - select_parts.push(format!( - "{expr} AS {}", - quote_identifier(&format!("__upd_{col}")) - )); - } - } - - // Add source columns for INSERT path (all source columns via s.*) - // We also need insert expressions if they differ from source columns - if !parsed.inserts.is_empty() { - select_parts.push(format!("{s_alias}.*")); - } - - let select_clause = select_parts.join(", "); - // Safety: all interpolated values (select_clause, source_ref, s_alias, t_alias, on_condition) - // originate from sqlparser AST's Display impl, so they are well-formed SQL fragments. - let join_sql = format!( - "SELECT {select_clause} FROM {source_ref} AS {s_alias} \ - LEFT JOIN {target_ref} AS {t_alias} ON {on_condition}" - ); - - let join_result = ctx.ctx().sql(&join_sql).await?.collect().await?; - - // 3. Split by _ROW_ID null/not-null - let mut all_messages = Vec::new(); - let mut total_count: u64 = 0; - - // Separate matched and not-matched rows - let (matched_batches, not_matched_batches) = split_by_row_id(&join_result)?; - - // 4. Handle matched rows (UPDATE or DELETE) - if let Some(mut writer) = update_writer { - let upd = parsed.update.as_ref().unwrap(); - let matched_count: usize = matched_batches.iter().map(|b| b.num_rows()).sum(); - if matched_count > 0 { - // Extract _ROW_ID + update columns (rename __upd_X → X) - let update_batches = project_update_columns(&matched_batches, &upd.columns)?; - for batch in update_batches { - writer - .add_matched_batch(batch) - .map_err(to_datafusion_error)?; - } - let update_messages = writer.prepare_commit().await.map_err(to_datafusion_error)?; - all_messages.extend(update_messages); - total_count += matched_count as u64; - } - } - if let Some(mut writer) = delete_writer { - let matched_count: usize = matched_batches.iter().map(|b| b.num_rows()).sum(); - if matched_count > 0 { - for batch in &matched_batches { - writer - .add_matched_batch(batch.clone()) - .map_err(to_datafusion_error)?; - } - let delete_messages = writer.prepare_commit().await.map_err(to_datafusion_error)?; - all_messages.extend(delete_messages); - total_count += matched_count as u64; - } - } - - // 5. Handle not-matched rows (INSERT) - if !parsed.inserts.is_empty() { - // Collect the exact set of injected column names to strip from JOIN result - let mut injected_columns: Vec = vec!["_ROW_ID".to_string()]; - if let Some(ref upd) = parsed.update { - for col in &upd.columns { - injected_columns.push(format!("__upd_{col}")); - } - } - let mut temp_tracker = TempTableTracker::new(ctx); - let insert_batches = build_insert_batches( - ctx, - ¬_matched_batches, - &parsed.inserts, - s_alias, - &injected_columns, - table.schema().fields(), - &mut temp_tracker, - ) - .await?; - let insert_count: usize = insert_batches.iter().map(|b| b.num_rows()).sum(); - if insert_count > 0 { - let mut table_write = wb.new_write().map_err(to_datafusion_error)?; - for batch in &insert_batches { - table_write - .write_arrow_batch(batch) - .await - .map_err(to_datafusion_error)?; - } - let insert_messages = table_write - .prepare_commit() - .await - .map_err(to_datafusion_error)?; - all_messages.extend(insert_messages); - total_count += insert_count as u64; - } - } - - // 6. Commit all messages atomically - if !all_messages.is_empty() { - wb.try_new_commit() - .map_err(to_datafusion_error)? - .commit(all_messages) - .await - .map_err(to_datafusion_error)?; - } - - ok_result(ctx.ctx(), total_count) -} - -/// Split join result into matched (_ROW_ID not null) and not-matched (_ROW_ID null) batches. -fn split_by_row_id(batches: &[RecordBatch]) -> DFResult<(Vec, Vec)> { - let mut matched = Vec::new(); - let mut not_matched = Vec::new(); - - for batch in batches { - if batch.num_rows() == 0 { - continue; - } - let row_id_col = batch.column_by_name("_ROW_ID").ok_or_else(|| { - DataFusionError::Internal("_ROW_ID column not found in join result".to_string()) - })?; - - let is_not_null = compute::is_not_null(row_id_col)?; - let is_null = compute::is_null(row_id_col)?; - - let matched_batch = compute::filter_record_batch(batch, &is_not_null)?; - if matched_batch.num_rows() > 0 { - matched.push(matched_batch); - } - - let not_matched_batch = compute::filter_record_batch(batch, &is_null)?; - if not_matched_batch.num_rows() > 0 { - not_matched.push(not_matched_batch); - } - } - - Ok((matched, not_matched)) -} - -/// Extract _ROW_ID + __upd_X columns from matched batches, renaming __upd_X → X. -pub(crate) fn project_update_columns( - batches: &[RecordBatch], - update_columns: &[String], -) -> DFResult> { - let mut result = Vec::new(); - for batch in batches { - let row_id_idx = batch - .schema() - .index_of("_ROW_ID") - .map_err(|e| DataFusionError::Internal(format!("_ROW_ID not found: {e}")))?; - - let mut columns = vec![batch.column(row_id_idx).clone()]; - let mut fields = vec![batch.schema().field(row_id_idx).clone()]; - - for col in update_columns { - let prefixed = format!("__upd_{col}"); - let idx = batch.schema().index_of(&prefixed).map_err(|e| { - DataFusionError::Internal(format!("Column {prefixed} not found: {e}")) - })?; - columns.push(batch.column(idx).clone()); - fields.push(Field::new( - col, - batch.schema().field(idx).data_type().clone(), - true, - )); - } - - let schema = Arc::new(Schema::new(fields)); - let projected = RecordBatch::try_new(schema, columns)?; - result.push(projected); - } - Ok(result) -} - -/// Build insert batches from not-matched rows, applying INSERT clause projections and predicates. -async fn build_insert_batches( - ctx: &SQLContext, - not_matched_batches: &[RecordBatch], - inserts: &[MergeInsertClause], - s_alias: &str, - injected_columns: &[String], - table_fields: &[DataField], - temp_tracker: &mut TempTableTracker<'_>, -) -> DFResult> { - if not_matched_batches.is_empty() || not_matched_batches.iter().all(|b| b.num_rows() == 0) { - return Ok(Vec::new()); - } - - // Strip injected columns (_ROW_ID, __upd_*) — keep only source columns - let source_batches = strip_non_source_columns(not_matched_batches, injected_columns)?; - - // Register as temp table for SQL-based projection/filtering - let first_schema = source_batches[0].schema(); - let tmp_name = next_cow_table_name("__merge_not_matched"); - - let mem_table = datafusion::datasource::MemTable::try_new(first_schema, vec![source_batches])?; - ctx.register_temp_table(&tmp_name, Arc::new(mem_table))?; - temp_tracker.register(&tmp_name); - - let result = build_insert_batches_inner(ctx, inserts, s_alias, &tmp_name, table_fields).await; - - result -} - -/// Execute INSERT clause queries against the registered temp table. -async fn build_insert_batches_inner( - ctx: &SQLContext, - inserts: &[MergeInsertClause], - s_alias: &str, - tmp_name: &str, - table_fields: &[DataField], -) -> DFResult> { - let mut all_batches = Vec::new(); - let mut consumed_predicates: Vec = Vec::new(); - - for ins in inserts { - let mut conditions = Vec::new(); - for prev in &consumed_predicates { - // Only TRUE consumes a row; FALSE and NULL fall through. - conditions.push(format!("({prev}) IS NOT TRUE")); - } - if let Some(ref pred) = ins.predicate { - // Preserve OR grouping when combined with previous-clause exclusions. - conditions.push(format!("({pred})")); - consumed_predicates.push(pred.clone()); - } else { - consumed_predicates.push("TRUE".to_string()); - } - - let where_clause = if conditions.is_empty() { - String::new() - } else { - format!(" WHERE {}", conditions.join(" AND ")) - }; - - let select_clause = insert_select_clause(ins, table_fields)?; - let sql = format!("SELECT {select_clause} FROM {tmp_name} AS {s_alias}{where_clause}"); - - let batches = ctx.ctx().sql(&sql).await?.collect().await?; - for batch in batches { - all_batches.push(normalize_insert_batch_to_table_schema( - &batch, - table_fields, - )?); - } - } - - Ok(all_batches) -} - -fn normalize_insert_batch_to_table_schema( - batch: &RecordBatch, - table_fields: &[DataField], -) -> DFResult { - if batch.num_columns() != table_fields.len() { - return Err(DataFusionError::Plan(format!( - "MERGE INSERT output has {} columns but target table has {}", - batch.num_columns(), - table_fields.len() - ))); - } - - let target_schema = - paimon::arrow::build_target_arrow_schema(table_fields).map_err(to_datafusion_error)?; - let mut columns = Vec::with_capacity(table_fields.len()); - - for (target_idx, field) in table_fields.iter().enumerate() { - let column = batch.column(target_idx).clone(); - let target_type = target_schema.field(target_idx).data_type(); - let column = cast_insert_column(field.name(), column, target_type)?; - columns.push(column); - } - - RecordBatch::try_new(target_schema, columns).map_err(DataFusionError::from) -} - -fn cast_insert_column( - name: &str, - column: ArrayRef, - target_type: &ArrowDataType, -) -> DFResult { - if column.data_type() == target_type { - return Ok(column); - } - - compute::cast(column.as_ref(), target_type).map_err(|e| { - DataFusionError::Plan(format!( - "Cannot cast MERGE INSERT column '{name}' from {:?} to {:?}: {e}", - column.data_type(), - target_type - )) - }) -} - -/// Remove injected columns from batches, keeping only source columns. -fn strip_non_source_columns( - batches: &[RecordBatch], - injected_columns: &[String], -) -> DFResult> { - let mut result = Vec::new(); - for batch in batches { - let schema = batch.schema(); - let mut indices = Vec::new(); - let mut fields = Vec::new(); - for (i, field) in schema.fields().iter().enumerate() { - if injected_columns.contains(field.name()) { - continue; - } - indices.push(i); - fields.push(field.as_ref().clone()); - } - let new_schema = Arc::new(Schema::new(fields)); - let columns: Vec<_> = indices.iter().map(|&i| batch.column(i).clone()).collect(); - let projected = RecordBatch::try_new(new_schema, columns)?; - result.push(projected); - } - Ok(result) -} - -/// Build the SELECT clause for an INSERT clause, ordered by table schema fields. -/// -/// When the INSERT specifies explicit columns (`INSERT (col2, col1) VALUES (expr2, expr1)`), -/// the output must be reordered to match the table schema so that `write_arrow_batch` -/// (which reads columns by positional index) maps them correctly. When the INSERT omits -/// the column list (`INSERT VALUES (expr1, expr2, ...)`), the values map to the table's -/// columns by position, following standard SQL, where the MERGE INSERT column list is -/// optional (unlike Spark, whose grammar requires an explicit list or `INSERT *`). -fn insert_select_clause(ins: &MergeInsertClause, table_fields: &[DataField]) -> DFResult { - if ins.columns.is_empty() && ins.value_exprs.is_empty() { - return Ok("*".to_string()); - } - - if ins.columns.is_empty() { - // No explicit column list: the values fill the table's columns by - // position. Without this branch the empty column->expr map below would - // emit `NULL` for every column and silently drop the source values, - // inserting all-NULL rows (or failing on a non-null column). - if ins.value_exprs.len() != table_fields.len() { - return Err(DataFusionError::Plan(format!( - "MERGE INSERT has {} value(s) but the table has {} column(s); \ - list the target columns explicitly or provide one value per column", - ins.value_exprs.len(), - table_fields.len() - ))); - } - return Ok(table_fields - .iter() - .zip(ins.value_exprs.iter()) - .map(|(field, expr)| format!("{expr} AS {}", quote_identifier(field.name()))) - .collect::>() - .join(", ")); - } - - // Build column_name -> expression mapping from the INSERT clause - let col_expr_map: HashMap<&str, &str> = ins - .columns - .iter() - .zip(ins.value_exprs.iter()) - .map(|(col, expr)| (col.as_str(), expr.as_str())) - .collect(); - - // Emit SELECT in table schema order - Ok(table_fields - .iter() - .map(|field| { - match col_expr_map.get(field.name()) { - Some(expr) => format!("{expr} AS {}", quote_identifier(field.name())), - // Column not in INSERT list — fill with NULL - None => format!("NULL AS {}", quote_identifier(field.name())), - } - }) - .collect::>() - .join(", ")) -} - -/// Parsed WHEN NOT MATCHED THEN INSERT clause. -struct MergeInsertClause { - /// Column names from INSERT (col1, col2). Empty means INSERT *. - columns: Vec, - /// SQL expressions from VALUES(...). - value_exprs: Vec, - /// Optional AND predicate (SQL string). - predicate: Option, -} - -fn extract_merge_insert_clause( - insert_expr: &MergeInsertExpr, - predicate: Option<&SqlExpr>, - enable_ident_normalization: bool, -) -> DFResult { - let columns = insert_expr - .columns - .iter() - .map(|name| match name.0.as_slice() { - [part] => part - .as_ident() - .map(|ident| normalize_schema_identifier(ident, enable_ident_normalization)) - .ok_or_else(|| { - DataFusionError::Plan(format!("Invalid column name in MERGE INSERT: {name}")) - }), - _ => Err(DataFusionError::Plan(format!( - "Expected a simple column name in MERGE INSERT, got: {name}" - ))), - }) - .collect::>>()?; - - let value_exprs = match &insert_expr.kind { - MergeInsertKind::Values(values) => { - let row = values.rows.first().ok_or_else(|| { - DataFusionError::Plan("INSERT VALUES must have at least one row".to_string()) - })?; - row.iter().map(ToString::to_string).collect() - } - // INSERT ROW — BigQuery syntax, treat as INSERT *. - MergeInsertKind::Row => Vec::new(), - }; - - Ok(MergeInsertClause { - columns, - value_exprs, - predicate: predicate.map(ToString::to_string), - }) -} - -fn validate_merge_insert_columns( - inserts: &[MergeInsertClause], - table_fields: &[DataField], -) -> DFResult<()> { - for insert in inserts { - validate_target_columns(&insert.columns, table_fields, "MERGE INSERT")?; - // A column-less positional INSERT maps its VALUES to the table columns by - // position, so the value count must match the column count. Validate it - // here, upfront: this runs in both CoW and data-evolution modes before any - // batch is built, whereas `insert_select_clause` only runs when unmatched - // rows exist — an all-matched MERGE would otherwise skip the arity check - // and commit. `INSERT *`/`INSERT ROW` (no columns and no values) is left - // untouched. - if insert.columns.is_empty() - && !insert.value_exprs.is_empty() - && insert.value_exprs.len() != table_fields.len() - { - return Err(DataFusionError::Plan(format!( - "MERGE INSERT has {} value(s) but the table has {} column(s); \ - list the target columns explicitly or provide one value per column", - insert.value_exprs.len(), - table_fields.len() - ))); - } - } - - Ok(()) -} - -pub(crate) fn validate_update_columns( - columns: &[String], - table_fields: &[DataField], -) -> DFResult<()> { - validate_target_columns(columns, table_fields, "UPDATE") -} - -fn validate_target_columns( - columns: &[String], - table_fields: &[DataField], - operation: &str, -) -> DFResult<()> { - let target_columns: HashSet<&str> = table_fields.iter().map(|field| field.name()).collect(); - let mut seen = HashSet::new(); - - for column in columns { - if !seen.insert(column.as_str()) { - return Err(DataFusionError::Plan(format!( - "Duplicate column '{column}' in {operation}" - ))); - } - if !target_columns.contains(column.as_str()) { - return Err(DataFusionError::Plan(format!( - "Unknown column '{column}' in {operation}" - ))); - } - } - - Ok(()) -} - -/// Parsed WHEN MATCHED THEN UPDATE clause. -struct MergeUpdateClause { - columns: Vec, - exprs: Vec, -} - -/// Parsed merge clauses. -struct ParsedMergeClauses { - update: Option, - delete: bool, - inserts: Vec, -} - -/// Extract UPDATE and INSERT clauses from the MERGE AST. -fn extract_merge_clauses( - merge: &Merge, - enable_ident_normalization: bool, -) -> DFResult { - let mut update: Option = None; - let mut delete = false; - let mut inserts: Vec = Vec::new(); - - for clause in &merge.clauses { - match clause.clause_kind { - MergeClauseKind::Matched => { - if update.is_some() || delete { - return Err(DataFusionError::Plan( - "Multiple WHEN MATCHED clauses are not yet supported".to_string(), - )); - } - if clause.predicate.is_some() { - return Err(DataFusionError::Plan( - "WHEN MATCHED AND is not yet supported".to_string(), - )); - } - match &clause.action { - MergeAction::Update(update_expr) => { - let mut columns = Vec::new(); - let mut exprs = Vec::new(); - for assignment in &update_expr.assignments { - let col_name = match &assignment.target { - AssignmentTarget::ColumnName(name) => name - .0 - .last() - .and_then(|p| p.as_ident()) - .map(|ident| { - normalize_schema_identifier( - ident, - enable_ident_normalization, - ) - }) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Invalid column name in SET: {name}" - )) - })?, - AssignmentTarget::Tuple(_) => { - return Err(DataFusionError::Plan( - "Tuple assignment in MERGE INTO SET is not supported" - .to_string(), - )); - } - }; - columns.push(col_name); - exprs.push(assignment.value.to_string()); - } - update = Some(MergeUpdateClause { columns, exprs }); - } - MergeAction::Delete { .. } => { - delete = true; - } - MergeAction::Insert(_) => { - return Err(DataFusionError::Plan( - "WHEN MATCHED THEN INSERT is not valid SQL".to_string(), - )); - } - } - } - MergeClauseKind::NotMatched | MergeClauseKind::NotMatchedByTarget => { - match &clause.action { - MergeAction::Insert(insert_expr) => { - inserts.push(extract_merge_insert_clause( - insert_expr, - clause.predicate.as_ref(), - enable_ident_normalization, - )?); - } - _ => { - return Err(DataFusionError::Plan( - "WHEN NOT MATCHED only supports INSERT".to_string(), - )); - } - } - } - MergeClauseKind::NotMatchedBySource => { - return Err(DataFusionError::Plan( - "WHEN NOT MATCHED BY SOURCE is not yet supported for data evolution MERGE INTO" - .to_string(), - )); - } - } - } - - if update.is_none() && !delete && inserts.is_empty() { - return Err(DataFusionError::Plan( - "MERGE INTO requires at least one WHEN MATCHED or WHEN NOT MATCHED clause".to_string(), - )); - } - - Ok(ParsedMergeClauses { - update, - delete, - inserts, - }) -} - -/// Extract table name and optional alias from a TableFactor. -fn extract_table_ref(table: &TableFactor) -> DFResult<(String, Option)> { - match table { - TableFactor::Table { name, alias, .. } => { - let table_name = name.to_string(); - let alias_name = alias.as_ref().map(|a| a.name.value.clone()); - Ok((table_name, alias_name)) - } - other => Err(DataFusionError::Plan(format!( - "Unsupported table reference in MERGE INTO: {other}" - ))), - } -} - -/// Extract source reference (table or subquery) from a TableFactor. -fn extract_source_ref(source: &TableFactor) -> DFResult<(String, Option)> { - match source { - TableFactor::Table { name, alias, .. } => { - let table_name = name.to_string(); - let alias_name = alias.as_ref().map(|a| a.name.value.clone()); - Ok((table_name, alias_name)) - } - TableFactor::Derived { - subquery, alias, .. - } => { - let subquery_sql = format!("({subquery})"); - let alias_name = alias.as_ref().map(|a| a.name.value.clone()); - if alias_name.is_none() { - return Err(DataFusionError::Plan( - "Subquery source in MERGE INTO must have an alias".to_string(), - )); - } - Ok((subquery_sql, alias_name)) - } - other => Err(DataFusionError::Plan(format!( - "Unsupported source in MERGE INTO: {other}" - ))), - } -} - -/// Extract __paimon_file_idx and __paimon_row_offset columns from a JOIN result batch. -pub(crate) fn extract_tracking_columns( - batch: &RecordBatch, -) -> DFResult<(&Int32Array, &UInt32Array)> { - let file_idx_col = batch - .column_by_name("__paimon_file_idx") - .ok_or_else(|| DataFusionError::Internal("__paimon_file_idx not found".to_string()))? - .as_any() - .downcast_ref::() - .ok_or_else(|| DataFusionError::Internal("__paimon_file_idx is not Int32".to_string()))?; - - let row_offset_col = batch - .column_by_name("__paimon_row_offset") - .ok_or_else(|| DataFusionError::Internal("__paimon_row_offset not found".to_string()))? - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal("__paimon_row_offset is not UInt32".to_string()) - })?; - - Ok((file_idx_col, row_offset_col)) -} - -/// Read all files from a table via the CoW writer's file index, attach `__paimon_file_idx` -/// and `__paimon_row_offset` tracking columns, and register the result as a MemTable. -/// -/// Returns `(has_data, cow_table_name)`. The caller is responsible for deregistering -/// via `TempTableTracker`. -/// -/// Note: all matching partition files are loaded into memory at once. For partitions -/// with many large files this may cause significant memory pressure. A future -/// optimization could stream or batch-process files instead of materializing everything. -pub(crate) async fn register_cow_target_table( - ctx: &SQLContext, - table: &Table, - writer: &CopyOnWriteMergeWriter, - temp_tracker: &mut TempTableTracker<'_>, -) -> DFResult<(bool, String)> { - let file_index = writer.file_index(); - if file_index.is_empty() { - let table_name = next_cow_table_name("__cow_target"); - return Ok((false, table_name)); - } - - // Read all files in parallel - let read_futures: Vec<_> = file_index - .iter() - .enumerate() - .map(|(file_idx, file_info)| async move { - let single_split = DataSplitBuilder::new() - .with_snapshot(file_info.snapshot_id) - .with_partition( - paimon::spec::BinaryRow::from_serialized_bytes(&file_info.partition) - .map_err(to_datafusion_error)?, - ) - .with_bucket(file_info.bucket) - .with_bucket_path(file_info.bucket_path.clone()) - .with_total_buckets(file_info.total_buckets) - .with_data_files(vec![file_info.file_meta.clone()]) - .build() - .map_err(to_datafusion_error)?; - - let read = table - .new_read_builder() - .new_read() - .map_err(to_datafusion_error)?; - let batches: Vec = read - .to_arrow(&[single_split]) - .map_err(to_datafusion_error)? - .try_collect() - .await - .map_err(to_datafusion_error)?; - - Ok::<_, DataFusionError>((file_idx, batches)) - }) - .collect(); - - let file_results = futures::future::try_join_all(read_futures).await?; - - let mut all_batches: Vec = Vec::new(); - let mut schema: Option> = None; - - for (file_idx, batches) in file_results { - let mut row_offset = 0u32; - for batch in batches { - let num_rows = batch.num_rows(); - if num_rows == 0 { - continue; - } - - let file_idx_i32 = i32::try_from(file_idx).map_err(|_| { - DataFusionError::Internal(format!("file_idx {file_idx} exceeds i32 range")) - })?; - let num_rows_u32 = u32::try_from(num_rows).map_err(|_| { - DataFusionError::Internal(format!("batch num_rows {num_rows} exceeds u32 range")) - })?; - let file_idx_col = Arc::new(Int32Array::from(vec![file_idx_i32; num_rows])); - let end_offset = row_offset.checked_add(num_rows_u32).ok_or_else(|| { - DataFusionError::Internal(format!( - "row_offset overflow: {row_offset} + {num_rows_u32}" - )) - })?; - let row_offset_col = Arc::new(UInt32Array::from( - (row_offset..end_offset).collect::>(), - )); - - let mut fields: Vec = batch - .schema() - .fields() - .iter() - .map(|f| f.as_ref().clone()) - .collect(); - fields.push(Field::new("__paimon_file_idx", ArrowDataType::Int32, false)); - fields.push(Field::new( - "__paimon_row_offset", - ArrowDataType::UInt32, - false, - )); - let augmented_schema = Arc::new(Schema::new(fields)); - - let mut columns: Vec> = batch.columns().to_vec(); - columns.push(file_idx_col); - columns.push(row_offset_col); - - let augmented = RecordBatch::try_new(augmented_schema.clone(), columns) - .map_err(|e| DataFusionError::Internal(format!("Failed to augment batch: {e}")))?; - - if schema.is_none() { - schema = Some(augmented.schema()); - } - all_batches.push(augmented); - row_offset = end_offset; - } - } - - let has_data = !all_batches.is_empty(); - let table_name = next_cow_table_name("__cow_target"); - - if has_data { - let s = schema.unwrap(); - let mem_table = datafusion::datasource::MemTable::try_new(s, vec![all_batches])?; - ctx.register_temp_table(&table_name, Arc::new(mem_table))?; - temp_tracker.register(&table_name); - } - - Ok((has_data, table_name)) -} - -/// Build a partition set from Arrow batches containing partition column values. -/// -/// Converts each row's partition columns into serialized `BinaryRow` bytes. -/// Returns `None` for non-partitioned tables. -pub(crate) fn build_partition_set_from_batches( - table: &Table, - batches: &[RecordBatch], -) -> DFResult>>> { - let partition_keys = table.schema().partition_keys(); - if partition_keys.is_empty() { - return Ok(None); - } - - let partition_fields = table.schema().partition_fields(); - let mut partition_set = HashSet::new(); - - for batch in batches { - for row in 0..batch.num_rows() { - let datums: Vec<(Option, paimon::spec::DataType)> = - partition_fields - .iter() - .enumerate() - .map(|(col_idx, field)| { - let datum = - extract_datum_from_arrow(batch, row, col_idx, field.data_type()) - .map_err(to_datafusion_error)?; - Ok((datum, field.data_type().clone())) - }) - .collect::>()?; - let refs: Vec<(&Option, &paimon::spec::DataType)> = - datums.iter().map(|(d, t)| (d, t)).collect(); - partition_set.insert(datums_to_binary_row(&refs)); - } - } - - Ok(Some(partition_set)) -} - -/// Query a table for distinct partition values matching an optional WHERE clause. -/// -/// Returns `None` for non-partitioned tables. -pub(crate) async fn build_partition_set_from_where( - ctx: &SQLContext, - table: &Table, - table_ref: &str, - where_clause: Option<&str>, -) -> DFResult>>> { - let partition_keys = table.schema().partition_keys(); - if partition_keys.is_empty() { - return Ok(None); - } - - let cols = partition_keys - .iter() - .map(|k| quote_identifier(k)) - .collect::>() - .join(", "); - let where_part = match where_clause { - Some(w) => format!(" WHERE {w}"), - None => String::new(), - }; - let sql = format!("SELECT DISTINCT {cols} FROM {table_ref}{where_part}"); - let batches = ctx.ctx().sql(&sql).await?.collect().await?; - - build_partition_set_from_batches(table, &batches) -} - -/// Query source table for distinct partition values and build a partition set. -/// -/// Returns `None` for non-partitioned tables or when the source lacks matching -/// partition key columns (falls back to full-partition scan). The source values -/// are only safe to apply to the target when the MERGE condition equates every -/// target partition column with the same source column. -async fn build_source_partition_set( - ctx: &SQLContext, - table: &Table, - source_ref: &str, - s_alias: &str, - t_alias: &str, - on_expr: &SqlExpr, -) -> DFResult>>> { - let partition_keys = table.schema().partition_keys(); - if partition_keys.is_empty() { - return Ok(None); - } - if !can_prune_target_by_source_partitions(partition_keys, t_alias, s_alias, on_expr) { - return Ok(None); - } - - let cols = partition_keys - .iter() - .map(|k| format!("{s_alias}.{}", quote_identifier(k))) - .collect::>() - .join(", "); - let sql = format!("SELECT DISTINCT {cols} FROM {source_ref} AS {s_alias}"); - match ctx.ctx().sql(&sql).await { - Ok(df) => { - let batches = df.collect().await?; - build_partition_set_from_batches(table, &batches) - } - Err(_) => Ok(None), - } -} - -fn can_prune_target_by_source_partitions( - partition_keys: &[String], - t_alias: &str, - s_alias: &str, - on_expr: &SqlExpr, -) -> bool { - partition_keys - .iter() - .all(|key| on_expr_contains_alias_column_eq(on_expr, t_alias, key, s_alias, key)) -} - -fn on_expr_contains_alias_column_eq( - expr: &SqlExpr, - left_alias: &str, - left_column: &str, - right_alias: &str, - right_column: &str, -) -> bool { - match expr { - SqlExpr::BinaryOp { - left, - op: BinaryOperator::And, - right, - } => { - on_expr_contains_alias_column_eq( - left, - left_alias, - left_column, - right_alias, - right_column, - ) || on_expr_contains_alias_column_eq( - right, - left_alias, - left_column, - right_alias, - right_column, - ) - } - SqlExpr::BinaryOp { - left, - op: BinaryOperator::Eq, - right, - } => { - is_alias_column(left, left_alias, left_column) - && is_alias_column(right, right_alias, right_column) - || is_alias_column(left, right_alias, right_column) - && is_alias_column(right, left_alias, left_column) - } - _ => false, - } -} - -fn is_alias_column(expr: &SqlExpr, alias: &str, column: &str) -> bool { - match expr { - SqlExpr::CompoundIdentifier(parts) => is_qualified_column(parts, alias, column), - _ => false, - } -} - -fn is_qualified_column(parts: &[Ident], alias: &str, column: &str) -> bool { - parts.len() == 2 - && parts[0].value.eq_ignore_ascii_case(alias) - && parts[1].value.eq_ignore_ascii_case(column) -} - -/// Replace table qualifiers in column references, leaving literals and column names intact. -/// -/// For example, `paimon.test_db.target.a = source.a` becomes `t.a = s.a` -/// when target_ref="paimon.test_db.target", t_alias="t", source_ref="source", s_alias="s". -/// -/// Match the longest qualifier first, and rewrite it only once so an alias that -/// equals another table name cannot be rewritten again. -fn rewrite_table_references( - node: &mut T, - target_ref: &str, - t_alias: &str, - source_ref: &str, - s_alias: &str, -) { - let _: ControlFlow<()> = visit_expressions_mut(node, |expr| { - if let SqlExpr::CompoundIdentifier(parts) = expr { - for prefix_len in (1..parts.len()).rev() { - let qualifier = parts[..prefix_len] - .iter() - .map(ToString::to_string) - .collect::>() - .join("."); - let alias = if qualifier == target_ref { - Some(t_alias) - } else if qualifier == source_ref { - Some(s_alias) - } else { - None - }; - if let Some(alias) = alias { - parts.splice(..prefix_len, [Ident::new(alias)]); - break; - } - } - } - ControlFlow::Continue(()) - }); + build_partition_set_from_batches(table, &batches) } /// Return a DataFrame with a single "count" column. @@ -1684,1029 +1207,4 @@ pub(crate) fn ok_result(ctx: &SessionContext, count: u64) -> DFResult } #[cfg(test)] -mod tests { - use super::*; - use datafusion::sql::sqlparser::dialect::GenericDialect; - use datafusion::sql::sqlparser::parser::Parser; - use paimon::catalog::{Catalog, Identifier}; - use paimon::io::FileIOBuilder; - use paimon::spec::{DataField, DataType, IntType, Schema as PaimonSchema, TableSchema}; - use paimon::{CatalogOptions, FileSystemCatalog, Options}; - use tempfile::TempDir; - - use crate::SQLContext; - - async fn setup_sql_context() -> (TempDir, SQLContext, Arc) { - let temp_dir = TempDir::new().unwrap(); - let warehouse = format!("file://{}", temp_dir.path().display()); - let mut options = Options::new(); - options.set(CatalogOptions::WAREHOUSE, warehouse); - let catalog = Arc::new(FileSystemCatalog::new(options).unwrap()); - - let mut sql_context = SQLContext::new(); - sql_context - .register_catalog("paimon", catalog.clone()) - .await - .unwrap(); - sql_context - .sql("CREATE SCHEMA paimon.test_db") - .await - .unwrap(); - - (temp_dir, sql_context, catalog) - } - - async fn setup_data_evolution_table(name: &str) -> (TempDir, SQLContext, Table) { - let (tmp, sql_context, catalog) = setup_sql_context().await; - - sql_context - .sql(&format!( - "CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT) WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true')" - )) - .await - .unwrap(); - - sql_context - .sql(&format!( - "INSERT INTO paimon.test_db.{name} (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)" - )) - .await - .unwrap() - .collect() - .await - .unwrap(); - - let table = catalog - .get_table(&Identifier::new("test_db", name)) - .await - .unwrap(); - let mut extra = std::collections::HashMap::new(); - extra.insert("data-evolution.enabled".to_string(), "true".to_string()); - extra.insert("row-tracking.enabled".to_string(), "true".to_string()); - let de_table = table.copy_with_options(extra); - - (tmp, sql_context, de_table) - } - - fn parse_merge(sql: &str) -> Merge { - let dialect = GenericDialect {}; - let stmts = Parser::parse_sql(&dialect, sql).unwrap(); - match stmts.into_iter().next().unwrap() { - datafusion::sql::sqlparser::ast::Statement::Merge(m) => m, - _ => panic!("Expected MERGE statement"), - } - } - - #[test] - fn test_merge_reference_rewrite_preserves_literal_contents() { - for (input, expected) in [ - ("target.id = source.id", "t.id = s.id"), - ("'source.payload'", "'source.payload'"), - ("'it''s target.payload'", "'it''s target.payload'"), - ( - "'{\"path\":\"source.payload\"}'", - "'{\"path\":\"source.payload\"}'", - ), - ("source.name LIKE 'target.%'", "s.name LIKE 'target.%'"), - ("source.\"source.payload\"", "s.\"source.payload\""), - ("my_source.id = source.id", "my_source.id = s.id"), - ( - "target.nested.value = source.value", - "t.nested.value = s.value", - ), - ] { - let parse_expr = |sql| { - Parser::new(&GenericDialect {}) - .try_with_sql(sql) - .unwrap() - .parse_expr() - .unwrap() - }; - let mut actual = parse_expr(input); - rewrite_table_references(&mut actual, "target", "t", "source", "s"); - assert_eq!(actual, parse_expr(expected), "{input}"); - } - } - - #[test] - fn test_merge_update_assignments_normalize_identifiers() { - for (target, enabled, expected) in [ - ("VALUE", true, "value"), - ("\"VALUE\"", true, "VALUE"), - ("VALUE", false, "VALUE"), - ] { - let merge = parse_merge(&format!( - "MERGE INTO target t USING source s ON t.id = s.id \ - WHEN MATCHED THEN UPDATE SET {target} = s.value" - )); - - let parsed = extract_merge_clauses(&merge, enabled).unwrap(); - assert_eq!(parsed.update.unwrap().columns, [expected]); - - let cow = extract_cow_merge_clauses(&merge, enabled).unwrap(); - let CowMatchedAction::Update(update) = &cow.matched[0].action else { - panic!("expected update action"); - }; - assert_eq!(update.columns, [expected]); - } - } - - #[test] - fn test_merge_insert_columns_normalize_identifiers() { - for (target, enabled, expected) in [ - ("MIXEDCASE", true, "mixedcase"), - ("\"MixedCase\"", true, "MixedCase"), - ("MixedCase", false, "MixedCase"), - ] { - let merge = parse_merge(&format!( - "MERGE INTO target t USING source s ON t.id = s.id \ - WHEN NOT MATCHED THEN INSERT ({target}) VALUES (s.value)" - )); - - let parsed = extract_merge_clauses(&merge, enabled).unwrap(); - assert_eq!(parsed.inserts[0].columns, [expected]); - - let cow = extract_cow_merge_clauses(&merge, enabled).unwrap(); - assert_eq!(cow.inserts[0].columns, [expected]); - } - } - - #[test] - fn test_normalize_merge_insert_batch_uses_position() { - let table_fields = vec![ - DataField::new(0, "a".to_string(), DataType::Int(IntType::new())), - DataField::new(1, "b".to_string(), DataType::Int(IntType::new())), - ]; - let batch = RecordBatch::try_new( - Arc::new(Schema::new(vec![ - Field::new("b", ArrowDataType::Int32, false), - Field::new("x", ArrowDataType::Int32, false), - ])), - vec![ - Arc::new(Int32Array::from(vec![100])), - Arc::new(Int32Array::from(vec![7])), - ], - ) - .unwrap(); - - let normalized = normalize_insert_batch_to_table_schema(&batch, &table_fields).unwrap(); - let first = normalized - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let second = normalized - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - - assert_eq!(normalized.schema().field(0).name(), "a"); - assert_eq!(normalized.schema().field(1).name(), "b"); - assert_eq!(first.value(0), 100); - assert_eq!(second.value(0), 7); - } - - #[test] - fn test_source_partition_pruning_requires_partition_equality() { - let merge = parse_merge( - "MERGE INTO target t USING source s ON t.a = s.a \ - WHEN MATCHED THEN UPDATE SET b = s.b", - ); - let partition_keys = vec!["pt".to_string()]; - - assert!(!can_prune_target_by_source_partitions( - &partition_keys, - "t", - "s", - &merge.on, - )); - } - - #[test] - fn test_source_partition_pruning_accepts_partition_equality() { - let merge = parse_merge( - "MERGE INTO target t USING source s ON t.a = s.a AND t.pt = s.pt \ - WHEN MATCHED THEN UPDATE SET b = s.b", - ); - let partition_keys = vec!["pt".to_string()]; - - assert!(can_prune_target_by_source_partitions( - &partition_keys, - "t", - "s", - &merge.on, - )); - } - - #[test] - fn test_source_partition_pruning_accepts_reversed_partition_equality() { - let merge = parse_merge( - "MERGE INTO target t USING source s ON s.pt = t.pt AND t.a = s.a \ - WHEN MATCHED THEN UPDATE SET b = s.b", - ); - let partition_keys = vec!["pt".to_string()]; - - assert!(can_prune_target_by_source_partitions( - &partition_keys, - "t", - "s", - &merge.on, - )); - } - - #[tokio::test] - async fn test_merge_into_updates_matched_rows() { - let (_tmp, sql_context, table) = setup_data_evolution_table("t_merge").await; - - // Create source table with updates - sql_context - .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") - .await - .unwrap() - .collect() - .await - .unwrap(); - - sql_context - .sql("INSERT INTO paimon.test_db.source VALUES (1, 'ALICE'), (3, 'CHARLIE')") - .await - .unwrap() - .collect() - .await - .unwrap(); - - // Execute MERGE INTO - let merge = parse_merge( - "MERGE INTO paimon.test_db.t_merge t USING paimon.test_db.source s ON t.id = s.id \ - WHEN MATCHED THEN UPDATE SET NAME = s.name", - ); - execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap(); - - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.t_merge ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let mut rows = Vec::new(); - for batch in &batches { - let ids = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let names = batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - let values = batch - .column(2) - .as_any() - .downcast_ref::() - .unwrap(); - for i in 0..batch.num_rows() { - rows.push((ids.value(i), names.value(i).to_string(), values.value(i))); - } - } - - assert_eq!( - rows, - vec![ - (1, "ALICE".to_string(), 10), - (2, "bob".to_string(), 20), - (3, "CHARLIE".to_string(), 30), - ] - ); - } - - #[tokio::test] - async fn test_merge_into_no_matches() { - let (_tmp, sql_context, table) = setup_data_evolution_table("t_merge2").await; - - sql_context - .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") - .await - .unwrap() - .collect() - .await - .unwrap(); - - sql_context - .sql("INSERT INTO paimon.test_db.source VALUES (99, 'nobody')") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let merge = parse_merge( - "MERGE INTO paimon.test_db.t_merge2 t USING paimon.test_db.source s ON t.id = s.id \ - WHEN MATCHED THEN UPDATE SET name = s.name", - ); - let result = execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap(); - let batches = result.collect().await.unwrap(); - let count = batches[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, 0); - } - - #[tokio::test] - async fn test_merge_into_rejects_pk_table_without_data_evolution() { - let file_io = FileIOBuilder::new("memory").build().unwrap(); - let table_path = "memory:/test_merge_reject"; - file_io - .mkdirs(&format!("{table_path}/snapshot/")) - .await - .unwrap(); - file_io - .mkdirs(&format!("{table_path}/manifest/")) - .await - .unwrap(); - - let schema = PaimonSchema::builder() - .column("id", DataType::Int(IntType::new())) - .primary_key(["id"]) - .option("bucket", "1") - .build() - .unwrap(); - let table_schema = TableSchema::new(0, &schema); - let table = Table::new( - file_io, - Identifier::new("default", "t"), - table_path.to_string(), - table_schema, - None, - ); - - let sql_context = SQLContext::new(); - let merge = parse_merge( - "MERGE INTO t USING s ON t.id = s.id \ - WHEN MATCHED THEN UPDATE SET id = s.id", - ); - let result = execute_merge_into(&sql_context, &merge, table, true).await; - assert!(result.is_err()); - assert!(result - .unwrap_err() - .to_string() - .contains("primary-key tables without data-evolution")); - } - - // ----------------------------------------------------------------------- - // CoW MERGE INTO tests (append-only tables) - // ----------------------------------------------------------------------- - - async fn setup_append_only_table(name: &str) -> (TempDir, SQLContext, Table) { - let (tmp, sql_context, catalog) = setup_sql_context().await; - - sql_context - .sql(&format!( - "CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT)" - )) - .await - .unwrap(); - - sql_context - .sql(&format!( - "INSERT INTO paimon.test_db.{name} (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)" - )) - .await - .unwrap() - .collect() - .await - .unwrap(); - - let table = catalog - .get_table(&Identifier::new("test_db", name)) - .await - .unwrap(); - - (tmp, sql_context, table) - } - - fn collect_rows(batches: &[RecordBatch]) -> Vec<(i32, String, i32)> { - let mut rows = Vec::new(); - for batch in batches { - let ids = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let names = batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - let values = batch - .column(2) - .as_any() - .downcast_ref::() - .unwrap(); - for i in 0..batch.num_rows() { - rows.push((ids.value(i), names.value(i).to_string(), values.value(i))); - } - } - rows.sort_by_key(|r| r.0); - rows - } - - #[tokio::test] - async fn test_cow_merge_update_matched_rows() { - let (_tmp, sql_context, table) = setup_append_only_table("t_cow_upd").await; - - sql_context - .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") - .await - .unwrap(); - sql_context - .sql("INSERT INTO paimon.test_db.source (id, name) VALUES (1, 'ALICE'), (3, 'CHARLIE')") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let merge = parse_merge( - "MERGE INTO paimon.test_db.t_cow_upd t USING paimon.test_db.source s ON t.id = s.id \ - WHEN MATCHED THEN UPDATE SET NAME = s.name", - ); - execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap(); - - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.t_cow_upd ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let rows = collect_rows(&batches); - assert_eq!( - rows, - vec![ - (1, "ALICE".to_string(), 10), - (2, "bob".to_string(), 20), - (3, "CHARLIE".to_string(), 30), - ] - ); - } - - #[tokio::test] - async fn test_cow_merge_delete_matched_rows() { - let (_tmp, sql_context, table) = setup_append_only_table("t_cow_del").await; - - sql_context - .sql("CREATE TABLE paimon.test_db.source (id INT)") - .await - .unwrap(); - sql_context - .sql("INSERT INTO paimon.test_db.source (id) VALUES (2)") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let merge = parse_merge( - "MERGE INTO paimon.test_db.t_cow_del t USING paimon.test_db.source s ON t.id = s.id \ - WHEN MATCHED THEN DELETE", - ); - execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap(); - - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.t_cow_del ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let rows = collect_rows(&batches); - assert_eq!( - rows, - vec![(1, "alice".to_string(), 10), (3, "charlie".to_string(), 30),] - ); - } - - #[tokio::test] - async fn test_cow_merge_insert_not_matched() { - let (_tmp, sql_context, table) = setup_append_only_table("t_cow_ins").await; - - sql_context - .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR, value INT)") - .await - .unwrap(); - sql_context - .sql("INSERT INTO paimon.test_db.source VALUES (4, 'dave', 40), (5, 'eve', 50)") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let merge = parse_merge( - "MERGE INTO paimon.test_db.t_cow_ins t USING paimon.test_db.source s ON t.id = s.id \ - WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)", - ); - execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap(); - - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.t_cow_ins ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let rows = collect_rows(&batches); - assert_eq!( - rows, - vec![ - (1, "alice".to_string(), 10), - (2, "bob".to_string(), 20), - (3, "charlie".to_string(), 30), - (4, "dave".to_string(), 40), - (5, "eve".to_string(), 50), - ] - ); - } - - #[tokio::test] - async fn test_cow_merge_insert_not_matched_without_columns() { - let (_tmp, sql_context, table) = setup_append_only_table("t_cow_ins_nocols").await; - - sql_context - .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR, value INT)") - .await - .unwrap(); - sql_context - .sql("INSERT INTO paimon.test_db.source VALUES (4, 'dave', 40), (5, 'eve', 50)") - .await - .unwrap() - .collect() - .await - .unwrap(); - - // No column list: the VALUES map to the table's columns by position. - // Before the fix this dropped the source values and inserted all-NULL rows. - let merge = parse_merge( - "MERGE INTO paimon.test_db.t_cow_ins_nocols t USING paimon.test_db.source s \ - ON t.id = s.id \ - WHEN NOT MATCHED THEN INSERT VALUES (s.id, s.name, s.value)", - ); - execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap(); - - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.t_cow_ins_nocols ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let rows = collect_rows(&batches); - assert_eq!( - rows, - vec![ - (1, "alice".to_string(), 10), - (2, "bob".to_string(), 20), - (3, "charlie".to_string(), 30), - (4, "dave".to_string(), 40), - (5, "eve".to_string(), 50), - ] - ); - } - - #[tokio::test] - async fn test_cow_merge_insert_arity_validated_even_when_all_matched() { - let (_tmp, sql_context, table) = setup_append_only_table("t_cow_ins_arity").await; - - sql_context - .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR, value INT)") - .await - .unwrap(); - // id = 1 matches the target, so the NOT MATCHED INSERT branch is never - // built — the arity check must still run upfront. - sql_context - .sql("INSERT INTO paimon.test_db.source VALUES (1, 'ALICE', 99)") - .await - .unwrap() - .collect() - .await - .unwrap(); - - // The column-less INSERT lists 2 values for a 3-column table. Even with no - // row to insert, the invalid arity must be rejected before execution so the - // matched UPDATE is not applied. - let merge = parse_merge( - "MERGE INTO paimon.test_db.t_cow_ins_arity t USING paimon.test_db.source s \ - ON t.id = s.id \ - WHEN MATCHED THEN UPDATE SET value = s.value \ - WHEN NOT MATCHED THEN INSERT VALUES (s.id, s.name)", - ); - let err = execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap_err(); - assert!( - err.to_string().contains("value(s) but the table has"), - "got {err}" - ); - - // The target is unchanged: the matched UPDATE did not apply. - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.t_cow_ins_arity ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_eq!( - collect_rows(&batches), - vec![ - (1, "alice".to_string(), 10), - (2, "bob".to_string(), 20), - (3, "charlie".to_string(), 30), - ] - ); - } - - #[tokio::test] - async fn test_cow_merge_update_and_insert() { - let (_tmp, sql_context, table) = setup_append_only_table("t_cow_upsert").await; - - sql_context - .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR, value INT)") - .await - .unwrap(); - sql_context - .sql("INSERT INTO paimon.test_db.source VALUES (2, 'BOB', 200), (4, 'dave', 40)") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let merge = parse_merge( - "MERGE INTO paimon.test_db.t_cow_upsert t USING paimon.test_db.source s ON t.id = s.id \ - WHEN MATCHED THEN UPDATE SET name = s.name, value = s.value \ - WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)", - ); - execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap(); - - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.t_cow_upsert ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let rows = collect_rows(&batches); - assert_eq!( - rows, - vec![ - (1, "alice".to_string(), 10), - (2, "BOB".to_string(), 200), - (3, "charlie".to_string(), 30), - (4, "dave".to_string(), 40), - ] - ); - } - - #[tokio::test] - async fn test_cow_merge_matched_null_predicates_fall_through() { - for delete_fallback in [false, true] { - let (_tmp, sql_context, table) = setup_append_only_table("t_null_matched").await; - for sql in [ - "CREATE TABLE paimon.test_db.source (id INT, flag INT)", - "INSERT INTO paimon.test_db.source VALUES (1, NULL), (2, 1), (3, 0)", - ] { - sql_context.sql(sql).await.unwrap().collect().await.unwrap(); - } - - let fallback = if delete_fallback { - "DELETE" - } else { - "UPDATE SET name = 'fallback'" - }; - let merge = parse_merge(&format!( - "MERGE INTO paimon.test_db.t_null_matched t USING paimon.test_db.source s \ - ON t.id = s.id \ - WHEN MATCHED AND s.flag = 1 THEN UPDATE SET name = 'first' \ - WHEN MATCHED AND s.flag = 0 THEN UPDATE SET name = 'second' \ - WHEN MATCHED THEN {fallback}" - )); - let result = execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_eq!( - result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 3, - ); - - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.t_null_matched ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - let mut expected = vec![(2, "first".to_string(), 20), (3, "second".to_string(), 30)]; - if !delete_fallback { - expected.insert(0, (1, "fallback".to_string(), 10)); - } - assert_eq!(collect_rows(&batches), expected); - } - } - - #[tokio::test] - async fn test_merge_insert_null_predicates_fall_through() { - for data_evolution in [false, true] { - let (_tmp, sql_context, table) = if data_evolution { - setup_data_evolution_table("t_null_insert").await - } else { - setup_append_only_table("t_null_insert").await - }; - for sql in [ - "CREATE TABLE paimon.test_db.source (id INT, flag INT)", - "INSERT INTO paimon.test_db.source VALUES (4, NULL), (5, 1), (6, 0)", - ] { - sql_context.sql(sql).await.unwrap().collect().await.unwrap(); - } - - let merge = parse_merge( - "MERGE INTO paimon.test_db.t_null_insert t USING paimon.test_db.source s \ - ON t.id = s.id \ - WHEN NOT MATCHED AND s.flag = 1 THEN INSERT (id, name, value) VALUES (s.id, 'first', 100) \ - WHEN NOT MATCHED AND s.flag = 0 THEN INSERT (id, name, value) VALUES (s.id, 'second', 100) \ - WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, 'fallback', 100)", - ); - let result = execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_eq!( - result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 3, - ); - - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.t_null_insert ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_eq!( - collect_rows(&batches), - vec![ - (1, "alice".to_string(), 10), - (2, "bob".to_string(), 20), - (3, "charlie".to_string(), 30), - (4, "fallback".to_string(), 100), - (5, "first".to_string(), 100), - (6, "second".to_string(), 100), - ], - ); - } - } - - #[tokio::test] - async fn test_cow_merge_matched_or_predicates_preserve_clause_order() { - for delete_second in [false, true] { - let (_tmp, sql_context, table) = setup_append_only_table("t_or_matched").await; - for sql in [ - "CREATE TABLE paimon.test_db.source (id INT, flag INT)", - "INSERT INTO paimon.test_db.source VALUES (1, 1), (2, 0), (3, NULL)", - ] { - sql_context.sql(sql).await.unwrap().collect().await.unwrap(); - } - - let second_action = if delete_second { - "DELETE" - } else { - "UPDATE SET name = 'second'" - }; - let merge = parse_merge(&format!( - "MERGE INTO paimon.test_db.t_or_matched t USING paimon.test_db.source s \ - ON t.id = s.id \ - WHEN MATCHED AND s.flag = 1 THEN UPDATE SET name = 'first' \ - WHEN MATCHED AND s.flag = 0 OR s.flag = 1 THEN {second_action} \ - WHEN MATCHED THEN UPDATE SET name = 'fallback'" - )); - let result = execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_eq!( - result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 3, - ); - - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.t_or_matched ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - let mut expected = vec![ - (1, "first".to_string(), 10), - (3, "fallback".to_string(), 30), - ]; - if !delete_second { - expected.insert(1, (2, "second".to_string(), 20)); - } - assert_eq!(collect_rows(&batches), expected); - } - } - - #[tokio::test] - async fn test_merge_insert_or_predicates_preserve_clause_order() { - for data_evolution in [false, true] { - let (_tmp, sql_context, table) = if data_evolution { - setup_data_evolution_table("t_or_insert").await - } else { - setup_append_only_table("t_or_insert").await - }; - for sql in [ - "CREATE TABLE paimon.test_db.source (id INT, flag INT)", - "INSERT INTO paimon.test_db.source VALUES (4, 1), (5, 0), (6, NULL)", - ] { - sql_context.sql(sql).await.unwrap().collect().await.unwrap(); - } - - let merge = parse_merge( - "MERGE INTO paimon.test_db.t_or_insert t USING paimon.test_db.source s \ - ON t.id = s.id \ - WHEN NOT MATCHED AND s.flag = 1 THEN INSERT (id, name, value) VALUES (s.id, 'first', 100) \ - WHEN NOT MATCHED AND s.flag = 0 OR s.flag = 1 THEN INSERT (id, name, value) VALUES (s.id, 'second', 100) \ - WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, 'fallback', 100)", - ); - let result = execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_eq!( - result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 3, - ); - - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.t_or_insert ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_eq!( - collect_rows(&batches), - vec![ - (1, "alice".to_string(), 10), - (2, "bob".to_string(), 20), - (3, "charlie".to_string(), 30), - (4, "first".to_string(), 100), - (5, "second".to_string(), 100), - (6, "fallback".to_string(), 100), - ], - ); - } - } - - #[tokio::test] - async fn test_cow_merge_preserves_string_literals() { - for (on_predicate, matched_predicate, insert_predicate) in [ - ("", "", ""), - (" AND s.name = 'paimon.test_db.source.payload'", "", ""), - ("", " AND s.name = 'paimon.test_db.source.payload'", ""), - ("", "", " AND s.name = 'paimon.test_db.source.payload'"), - ] { - let (_tmp, sql_context, table) = setup_append_only_table("target").await; - for sql in [ - "CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)", - "INSERT INTO paimon.test_db.source VALUES (1, 'paimon.test_db.source.payload'), (4, 'paimon.test_db.source.payload')", - ] { - sql_context.sql(sql).await.unwrap().collect().await.unwrap(); - } - - let merge = parse_merge(&format!( - "MERGE INTO paimon.test_db.target t USING paimon.test_db.source s \ - ON paimon.test_db.target.id = paimon.test_db.source.id{on_predicate} \ - WHEN MATCHED{matched_predicate} THEN UPDATE SET name = 'paimon.test_db.target.payload' \ - WHEN NOT MATCHED{insert_predicate} THEN INSERT (id, name, value) VALUES (s.id, 'paimon.test_db.source.payload', 40)" - )); - let result = execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_eq!( - result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 2, - ); - - let batches = sql_context - .sql("SELECT id, name, value FROM paimon.test_db.target ORDER BY id") - .await - .unwrap() - .collect() - .await - .unwrap(); - assert_eq!( - collect_rows(&batches), - vec![ - (1, "paimon.test_db.target.payload".to_string(), 10), - (2, "bob".to_string(), 20), - (3, "charlie".to_string(), 30), - (4, "paimon.test_db.source.payload".to_string(), 40), - ], - "ON: {on_predicate}, MATCHED: {matched_predicate}, INSERT: {insert_predicate}", - ); - } - } - - #[tokio::test] - async fn test_cow_merge_no_matches() { - let (_tmp, sql_context, table) = setup_append_only_table("t_cow_nomatch").await; - - sql_context - .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let merge = parse_merge( - "MERGE INTO paimon.test_db.t_cow_nomatch t USING paimon.test_db.source s ON t.id = s.id \ - WHEN MATCHED THEN UPDATE SET name = s.name", - ); - let result = execute_merge_into(&sql_context, &merge, table, true) - .await - .unwrap(); - let batches = result.collect().await.unwrap(); - let count = batches[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - assert_eq!(count, 0); - } -} +mod tests; diff --git a/crates/integrations/datafusion/src/merge_into/normalize.rs b/crates/integrations/datafusion/src/merge_into/normalize.rs new file mode 100644 index 000000000..68a60751e --- /dev/null +++ b/crates/integrations/datafusion/src/merge_into/normalize.rs @@ -0,0 +1,130 @@ +// 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. + +use std::ops::ControlFlow; + +use datafusion::sql::sqlparser::ast::{ + Expr, Ident, Merge, MergeAction, TableFactor, VisitMut, VisitorMut, +}; + +use crate::sql_context::normalize_schema_identifier; + +fn qualifier(table: &TableFactor) -> Vec { + match table { + TableFactor::Table { + alias: Some(alias), .. + } + | TableFactor::Derived { + alias: Some(alias), .. + } => vec![alias.name.clone()], + TableFactor::Table { name, .. } => name + .0 + .iter() + .filter_map(|p| p.as_ident().cloned()) + .collect(), + _ => vec![], + } +} + +/// Restore source-only resolution for NOT MATCHED without changing variables, +/// lambda bindings, or already-qualified target/source references. +pub(crate) fn qualify_not_matched_exprs(merge: &mut Merge, normalize: bool) { + let source = qualifier(&merge.source); + if source.is_empty() { + return; + } + let mut scope = SourceScope { + source, + target: qualifier(&merge.table), + normalize, + locals: vec![], + }; + for clause in &mut merge.clauses { + if let MergeAction::Insert(insert) = &mut clause.action { + let _ = insert.kind.visit(&mut scope); + let _ = clause.predicate.visit(&mut scope); + } + } +} + +struct SourceScope { + source: Vec, + target: Vec, + normalize: bool, + locals: Vec>, +} + +impl SourceScope { + fn name(&self, ident: &Ident) -> String { + normalize_schema_identifier(ident, self.normalize) + } + + fn is_variable(&self, ident: &Ident) -> bool { + (ident.quote_style.is_none() && ident.value.starts_with('@')) + || self + .locals + .iter() + .any(|scope| scope.contains(&self.name(ident))) + } + + fn is_qualified(&self, parts: &[Ident]) -> bool { + [&self.source, &self.target].into_iter().any(|qualifier| { + (0..qualifier.len()).any(|start| { + let qualifier = &qualifier[start..]; + parts.len() > qualifier.len() + && parts + .iter() + .zip(qualifier) + .all(|(part, name)| self.name(part) == self.name(name)) + }) + }) + } +} + +impl VisitorMut for SourceScope { + type Break = (); + + fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow { + match expr { + Expr::Lambda(lambda) => { + self.locals + .push(lambda.params.iter().map(|p| self.name(&p.name)).collect()); + } + Expr::Identifier(ident) if !self.is_variable(ident) => { + let mut parts = self.source.clone(); + parts.push(ident.clone()); + *expr = Expr::CompoundIdentifier(parts); + } + Expr::CompoundIdentifier(parts) + if !self.is_variable(&parts[0]) && !self.is_qualified(parts) => + { + let mut qualified = self.source.clone(); + qualified.extend(parts.iter().cloned()); + *parts = qualified; + } + _ => {} + } + ControlFlow::Continue(()) + } + + fn post_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow { + if matches!(expr, Expr::Lambda(_)) { + self.locals.pop(); + } + ControlFlow::Continue(()) + } +} diff --git a/crates/integrations/datafusion/src/merge_into/tests.rs b/crates/integrations/datafusion/src/merge_into/tests.rs new file mode 100644 index 000000000..af217bfb5 --- /dev/null +++ b/crates/integrations/datafusion/src/merge_into/tests.rs @@ -0,0 +1,882 @@ +// 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. + +//! MERGE INTO unit and behavioral tests. + +use super::*; +use datafusion::sql::sqlparser::ast::Merge; +use datafusion::sql::sqlparser::dialect::GenericDialect; +use datafusion::sql::sqlparser::parser::Parser; +use paimon::{CatalogOptions, FileSystemCatalog, Options}; +use tempfile::TempDir; + +use crate::SQLContext; + +#[tokio::test] +async fn test_retry_limits_and_non_conflict_errors() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + for (message, classify) in [ + ( + "Row ID conflict", + is_row_id_conflict as fn(&DataFusionError) -> bool, + ), + ( + "Delete conflict", + is_delete_conflict as fn(&DataFusionError) -> bool, + ), + ] { + let attempts = AtomicUsize::new(0); + let result: DFResult = retry_on_conflict("MERGE INTO", classify, || async { + attempts.fetch_add(1, Ordering::SeqCst); + Err(DataFusionError::External(Box::new(std::io::Error::other( + message, + )))) + }) + .await; + assert_eq!(attempts.load(Ordering::SeqCst), 5); + let err = result.unwrap_err().to_string(); + assert!(err.contains("after 5 retries"), "{err}"); + assert!(err.contains(message), "{err}"); + + let attempts = AtomicUsize::new(0); + let result: DFResult = retry_on_conflict("MERGE INTO", classify, || async { + attempts.fetch_add(1, Ordering::SeqCst); + Err(DataFusionError::Execution("unrelated failure".into())) + }) + .await; + assert_eq!(attempts.load(Ordering::SeqCst), 1); + assert!( + matches!(result, Err(DataFusionError::Execution(message)) if message == "unrelated failure") + ); + let overlapping_update = to_datafusion_error(paimon::Error::DataInvalid { + message: "For Data Evolution table, multiple MERGE INTO operations have encountered conflicts, updating the same file".into(), + source: None, + }); + assert!(!classify(&overlapping_update)); + } +} + +async fn setup_sql_context() -> (TempDir, SQLContext) { + let temp_dir = TempDir::new().unwrap(); + let warehouse = format!("file://{}", temp_dir.path().display()); + let mut options = Options::new(); + options.set(CatalogOptions::WAREHOUSE, warehouse); + let catalog = Arc::new(FileSystemCatalog::new(options).unwrap()); + + let mut sql_context = SQLContext::new(); + sql_context + .register_catalog("paimon", catalog) + .await + .unwrap(); + sql_context + .sql("CREATE SCHEMA paimon.test_db") + .await + .unwrap(); + + (temp_dir, sql_context) +} + +async fn setup_data_evolution_table(name: &str) -> (TempDir, SQLContext) { + let (tmp, sql_context) = setup_sql_context().await; + + sql_context + .sql(&format!( + "CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT) WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true')" + )) + .await + .unwrap(); + + sql_context + .sql(&format!( + "INSERT INTO paimon.test_db.{name} (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + + (tmp, sql_context) +} + +// Exercise the native planner and provider while retaining these behavioral fixtures. +async fn execute_merge_into(ctx: &SQLContext, merge: &Merge) -> DFResult { + let batches = ctx.sql(&merge.to_string()).await?.collect().await?; + ctx.ctx().read_batches(batches) +} + +fn parse_merge(sql: &str) -> Merge { + let dialect = GenericDialect {}; + let stmts = Parser::parse_sql(&dialect, sql).unwrap(); + match stmts.into_iter().next().unwrap() { + datafusion::sql::sqlparser::ast::Statement::Merge(m) => m, + _ => panic!("Expected MERGE statement"), + } +} + +#[tokio::test] +async fn test_merge_into_updates_matched_rows() { + let (_tmp, sql_context) = setup_data_evolution_table("t_merge").await; + + // Create source table with updates + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + sql_context + .sql("INSERT INTO paimon.test_db.source VALUES (1, 'ALICE'), (3, 'CHARLIE')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + // Execute MERGE INTO + let merge = parse_merge( + "MERGE INTO paimon.test_db.t_merge t USING paimon.test_db.source s ON t.id = s.id \ + WHEN MATCHED THEN UPDATE SET NAME = s.name", + ); + execute_merge_into(&sql_context, &merge).await.unwrap(); + + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_merge ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let mut rows = Vec::new(); + for batch in &batches { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let names = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + rows.push((ids.value(i), names.value(i).to_string(), values.value(i))); + } + } + + assert_eq!( + rows, + vec![ + (1, "ALICE".to_string(), 10), + (2, "bob".to_string(), 20), + (3, "CHARLIE".to_string(), 30), + ] + ); +} + +#[tokio::test] +async fn test_merge_into_no_matches() { + let (_tmp, sql_context) = setup_data_evolution_table("t_merge2").await; + + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + sql_context + .sql("INSERT INTO paimon.test_db.source VALUES (99, 'nobody')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let merge = parse_merge( + "MERGE INTO paimon.test_db.t_merge2 t USING paimon.test_db.source s ON t.id = s.id \ + WHEN MATCHED THEN UPDATE SET name = s.name", + ); + let result = execute_merge_into(&sql_context, &merge).await.unwrap(); + let batches = result.collect().await.unwrap(); + let count = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, 0); +} + +// ----------------------------------------------------------------------- +// CoW MERGE INTO tests (append-only tables) +// ----------------------------------------------------------------------- + +async fn setup_append_only_table(name: &str) -> (TempDir, SQLContext) { + let (tmp, sql_context) = setup_sql_context().await; + + sql_context + .sql(&format!( + "CREATE TABLE paimon.test_db.{name} (id INT, name VARCHAR, value INT)" + )) + .await + .unwrap(); + + sql_context + .sql(&format!( + "INSERT INTO paimon.test_db.{name} (id, name, value) VALUES (1, 'alice', 10), (2, 'bob', 20), (3, 'charlie', 30)" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + + (tmp, sql_context) +} + +fn collect_rows(batches: &[RecordBatch]) -> Vec<(i32, String, i32)> { + let mut rows = Vec::new(); + for batch in batches { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let names = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + rows.push((ids.value(i), names.value(i).to_string(), values.value(i))); + } + } + rows.sort_by_key(|r| r.0); + rows +} + +#[tokio::test] +async fn test_cow_merge_update_matched_rows() { + let (_tmp, sql_context) = setup_append_only_table("t_cow_upd").await; + + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.test_db.source (id, name) VALUES (1, 'ALICE'), (3, 'CHARLIE')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let merge = parse_merge( + "MERGE INTO paimon.test_db.t_cow_upd t USING paimon.test_db.source s ON t.id = s.id \ + WHEN MATCHED THEN UPDATE SET NAME = s.name", + ); + execute_merge_into(&sql_context, &merge).await.unwrap(); + + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_cow_upd ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let rows = collect_rows(&batches); + assert_eq!( + rows, + vec![ + (1, "ALICE".to_string(), 10), + (2, "bob".to_string(), 20), + (3, "CHARLIE".to_string(), 30), + ] + ); +} + +#[tokio::test] +async fn test_cow_merge_delete_matched_rows() { + let (_tmp, sql_context) = setup_append_only_table("t_cow_del").await; + + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT)") + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.test_db.source (id) VALUES (2)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let merge = parse_merge( + "MERGE INTO paimon.test_db.t_cow_del t USING paimon.test_db.source s ON t.id = s.id \ + WHEN MATCHED THEN DELETE", + ); + execute_merge_into(&sql_context, &merge).await.unwrap(); + + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_cow_del ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let rows = collect_rows(&batches); + assert_eq!( + rows, + vec![(1, "alice".to_string(), 10), (3, "charlie".to_string(), 30),] + ); +} + +#[tokio::test] +async fn test_cow_merge_insert_not_matched() { + let (_tmp, sql_context) = setup_append_only_table("t_cow_ins").await; + + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR, value INT)") + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.test_db.source VALUES (4, 'dave', 40), (5, 'eve', 50)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let merge = parse_merge( + "MERGE INTO paimon.test_db.t_cow_ins t USING paimon.test_db.source s ON t.id = s.id \ + WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)", + ); + execute_merge_into(&sql_context, &merge).await.unwrap(); + + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_cow_ins ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let rows = collect_rows(&batches); + assert_eq!( + rows, + vec![ + (1, "alice".to_string(), 10), + (2, "bob".to_string(), 20), + (3, "charlie".to_string(), 30), + (4, "dave".to_string(), 40), + (5, "eve".to_string(), 50), + ] + ); +} + +#[tokio::test] +async fn test_cow_merge_insert_not_matched_without_columns() { + let (_tmp, sql_context) = setup_append_only_table("t_cow_ins_nocols").await; + + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR, value INT)") + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.test_db.source VALUES (4, 'dave', 40), (5, 'eve', 50)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + // No column list: the VALUES map to the table's columns by position. + // Before the fix this dropped the source values and inserted all-NULL rows. + let merge = parse_merge( + "MERGE INTO paimon.test_db.t_cow_ins_nocols t USING paimon.test_db.source s \ + ON t.id = s.id \ + WHEN NOT MATCHED THEN INSERT VALUES (s.id, s.name, s.value)", + ); + execute_merge_into(&sql_context, &merge).await.unwrap(); + + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_cow_ins_nocols ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let rows = collect_rows(&batches); + assert_eq!( + rows, + vec![ + (1, "alice".to_string(), 10), + (2, "bob".to_string(), 20), + (3, "charlie".to_string(), 30), + (4, "dave".to_string(), 40), + (5, "eve".to_string(), 50), + ] + ); +} + +#[tokio::test] +async fn test_cow_merge_insert_arity_validated_even_when_all_matched() { + let (_tmp, sql_context) = setup_append_only_table("t_cow_ins_arity").await; + + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR, value INT)") + .await + .unwrap(); + // id = 1 matches the target, so the NOT MATCHED INSERT branch is never + // built — the arity check must still run upfront. + sql_context + .sql("INSERT INTO paimon.test_db.source VALUES (1, 'ALICE', 99)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + // The column-less INSERT lists 2 values for a 3-column table. Even with no + // row to insert, the invalid arity must be rejected before execution so the + // matched UPDATE is not applied. + let merge = parse_merge( + "MERGE INTO paimon.test_db.t_cow_ins_arity t USING paimon.test_db.source s \ + ON t.id = s.id \ + WHEN MATCHED THEN UPDATE SET value = s.value \ + WHEN NOT MATCHED THEN INSERT VALUES (s.id, s.name)", + ); + let err = execute_merge_into(&sql_context, &merge).await.unwrap_err(); + assert!(err.to_string().contains("column(s) but"), "got {err}"); + + // The target is unchanged: the matched UPDATE did not apply. + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_cow_ins_arity ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + collect_rows(&batches), + vec![ + (1, "alice".to_string(), 10), + (2, "bob".to_string(), 20), + (3, "charlie".to_string(), 30), + ] + ); +} + +#[tokio::test] +async fn test_cow_merge_update_and_insert() { + let (_tmp, sql_context) = setup_append_only_table("t_cow_upsert").await; + + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR, value INT)") + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.test_db.source VALUES (2, 'BOB', 200), (4, 'dave', 40)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let merge = parse_merge( + "MERGE INTO paimon.test_db.t_cow_upsert t USING paimon.test_db.source s ON t.id = s.id \ + WHEN MATCHED THEN UPDATE SET name = s.name, value = s.value \ + WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, s.name, s.value)", + ); + execute_merge_into(&sql_context, &merge).await.unwrap(); + + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_cow_upsert ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let rows = collect_rows(&batches); + assert_eq!( + rows, + vec![ + (1, "alice".to_string(), 10), + (2, "BOB".to_string(), 200), + (3, "charlie".to_string(), 30), + (4, "dave".to_string(), 40), + ] + ); +} + +#[tokio::test] +async fn test_cow_merge_matched_null_predicates_fall_through() { + for delete_fallback in [false, true] { + let (_tmp, sql_context) = setup_append_only_table("t_null_matched").await; + for sql in [ + "CREATE TABLE paimon.test_db.source (id INT, flag INT)", + "INSERT INTO paimon.test_db.source VALUES (1, NULL), (2, 1), (3, 0)", + ] { + sql_context.sql(sql).await.unwrap().collect().await.unwrap(); + } + + let fallback = if delete_fallback { + "DELETE" + } else { + "UPDATE SET name = 'fallback'" + }; + let merge = parse_merge(&format!( + "MERGE INTO paimon.test_db.t_null_matched t USING paimon.test_db.source s \ + ON t.id = s.id \ + WHEN MATCHED AND s.flag = 1 THEN UPDATE SET name = 'first' \ + WHEN MATCHED AND s.flag = 0 THEN UPDATE SET name = 'second' \ + WHEN MATCHED THEN {fallback}" + )); + let result = execute_merge_into(&sql_context, &merge) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 3, + ); + + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_null_matched ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + let mut expected = vec![(2, "first".to_string(), 20), (3, "second".to_string(), 30)]; + if !delete_fallback { + expected.insert(0, (1, "fallback".to_string(), 10)); + } + assert_eq!(collect_rows(&batches), expected); + } +} + +#[tokio::test] +async fn test_merge_insert_null_predicates_fall_through() { + for data_evolution in [false, true] { + let (_tmp, sql_context) = if data_evolution { + setup_data_evolution_table("t_null_insert").await + } else { + setup_append_only_table("t_null_insert").await + }; + for sql in [ + "CREATE TABLE paimon.test_db.source (id INT, flag INT)", + "INSERT INTO paimon.test_db.source VALUES (4, NULL), (5, 1), (6, 0)", + ] { + sql_context.sql(sql).await.unwrap().collect().await.unwrap(); + } + + let merge = parse_merge( + "MERGE INTO paimon.test_db.t_null_insert t USING paimon.test_db.source s \ + ON t.id = s.id \ + WHEN NOT MATCHED AND s.flag = 1 THEN INSERT (id, name, value) VALUES (s.id, 'first', 100) \ + WHEN NOT MATCHED AND s.flag = 0 THEN INSERT (id, name, value) VALUES (s.id, 'second', 100) \ + WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, 'fallback', 100)", + ); + let result = execute_merge_into(&sql_context, &merge) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 3, + ); + + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_null_insert ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + collect_rows(&batches), + vec![ + (1, "alice".to_string(), 10), + (2, "bob".to_string(), 20), + (3, "charlie".to_string(), 30), + (4, "fallback".to_string(), 100), + (5, "first".to_string(), 100), + (6, "second".to_string(), 100), + ], + ); + } +} + +#[tokio::test] +async fn test_cow_merge_matched_or_predicates_preserve_clause_order() { + for delete_second in [false, true] { + let (_tmp, sql_context) = setup_append_only_table("t_or_matched").await; + for sql in [ + "CREATE TABLE paimon.test_db.source (id INT, flag INT)", + "INSERT INTO paimon.test_db.source VALUES (1, 1), (2, 0), (3, NULL)", + ] { + sql_context.sql(sql).await.unwrap().collect().await.unwrap(); + } + + let second_action = if delete_second { + "DELETE" + } else { + "UPDATE SET name = 'second'" + }; + let merge = parse_merge(&format!( + "MERGE INTO paimon.test_db.t_or_matched t USING paimon.test_db.source s \ + ON t.id = s.id \ + WHEN MATCHED AND s.flag = 1 THEN UPDATE SET name = 'first' \ + WHEN MATCHED AND s.flag = 0 OR s.flag = 1 THEN {second_action} \ + WHEN MATCHED THEN UPDATE SET name = 'fallback'" + )); + let result = execute_merge_into(&sql_context, &merge) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 3, + ); + + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_or_matched ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + let mut expected = vec![ + (1, "first".to_string(), 10), + (3, "fallback".to_string(), 30), + ]; + if !delete_second { + expected.insert(1, (2, "second".to_string(), 20)); + } + assert_eq!(collect_rows(&batches), expected); + } +} + +#[tokio::test] +async fn test_merge_insert_or_predicates_preserve_clause_order() { + for data_evolution in [false, true] { + let (_tmp, sql_context) = if data_evolution { + setup_data_evolution_table("t_or_insert").await + } else { + setup_append_only_table("t_or_insert").await + }; + for sql in [ + "CREATE TABLE paimon.test_db.source (id INT, flag INT)", + "INSERT INTO paimon.test_db.source VALUES (4, 1), (5, 0), (6, NULL)", + ] { + sql_context.sql(sql).await.unwrap().collect().await.unwrap(); + } + + let merge = parse_merge( + "MERGE INTO paimon.test_db.t_or_insert t USING paimon.test_db.source s \ + ON t.id = s.id \ + WHEN NOT MATCHED AND s.flag = 1 THEN INSERT (id, name, value) VALUES (s.id, 'first', 100) \ + WHEN NOT MATCHED AND s.flag = 0 OR s.flag = 1 THEN INSERT (id, name, value) VALUES (s.id, 'second', 100) \ + WHEN NOT MATCHED THEN INSERT (id, name, value) VALUES (s.id, 'fallback', 100)", + ); + let result = execute_merge_into(&sql_context, &merge) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 3, + ); + + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.t_or_insert ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + collect_rows(&batches), + vec![ + (1, "alice".to_string(), 10), + (2, "bob".to_string(), 20), + (3, "charlie".to_string(), 30), + (4, "first".to_string(), 100), + (5, "second".to_string(), 100), + (6, "fallback".to_string(), 100), + ], + ); + } +} + +#[tokio::test] +async fn test_cow_merge_preserves_string_literals() { + for (on_predicate, matched_predicate, insert_predicate) in [ + ("", "", ""), + (" AND s.name = 'paimon.test_db.source.payload'", "", ""), + ("", " AND s.name = 'paimon.test_db.source.payload'", ""), + ("", "", " AND s.name = 'paimon.test_db.source.payload'"), + ] { + let (_tmp, sql_context) = setup_append_only_table("target").await; + for sql in [ + "CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)", + "INSERT INTO paimon.test_db.source VALUES (1, 'paimon.test_db.source.payload'), (4, 'paimon.test_db.source.payload')", + ] { + sql_context.sql(sql).await.unwrap().collect().await.unwrap(); + } + + let merge = parse_merge(&format!( + "MERGE INTO paimon.test_db.target t USING paimon.test_db.source s \ + ON t.id = s.id{on_predicate} \ + WHEN MATCHED{matched_predicate} THEN UPDATE SET name = 'paimon.test_db.target.payload' \ + WHEN NOT MATCHED{insert_predicate} THEN INSERT (id, name, value) VALUES (s.id, 'paimon.test_db.source.payload', 40)" + )); + let result = execute_merge_into(&sql_context, &merge) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2, + ); + + let batches = sql_context + .sql("SELECT id, name, value FROM paimon.test_db.target ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + collect_rows(&batches), + vec![ + (1, "paimon.test_db.target.payload".to_string(), 10), + (2, "bob".to_string(), 20), + (3, "charlie".to_string(), 30), + (4, "paimon.test_db.source.payload".to_string(), 40), + ], + "ON: {on_predicate}, MATCHED: {matched_predicate}, INSERT: {insert_predicate}", + ); + } +} + +#[tokio::test] +async fn test_cow_merge_no_matches() { + let (_tmp, sql_context) = setup_append_only_table("t_cow_nomatch").await; + + sql_context + .sql("CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let merge = parse_merge( + "MERGE INTO paimon.test_db.t_cow_nomatch t USING paimon.test_db.source s ON t.id = s.id \ + WHEN MATCHED THEN UPDATE SET name = s.name", + ); + let result = execute_merge_into(&sql_context, &merge).await.unwrap(); + let batches = result.collect().await.unwrap(); + let count = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, 0); +} + +#[test] +fn source_partition_pruning_requires_a_conjunctive_target_source_equality() { + let target = column_name(1, 2); + let source = column_name(2, 2); + for on in [ + col(&target).eq(col(&source)), + col(&source).eq(col(&target)), + col(&target) + .eq(col(&source)) + .and(col(column_name(0, 2)).gt(lit(0))), + ] { + assert_eq!(partition_source_index(&on, &target, 2, 4), Some(0)); + } + for on in [ + col(&target) + .eq(col(&source)) + .or(col(column_name(0, 2)).eq(lit(0))), + col(&target).eq(col(column_name(0, 2))), + col(column_name(0, 2)).eq(col(&source)), + ] { + assert_eq!(partition_source_index(&on, &target, 2, 4), None); + } +} diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index 2ee74f597..44fd89308 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -1182,6 +1182,83 @@ impl ExecutionPlan for PaimonTableScan { vec![] } + fn try_swapping_with_projection( + &self, + projection: &datafusion::physical_plan::projection::ProjectionExec, + ) -> DFResult>> { + use datafusion::physical_plan::projection::ProjectionExec; + let Some(indices) = projection + .expr() + .iter() + .map(|expr| { + expr.expr + .downcast_ref::() + .map(|column| column.index()) + }) + .collect::>>() + else { + return Ok(None); + }; + let mut required = indices.clone(); + for filter in self.runtime_filters.iter().chain(&self.decoder_filters) { + required.extend(collect_columns(filter).iter().map(|column| column.index())); + } + required.sort_unstable(); + required.dedup(); + if required.len() == self.read_type.len() { + return Ok(None); + } + + // Keep predicate dependencies in the reader schema. A projection above + // the scan removes them only after filtering, and applies any aliases. + let schema = self.schema(); + let reader_projection = ProjectionExec::try_new( + required.iter().map(|index| { + ( + Arc::new(Column::new(schema.field(*index).name(), *index)) + as Arc, + schema.field(*index).name().clone(), + ) + }), + Arc::new(self.clone()), + )?; + let mut scan = self.clone(); + scan.read_type = required + .iter() + .map(|index| self.read_type[*index].clone()) + .collect(); + scan.plan_properties = reader_projection.properties().clone(); + scan.runtime_filters = self + .runtime_filters + .iter() + .map(|filter| reassign_expr_columns(filter.clone(), scan.schema().as_ref())) + .collect::>>()?; + scan.decoder_filters = self + .decoder_filters + .iter() + .map(|filter| reassign_expr_columns(filter.clone(), scan.schema().as_ref())) + .collect::>>()?; + let output = projection + .expr() + .iter() + .zip(indices) + .map(|(expr, index)| { + let index = required + .binary_search(&index) + .expect("projection field retained"); + ( + Arc::new(Column::new(scan.schema().field(index).name(), index)) + as Arc, + expr.alias.clone(), + ) + }) + .collect::>(); + Ok(Some(Arc::new(ProjectionExec::try_new( + output, + Arc::new(scan), + )?))) + } + fn apply_expressions( &self, f: &mut dyn FnMut(&Arc) -> DFResult, diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 0adf576eb..99b92af7c 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -659,10 +659,7 @@ impl SQLContext { ) .await } - Statement::Merge(merge) => { - self.handle_merge_into(merge, enable_ident_normalization) - .await - } + Statement::Merge(merge) => self.handle_merge_into(merge, options).await, Statement::Update(update) => { self.handle_update(update, enable_ident_normalization).await } @@ -836,6 +833,15 @@ impl SQLContext { .await } Statement::Query(_) | Statement::Explain { .. } => { + let mut statement = statements[0].clone(); + if let Statement::Explain { + statement: inner, .. + } = &mut statement + { + if let Statement::Merge(merge) = inner.as_mut() { + *merge = self.normalize_merge_into(merge).await?; + } + } let current_catalog = self.current_catalog_name(); let current_database = self .ctx @@ -846,7 +852,7 @@ impl SQLContext { .clone(); let expanded = crate::sql_function::expand_statement( &self.ctx.state(), - statements[0].clone(), + statement, &self.catalogs, ¤t_catalog, ¤t_database, @@ -1799,11 +1805,12 @@ impl SQLContext { Ok(()) } - async fn handle_merge_into( - &self, - merge: &Merge, - enable_ident_normalization: bool, - ) -> DFResult { + async fn handle_merge_into(&self, merge: &Merge, options: SQLOptions) -> DFResult { + let merge = self.normalize_merge_into(merge).await?; + self.ctx.sql_with_options(&merge.to_string(), options).await + } + + async fn normalize_merge_into(&self, merge: &Merge) -> DFResult { self.ensure_no_time_travel_for_write("MERGE INTO")?; let table_name = match &merge.table { TableFactor::Table { name, .. } => name.clone(), @@ -1814,14 +1821,49 @@ impl SQLContext { } }; Self::ensure_main_branch_write_target(&table_name, "MERGE INTO")?; - let (catalog, _catalog_name, identifier) = self.resolve_catalog_and_table(&table_name)?; - - let table = catalog - .get_table(&identifier) - .await - .map_err(to_datafusion_error)?; - - crate::merge_into::execute_merge_into(self, merge, table, enable_ident_normalization).await + // Use the configured native dialect for local expression scopes; the + // front parser treats lambda arrows as ordinary binary operators. + let state = self.ctx.state(); + let parsed = state.sql_to_statement( + &merge.to_string(), + &state.config_options().sql_parser.dialect, + )?; + let datafusion::sql::parser::Statement::Statement(statement) = parsed else { + return Err(DataFusionError::Plan( + "Expected a native MERGE statement".into(), + )); + }; + let Statement::Merge(mut merge) = *statement else { + return Err(DataFusionError::Plan( + "Expected a native MERGE statement".into(), + )); + }; + crate::merge_into::qualify_not_matched_exprs( + &mut merge, + self.ctx.enable_ident_normalization(), + ); + // `_ROW_ID` is a query-visible virtual column, but positional INSERTs + // still map to the persisted user columns. Make that list explicit + // before handing the statement to DataFusion's native planner. + if merge.clauses.iter().any(|clause| matches!( + &clause.action, + datafusion::sql::sqlparser::ast::MergeAction::Insert(insert) if insert.columns.is_empty() + )) { + let (catalog, _, identifier) = self.resolve_catalog_and_table(&table_name)?; + let table = catalog.get_table(&identifier).await.map_err(to_datafusion_error)?; + if table.schema().core_options().data_evolution_enabled() { + for clause in &mut merge.clauses { + if let datafusion::sql::sqlparser::ast::MergeAction::Insert(insert) = &mut clause.action { + if insert.columns.is_empty() { + insert.columns = table.schema().fields().iter().map(|field| { + ObjectName::from(vec![Ident::with_quote('"', field.name())]) + }).collect(); + } + } + } + } + } + Ok(merge) } async fn handle_update( diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 730cd30bb..6d9219d10 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -25,10 +25,11 @@ use datafusion::arrow::datatypes::{ DataType as ArrowDataType, Field, Schema, SchemaRef as ArrowSchemaRef, }; use datafusion::catalog::Session; +use datafusion::common::DFSchemaRef; use datafusion::datasource::sink::DataSinkExec; use datafusion::datasource::{TableProvider, TableType}; use datafusion::error::{DataFusionError, Result as DFResult}; -use datafusion::logical_expr::dml::InsertOp; +use datafusion::logical_expr::dml::{InsertOp, MergeIntoClause}; use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; use datafusion::physical_plan::ExecutionPlan; use paimon::spec::{ @@ -113,7 +114,7 @@ fn query_auth_output_schema(table: &Table, schema: ArrowSchemaRef) -> ArrowSchem )) } -/// Read-only table provider for a Paimon table. +/// Table provider for Paimon scans, INSERT, and native MERGE execution. /// /// Supports full table scan, column projection, and predicate pushdown for /// planning. Partition predicates prune splits eagerly, while supported @@ -533,6 +534,17 @@ impl TableProvider for PaimonTableProvider { Ok(Arc::new(DataSinkExec::new(input, Arc::new(sink), None))) } + async fn merge_into( + &self, + state: &dyn Session, + source: Arc, + merge_schema: DFSchemaRef, + on: Expr, + clauses: Vec, + ) -> DFResult> { + crate::merge_into::plan_merge(state, self.table.clone(), source, merge_schema, on, clauses) + } + fn supports_filters_pushdown( &self, filters: &[&Expr], diff --git a/crates/integrations/datafusion/tests/blob_tests.rs b/crates/integrations/datafusion/tests/blob_tests.rs index 3c5f44a30..f9d3bb4c1 100644 --- a/crates/integrations/datafusion/tests/blob_tests.rs +++ b/crates/integrations/datafusion/tests/blob_tests.rs @@ -565,6 +565,10 @@ async fn test_merge_into_updates_non_blob_on_descriptor_table() { ) .await; + // Updating name must not resolve unrelated external BLOB payloads. + std::fs::remove_file(tmp.path().join("first.bin")).unwrap(); + std::fs::remove_file(tmp.path().join("second.bin")).unwrap(); + exec( &sql_context, "CREATE TEMPORARY TABLE paimon.test_db.src AS SELECT * FROM (VALUES (1, 'Updated')) AS t(id, name)", diff --git a/crates/integrations/datafusion/tests/merge_into_tests.rs b/crates/integrations/datafusion/tests/merge_into_tests.rs index 6a21c0d57..1a79cb3b0 100644 --- a/crates/integrations/datafusion/tests/merge_into_tests.rs +++ b/crates/integrations/datafusion/tests/merge_into_tests.rs @@ -133,12 +133,11 @@ async fn collect_row_ids(sql_context: &SQLContext, sql: &str) -> Vec<(i64, i32)> } async fn assert_merge_error(sql_context: &SQLContext, sql: &str, expected_substring: &str) { - let result = sql_context.sql(sql).await; - assert!( - result.is_err(), - "Expected error containing '{expected_substring}', but got Ok" - ); - let err_msg = result.unwrap_err().to_string(); + let result = match sql_context.sql(sql).await { + Ok(df) => df.collect().await, + Err(err) => Err(err), + }; + let err_msg = result.expect_err("Expected MERGE to fail").to_string(); assert!( err_msg.contains(expected_substring), "Error message '{err_msg}' does not contain '{expected_substring}'" @@ -269,7 +268,7 @@ async fn test_merge_insert_rejects_unknown_and_duplicate_columns() { .unwrap(); for (columns, expected_error) in [ - ("id, missing", "Unknown column 'missing' in MERGE INSERT"), + ("id, missing", "No field named missing"), ("id, ID", "Duplicate column 'id' in MERGE INSERT"), ] { let sql = format!( @@ -320,8 +319,11 @@ async fn test_merge_update_rejects_unknown_and_duplicate_target_columns() { .unwrap(); for (assignments, expected_error) in [ - ("\"VALUE\" = 3", "Unknown column 'VALUE' in UPDATE"), - ("value = 3, VALUE = 4", "Duplicate column 'value' in UPDATE"), + ("\"VALUE\" = 3", "No field named \"VALUE\""), + ( + "value = 3, VALUE = 4", + "Duplicate column 'value' in MERGE UPDATE", + ), ] { let sql = format!( "MERGE INTO paimon.test_db.{table_name} t diff --git a/crates/integrations/datafusion/tests/native_merge.rs b/crates/integrations/datafusion/tests/native_merge.rs new file mode 100644 index 000000000..bdd0ac4b0 --- /dev/null +++ b/crates/integrations/datafusion/tests/native_merge.rs @@ -0,0 +1,709 @@ +// 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. + +mod common; + +use arrow_array::UInt64Array; +use common::{collect_id_name, create_sql_context, create_test_env, exec}; +use datafusion::logical_expr::{dml::WriteOp, LogicalPlan}; + +#[tokio::test] +async fn native_merge_defers_mutation_until_execution() { + for options in [ + "", + "WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true')", + ] { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + &format!("CREATE TABLE paimon.test_db.target (id INT, name STRING) {options}"), + ) + .await; + exec( + &ctx, + "INSERT INTO paimon.test_db.target (id, name) VALUES (1, 'old'), (2, 'keep')", + ) + .await; + let sql = "MERGE INTO paimon.test_db.target AS t \ + USING (VALUES (1, 'new'), (3, 'insert')) AS s(id, name) ON t.id = s.id \ + WHEN MATCHED THEN UPDATE SET name = s.name \ + WHEN NOT MATCHED THEN INSERT (id, name) VALUES (s.id, s.name)"; + let before = vec![(1, "old".to_string()), (2, "keep".to_string())]; + let query = "SELECT id, name FROM paimon.test_db.target ORDER BY id"; + + let df = ctx.sql(sql).await.unwrap(); + assert!( + matches!(df.logical_plan(), LogicalPlan::Dml(dml) if matches!(dml.op, WriteOp::MergeInto(_))) + ); + assert_eq!(collect_id_name(&ctx, query).await, before); + df.clone().create_physical_plan().await.unwrap(); + assert_eq!(collect_id_name(&ctx, query).await, before); + exec(&ctx, &format!("EXPLAIN {sql}")).await; + assert_eq!(collect_id_name(&ctx, query).await, before); + + let result = df.collect().await.unwrap(); + assert_eq!( + result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2 + ); + assert_eq!( + collect_id_name(&ctx, query).await, + vec![(1, "new".into()), (2, "keep".into()), (3, "insert".into())] + ); + } +} + +#[tokio::test] +async fn native_session_can_merge_registered_paimon_provider() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + "CREATE TABLE paimon.test_db.target (id INT, name STRING)", + ) + .await; + exec(&ctx, "INSERT INTO paimon.test_db.target VALUES (1, 'old')").await; + // Bypass SQLContext: execution must be implemented by the table provider. + ctx.ctx() + .sql( + "MERGE INTO paimon.test_db.target AS t \ + USING (VALUES (1, 'new')) AS s(id, name) ON t.id = s.id \ + WHEN MATCHED THEN UPDATE SET name = s.name", + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + collect_id_name(&ctx, "SELECT id, name FROM paimon.test_db.target").await, + vec![(1, "new".into())] + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn native_merge_retries_with_fresh_target_and_one_source_capture() { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + use arrow_schema::DataType; + use datafusion::logical_expr::{create_udf, Volatility}; + + for (options, competing_sql, expected_value) in [ + ("", "UPDATE paimon.test_db.target SET value = 100 WHERE id = 1", 101), + ("WITH ('row-tracking.enabled' = 'true')", "UPDATE paimon.test_db.target SET value = 100 WHERE id = 1", 101), + ("WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true', 'write-only' = 'true', 'compaction.min.file-num' = '2')", "CALL sys.compact(table => 'test_db.target')", 11), + ] { + let (_tmp, catalog) = create_test_env(); + let ctx = Arc::new(create_sql_context(catalog).await); + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + &format!("CREATE TABLE paimon.test_db.target (id INT, value INT) {options}"), + ) + .await; + exec( + &ctx, + "INSERT INTO paimon.test_db.target (id, value) VALUES (1, 10)", + ) + .await; + exec(&ctx, "INSERT INTO paimon.test_db.target (id, value) VALUES (3, 30)").await; + + let source_calls = Arc::new(AtomicUsize::new(0)); + let calls = source_calls.clone(); + ctx.ctx().register_udf(create_udf( + "capture_source", + vec![DataType::Int32], + DataType::Int32, + Volatility::Volatile, + Arc::new(move |args| { + calls.fetch_add(1, Ordering::SeqCst); + Ok(args[0].clone()) + }), + )); + let target_calls = Arc::new(AtomicUsize::new(0)); + let calls = target_calls.clone(); + let weak = Arc::downgrade(&ctx); + ctx.ctx().register_udf(create_udf( + "interfere", + vec![DataType::Int32], + DataType::Int32, + Volatility::Volatile, + Arc::new(move |args| { + if calls.fetch_add(1, Ordering::SeqCst) == 0 { + let ctx = weak.upgrade().unwrap(); + // Publish a competing change after matching and before the MERGE + // commit. No sleeps or scheduler-dependent race is needed. + tokio::task::block_in_place(|| { + tokio::runtime::Handle::current().block_on(async { + ctx.sql(competing_sql) + .await? + .collect() + .await + }) + })?; + } + Ok(args[0].clone()) + }), + )); + + let count = common::dml_count(&ctx, + "MERGE INTO paimon.test_db.target AS t \ + USING (SELECT capture_source(id) AS id FROM (VALUES (CAST(1 AS INT)), (CAST(2 AS INT))) AS v(id)) AS s \ + ON t.id = s.id \ + WHEN MATCHED THEN UPDATE SET value = interfere(CAST(t.value + 1 AS INT)) \ + WHEN NOT MATCHED THEN INSERT (id, value) VALUES (s.id, 20)").await; + assert_eq!(count, 2); + assert_eq!( + common::collect_id_value(&ctx, "SELECT id, value FROM paimon.test_db.target").await, + vec![(1, expected_value), (2, 20), (3, 30)] + ); + assert_eq!( + source_calls.load(Ordering::SeqCst), + 1, + "source must be replayed on retry" + ); + assert_eq!( + target_calls.load(Ordering::SeqCst), + 2, + "the conflict must cause exactly one retry" + ); + } +} + +#[tokio::test] +async fn native_merge_duplicate_matches_do_not_publish_any_actions() { + for options in [ + "", + "WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true')", + ] { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + &format!("CREATE TABLE paimon.test_db.target (id INT, name STRING) {options}"), + ) + .await; + exec( + &ctx, + "INSERT INTO paimon.test_db.target (id, name) VALUES (1, 'old')", + ) + .await; + common::assert_sql_error( + &ctx, + "MERGE INTO paimon.test_db.target AS t \ + USING (VALUES (1, 'a'), (1, 'b'), (2, 'insert')) AS s(id, name) ON t.id = s.id \ + WHEN MATCHED THEN UPDATE SET name = s.name \ + WHEN NOT MATCHED THEN INSERT (id, name) VALUES (s.id, s.name)", + "duplicate UPDATE", + ) + .await; + assert_eq!( + collect_id_name(&ctx, "SELECT id, name FROM paimon.test_db.target").await, + vec![(1, "old".into())] + ); + } +} + +#[tokio::test] +async fn native_merge_respects_dml_options() { + use datafusion::execution::context::SQLOptions; + + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + "CREATE TABLE paimon.test_db.target (id INT, name STRING)", + ) + .await; + let result = ctx + .sql_with_options( + "MERGE INTO paimon.test_db.target AS t \ + USING (VALUES (1, 'new')) AS s(id, name) ON t.id = s.id \ + WHEN NOT MATCHED THEN INSERT (id, name) VALUES (s.id, s.name)", + SQLOptions::new().with_allow_dml(false), + ) + .await; + assert!(result.is_err()); + assert!( + collect_id_name(&ctx, "SELECT id, name FROM paimon.test_db.target") + .await + .is_empty() + ); +} + +#[tokio::test] +async fn data_evolution_positional_insert_and_explain_exclude_virtual_row_id() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec(&ctx, "CREATE TABLE paimon.test_db.target (id INT, name STRING) WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true')").await; + let sql = "MERGE INTO paimon.test_db.target AS t \ + USING (VALUES (1, 's.name')) AS s(id, name) ON t.id = s.id \ + WHEN NOT MATCHED AND id = 1 THEN INSERT VALUES (id, name)"; + exec(&ctx, &format!("EXPLAIN {sql}")).await; + assert!( + collect_id_name(&ctx, "SELECT id, name FROM paimon.test_db.target") + .await + .is_empty() + ); + assert_eq!(common::dml_count(&ctx, sql).await, 1); + assert_eq!( + collect_id_name(&ctx, "SELECT id, name FROM paimon.test_db.target").await, + vec![(1, "s.name".into())] + ); + let batches = ctx + .sql("SELECT \"_ROW_ID\" FROM paimon.test_db.target") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 0 + ); +} + +#[tokio::test] +async fn cow_tracking_columns_do_not_shadow_user_columns() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec(&ctx, "CREATE TABLE paimon.test_db.target (id INT, value INT, \"__paimon_file_idx\" INT, \"__paimon_row_offset\" INT) WITH ('write-only' = 'true')").await; + exec( + &ctx, + "INSERT INTO paimon.test_db.target VALUES (1, 10, 99, 98)", + ) + .await; + exec( + &ctx, + "INSERT INTO paimon.test_db.target VALUES (2, 20, 97, 96)", + ) + .await; + assert_eq!( + common::dml_count( + &ctx, + "MERGE INTO paimon.test_db.target t USING (VALUES (1, 100)) AS s(id, value) \ + ON t.id = s.id WHEN MATCHED THEN UPDATE SET value = s.value" + ) + .await, + 1 + ); + assert_eq!( + common::collect_id_value(&ctx, "SELECT id, value FROM paimon.test_db.target").await, + vec![(1, 100), (2, 20)] + ); + let batches = ctx.sql("SELECT id, \"__paimon_file_idx\", \"__paimon_row_offset\" FROM paimon.test_db.target ORDER BY id") + .await.unwrap().collect().await.unwrap(); + let mut rows = common::collect_three_ints(&batches); + rows.sort_unstable(); + assert_eq!(rows, vec![(1, 99, 98), (2, 97, 96)]); +} + +#[tokio::test] +async fn native_merge_retains_session_variable_providers() { + use datafusion::arrow::datatypes::DataType; + use datafusion::common::ScalarValue; + use datafusion::variable::{VarProvider, VarType}; + use std::sync::Arc; + + #[derive(Debug)] + struct Variables; + impl VarProvider for Variables { + fn get_value(&self, _names: Vec) -> datafusion::error::Result { + Ok(ScalarValue::Int32(Some(7))) + } + fn get_type(&self, _names: &[String]) -> Option { + Some(DataType::Int32) + } + } + for options in [ + "", + "WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true')", + ] { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + ctx.ctx() + .register_variable(VarType::UserDefined, Arc::new(Variables)); + ctx.ctx() + .register_variable(VarType::System, Arc::new(Variables)); + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + &format!("CREATE TABLE paimon.test_db.target (id INT, value INT) {options}"), + ) + .await; + exec( + &ctx, + "INSERT INTO paimon.test_db.target (id, value) VALUES (1, 10)", + ) + .await; + // Data-evolution tables only support unconditional matched clauses. + let matched_predicate = if options.is_empty() { + " AND @delta = 7" + } else { + "" + }; + assert_eq!( + common::dml_count( + &ctx, + &format!( + "MERGE INTO paimon.test_db.target t USING (VALUES (1)) AS s(id) \ + ON t.id = s.id AND @@flag = 7 \ + WHEN MATCHED{matched_predicate} THEN UPDATE SET value = @delta" + ) + ) + .await, + 1 + ); + assert_eq!( + common::dml_count( + &ctx, + "MERGE INTO paimon.test_db.target t USING (VALUES (1), (2)) AS s(id) \ + ON t.id = s.id AND @delta = 7 \ + WHEN MATCHED THEN UPDATE SET value = @delta \ + WHEN NOT MATCHED AND @delta = 7 THEN INSERT (id, value) VALUES (s.id, @delta)" + ) + .await, + 2 + ); + assert_eq!( + common::collect_id_value(&ctx, "SELECT id, value FROM paimon.test_db.target").await, + vec![(1, 7), (2, 7)] + ); + if options.is_empty() { + assert_eq!( + common::dml_count( + &ctx, + "MERGE INTO paimon.test_db.target t USING (VALUES (1)) AS s(id) \ + ON t.id = s.id WHEN MATCHED AND @delta <> 7 THEN UPDATE SET value = 0" + ) + .await, + 0 + ); + assert_eq!( + common::collect_id_value(&ctx, "SELECT id, value FROM paimon.test_db.target").await, + vec![(1, 7), (2, 7)] + ); + } + } +} + +#[tokio::test] +async fn merge_insert_resolves_unqualified_nested_fields_against_source() { + for options in [ + "", + "WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true')", + ] { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec(&ctx, &format!("CREATE TABLE paimon.test_db.target (id INT, value INT, payload STRUCT) {options}")).await; + assert_eq!( + common::dml_count( + &ctx, + "MERGE INTO paimon.test_db.target t \ + USING (SELECT 1 AS id, named_struct('foo', 42) AS payload) s ON t.id = s.id \ + WHEN NOT MATCHED THEN INSERT (id, value) VALUES (id, payload.foo)" + ) + .await, + 1 + ); + assert_eq!( + common::collect_id_value(&ctx, "SELECT id, value FROM paimon.test_db.target").await, + vec![(1, 42)] + ); + } +} + +#[tokio::test] +async fn failed_merge_insert_discards_prepared_cow_files() { + fn parquet_files(path: &std::path::Path) -> std::collections::BTreeSet { + let mut files = std::collections::BTreeSet::new(); + for entry in std::fs::read_dir(path).unwrap() { + let path = entry.unwrap().path(); + if path.is_dir() { + files.extend(parquet_files(&path)); + } else if path.extension().is_some_and(|ext| ext == "parquet") { + files.insert(path); + } + } + files + } + let (tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + "CREATE TABLE paimon.test_db.target (id INT, value INT, pt STRING) PARTITIONED BY (pt)", + ) + .await; + exec( + &ctx, + "INSERT INTO paimon.test_db.target VALUES (1, 10, 'a')", + ) + .await; + // An unavailable destination affects only the new INSERT partition, after + // rewriting the existing partition has succeeded. + std::fs::write( + tmp.path().join("test_db.db/target/pt=b"), + b"blocked directory", + ) + .unwrap(); + let before = parquet_files(tmp.path()); + common::assert_sql_error( + &ctx, + "MERGE INTO paimon.test_db.target t \ + USING (VALUES (1, 100, 'a'), (2, 200, 'b')) AS s(id, value, pt) \ + ON t.id = s.id WHEN MATCHED THEN UPDATE SET value = s.value \ + WHEN NOT MATCHED THEN INSERT (id, value, pt) VALUES (s.id, s.value, s.pt)", + "Failed to create directory", + ) + .await; + assert_eq!( + common::collect_id_value(&ctx, "SELECT id, value FROM paimon.test_db.target").await, + vec![(1, 10)] + ); + assert_eq!(parquet_files(tmp.path()), before); +} + +#[tokio::test] +async fn merge_insert_preserves_lambda_scopes() { + for (dialect, expression, expected) in [ + ( + "spark", + "array_length(array_filter([1, 2, 3], x -> x > 1))", + 2, + ), + ( + "duckdb", + "array_element(array_transform(s.arr, x -> x + 1), 1)", + 2, + ), + ] { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec( + &ctx, + &format!("SET datafusion.sql_parser.dialect = '{dialect}'"), + ) + .await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + "CREATE TABLE paimon.test_db.target (id INT, value INT)", + ) + .await; + assert_eq!( + common::dml_count( + &ctx, + &format!( + "MERGE INTO paimon.test_db.target t \ + USING (SELECT 1 AS id, [1, 2, 3] AS arr) s ON t.id = s.id \ + WHEN NOT MATCHED THEN INSERT (id, value) VALUES (id, {expression})" + ) + ) + .await, + 1 + ); + assert_eq!( + common::collect_id_value(&ctx, "SELECT id, value FROM paimon.test_db.target").await, + vec![(1, expected)] + ); + } +} + +#[tokio::test] +async fn merge_unconditional_delete_collapses_duplicate_source_matches() { + for options in ["", "WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true', 'deletion-vectors.enabled' = 'true')"] { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec(&ctx, &format!("CREATE TABLE paimon.test_db.target (id INT, value INT) {options}")).await; + exec(&ctx, "INSERT INTO paimon.test_db.target (id, value) VALUES (1, 10), (3, 30)").await; + assert_eq!(common::dml_count(&ctx, + "MERGE INTO paimon.test_db.target t USING (VALUES (1), (1), (2)) AS s(id) \ + ON t.id = s.id WHEN MATCHED THEN DELETE \ + WHEN NOT MATCHED THEN INSERT (id, value) VALUES (s.id, 20)").await, 2); + assert_eq!(common::collect_id_value(&ctx, "SELECT id, value FROM paimon.test_db.target").await, vec![(2, 20), (3, 30)]); + } +} + +#[tokio::test] +async fn merge_does_not_evaluate_not_matched_predicate_on_matched_rows() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + "CREATE TABLE paimon.test_db.target (id INT, value INT)", + ) + .await; + exec(&ctx, "INSERT INTO paimon.test_db.target VALUES (1, 10)").await; + assert_eq!( + common::dml_count( + &ctx, + "MERGE INTO paimon.test_db.target t USING (VALUES (1, 0), (2, 1)) AS s(id, divisor) \ + ON t.id = s.id WHEN NOT MATCHED AND 10 / s.divisor > 1 \ + THEN INSERT (id, value) VALUES (s.id, 20)" + ) + .await, + 1 + ); + assert_eq!( + common::collect_id_value(&ctx, "SELECT id, value FROM paimon.test_db.target").await, + vec![(1, 10), (2, 20)] + ); +} + +#[tokio::test] +async fn merge_not_matched_rejects_target_column_dependencies() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + "CREATE TABLE paimon.test_db.target (id INT, value INT)", + ) + .await; + common::assert_sql_error( + &ctx, + "MERGE INTO paimon.test_db.target t USING (VALUES (1)) AS s(id) ON t.id = s.id \ + WHEN NOT MATCHED THEN INSERT (id, value) VALUES (s.id, t.value)", + "source columns", + ) + .await; +} + +#[tokio::test] +async fn native_merge_does_not_read_unused_source_blob_columns() { + use paimon::spec::BlobDescriptor; + for options in [ + "", + "WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true')", + ] { + let (tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + &format!("CREATE TABLE paimon.test_db.target (id INT, name STRING) {options}"), + ) + .await; + exec( + &ctx, + "INSERT INTO paimon.test_db.target (id, name) VALUES (1, 'old')", + ) + .await; + exec(&ctx, "CREATE TABLE paimon.test_db.source (id INT, name STRING, picture BLOB, marker INT) WITH ('data-evolution.enabled' = 'true', 'row-tracking.enabled' = 'true', 'blob-descriptor-field' = 'picture')").await; + let path = tmp.path().join("payload.bin"); + std::fs::write(&path, b"data").unwrap(); + let descriptor = BlobDescriptor::new(path.to_str().unwrap().to_string(), 0, 4).serialize(); + let hex = descriptor + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + exec( + &ctx, + &format!( + "INSERT INTO paimon.test_db.source (id, name, picture, marker) VALUES (1, 'new', X'{hex}', 1)" + ), + ) + .await; + std::fs::remove_file(path).unwrap(); + for (source, name) in [ + ("paimon.test_db.source", "name"), + ("(SELECT * FROM paimon.test_db.source WHERE id = 1)", "name"), + ( + "(SELECT * FROM paimon.test_db.source WHERE marker = 1)", + "name", + ), + ( + "(SELECT name AS renamed, id, picture FROM paimon.test_db.source WHERE marker = 1)", + "renamed", + ), + ] { + let result = ctx + .ctx() + .sql(&format!( + "MERGE INTO paimon.test_db.target t USING {source} s \ + ON t.id = s.id WHEN MATCHED THEN UPDATE SET name = s.{name}" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + result[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 1 + ); + } + assert_eq!( + collect_id_name(&ctx, "SELECT id, name FROM paimon.test_db.target").await, + vec![(1, "new".into())] + ); + } +} + +#[tokio::test] +async fn native_merge_preserves_source_cardinality_with_zero_column_projection() { + let (_tmp, catalog) = create_test_env(); + let ctx = create_sql_context(catalog).await; + exec(&ctx, "CREATE SCHEMA paimon.test_db").await; + exec( + &ctx, + "CREATE TABLE paimon.test_db.target (id INT, value INT)", + ) + .await; + exec(&ctx, "CREATE TABLE paimon.test_db.source (unused INT)").await; + exec(&ctx, "INSERT INTO paimon.test_db.source VALUES (1), (2)").await; + assert_eq!( + common::dml_count( + &ctx, + "MERGE INTO paimon.test_db.target t USING paimon.test_db.source s ON FALSE \ + WHEN NOT MATCHED THEN INSERT (id, value) VALUES (42, 7)" + ) + .await, + 2 + ); + assert_eq!( + common::collect_id_value(&ctx, "SELECT id, value FROM paimon.test_db.target").await, + vec![(42, 7), (42, 7)] + ); +} diff --git a/crates/integrations/datafusion/tests/procedures.rs b/crates/integrations/datafusion/tests/procedures.rs index d149d9712..e648aa7aa 100644 --- a/crates/integrations/datafusion/tests/procedures.rs +++ b/crates/integrations/datafusion/tests/procedures.rs @@ -59,6 +59,25 @@ async fn setup_table_with_snapshots() -> (tempfile::TempDir, paimon_datafusion:: (tmp, sql_context) } +const SNAPSHOT_TIMES: [u64; 3] = [1_700_000_000_000, 1_700_000_001_000, 1_700_000_002_000]; + +async fn setup_table_with_snapshot_times() -> (tempfile::TempDir, paimon_datafusion::SQLContext) { + let (tmp, sql_context) = setup_table_with_snapshots().await; + // Procedure arguments are epoch milliseconds. `$snapshots.commit_time` is + // a local TIMESTAMP, whose Arrow values include the system timezone offset. + // Explicit metadata times also keep boundary tests independent of the clock. + for (index, time_millis) in SNAPSHOT_TIMES.into_iter().enumerate() { + let path = tmp + .path() + .join(format!("test_db.db/t1/snapshot/snapshot-{}", index + 1)); + let mut snapshot: serde_json::Value = + serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + snapshot["timeMillis"] = serde_json::json!(time_millis); + std::fs::write(path, serde_json::to_vec(&snapshot).unwrap()).unwrap(); + } + (tmp, sql_context) +} + async fn setup_table_with_watermarks( watermarks: [Option; 3], ) -> (tempfile::TempDir, paimon_datafusion::SQLContext) { @@ -2693,73 +2712,83 @@ async fn test_rollback_to_tag() { #[tokio::test] async fn test_rollback_to_timestamp() { - let (_tmp, sql_context) = setup_table_with_snapshots().await; - - // Get the timestamp of snapshot 1 from $snapshots system table - let batches = sql_context - .sql("SELECT snapshot_id, commit_time FROM paimon.test_db.`t1$snapshots` ORDER BY snapshot_id") - .await - .unwrap() - .collect() - .await - .unwrap(); - - // Use a timestamp between snapshot 1 and snapshot 2 - let snap1_time = batches[0] - .column_by_name("commit_time") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - - exec( - &sql_context, - &format!( - "CALL sys.rollback_to_timestamp(table => 'test_db.t1', timestamp => '{snap1_time}')" - ), - ) - .await; + for (timestamp, expected_snapshot) in [ + (SNAPSHOT_TIMES[0], 1), + (SNAPSHOT_TIMES[1] - 1, 1), + (SNAPSHOT_TIMES[1], 2), + ] { + let (_tmp, sql_context) = setup_table_with_snapshot_times().await; + exec( + &sql_context, + &format!( + "CALL sys.rollback_to_timestamp(table => 'test_db.t1', timestamp => '{timestamp}')" + ), + ) + .await; - let count = row_count(&sql_context, "SELECT * FROM paimon.test_db.t1").await; - assert_eq!(count, 1); + let rows = collect_id_name(&sql_context, "SELECT id, name FROM paimon.test_db.t1").await; + let expected = [(1, "alice".to_string()), (2, "bob".to_string())]; + assert_eq!(rows, expected[..expected_snapshot], "timestamp {timestamp}"); + let snapshots = sql_context + .sql("SELECT snapshot_id FROM paimon.test_db.`t1$snapshots` ORDER BY snapshot_id") + .await + .unwrap() + .collect() + .await + .unwrap(); + let ids = snapshots + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!(ids, (1..=expected_snapshot as i64).collect::>()); + } } #[tokio::test] async fn test_create_tag_from_timestamp() { - let (_tmp, sql_context) = setup_table_with_snapshots().await; - - // Get the timestamp of snapshot 2 - let batches = sql_context - .sql("SELECT snapshot_id, commit_time FROM paimon.test_db.`t1$snapshots` ORDER BY snapshot_id") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let snap2_time = batches[0] - .column_by_name("commit_time") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .value(1); - - exec( - &sql_context, - &format!( - "CALL sys.create_tag_from_timestamp(table => 'test_db.t1', tag => 'ts_tag', timestamp => '{snap2_time}')" - ), - ) - .await; - - let count = row_count( - &sql_context, - "SELECT * FROM paimon.test_db.`t1$tags` WHERE tag_name = 'ts_tag'", - ) - .await; - assert_eq!(count, 1); + let (tmp, sql_context) = setup_table_with_snapshot_times().await; + for (tag, timestamp, expected_snapshot) in [ + ("before_first", SNAPSHOT_TIMES[0] - 1, 1), + ("between", SNAPSHOT_TIMES[1] - 1, 2), + ("exact", SNAPSHOT_TIMES[1], 2), + ] { + exec( + &sql_context, + &format!( + "CALL sys.create_tag_from_timestamp(table => 'test_db.t1', tag => '{tag}', timestamp => '{timestamp}')" + ), + ) + .await; + assert_eq!(read_tag(tmp.path(), tag)["id"], expected_snapshot); + let batches = sql_context + .sql(&format!( + "SELECT snapshot_id FROM paimon.test_db.`t1$tags` WHERE tag_name = '{tag}'" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 1); + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + expected_snapshot + ); + } } #[tokio::test] diff --git a/docs/src/sql.md b/docs/src/sql.md index e15ac8f75..3cec9a760 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -41,7 +41,7 @@ and writes to ordinary Paimon tables can create them with `file.format=mosaic`. SQL support has two layers: - DataFusion provides the parser, query planner, optimizer, execution engine, expressions, scalar functions, aggregate functions, and window functions. SQL statements that `SQLContext` does not intercept are delegated to DataFusion. This includes the DataFusion SQL surface for `SELECT` queries, CTEs (including recursive CTEs), subqueries, joins including `LATERAL` joins, SQL lambda functions, grouping, `HAVING`, window clauses, `QUALIFY`, set operations, `ORDER BY`, `LIMIT`/`OFFSET`, `EXPLAIN`, information-schema commands such as `SHOW TABLES`, `DESCRIBE`, `COPY`, and ordinary `INSERT`. -- Paimon-specific table management and row-level writes are implemented by `SQLContext`. This includes Paimon `CREATE TABLE`, `ALTER TABLE`, `DROP TABLE`, `CREATE TEMPORARY TABLE`, `CREATE TEMPORARY VIEW`, REST Catalog persistent `CREATE VIEW`, `DROP VIEW`, and `CREATE FUNCTION`, `DROP TEMPORARY TABLE` / `VIEW`, `INSERT OVERWRITE ... PARTITION`, `UPDATE`, `DELETE`, `MERGE INTO`, `TRUNCATE TABLE`, `ALTER TABLE ... ADD PARTITION`, `ALTER TABLE ... DROP PARTITION`, `SHOW PARTITIONS`, `MSCK REPAIR TABLE`, `ANALYZE TABLE`, `CALL sys.*`, Paimon time travel, and `SET` / `RESET 'paimon.*'`. +- Paimon-specific table management and row-level writes are implemented by `SQLContext`. This includes Paimon `CREATE TABLE`, `ALTER TABLE`, `DROP TABLE`, `CREATE TEMPORARY TABLE`, `CREATE TEMPORARY VIEW`, REST Catalog persistent `CREATE VIEW`, `DROP VIEW`, and `CREATE FUNCTION`, `DROP TEMPORARY TABLE` / `VIEW`, `INSERT OVERWRITE ... PARTITION`, `UPDATE`, `DELETE`, `TRUNCATE TABLE`, `ALTER TABLE ... ADD PARTITION`, `ALTER TABLE ... DROP PARTITION`, `SHOW PARTITIONS`, `MSCK REPAIR TABLE`, `ANALYZE TABLE`, `CALL sys.*`, Paimon time travel, and `SET` / `RESET 'paimon.*'`. `MERGE INTO` uses DataFusion's native planner and the Paimon table provider for execution. Not every DataFusion DDL/DML statement maps to a Paimon table operation. For Paimon catalogs, `CREATE EXTERNAL TABLE`, `LOCATION`, `CREATE MATERIALIZED VIEW`, and persistent `CREATE TABLE AS SELECT` are rejected or not implemented. Persistent `CREATE FUNCTION` is supported only for the REST Catalog SQL scalar form documented below. DataFusion `COPY` can export query results to files; it does not create or commit Paimon table files. @@ -76,11 +76,13 @@ and `unnest_outer`) and its planner, aggregation, and spill improvements are available through the normal session. Custom spill backends can be configured on the `RuntimeEnv` passed to `SQLContext::builder`. -Paimon `MERGE INTO` continues to use `SQLContext`'s Paimon write implementation. -DataFusion's generic `TableProvider::merge_into` hook, `file_row_index()` and -`input_file_name()` for Paimon scans, and custom query-planner FFI are not -provided by this upgrade. Paimon's logical row IDs are not physical file row -offsets. +Paimon `MERGE INTO` uses DataFusion's native DML plan and +`TableProvider::merge_into` hook. Planning and `EXPLAIN` do not modify data; +collecting or executing the plan performs the write and returns an affected-row +`count`. `SQLOptions` DML restrictions apply to MERGE statements. +`file_row_index()` and `input_file_name()` for Paimon scans and custom +query-planner FFI are not provided. Paimon's logical row IDs are not physical +file row offsets. Scans declare DataFusion `Range` partitioning when every table partition key is present unchanged in the scan output and uses a supported type: Boolean, @@ -1047,6 +1049,26 @@ WHEN MATCHED THEN UPDATE SET name = source.name; For append-only tables, `MERGE INTO` uses Copy-on-Write file rewriting and supports matched `UPDATE`, matched `DELETE`, and not-matched `INSERT`. For data-evolution row-tracking tables without primary keys, `MERGE INTO` uses the `_ROW_ID` virtual column for row-level tracking and supports matched `UPDATE`, matched `DELETE` when deletion vectors are enabled, and not-matched `INSERT`. Primary-key tables are not supported for `MERGE INTO`. +The source is captured once per execution. On row-ID or CoW delete conflicts, +the operation retries up to five attempts, refreshing the target scan and +writers while replaying the captured source. Update, delete, and insert messages +are published in one commit. Multiple source rows selecting conflicting actions +for the same target row fail without publishing any actions. A sole unconditional +matched DELETE collapses duplicate matches and counts each deleted target once. + +Append-only tables support multiple conditional matched clauses with the first +true clause taking effect; NULL predicates fall through. Data-evolution tables +retain the restriction of one unconditional matched clause. Both support +conditional not-matched inserts. `NOT MATCHED BY SOURCE` and subqueries inside +`ON` or action expressions are not supported by the Paimon executor. Native +planner restrictions also apply, including `UPDATE SET *`, `INSERT *`, `OUTPUT`, +and action-level `WHERE` predicates. After declaring a table alias, use that +alias in column references. + +`SQLContext` expands positional MERGE inserts to persisted column names for +data-evolution tables, excluding the virtual `_ROW_ID`. When using a raw +DataFusion `SessionContext`, list the insert columns explicitly for these tables. + ### TRUNCATE TABLE Truncate an entire table or specific partitions: @@ -1443,7 +1465,7 @@ by a conditional update or transaction. ### create_tag_from_timestamp -Create a named tag from a commit timestamp (finds the earliest retained snapshot at or after the given time): +Create a named tag from an epoch timestamp in milliseconds (finds the earliest retained snapshot at or after the given time): ```sql CALL sys.create_tag_from_timestamp(table => 'paimon.my_db.my_table', tag => 'my_tag', timestamp => 1234567890000); @@ -1504,12 +1526,16 @@ CALL sys.rollback_to(table => 'paimon.my_db.my_table', tag => 'my_tag'); ### rollback_to_timestamp -Rollback a table to a specific timestamp: +Rollback a table to an epoch timestamp in milliseconds (selects the latest retained snapshot at or before that time): ```sql CALL sys.rollback_to_timestamp(table => 'paimon.my_db.my_table', timestamp => 1234567890000); ``` +Both timestamp procedures take epoch milliseconds. The `$snapshots.commit_time` +column is a local `TIMESTAMP`; its raw Arrow value includes the local timezone +offset and must not be passed directly as an epoch timestamp. + ### create_global_index Build and commit a global index for a table column.