wirybeaver commented on code in PR #24195:
URL: https://github.com/apache/datafusion/pull/24195#discussion_r4086837768
##########
datafusion/catalog/src/memory/table.rs:
##########
@@ -572,6 +596,735 @@ impl MemTable {
Ok(Arc::new(DmlResultExec::new(total_updated)))
}
+
+ #[expect(
+ clippy::needless_pass_by_value,
+ reason = "matches the TableProvider::merge_into signature"
+ )]
+ fn merge_into_boxed<'a>(
+ &'a self,
+ state: &'a dyn Session,
+ source: Arc<dyn ExecutionPlan>,
+ merge_schema: DFSchemaRef,
+ on: Expr,
+ clauses: Vec<MergeIntoClause>,
+ ) -> BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {
+ Box::pin(ready(self.merge_into_inner(
+ state,
+ source,
+ merge_schema.as_ref(),
+ on,
+ clauses,
+ )))
+ }
+
+ fn merge_into_inner(
+ &self,
+ state: &dyn Session,
+ source: Arc<dyn ExecutionPlan>,
+ merge_schema: &DFSchema,
+ on: Expr,
+ clauses: Vec<MergeIntoClause>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if self.batches.is_empty() {
+ return plan_err!("No partitions provided, expected at least one
partition");
+ }
+
+ let source_schema = source.schema();
+ let target_width = self.schema.fields().len();
+ let source_width = source_schema.fields().len();
+ if merge_schema.fields().len() != target_width + source_width {
+ return plan_err!(
+ "MERGE INTO schema mismatch: expected {} target/source
columns, got {}",
+ target_width + source_width,
+ merge_schema.fields().len()
+ );
+ }
+
+ let merge_arrow_schema = Arc::new(merge_schema.as_arrow().clone());
+ let on = on.cast_to(&DataType::Boolean, merge_schema)?;
+ let on = state.create_physical_expr(on, merge_schema)?;
+ let clauses = compile_merge_clauses(self, state, merge_schema,
clauses)?;
+
+ Ok(Arc::new(MergeIntoExec::new(
+ self.batches.clone(),
+ Arc::clone(&self.schema),
+ Arc::clone(&self.sort_order),
+ source,
+ source_schema,
+ merge_arrow_schema,
+ on,
+ clauses,
+ )))
+ }
+}
+
+#[derive(Debug, Clone)]
+struct CompiledMergeClause {
+ kind: MergeIntoClauseKind,
+ predicate: Option<Arc<dyn PhysicalExpr>>,
+ action: CompiledMergeAction,
+}
+
+#[derive(Debug, Clone)]
+enum CompiledMergeAction {
+ Update(Vec<CompiledMergeAssignment>),
+ Insert(Vec<CompiledInsertValue>),
+ Delete,
+}
+
+#[derive(Debug, Clone)]
+struct CompiledMergeAssignment {
+ target_index: usize,
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+}
+
+#[derive(Debug, Clone)]
+enum CompiledInsertValue {
+ MergeExpr {
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+ },
+ DefaultExpr {
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+ },
+ Null(ScalarValue),
+}
+
+#[derive(Debug, Clone)]
+struct MergeIntoExec {
+ batches: Vec<PartitionData>,
+ target_schema: SchemaRef,
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ source: Arc<dyn ExecutionPlan>,
+ source_schema: SchemaRef,
+ merge_schema: SchemaRef,
+ on: Arc<dyn PhysicalExpr>,
+ clauses: Vec<CompiledMergeClause>,
+ schema: SchemaRef,
+ properties: Arc<PlanProperties>,
+}
+
+impl MergeIntoExec {
+ #[expect(clippy::too_many_arguments)]
+ fn new(
+ batches: Vec<PartitionData>,
+ target_schema: SchemaRef,
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ source: Arc<dyn ExecutionPlan>,
+ source_schema: SchemaRef,
+ merge_schema: SchemaRef,
+ on: Arc<dyn PhysicalExpr>,
+ clauses: Vec<CompiledMergeClause>,
+ ) -> Self {
+ let schema = dml_result_schema();
+ let properties = dml_result_properties(&schema);
+ Self {
+ batches,
+ target_schema,
+ sort_order,
+ source,
+ source_schema,
+ merge_schema,
+ on,
+ clauses,
+ schema,
+ properties,
+ }
+ }
+
+ async fn execute_merge(
+ &self,
+ context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<ArrowRecordBatch> {
+ let source_partitions =
+ collect_partitioned(Arc::clone(&self.source), context).await?;
+ let source_rows = partitioned_batches_to_rows(&source_partitions)?;
+
+ let mut target_batches = vec![];
+ for partition_data in &self.batches {
+ let partition = partition_data.read().await;
+ target_batches.extend(partition.iter().cloned());
+ }
+ let target_rows = batches_to_rows(&target_batches)?;
+
+ let null_target = null_row_for_schema(&self.target_schema)?;
+ let null_source = null_row_for_schema(&self.source_schema)?;
+
+ let mut target_matches: Vec<Option<usize>> = vec![None;
target_rows.len()];
+ let mut source_matched = vec![false; source_rows.len()];
+
+ for (target_idx, target_row) in target_rows.iter().enumerate() {
+ for (source_idx, source_row) in source_rows.iter().enumerate() {
+ let combined = combined_row_batch(
Review Comment:
Partially addressed in 062af76a3. Generic ON predicates still require
considering target/source pairs, but evaluation now runs in bounded 8K-pair
Arrow batches instead of allocating one RecordBatch per pair. I am leaving this
thread open because the generic algorithm remains O(n*m).
##########
datafusion/expr/src/logical_plan/dml.rs:
##########
@@ -479,6 +479,23 @@ impl MergeIntoOp {
}
}
+impl Display for MergeIntoOp {
+ fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+ write!(f, "on=[{}]", self.on)?;
+ if !self.clauses.is_empty() {
Review Comment:
Addressed in 062af76a3. MergeIntoOp now always renders clauses=[...],
including clauses=[] when empty.
##########
datafusion/catalog/src/memory/table.rs:
##########
@@ -572,6 +596,735 @@ impl MemTable {
Ok(Arc::new(DmlResultExec::new(total_updated)))
}
+
+ #[expect(
+ clippy::needless_pass_by_value,
+ reason = "matches the TableProvider::merge_into signature"
+ )]
+ fn merge_into_boxed<'a>(
+ &'a self,
+ state: &'a dyn Session,
+ source: Arc<dyn ExecutionPlan>,
+ merge_schema: DFSchemaRef,
+ on: Expr,
+ clauses: Vec<MergeIntoClause>,
+ ) -> BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {
+ Box::pin(ready(self.merge_into_inner(
+ state,
+ source,
+ merge_schema.as_ref(),
+ on,
+ clauses,
+ )))
+ }
+
+ fn merge_into_inner(
+ &self,
+ state: &dyn Session,
+ source: Arc<dyn ExecutionPlan>,
+ merge_schema: &DFSchema,
+ on: Expr,
+ clauses: Vec<MergeIntoClause>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if self.batches.is_empty() {
+ return plan_err!("No partitions provided, expected at least one
partition");
+ }
+
+ let source_schema = source.schema();
+ let target_width = self.schema.fields().len();
+ let source_width = source_schema.fields().len();
+ if merge_schema.fields().len() != target_width + source_width {
+ return plan_err!(
+ "MERGE INTO schema mismatch: expected {} target/source
columns, got {}",
+ target_width + source_width,
+ merge_schema.fields().len()
+ );
+ }
+
+ let merge_arrow_schema = Arc::new(merge_schema.as_arrow().clone());
+ let on = on.cast_to(&DataType::Boolean, merge_schema)?;
+ let on = state.create_physical_expr(on, merge_schema)?;
+ let clauses = compile_merge_clauses(self, state, merge_schema,
clauses)?;
+
+ Ok(Arc::new(MergeIntoExec::new(
+ self.batches.clone(),
+ Arc::clone(&self.schema),
+ Arc::clone(&self.sort_order),
+ source,
+ source_schema,
+ merge_arrow_schema,
+ on,
+ clauses,
+ )))
+ }
+}
+
+#[derive(Debug, Clone)]
+struct CompiledMergeClause {
+ kind: MergeIntoClauseKind,
+ predicate: Option<Arc<dyn PhysicalExpr>>,
+ action: CompiledMergeAction,
+}
+
+#[derive(Debug, Clone)]
+enum CompiledMergeAction {
+ Update(Vec<CompiledMergeAssignment>),
+ Insert(Vec<CompiledInsertValue>),
+ Delete,
+}
+
+#[derive(Debug, Clone)]
+struct CompiledMergeAssignment {
+ target_index: usize,
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+}
+
+#[derive(Debug, Clone)]
+enum CompiledInsertValue {
+ MergeExpr {
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+ },
+ DefaultExpr {
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+ },
+ Null(ScalarValue),
+}
+
+#[derive(Debug, Clone)]
+struct MergeIntoExec {
+ batches: Vec<PartitionData>,
+ target_schema: SchemaRef,
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ source: Arc<dyn ExecutionPlan>,
+ source_schema: SchemaRef,
+ merge_schema: SchemaRef,
+ on: Arc<dyn PhysicalExpr>,
+ clauses: Vec<CompiledMergeClause>,
+ schema: SchemaRef,
+ properties: Arc<PlanProperties>,
+}
+
+impl MergeIntoExec {
+ #[expect(clippy::too_many_arguments)]
+ fn new(
+ batches: Vec<PartitionData>,
+ target_schema: SchemaRef,
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ source: Arc<dyn ExecutionPlan>,
+ source_schema: SchemaRef,
+ merge_schema: SchemaRef,
+ on: Arc<dyn PhysicalExpr>,
+ clauses: Vec<CompiledMergeClause>,
+ ) -> Self {
+ let schema = dml_result_schema();
+ let properties = dml_result_properties(&schema);
+ Self {
+ batches,
+ target_schema,
+ sort_order,
+ source,
+ source_schema,
+ merge_schema,
+ on,
+ clauses,
+ schema,
+ properties,
+ }
+ }
+
+ async fn execute_merge(
+ &self,
+ context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<ArrowRecordBatch> {
+ let source_partitions =
+ collect_partitioned(Arc::clone(&self.source), context).await?;
+ let source_rows = partitioned_batches_to_rows(&source_partitions)?;
+
+ let mut target_batches = vec![];
+ for partition_data in &self.batches {
+ let partition = partition_data.read().await;
+ target_batches.extend(partition.iter().cloned());
+ }
+ let target_rows = batches_to_rows(&target_batches)?;
+
+ let null_target = null_row_for_schema(&self.target_schema)?;
+ let null_source = null_row_for_schema(&self.source_schema)?;
+
+ let mut target_matches: Vec<Option<usize>> = vec![None;
target_rows.len()];
+ let mut source_matched = vec![false; source_rows.len()];
+
+ for (target_idx, target_row) in target_rows.iter().enumerate() {
+ for (source_idx, source_row) in source_rows.iter().enumerate() {
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ target_row,
+ source_row,
+ )?;
+ if evaluate_merge_predicate(&self.on, &combined)? {
+ if let Some(first_source_idx) = target_matches[target_idx]
{
+ return plan_err!(
+ "MERGE INTO matched target row {target_idx} with
more than one source row ({first_source_idx} and {source_idx})"
+ );
+ }
+ target_matches[target_idx] = Some(source_idx);
+ source_matched[source_idx] = true;
+ }
+ }
+ }
+
+ let default_batch = one_row_empty_batch()?;
+ let mut merged_rows =
+
Vec::with_capacity(target_rows.len().saturating_add(source_rows.len()));
+ let mut rows_affected = 0_u64;
+
+ for (target_idx, target_row) in target_rows.iter().enumerate() {
+ let (source_row, clause_kind) =
+ if let Some(source_idx) = target_matches[target_idx] {
+ (&source_rows[source_idx], MergeIntoClauseKind::Matched)
+ } else {
+ (&null_source, MergeIntoClauseKind::NotMatchedBySource)
+ };
+
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ target_row,
+ source_row,
+ )?;
+ let application = apply_first_merge_clause(
+ &self.clauses,
+ clause_kind,
+ &combined,
+ &default_batch,
+ Some(target_row),
+ )?;
+ if application.affected {
+ rows_affected += 1;
+ }
+ if let Some(row) = application.row {
+ merged_rows.push(row);
+ }
+ }
+
+ for (source_idx, source_row) in source_rows.iter().enumerate() {
+ if source_matched[source_idx] {
+ continue;
+ }
+
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ &null_target,
+ source_row,
+ )?;
+ let application = apply_first_merge_clause(
+ &self.clauses,
+ MergeIntoClauseKind::NotMatchedByTarget,
+ &combined,
+ &default_batch,
+ None,
+ )?;
+ if application.affected {
+ rows_affected += 1;
+ }
+ if let Some(row) = application.row {
+ merged_rows.push(row);
+ }
+ }
+
+ let merged_batch = rows_to_batch(Arc::clone(&self.target_schema),
&merged_rows)?;
+
+ *self.sort_order.lock() = vec![];
+ let mut wrote_first_partition = false;
+ for partition_data in &self.batches {
+ let mut partition = partition_data.write().await;
+ if !wrote_first_partition {
+ if merged_batch.num_rows() == 0 {
+ partition.clear();
+ } else {
+ *partition = vec![merged_batch.clone()];
+ }
+ wrote_first_partition = true;
+ } else {
+ partition.clear();
+ }
+ }
+
+ dml_result_batch(Arc::clone(&self.schema), rows_affected)
+ }
+
+ fn expressions(&self) -> impl Iterator<Item = &Arc<dyn PhysicalExpr>> {
+ std::iter::once(&self.on).chain(self.clauses.iter().flat_map(|clause| {
+ clause.predicate.iter().chain(clause.action.expressions())
+ }))
+ }
+}
+
+impl CompiledMergeAction {
+ fn expressions(&self) -> Box<dyn Iterator<Item = &Arc<dyn PhysicalExpr>> +
'_> {
+ match self {
+ Self::Update(assignments) => {
+ Box::new(assignments.iter().map(|assignment| &assignment.expr))
+ }
+ Self::Insert(values) => {
+ Box::new(values.iter().filter_map(|value| match value {
+ CompiledInsertValue::MergeExpr { expr, .. }
+ | CompiledInsertValue::DefaultExpr { expr, .. } =>
Some(expr),
+ CompiledInsertValue::Null(_) => None,
+ }))
+ }
+ Self::Delete => Box::new(std::iter::empty()),
+ }
+ }
+}
+
+impl DisplayAs for MergeIntoExec {
+ fn fmt_as(
+ &self,
+ t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default
+ | DisplayFormatType::Verbose
+ | DisplayFormatType::TreeRender => write!(f, "MergeIntoExec"),
+ }
+ }
+}
+
+impl ExecutionPlan for MergeIntoExec {
+ fn name(&self) -> &str {
+ "MergeIntoExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.source]
+ }
+
+ fn replace_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ _: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ validate_child_count!(self, children);
+ Ok(Arc::new(Self {
+ source: children.swap_remove(0),
+ ..Self::clone(&self)
+ }))
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ fn execute(
+ &self,
+ _partition: usize,
+ context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<datafusion_execution::SendableRecordBatchStream> {
+ let exec = self.clone();
+ Ok(Box::pin(RecordBatchStreamAdapter::new(
+ Arc::clone(&self.schema),
+ futures::stream::once(async move {
exec.execute_merge(context).await }),
+ )))
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ apply_expression_roots(self.expressions(), f)
+ }
+}
+
+struct MergeApplication {
+ row: Option<Vec<ScalarValue>>,
+ affected: bool,
+}
+
+fn compile_merge_clauses(
+ table: &MemTable,
+ state: &dyn Session,
+ merge_schema: &DFSchema,
+ clauses: Vec<MergeIntoClause>,
+) -> Result<Vec<CompiledMergeClause>> {
+ let empty_schema = DFSchema::empty();
+ clauses
+ .into_iter()
+ .map(|clause| {
+ let predicate = clause
+ .predicate
+ .map(|predicate| {
+ let predicate =
+ predicate.cast_to(&DataType::Boolean, merge_schema)?;
+ state.create_physical_expr(predicate, merge_schema)
+ })
+ .transpose()?;
+
+ let action = match (clause.kind.canonical(), clause.action) {
+ (MergeIntoClauseKind::Matched,
MergeIntoAction::Update(assignments))
+ | (
+ MergeIntoClauseKind::NotMatchedBySource,
+ MergeIntoAction::Update(assignments),
+ ) => CompiledMergeAction::Update(compile_merge_assignments(
+ table,
+ state,
+ merge_schema,
+ assignments,
+ )?),
+ (MergeIntoClauseKind::Matched, MergeIntoAction::Delete)
+ | (MergeIntoClauseKind::NotMatchedBySource,
MergeIntoAction::Delete) => {
+ CompiledMergeAction::Delete
+ }
+ (
+ MergeIntoClauseKind::NotMatchedByTarget,
+ MergeIntoAction::Insert { columns, values },
+ ) => CompiledMergeAction::Insert(compile_merge_insert_values(
+ table,
+ state,
+ merge_schema,
+ &empty_schema,
+ columns,
+ values,
+ )?),
+ (MergeIntoClauseKind::Matched, MergeIntoAction::Insert { .. })
=> {
+ return plan_err!("MERGE MATCHED INSERT is not supported");
+ }
+ (MergeIntoClauseKind::NotMatchedByTarget,
MergeIntoAction::Update(_)) => {
+ return plan_err!("MERGE NOT MATCHED UPDATE is not
supported");
+ }
+ (MergeIntoClauseKind::NotMatchedByTarget,
MergeIntoAction::Delete) => {
+ return plan_err!("MERGE NOT MATCHED DELETE is not
supported");
+ }
+ (
+ MergeIntoClauseKind::NotMatchedBySource,
+ MergeIntoAction::Insert { .. },
+ ) => {
+ return plan_err!(
+ "MERGE NOT MATCHED BY SOURCE INSERT is not supported"
+ );
+ }
+ (MergeIntoClauseKind::NotMatched, _) => {
+ unreachable!("canonical() never returns NotMatched")
+ }
+ };
+
+ Ok(CompiledMergeClause {
+ kind: clause.kind,
+ predicate,
+ action,
+ })
+ })
+ .collect()
+}
+
+fn compile_merge_assignments(
+ table: &MemTable,
+ state: &dyn Session,
+ merge_schema: &DFSchema,
+ assignments: Vec<(String, Expr)>,
+) -> Result<Vec<CompiledMergeAssignment>> {
+ let available_columns = table.available_column_names();
Review Comment:
Addressed in 062af76a3. Available column names are now collected lazily only
in the missing-column error path.
##########
datafusion/catalog/src/memory/table.rs:
##########
@@ -572,6 +596,735 @@ impl MemTable {
Ok(Arc::new(DmlResultExec::new(total_updated)))
}
+
+ #[expect(
+ clippy::needless_pass_by_value,
+ reason = "matches the TableProvider::merge_into signature"
+ )]
+ fn merge_into_boxed<'a>(
+ &'a self,
+ state: &'a dyn Session,
+ source: Arc<dyn ExecutionPlan>,
+ merge_schema: DFSchemaRef,
+ on: Expr,
+ clauses: Vec<MergeIntoClause>,
+ ) -> BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {
+ Box::pin(ready(self.merge_into_inner(
+ state,
+ source,
+ merge_schema.as_ref(),
+ on,
+ clauses,
+ )))
+ }
+
+ fn merge_into_inner(
+ &self,
+ state: &dyn Session,
+ source: Arc<dyn ExecutionPlan>,
+ merge_schema: &DFSchema,
+ on: Expr,
+ clauses: Vec<MergeIntoClause>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if self.batches.is_empty() {
+ return plan_err!("No partitions provided, expected at least one
partition");
+ }
+
+ let source_schema = source.schema();
+ let target_width = self.schema.fields().len();
+ let source_width = source_schema.fields().len();
+ if merge_schema.fields().len() != target_width + source_width {
+ return plan_err!(
+ "MERGE INTO schema mismatch: expected {} target/source
columns, got {}",
+ target_width + source_width,
+ merge_schema.fields().len()
+ );
+ }
+
+ let merge_arrow_schema = Arc::new(merge_schema.as_arrow().clone());
+ let on = on.cast_to(&DataType::Boolean, merge_schema)?;
+ let on = state.create_physical_expr(on, merge_schema)?;
+ let clauses = compile_merge_clauses(self, state, merge_schema,
clauses)?;
+
+ Ok(Arc::new(MergeIntoExec::new(
+ self.batches.clone(),
+ Arc::clone(&self.schema),
+ Arc::clone(&self.sort_order),
+ source,
+ source_schema,
+ merge_arrow_schema,
+ on,
+ clauses,
+ )))
+ }
+}
+
+#[derive(Debug, Clone)]
+struct CompiledMergeClause {
+ kind: MergeIntoClauseKind,
+ predicate: Option<Arc<dyn PhysicalExpr>>,
+ action: CompiledMergeAction,
+}
+
+#[derive(Debug, Clone)]
+enum CompiledMergeAction {
+ Update(Vec<CompiledMergeAssignment>),
+ Insert(Vec<CompiledInsertValue>),
+ Delete,
+}
+
+#[derive(Debug, Clone)]
+struct CompiledMergeAssignment {
+ target_index: usize,
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+}
+
+#[derive(Debug, Clone)]
+enum CompiledInsertValue {
+ MergeExpr {
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+ },
+ DefaultExpr {
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+ },
+ Null(ScalarValue),
+}
+
+#[derive(Debug, Clone)]
+struct MergeIntoExec {
+ batches: Vec<PartitionData>,
+ target_schema: SchemaRef,
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ source: Arc<dyn ExecutionPlan>,
+ source_schema: SchemaRef,
+ merge_schema: SchemaRef,
+ on: Arc<dyn PhysicalExpr>,
+ clauses: Vec<CompiledMergeClause>,
+ schema: SchemaRef,
+ properties: Arc<PlanProperties>,
+}
+
+impl MergeIntoExec {
+ #[expect(clippy::too_many_arguments)]
+ fn new(
+ batches: Vec<PartitionData>,
+ target_schema: SchemaRef,
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ source: Arc<dyn ExecutionPlan>,
+ source_schema: SchemaRef,
+ merge_schema: SchemaRef,
+ on: Arc<dyn PhysicalExpr>,
+ clauses: Vec<CompiledMergeClause>,
+ ) -> Self {
+ let schema = dml_result_schema();
+ let properties = dml_result_properties(&schema);
+ Self {
+ batches,
+ target_schema,
+ sort_order,
+ source,
+ source_schema,
+ merge_schema,
+ on,
+ clauses,
+ schema,
+ properties,
+ }
+ }
+
+ async fn execute_merge(
+ &self,
+ context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<ArrowRecordBatch> {
+ let source_partitions =
+ collect_partitioned(Arc::clone(&self.source), context).await?;
+ let source_rows = partitioned_batches_to_rows(&source_partitions)?;
+
+ let mut target_batches = vec![];
+ for partition_data in &self.batches {
+ let partition = partition_data.read().await;
+ target_batches.extend(partition.iter().cloned());
+ }
+ let target_rows = batches_to_rows(&target_batches)?;
+
+ let null_target = null_row_for_schema(&self.target_schema)?;
+ let null_source = null_row_for_schema(&self.source_schema)?;
+
+ let mut target_matches: Vec<Option<usize>> = vec![None;
target_rows.len()];
+ let mut source_matched = vec![false; source_rows.len()];
+
+ for (target_idx, target_row) in target_rows.iter().enumerate() {
+ for (source_idx, source_row) in source_rows.iter().enumerate() {
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ target_row,
+ source_row,
+ )?;
+ if evaluate_merge_predicate(&self.on, &combined)? {
+ if let Some(first_source_idx) = target_matches[target_idx]
{
+ return plan_err!(
+ "MERGE INTO matched target row {target_idx} with
more than one source row ({first_source_idx} and {source_idx})"
+ );
+ }
+ target_matches[target_idx] = Some(source_idx);
+ source_matched[source_idx] = true;
+ }
+ }
+ }
+
+ let default_batch = one_row_empty_batch()?;
+ let mut merged_rows =
+
Vec::with_capacity(target_rows.len().saturating_add(source_rows.len()));
+ let mut rows_affected = 0_u64;
+
+ for (target_idx, target_row) in target_rows.iter().enumerate() {
+ let (source_row, clause_kind) =
+ if let Some(source_idx) = target_matches[target_idx] {
+ (&source_rows[source_idx], MergeIntoClauseKind::Matched)
+ } else {
+ (&null_source, MergeIntoClauseKind::NotMatchedBySource)
+ };
+
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ target_row,
+ source_row,
+ )?;
+ let application = apply_first_merge_clause(
+ &self.clauses,
+ clause_kind,
+ &combined,
+ &default_batch,
+ Some(target_row),
+ )?;
+ if application.affected {
+ rows_affected += 1;
+ }
+ if let Some(row) = application.row {
+ merged_rows.push(row);
+ }
+ }
+
+ for (source_idx, source_row) in source_rows.iter().enumerate() {
+ if source_matched[source_idx] {
+ continue;
+ }
+
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ &null_target,
+ source_row,
+ )?;
+ let application = apply_first_merge_clause(
+ &self.clauses,
+ MergeIntoClauseKind::NotMatchedByTarget,
+ &combined,
+ &default_batch,
+ None,
+ )?;
+ if application.affected {
+ rows_affected += 1;
+ }
+ if let Some(row) = application.row {
+ merged_rows.push(row);
+ }
+ }
+
+ let merged_batch = rows_to_batch(Arc::clone(&self.target_schema),
&merged_rows)?;
+
+ *self.sort_order.lock() = vec![];
+ let mut wrote_first_partition = false;
+ for partition_data in &self.batches {
+ let mut partition = partition_data.write().await;
+ if !wrote_first_partition {
+ if merged_batch.num_rows() == 0 {
+ partition.clear();
+ } else {
+ *partition = vec![merged_batch.clone()];
+ }
+ wrote_first_partition = true;
+ } else {
+ partition.clear();
+ }
+ }
+
+ dml_result_batch(Arc::clone(&self.schema), rows_affected)
+ }
+
+ fn expressions(&self) -> impl Iterator<Item = &Arc<dyn PhysicalExpr>> {
+ std::iter::once(&self.on).chain(self.clauses.iter().flat_map(|clause| {
+ clause.predicate.iter().chain(clause.action.expressions())
+ }))
+ }
+}
+
+impl CompiledMergeAction {
+ fn expressions(&self) -> Box<dyn Iterator<Item = &Arc<dyn PhysicalExpr>> +
'_> {
+ match self {
+ Self::Update(assignments) => {
+ Box::new(assignments.iter().map(|assignment| &assignment.expr))
+ }
+ Self::Insert(values) => {
+ Box::new(values.iter().filter_map(|value| match value {
+ CompiledInsertValue::MergeExpr { expr, .. }
+ | CompiledInsertValue::DefaultExpr { expr, .. } =>
Some(expr),
+ CompiledInsertValue::Null(_) => None,
+ }))
+ }
+ Self::Delete => Box::new(std::iter::empty()),
+ }
+ }
+}
+
+impl DisplayAs for MergeIntoExec {
+ fn fmt_as(
+ &self,
+ t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default
+ | DisplayFormatType::Verbose
+ | DisplayFormatType::TreeRender => write!(f, "MergeIntoExec"),
+ }
+ }
+}
+
+impl ExecutionPlan for MergeIntoExec {
+ fn name(&self) -> &str {
+ "MergeIntoExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.source]
+ }
+
+ fn replace_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ _: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ validate_child_count!(self, children);
+ Ok(Arc::new(Self {
+ source: children.swap_remove(0),
+ ..Self::clone(&self)
+ }))
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ fn execute(
+ &self,
+ _partition: usize,
+ context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<datafusion_execution::SendableRecordBatchStream> {
+ let exec = self.clone();
+ Ok(Box::pin(RecordBatchStreamAdapter::new(
+ Arc::clone(&self.schema),
+ futures::stream::once(async move {
exec.execute_merge(context).await }),
+ )))
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ apply_expression_roots(self.expressions(), f)
+ }
+}
+
+struct MergeApplication {
+ row: Option<Vec<ScalarValue>>,
+ affected: bool,
+}
+
+fn compile_merge_clauses(
+ table: &MemTable,
+ state: &dyn Session,
+ merge_schema: &DFSchema,
+ clauses: Vec<MergeIntoClause>,
+) -> Result<Vec<CompiledMergeClause>> {
+ let empty_schema = DFSchema::empty();
+ clauses
+ .into_iter()
+ .map(|clause| {
+ let predicate = clause
+ .predicate
+ .map(|predicate| {
+ let predicate =
+ predicate.cast_to(&DataType::Boolean, merge_schema)?;
+ state.create_physical_expr(predicate, merge_schema)
+ })
+ .transpose()?;
+
+ let action = match (clause.kind.canonical(), clause.action) {
+ (MergeIntoClauseKind::Matched,
MergeIntoAction::Update(assignments))
+ | (
+ MergeIntoClauseKind::NotMatchedBySource,
+ MergeIntoAction::Update(assignments),
+ ) => CompiledMergeAction::Update(compile_merge_assignments(
+ table,
+ state,
+ merge_schema,
+ assignments,
+ )?),
+ (MergeIntoClauseKind::Matched, MergeIntoAction::Delete)
+ | (MergeIntoClauseKind::NotMatchedBySource,
MergeIntoAction::Delete) => {
+ CompiledMergeAction::Delete
+ }
+ (
+ MergeIntoClauseKind::NotMatchedByTarget,
+ MergeIntoAction::Insert { columns, values },
+ ) => CompiledMergeAction::Insert(compile_merge_insert_values(
+ table,
+ state,
+ merge_schema,
+ &empty_schema,
+ columns,
+ values,
+ )?),
+ (MergeIntoClauseKind::Matched, MergeIntoAction::Insert { .. })
=> {
+ return plan_err!("MERGE MATCHED INSERT is not supported");
+ }
+ (MergeIntoClauseKind::NotMatchedByTarget,
MergeIntoAction::Update(_)) => {
+ return plan_err!("MERGE NOT MATCHED UPDATE is not
supported");
+ }
+ (MergeIntoClauseKind::NotMatchedByTarget,
MergeIntoAction::Delete) => {
+ return plan_err!("MERGE NOT MATCHED DELETE is not
supported");
+ }
+ (
+ MergeIntoClauseKind::NotMatchedBySource,
+ MergeIntoAction::Insert { .. },
+ ) => {
+ return plan_err!(
+ "MERGE NOT MATCHED BY SOURCE INSERT is not supported"
+ );
+ }
+ (MergeIntoClauseKind::NotMatched, _) => {
+ unreachable!("canonical() never returns NotMatched")
+ }
+ };
+
+ Ok(CompiledMergeClause {
+ kind: clause.kind,
+ predicate,
+ action,
+ })
+ })
+ .collect()
+}
+
+fn compile_merge_assignments(
+ table: &MemTable,
+ state: &dyn Session,
+ merge_schema: &DFSchema,
+ assignments: Vec<(String, Expr)>,
+) -> Result<Vec<CompiledMergeAssignment>> {
+ let available_columns = table.available_column_names();
+ let mut seen = HashSet::new();
+ assignments
+ .into_iter()
+ .map(|(column, value)| {
+ if !seen.insert(column.clone()) {
+ return plan_err!("Duplicate column '{column}' in MERGE
UPDATE");
+ }
+ let (target_index, field) =
+ table.target_field(&column).ok_or_else(|| {
+ datafusion_common::DataFusionError::Plan(format!(
+ "MERGE UPDATE failed: column '{column}' does not
exist. Available columns: {}",
+ available_columns.join(", ")
+ ))
+ })?;
+ let value = value.cast_to(field.data_type(), merge_schema)?;
+ let expr = state.create_physical_expr(value, merge_schema)?;
+ Ok(CompiledMergeAssignment {
+ target_index,
+ data_type: field.data_type().clone(),
+ expr,
+ })
+ })
+ .collect()
+}
+
+fn compile_merge_insert_values(
+ table: &MemTable,
+ state: &dyn Session,
+ merge_schema: &DFSchema,
+ empty_schema: &DFSchema,
+ columns: Vec<String>,
+ values: Vec<Expr>,
+) -> Result<Vec<CompiledInsertValue>> {
+ let target_width = table.schema.fields().len();
+ if columns.is_empty() {
+ if values.len() != target_width {
+ return plan_err!(
+ "MERGE INSERT has {target_width} column(s) but {} value(s)",
+ values.len()
+ );
+ }
+ return values
+ .into_iter()
+ .zip(table.schema.fields())
+ .map(|(value, field)| {
+ let value = value.cast_to(field.data_type(), merge_schema)?;
+ let expr = state.create_physical_expr(value, merge_schema)?;
+ Ok(CompiledInsertValue::MergeExpr {
+ data_type: field.data_type().clone(),
+ expr,
+ })
+ })
+ .collect();
+ }
+
+ if columns.len() != values.len() {
+ return plan_err!(
+ "MERGE INSERT has {} column(s) but {} value(s)",
+ columns.len(),
+ values.len()
+ );
+ }
+
+ let mut insert_values = table
+ .schema
+ .fields()
+ .iter()
+ .map(|field| {
+ if let Some(default) = table.column_defaults.get(field.name()) {
+ let default = default.clone().cast_to(field.data_type(),
empty_schema)?;
+ let expr = state.create_physical_expr(default, empty_schema)?;
+ Ok(CompiledInsertValue::DefaultExpr {
+ data_type: field.data_type().clone(),
+ expr,
+ })
+ } else {
+ Ok(CompiledInsertValue::Null(ScalarValue::try_new_null(
+ field.data_type(),
+ )?))
+ }
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ let available_columns = table.available_column_names();
Review Comment:
Addressed in 062af76a3. Available column names are now collected only when
constructing the error.
##########
datafusion/catalog/src/memory/table.rs:
##########
@@ -572,6 +596,735 @@ impl MemTable {
Ok(Arc::new(DmlResultExec::new(total_updated)))
}
+
+ #[expect(
+ clippy::needless_pass_by_value,
+ reason = "matches the TableProvider::merge_into signature"
+ )]
+ fn merge_into_boxed<'a>(
+ &'a self,
+ state: &'a dyn Session,
+ source: Arc<dyn ExecutionPlan>,
+ merge_schema: DFSchemaRef,
+ on: Expr,
+ clauses: Vec<MergeIntoClause>,
+ ) -> BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {
+ Box::pin(ready(self.merge_into_inner(
+ state,
+ source,
+ merge_schema.as_ref(),
+ on,
+ clauses,
+ )))
+ }
+
+ fn merge_into_inner(
+ &self,
+ state: &dyn Session,
+ source: Arc<dyn ExecutionPlan>,
+ merge_schema: &DFSchema,
+ on: Expr,
+ clauses: Vec<MergeIntoClause>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if self.batches.is_empty() {
+ return plan_err!("No partitions provided, expected at least one
partition");
+ }
+
+ let source_schema = source.schema();
+ let target_width = self.schema.fields().len();
+ let source_width = source_schema.fields().len();
+ if merge_schema.fields().len() != target_width + source_width {
+ return plan_err!(
+ "MERGE INTO schema mismatch: expected {} target/source
columns, got {}",
+ target_width + source_width,
+ merge_schema.fields().len()
+ );
+ }
+
+ let merge_arrow_schema = Arc::new(merge_schema.as_arrow().clone());
+ let on = on.cast_to(&DataType::Boolean, merge_schema)?;
+ let on = state.create_physical_expr(on, merge_schema)?;
+ let clauses = compile_merge_clauses(self, state, merge_schema,
clauses)?;
+
+ Ok(Arc::new(MergeIntoExec::new(
+ self.batches.clone(),
+ Arc::clone(&self.schema),
+ Arc::clone(&self.sort_order),
+ source,
+ source_schema,
+ merge_arrow_schema,
+ on,
+ clauses,
+ )))
+ }
+}
+
+#[derive(Debug, Clone)]
+struct CompiledMergeClause {
+ kind: MergeIntoClauseKind,
+ predicate: Option<Arc<dyn PhysicalExpr>>,
+ action: CompiledMergeAction,
+}
+
+#[derive(Debug, Clone)]
+enum CompiledMergeAction {
+ Update(Vec<CompiledMergeAssignment>),
+ Insert(Vec<CompiledInsertValue>),
+ Delete,
+}
+
+#[derive(Debug, Clone)]
+struct CompiledMergeAssignment {
+ target_index: usize,
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+}
+
+#[derive(Debug, Clone)]
+enum CompiledInsertValue {
+ MergeExpr {
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+ },
+ DefaultExpr {
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+ },
+ Null(ScalarValue),
+}
+
+#[derive(Debug, Clone)]
+struct MergeIntoExec {
+ batches: Vec<PartitionData>,
+ target_schema: SchemaRef,
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ source: Arc<dyn ExecutionPlan>,
+ source_schema: SchemaRef,
+ merge_schema: SchemaRef,
+ on: Arc<dyn PhysicalExpr>,
+ clauses: Vec<CompiledMergeClause>,
+ schema: SchemaRef,
+ properties: Arc<PlanProperties>,
+}
+
+impl MergeIntoExec {
+ #[expect(clippy::too_many_arguments)]
+ fn new(
+ batches: Vec<PartitionData>,
+ target_schema: SchemaRef,
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ source: Arc<dyn ExecutionPlan>,
+ source_schema: SchemaRef,
+ merge_schema: SchemaRef,
+ on: Arc<dyn PhysicalExpr>,
+ clauses: Vec<CompiledMergeClause>,
+ ) -> Self {
+ let schema = dml_result_schema();
+ let properties = dml_result_properties(&schema);
+ Self {
+ batches,
+ target_schema,
+ sort_order,
+ source,
+ source_schema,
+ merge_schema,
+ on,
+ clauses,
+ schema,
+ properties,
+ }
+ }
+
+ async fn execute_merge(
+ &self,
+ context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<ArrowRecordBatch> {
+ let source_partitions =
+ collect_partitioned(Arc::clone(&self.source), context).await?;
+ let source_rows = partitioned_batches_to_rows(&source_partitions)?;
+
+ let mut target_batches = vec![];
+ for partition_data in &self.batches {
+ let partition = partition_data.read().await;
+ target_batches.extend(partition.iter().cloned());
+ }
+ let target_rows = batches_to_rows(&target_batches)?;
+
+ let null_target = null_row_for_schema(&self.target_schema)?;
+ let null_source = null_row_for_schema(&self.source_schema)?;
+
+ let mut target_matches: Vec<Option<usize>> = vec![None;
target_rows.len()];
+ let mut source_matched = vec![false; source_rows.len()];
+
+ for (target_idx, target_row) in target_rows.iter().enumerate() {
+ for (source_idx, source_row) in source_rows.iter().enumerate() {
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ target_row,
+ source_row,
+ )?;
+ if evaluate_merge_predicate(&self.on, &combined)? {
+ if let Some(first_source_idx) = target_matches[target_idx]
{
+ return plan_err!(
+ "MERGE INTO matched target row {target_idx} with
more than one source row ({first_source_idx} and {source_idx})"
+ );
+ }
+ target_matches[target_idx] = Some(source_idx);
+ source_matched[source_idx] = true;
+ }
+ }
+ }
+
+ let default_batch = one_row_empty_batch()?;
+ let mut merged_rows =
+
Vec::with_capacity(target_rows.len().saturating_add(source_rows.len()));
+ let mut rows_affected = 0_u64;
+
+ for (target_idx, target_row) in target_rows.iter().enumerate() {
+ let (source_row, clause_kind) =
+ if let Some(source_idx) = target_matches[target_idx] {
+ (&source_rows[source_idx], MergeIntoClauseKind::Matched)
+ } else {
+ (&null_source, MergeIntoClauseKind::NotMatchedBySource)
+ };
+
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ target_row,
+ source_row,
+ )?;
+ let application = apply_first_merge_clause(
+ &self.clauses,
+ clause_kind,
+ &combined,
+ &default_batch,
+ Some(target_row),
+ )?;
+ if application.affected {
+ rows_affected += 1;
+ }
+ if let Some(row) = application.row {
+ merged_rows.push(row);
+ }
+ }
+
+ for (source_idx, source_row) in source_rows.iter().enumerate() {
+ if source_matched[source_idx] {
+ continue;
+ }
+
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ &null_target,
+ source_row,
+ )?;
+ let application = apply_first_merge_clause(
+ &self.clauses,
+ MergeIntoClauseKind::NotMatchedByTarget,
+ &combined,
+ &default_batch,
+ None,
+ )?;
+ if application.affected {
+ rows_affected += 1;
+ }
+ if let Some(row) = application.row {
+ merged_rows.push(row);
+ }
+ }
+
+ let merged_batch = rows_to_batch(Arc::clone(&self.target_schema),
&merged_rows)?;
+
+ *self.sort_order.lock() = vec![];
+ let mut wrote_first_partition = false;
+ for partition_data in &self.batches {
+ let mut partition = partition_data.write().await;
+ if !wrote_first_partition {
+ if merged_batch.num_rows() == 0 {
+ partition.clear();
+ } else {
+ *partition = vec![merged_batch.clone()];
+ }
+ wrote_first_partition = true;
+ } else {
+ partition.clear();
+ }
+ }
+
+ dml_result_batch(Arc::clone(&self.schema), rows_affected)
+ }
+
+ fn expressions(&self) -> impl Iterator<Item = &Arc<dyn PhysicalExpr>> {
+ std::iter::once(&self.on).chain(self.clauses.iter().flat_map(|clause| {
+ clause.predicate.iter().chain(clause.action.expressions())
+ }))
+ }
+}
+
+impl CompiledMergeAction {
+ fn expressions(&self) -> Box<dyn Iterator<Item = &Arc<dyn PhysicalExpr>> +
'_> {
+ match self {
+ Self::Update(assignments) => {
+ Box::new(assignments.iter().map(|assignment| &assignment.expr))
+ }
+ Self::Insert(values) => {
+ Box::new(values.iter().filter_map(|value| match value {
+ CompiledInsertValue::MergeExpr { expr, .. }
+ | CompiledInsertValue::DefaultExpr { expr, .. } =>
Some(expr),
+ CompiledInsertValue::Null(_) => None,
+ }))
+ }
+ Self::Delete => Box::new(std::iter::empty()),
+ }
+ }
+}
+
+impl DisplayAs for MergeIntoExec {
+ fn fmt_as(
+ &self,
+ t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default
+ | DisplayFormatType::Verbose
+ | DisplayFormatType::TreeRender => write!(f, "MergeIntoExec"),
+ }
+ }
+}
+
+impl ExecutionPlan for MergeIntoExec {
+ fn name(&self) -> &str {
+ "MergeIntoExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.source]
+ }
+
+ fn replace_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ _: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ validate_child_count!(self, children);
+ Ok(Arc::new(Self {
+ source: children.swap_remove(0),
+ ..Self::clone(&self)
+ }))
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ fn execute(
+ &self,
+ _partition: usize,
+ context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<datafusion_execution::SendableRecordBatchStream> {
+ let exec = self.clone();
+ Ok(Box::pin(RecordBatchStreamAdapter::new(
+ Arc::clone(&self.schema),
+ futures::stream::once(async move {
exec.execute_merge(context).await }),
+ )))
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ apply_expression_roots(self.expressions(), f)
+ }
+}
+
+struct MergeApplication {
+ row: Option<Vec<ScalarValue>>,
+ affected: bool,
+}
+
+fn compile_merge_clauses(
+ table: &MemTable,
+ state: &dyn Session,
+ merge_schema: &DFSchema,
+ clauses: Vec<MergeIntoClause>,
+) -> Result<Vec<CompiledMergeClause>> {
+ let empty_schema = DFSchema::empty();
+ clauses
+ .into_iter()
+ .map(|clause| {
+ let predicate = clause
+ .predicate
+ .map(|predicate| {
+ let predicate =
+ predicate.cast_to(&DataType::Boolean, merge_schema)?;
+ state.create_physical_expr(predicate, merge_schema)
+ })
+ .transpose()?;
+
+ let action = match (clause.kind.canonical(), clause.action) {
+ (MergeIntoClauseKind::Matched,
MergeIntoAction::Update(assignments))
+ | (
+ MergeIntoClauseKind::NotMatchedBySource,
+ MergeIntoAction::Update(assignments),
+ ) => CompiledMergeAction::Update(compile_merge_assignments(
+ table,
+ state,
+ merge_schema,
+ assignments,
+ )?),
+ (MergeIntoClauseKind::Matched, MergeIntoAction::Delete)
+ | (MergeIntoClauseKind::NotMatchedBySource,
MergeIntoAction::Delete) => {
+ CompiledMergeAction::Delete
+ }
+ (
+ MergeIntoClauseKind::NotMatchedByTarget,
+ MergeIntoAction::Insert { columns, values },
+ ) => CompiledMergeAction::Insert(compile_merge_insert_values(
+ table,
+ state,
+ merge_schema,
+ &empty_schema,
+ columns,
+ values,
+ )?),
+ (MergeIntoClauseKind::Matched, MergeIntoAction::Insert { .. })
=> {
+ return plan_err!("MERGE MATCHED INSERT is not supported");
+ }
+ (MergeIntoClauseKind::NotMatchedByTarget,
MergeIntoAction::Update(_)) => {
+ return plan_err!("MERGE NOT MATCHED UPDATE is not
supported");
+ }
+ (MergeIntoClauseKind::NotMatchedByTarget,
MergeIntoAction::Delete) => {
+ return plan_err!("MERGE NOT MATCHED DELETE is not
supported");
+ }
+ (
+ MergeIntoClauseKind::NotMatchedBySource,
+ MergeIntoAction::Insert { .. },
+ ) => {
+ return plan_err!(
+ "MERGE NOT MATCHED BY SOURCE INSERT is not supported"
+ );
+ }
+ (MergeIntoClauseKind::NotMatched, _) => {
+ unreachable!("canonical() never returns NotMatched")
+ }
+ };
+
+ Ok(CompiledMergeClause {
+ kind: clause.kind,
+ predicate,
+ action,
+ })
+ })
+ .collect()
+}
+
+fn compile_merge_assignments(
+ table: &MemTable,
+ state: &dyn Session,
+ merge_schema: &DFSchema,
+ assignments: Vec<(String, Expr)>,
+) -> Result<Vec<CompiledMergeAssignment>> {
+ let available_columns = table.available_column_names();
+ let mut seen = HashSet::new();
+ assignments
+ .into_iter()
+ .map(|(column, value)| {
+ if !seen.insert(column.clone()) {
+ return plan_err!("Duplicate column '{column}' in MERGE
UPDATE");
+ }
+ let (target_index, field) =
+ table.target_field(&column).ok_or_else(|| {
Review Comment:
Addressed in 062af76a3. Target fields are indexed once in a HashMap and
reused for all assignments.
##########
datafusion/catalog/src/memory/table.rs:
##########
@@ -572,6 +596,735 @@ impl MemTable {
Ok(Arc::new(DmlResultExec::new(total_updated)))
}
+
+ #[expect(
+ clippy::needless_pass_by_value,
+ reason = "matches the TableProvider::merge_into signature"
+ )]
+ fn merge_into_boxed<'a>(
+ &'a self,
+ state: &'a dyn Session,
+ source: Arc<dyn ExecutionPlan>,
+ merge_schema: DFSchemaRef,
+ on: Expr,
+ clauses: Vec<MergeIntoClause>,
+ ) -> BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {
+ Box::pin(ready(self.merge_into_inner(
+ state,
+ source,
+ merge_schema.as_ref(),
+ on,
+ clauses,
+ )))
+ }
+
+ fn merge_into_inner(
+ &self,
+ state: &dyn Session,
+ source: Arc<dyn ExecutionPlan>,
+ merge_schema: &DFSchema,
+ on: Expr,
+ clauses: Vec<MergeIntoClause>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if self.batches.is_empty() {
+ return plan_err!("No partitions provided, expected at least one
partition");
+ }
+
+ let source_schema = source.schema();
+ let target_width = self.schema.fields().len();
+ let source_width = source_schema.fields().len();
+ if merge_schema.fields().len() != target_width + source_width {
+ return plan_err!(
+ "MERGE INTO schema mismatch: expected {} target/source
columns, got {}",
+ target_width + source_width,
+ merge_schema.fields().len()
+ );
+ }
+
+ let merge_arrow_schema = Arc::new(merge_schema.as_arrow().clone());
+ let on = on.cast_to(&DataType::Boolean, merge_schema)?;
+ let on = state.create_physical_expr(on, merge_schema)?;
+ let clauses = compile_merge_clauses(self, state, merge_schema,
clauses)?;
+
+ Ok(Arc::new(MergeIntoExec::new(
+ self.batches.clone(),
+ Arc::clone(&self.schema),
+ Arc::clone(&self.sort_order),
+ source,
+ source_schema,
+ merge_arrow_schema,
+ on,
+ clauses,
+ )))
+ }
+}
+
+#[derive(Debug, Clone)]
+struct CompiledMergeClause {
+ kind: MergeIntoClauseKind,
+ predicate: Option<Arc<dyn PhysicalExpr>>,
+ action: CompiledMergeAction,
+}
+
+#[derive(Debug, Clone)]
+enum CompiledMergeAction {
+ Update(Vec<CompiledMergeAssignment>),
+ Insert(Vec<CompiledInsertValue>),
+ Delete,
+}
+
+#[derive(Debug, Clone)]
+struct CompiledMergeAssignment {
+ target_index: usize,
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+}
+
+#[derive(Debug, Clone)]
+enum CompiledInsertValue {
+ MergeExpr {
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+ },
+ DefaultExpr {
+ data_type: DataType,
+ expr: Arc<dyn PhysicalExpr>,
+ },
+ Null(ScalarValue),
+}
+
+#[derive(Debug, Clone)]
+struct MergeIntoExec {
+ batches: Vec<PartitionData>,
+ target_schema: SchemaRef,
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ source: Arc<dyn ExecutionPlan>,
+ source_schema: SchemaRef,
+ merge_schema: SchemaRef,
+ on: Arc<dyn PhysicalExpr>,
+ clauses: Vec<CompiledMergeClause>,
+ schema: SchemaRef,
+ properties: Arc<PlanProperties>,
+}
+
+impl MergeIntoExec {
+ #[expect(clippy::too_many_arguments)]
+ fn new(
+ batches: Vec<PartitionData>,
+ target_schema: SchemaRef,
+ sort_order: Arc<Mutex<Vec<Vec<SortExpr>>>>,
+ source: Arc<dyn ExecutionPlan>,
+ source_schema: SchemaRef,
+ merge_schema: SchemaRef,
+ on: Arc<dyn PhysicalExpr>,
+ clauses: Vec<CompiledMergeClause>,
+ ) -> Self {
+ let schema = dml_result_schema();
+ let properties = dml_result_properties(&schema);
+ Self {
+ batches,
+ target_schema,
+ sort_order,
+ source,
+ source_schema,
+ merge_schema,
+ on,
+ clauses,
+ schema,
+ properties,
+ }
+ }
+
+ async fn execute_merge(
+ &self,
+ context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<ArrowRecordBatch> {
+ let source_partitions =
+ collect_partitioned(Arc::clone(&self.source), context).await?;
+ let source_rows = partitioned_batches_to_rows(&source_partitions)?;
+
+ let mut target_batches = vec![];
+ for partition_data in &self.batches {
+ let partition = partition_data.read().await;
+ target_batches.extend(partition.iter().cloned());
+ }
+ let target_rows = batches_to_rows(&target_batches)?;
+
+ let null_target = null_row_for_schema(&self.target_schema)?;
+ let null_source = null_row_for_schema(&self.source_schema)?;
+
+ let mut target_matches: Vec<Option<usize>> = vec![None;
target_rows.len()];
+ let mut source_matched = vec![false; source_rows.len()];
+
+ for (target_idx, target_row) in target_rows.iter().enumerate() {
+ for (source_idx, source_row) in source_rows.iter().enumerate() {
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ target_row,
+ source_row,
+ )?;
+ if evaluate_merge_predicate(&self.on, &combined)? {
+ if let Some(first_source_idx) = target_matches[target_idx]
{
+ return plan_err!(
+ "MERGE INTO matched target row {target_idx} with
more than one source row ({first_source_idx} and {source_idx})"
+ );
+ }
+ target_matches[target_idx] = Some(source_idx);
+ source_matched[source_idx] = true;
+ }
+ }
+ }
+
+ let default_batch = one_row_empty_batch()?;
+ let mut merged_rows =
+
Vec::with_capacity(target_rows.len().saturating_add(source_rows.len()));
+ let mut rows_affected = 0_u64;
+
+ for (target_idx, target_row) in target_rows.iter().enumerate() {
+ let (source_row, clause_kind) =
+ if let Some(source_idx) = target_matches[target_idx] {
+ (&source_rows[source_idx], MergeIntoClauseKind::Matched)
+ } else {
+ (&null_source, MergeIntoClauseKind::NotMatchedBySource)
+ };
+
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ target_row,
+ source_row,
+ )?;
+ let application = apply_first_merge_clause(
+ &self.clauses,
+ clause_kind,
+ &combined,
+ &default_batch,
+ Some(target_row),
+ )?;
+ if application.affected {
+ rows_affected += 1;
+ }
+ if let Some(row) = application.row {
+ merged_rows.push(row);
+ }
+ }
+
+ for (source_idx, source_row) in source_rows.iter().enumerate() {
+ if source_matched[source_idx] {
+ continue;
+ }
+
+ let combined = combined_row_batch(
+ Arc::clone(&self.merge_schema),
+ &null_target,
+ source_row,
+ )?;
+ let application = apply_first_merge_clause(
+ &self.clauses,
+ MergeIntoClauseKind::NotMatchedByTarget,
+ &combined,
+ &default_batch,
+ None,
+ )?;
+ if application.affected {
+ rows_affected += 1;
+ }
+ if let Some(row) = application.row {
+ merged_rows.push(row);
+ }
+ }
+
+ let merged_batch = rows_to_batch(Arc::clone(&self.target_schema),
&merged_rows)?;
+
+ *self.sort_order.lock() = vec![];
+ let mut wrote_first_partition = false;
+ for partition_data in &self.batches {
+ let mut partition = partition_data.write().await;
+ if !wrote_first_partition {
+ if merged_batch.num_rows() == 0 {
+ partition.clear();
+ } else {
+ *partition = vec![merged_batch.clone()];
+ }
+ wrote_first_partition = true;
+ } else {
+ partition.clear();
+ }
+ }
+
+ dml_result_batch(Arc::clone(&self.schema), rows_affected)
+ }
+
+ fn expressions(&self) -> impl Iterator<Item = &Arc<dyn PhysicalExpr>> {
+ std::iter::once(&self.on).chain(self.clauses.iter().flat_map(|clause| {
+ clause.predicate.iter().chain(clause.action.expressions())
+ }))
+ }
+}
+
+impl CompiledMergeAction {
+ fn expressions(&self) -> Box<dyn Iterator<Item = &Arc<dyn PhysicalExpr>> +
'_> {
+ match self {
+ Self::Update(assignments) => {
+ Box::new(assignments.iter().map(|assignment| &assignment.expr))
+ }
+ Self::Insert(values) => {
+ Box::new(values.iter().filter_map(|value| match value {
+ CompiledInsertValue::MergeExpr { expr, .. }
+ | CompiledInsertValue::DefaultExpr { expr, .. } =>
Some(expr),
+ CompiledInsertValue::Null(_) => None,
+ }))
+ }
+ Self::Delete => Box::new(std::iter::empty()),
+ }
+ }
+}
+
+impl DisplayAs for MergeIntoExec {
+ fn fmt_as(
+ &self,
+ t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default
+ | DisplayFormatType::Verbose
+ | DisplayFormatType::TreeRender => write!(f, "MergeIntoExec"),
+ }
+ }
+}
+
+impl ExecutionPlan for MergeIntoExec {
+ fn name(&self) -> &str {
+ "MergeIntoExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.source]
+ }
+
+ fn replace_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ _: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ validate_child_count!(self, children);
+ Ok(Arc::new(Self {
+ source: children.swap_remove(0),
+ ..Self::clone(&self)
+ }))
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ fn execute(
+ &self,
+ _partition: usize,
+ context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<datafusion_execution::SendableRecordBatchStream> {
+ let exec = self.clone();
+ Ok(Box::pin(RecordBatchStreamAdapter::new(
+ Arc::clone(&self.schema),
+ futures::stream::once(async move {
exec.execute_merge(context).await }),
+ )))
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ apply_expression_roots(self.expressions(), f)
+ }
+}
+
+struct MergeApplication {
+ row: Option<Vec<ScalarValue>>,
+ affected: bool,
+}
+
+fn compile_merge_clauses(
+ table: &MemTable,
+ state: &dyn Session,
+ merge_schema: &DFSchema,
+ clauses: Vec<MergeIntoClause>,
+) -> Result<Vec<CompiledMergeClause>> {
+ let empty_schema = DFSchema::empty();
+ clauses
+ .into_iter()
+ .map(|clause| {
+ let predicate = clause
+ .predicate
+ .map(|predicate| {
+ let predicate =
+ predicate.cast_to(&DataType::Boolean, merge_schema)?;
+ state.create_physical_expr(predicate, merge_schema)
+ })
+ .transpose()?;
+
+ let action = match (clause.kind.canonical(), clause.action) {
+ (MergeIntoClauseKind::Matched,
MergeIntoAction::Update(assignments))
+ | (
+ MergeIntoClauseKind::NotMatchedBySource,
+ MergeIntoAction::Update(assignments),
+ ) => CompiledMergeAction::Update(compile_merge_assignments(
+ table,
+ state,
+ merge_schema,
+ assignments,
+ )?),
+ (MergeIntoClauseKind::Matched, MergeIntoAction::Delete)
+ | (MergeIntoClauseKind::NotMatchedBySource,
MergeIntoAction::Delete) => {
+ CompiledMergeAction::Delete
+ }
+ (
+ MergeIntoClauseKind::NotMatchedByTarget,
+ MergeIntoAction::Insert { columns, values },
+ ) => CompiledMergeAction::Insert(compile_merge_insert_values(
+ table,
+ state,
+ merge_schema,
+ &empty_schema,
+ columns,
+ values,
+ )?),
+ (MergeIntoClauseKind::Matched, MergeIntoAction::Insert { .. })
=> {
+ return plan_err!("MERGE MATCHED INSERT is not supported");
+ }
+ (MergeIntoClauseKind::NotMatchedByTarget,
MergeIntoAction::Update(_)) => {
+ return plan_err!("MERGE NOT MATCHED UPDATE is not
supported");
+ }
+ (MergeIntoClauseKind::NotMatchedByTarget,
MergeIntoAction::Delete) => {
+ return plan_err!("MERGE NOT MATCHED DELETE is not
supported");
+ }
+ (
+ MergeIntoClauseKind::NotMatchedBySource,
+ MergeIntoAction::Insert { .. },
+ ) => {
+ return plan_err!(
+ "MERGE NOT MATCHED BY SOURCE INSERT is not supported"
+ );
+ }
+ (MergeIntoClauseKind::NotMatched, _) => {
+ unreachable!("canonical() never returns NotMatched")
+ }
+ };
+
+ Ok(CompiledMergeClause {
+ kind: clause.kind,
+ predicate,
+ action,
+ })
+ })
+ .collect()
+}
+
+fn compile_merge_assignments(
+ table: &MemTable,
+ state: &dyn Session,
+ merge_schema: &DFSchema,
+ assignments: Vec<(String, Expr)>,
+) -> Result<Vec<CompiledMergeAssignment>> {
+ let available_columns = table.available_column_names();
+ let mut seen = HashSet::new();
+ assignments
+ .into_iter()
+ .map(|(column, value)| {
+ if !seen.insert(column.clone()) {
+ return plan_err!("Duplicate column '{column}' in MERGE
UPDATE");
+ }
+ let (target_index, field) =
+ table.target_field(&column).ok_or_else(|| {
+ datafusion_common::DataFusionError::Plan(format!(
+ "MERGE UPDATE failed: column '{column}' does not
exist. Available columns: {}",
+ available_columns.join(", ")
+ ))
+ })?;
+ let value = value.cast_to(field.data_type(), merge_schema)?;
+ let expr = state.create_physical_expr(value, merge_schema)?;
+ Ok(CompiledMergeAssignment {
+ target_index,
+ data_type: field.data_type().clone(),
+ expr,
+ })
+ })
+ .collect()
+}
+
+fn compile_merge_insert_values(
+ table: &MemTable,
+ state: &dyn Session,
+ merge_schema: &DFSchema,
+ empty_schema: &DFSchema,
+ columns: Vec<String>,
+ values: Vec<Expr>,
+) -> Result<Vec<CompiledInsertValue>> {
+ let target_width = table.schema.fields().len();
+ if columns.is_empty() {
+ if values.len() != target_width {
+ return plan_err!(
+ "MERGE INSERT has {target_width} column(s) but {} value(s)",
+ values.len()
+ );
+ }
+ return values
+ .into_iter()
+ .zip(table.schema.fields())
+ .map(|(value, field)| {
+ let value = value.cast_to(field.data_type(), merge_schema)?;
+ let expr = state.create_physical_expr(value, merge_schema)?;
+ Ok(CompiledInsertValue::MergeExpr {
+ data_type: field.data_type().clone(),
+ expr,
+ })
+ })
+ .collect();
+ }
+
+ if columns.len() != values.len() {
+ return plan_err!(
+ "MERGE INSERT has {} column(s) but {} value(s)",
+ columns.len(),
+ values.len()
+ );
+ }
+
+ let mut insert_values = table
+ .schema
+ .fields()
+ .iter()
+ .map(|field| {
+ if let Some(default) = table.column_defaults.get(field.name()) {
+ let default = default.clone().cast_to(field.data_type(),
empty_schema)?;
+ let expr = state.create_physical_expr(default, empty_schema)?;
+ Ok(CompiledInsertValue::DefaultExpr {
+ data_type: field.data_type().clone(),
+ expr,
+ })
+ } else {
+ Ok(CompiledInsertValue::Null(ScalarValue::try_new_null(
+ field.data_type(),
+ )?))
+ }
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ let available_columns = table.available_column_names();
+ let mut seen = HashSet::new();
+ for (column, value) in columns.into_iter().zip(values) {
+ if !seen.insert(column.clone()) {
+ return plan_err!("Duplicate column '{column}' in MERGE INSERT");
+ }
+ let (target_index, field) = table.target_field(&column).ok_or_else(|| {
Review Comment:
Addressed in 062af76a3. Explicit INSERT columns now use one precomputed
target-field HashMap.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
