wirybeaver commented on code in PR #24195:
URL: https://github.com/apache/datafusion/pull/24195#discussion_r4086839415
##########
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;
Review Comment:
Addressed in 062af76a3. MERGE now acquires all target partition write locks
before reading target rows and retains them through replacement, eliminating
the read/write TOCTOU window.
##########
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!(
Review Comment:
Addressed in 062af76a3. Duplicate source matches now return exec_err! at
execution time.
##########
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()];
Review Comment:
Addressed in 062af76a3. The merged rows are now distributed evenly across
every existing target partition, with multipartition regression coverage.
##########
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 {
Review Comment:
Addressed in 062af76a3. Omitted non-nullable columns without defaults now
produce a planning error; supplied values and nullable/defaulted columns remain
supported.
--
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]
