leaves12138 commented on code in PR #1068: URL: https://github.com/apache/paimon-rust/pull/1068#discussion_r4212549694
########## crates/paimon/src/table/table_merge_into.rs: ########## @@ -0,0 +1,578 @@ +// 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. + +//! Data-evolution MERGE orchestration, independent of an expression engine. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use arrow_array::{Array, BooleanArray, Int64Array, RecordBatch, UInt64Array}; +use arrow_row::{RowConverter, SortField}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use arrow_select::{concat::concat_batches, take::take}; +use futures::{future::BoxFuture, TryStreamExt}; + +use super::data_evolution_writer::RowIdFileIndex; +use super::{CommitMessage, DataEvolutionWriter, Table, TableUpdateByRowId, UpdateAssignment}; + +const ROW_ID: &str = "_ROW_ID"; +type ConditionFunction = + dyn Fn(RecordBatch) -> BoxFuture<'static, crate::Result<BooleanArray>> + Send + Sync; + +/// An engine-compiled condition. NULL evaluates to false when choosing a clause. +/// Dependencies control target projection, including source fields for self MERGE. +/// Evaluation is asynchronous so expression engines can execute subquery plans. +#[derive(Clone)] +pub struct MergeCondition { + pub target_columns: Vec<String>, + pub source_columns: Vec<String>, + pub evaluate: Arc<ConditionFunction>, +} + +/// SET/INSERT values. Column references are evaluated by core, without callbacks. +#[derive(Clone)] +pub enum MergeAssignment { + SourceColumn(String), + TargetColumn(String), + Value(UpdateAssignment), +} + +/// Ordered matched clauses, corresponding to WHEN MATCHED UPDATE/DELETE. +#[derive(Clone)] +pub struct WhenMatched { + pub condition: Option<MergeCondition>, + pub delete: bool, + pub assignments: Vec<(String, MergeAssignment)>, +} + +/// Ordered WHEN NOT MATCHED INSERT clauses. +#[derive(Clone)] +pub struct WhenNotMatched { + pub condition: Option<MergeCondition>, + pub assignments: Vec<(String, MergeAssignment)>, +} + +/// Already materialized input, or the target itself joined on its row ID. +pub enum MergeSource { + Batches(Vec<RecordBatch>), + SelfTable, +} + +fn invalid(message: impl Into<String>) -> crate::Error { + crate::Error::DataInvalid { + message: message.into(), + source: None, + } +} + +fn gather(batch: &RecordBatch, rows: &[usize]) -> crate::Result<RecordBatch> { + let indices = UInt64Array::from(rows.iter().map(|&row| row as u64).collect::<Vec<_>>()); + let columns = batch + .columns() + .iter() + .map(|array| { + take(array.as_ref(), &indices, None).map_err(|error| invalid(error.to_string())) + }) + .collect::<crate::Result<_>>()?; + RecordBatch::try_new(batch.schema(), columns).map_err(|error| invalid(error.to_string())) +} + +fn aliases(target: Option<&RecordBatch>, source: &RecordBatch) -> crate::Result<RecordBatch> { + let mut fields = Vec::new(); + let mut columns = Vec::new(); + if let Some(target) = target { + fields.push(Arc::new(Field::new(ROW_ID, DataType::Int64, false))); + columns.push(target.column_by_name(ROW_ID).unwrap().clone()); + for (field, column) in target.schema().fields().iter().zip(target.columns()) { + fields.push(Arc::new( + field + .as_ref() + .clone() + .with_name(format!("t.{}", field.name())), + )); + columns.push(column.clone()); + } + } else { + fields.push(Arc::new(Field::new(ROW_ID, DataType::Int64, false))); + columns.push(Arc::new(Int64Array::from(vec![0; source.num_rows()]))); + } + for (field, column) in source.schema().fields().iter().zip(source.columns()) { + fields.push(Arc::new( + field + .as_ref() + .clone() + .with_name(format!("s.{}", field.name())), + )); + columns.push(column.clone()); + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns) + .map_err(|error| invalid(error.to_string())) +} + +async fn select_clause( + batch: RecordBatch, + condition: &Option<MergeCondition>, +) -> crate::Result<(RecordBatch, RecordBatch)> { + let mask = match condition { + Some(condition) => Some((condition.evaluate)(batch.clone()).await?), + None => None, + }; + if mask + .as_ref() + .is_some_and(|mask| mask.len() != batch.num_rows()) + { + return Err(invalid("MERGE condition length must match input row count")); + } + let mut selected = Vec::new(); + let mut remaining = Vec::new(); + for row in 0..batch.num_rows() { + if mask + .as_ref() + .is_none_or(|mask| !mask.is_null(row) && mask.value(row)) + { + selected.push(row); + } else { + remaining.push(row); + } + } + Ok((gather(&batch, &selected)?, gather(&batch, &remaining)?)) +} + +fn assigned( + batch: &RecordBatch, + assignments: &[(String, MergeAssignment)], + columns: &[String], + schema: SchemaRef, + insert: bool, +) -> crate::Result<RecordBatch> { + let mut values = Vec::new(); + for name in columns { + let value = assignments + .iter() + .find(|(column, _)| column == name) + .map(|(_, value)| value); + let assignment = match value { + Some(MergeAssignment::Value(value)) => value.clone(), + Some(MergeAssignment::SourceColumn(column)) + | Some(MergeAssignment::TargetColumn(column)) => { + let prefix = if matches!(value, Some(MergeAssignment::SourceColumn(_))) { + "s" + } else { + "t" + }; + let column = batch + .column_by_name(&format!("{prefix}.{column}")) + .ok_or_else(|| invalid(format!("Missing MERGE {prefix} column {column}")))?; + UpdateAssignment::Array(vec![column.clone()]) + } + None if insert => UpdateAssignment::Scalar(arrow_array::new_null_array( + schema.field_with_name(name).unwrap().data_type(), + 1, + )), + None => UpdateAssignment::Array(vec![batch + .column_by_name(&format!("t.{name}")) + .ok_or_else(|| invalid(format!("Missing target column {name}")))? + .clone()]), + }; + values.push((name.clone(), assignment)); + } + let batches = + super::update_assignment::assigned_batches(std::slice::from_ref(batch), values, schema)?; + let first = batches + .first() + .ok_or_else(|| invalid("MERGE assignment produced no rows"))?; + concat_batches(&first.schema(), &batches).map_err(|error| invalid(error.to_string())) +} + +pub(super) async fn merge_into( + table: &Table, + commit_user: &str, + source: MergeSource, + on: Vec<(String, String)>, + matched: Vec<WhenMatched>, + not_matched: Vec<WhenNotMatched>, +) -> crate::Result<Vec<CommitMessage>> { + if on.is_empty() || matched.is_empty() && not_matched.is_empty() { + return Err(invalid("MERGE requires ON keys and at least one action")); + } + let schema = crate::arrow::build_target_arrow_schema(table.schema().fields())?; + let mut update_columns = Vec::new(); + let mut projection: HashSet<String> = on.iter().map(|(target, _)| target.clone()).collect(); + projection.insert(ROW_ID.into()); + for clause in &matched { + if clause.delete && !clause.assignments.is_empty() { + return Err(invalid("MERGE DELETE cannot have assignments")); + } + for (name, value) in &clause.assignments { + schema + .field_with_name(name) + .map_err(|error| invalid(error.to_string()))?; + if table.schema().partition_keys().contains(name) { + return Err(invalid("MERGE does not support updating partition columns")); + } + if !update_columns.contains(name) { + update_columns.push(name.clone()); + } + if let MergeAssignment::TargetColumn(name) = value { + projection.insert(name.clone()); + } + if matches!(value, MergeAssignment::Value(UpdateAssignment::Function(_))) { + projection.extend(schema.fields().iter().map(|field| field.name().clone())); + } + } + if let Some(condition) = &clause.condition { + projection.extend(condition.target_columns.clone()); + } + } + projection.extend(update_columns.clone()); + for clause in ¬_matched { + if clause + .condition + .as_ref() + .is_some_and(|condition| !condition.target_columns.is_empty()) + || clause + .assignments + .iter() + .any(|(_, value)| matches!(value, MergeAssignment::TargetColumn(_))) + { + return Err(invalid("WHEN NOT MATCHED cannot reference target columns")); + } + } + for conditions in [ + matched + .iter() + .map(|clause| &clause.condition) + .collect::<Vec<_>>(), + not_matched.iter().map(|clause| &clause.condition).collect(), + ] { + if conditions + .iter() + .take(conditions.len().saturating_sub(1)) + .any(|condition| condition.is_none()) + { + return Err(invalid("Only the last MERGE clause may omit its condition")); + } + } + let _ = DataEvolutionWriter::new(table, update_columns.clone())?; + let snapshot = super::time_travel::resolve_snapshot(table).await?; + let scan_table = match &snapshot { + Some(snapshot) => table.copy_with_pinned_snapshot(snapshot), + None => table.clone(), + } + .copy_with_options(HashMap::from([ + ("scalar-index.search-mode".into(), "FULL".into()), + ("scan.mode".into(), "default".into()), + ])); + let self_merge = matches!(source, MergeSource::SelfTable); + if self_merge { + if on != [(ROW_ID.into(), ROW_ID.into())] || !not_matched.is_empty() { + return Err(invalid( + "Self MERGE requires ON _ROW_ID and no insert clauses", + )); + } + for clause in &matched { + if let Some(condition) = &clause.condition { + projection.extend(condition.source_columns.clone()); + } + for (_, assignment) in &clause.assignments { + if let MergeAssignment::SourceColumn(name) = assignment { + projection.insert(name.clone()); + } + } + } + } + let projection = projection.into_iter().collect::<Vec<_>>(); + let mut builder = scan_table.new_read_builder(); + builder.with_projection(&projection.iter().map(String::as_str).collect::<Vec<_>>())?; + let plan = if snapshot.is_some() { + builder.new_scan().with_scan_all_files().plan().await? + } else { + super::source::Plan::new(Vec::new()) + }; + let index = RowIdFileIndex::from_splits(scan_table.clone(), plan.splits())?; + let source = match source { + MergeSource::Batches(batches) => { + let first = batches + .first() + .ok_or_else(|| invalid("MERGE source needs an Arrow schema"))?; + if batches.iter().any(|batch| batch.schema() != first.schema()) { + return Err(invalid("MERGE source batches must share a schema")); + } + let names = first + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::<HashSet<_>>(); + if names.len() != first.num_columns() { + return Err(invalid("MERGE source has duplicate column names")); + } + Some( + concat_batches(&first.schema(), &batches) + .map_err(|error| invalid(error.to_string()))?, + ) + } + MergeSource::SelfTable => None, + }; + let source_schema = source + .as_ref() + .map_or_else(|| schema.clone(), RecordBatch::schema); + for (assignments, condition, insert) in matched + .iter() + .map(|clause| (&clause.assignments, &clause.condition, false)) + .chain( + not_matched + .iter() + .map(|clause| (&clause.assignments, &clause.condition, true)), + ) + { + let mut seen = HashSet::new(); + for (name, value) in assignments { + if !seen.insert(name) { + return Err(invalid(format!("Duplicate MERGE assignment {name}"))); + } + schema + .field_with_name(name) + .map_err(|error| invalid(error.to_string()))?; + match value { + MergeAssignment::SourceColumn(name) => { + if !(self_merge && name == ROW_ID) { + source_schema + .field_with_name(name) + .map_err(|error| invalid(error.to_string()))?; + } + } + MergeAssignment::TargetColumn(name) => { + if insert { + return Err(invalid("INSERT cannot reference target columns")); + } + if name != ROW_ID { + schema + .field_with_name(name) + .map_err(|error| invalid(error.to_string()))?; + } + } + MergeAssignment::Value(_) => {} + } + } + if let Some(condition) = condition { + for name in &condition.target_columns { + if name != ROW_ID { + schema + .field_with_name(name) + .map_err(|error| invalid(error.to_string()))?; + } + } + } + } + if matched.iter().any(|clause| clause.delete) + && !crate::spec::CoreOptions::new(table.schema().options()).deletion_vectors_enabled() + { + return Err(invalid( + "MERGE DELETE requires deletion-vectors.enabled=true", + )); + } + let mut converter = None; + let mut source_keys = HashMap::<Vec<u8>, Vec<usize>>::new(); + let mut seen_source = vec![false; source.as_ref().map_or(0, RecordBatch::num_rows)]; + if let Some(source) = &source { + let keys = on + .iter() + .map(|(target, source_name)| { + let source_column = source + .column_by_name(source_name) + .ok_or_else(|| invalid(format!("Missing source ON key {source_name}")))?; + let target_type = if target == ROW_ID { + &DataType::Int64 + } else { + schema + .field_with_name(target) + .map_err(|error| invalid(error.to_string()))? + .data_type() + }; + if source_column.data_type() != target_type + || !super::upsert_key_matcher::supported_key_type(target_type) + { + return Err(invalid( + "MERGE ON key types must match and support row encoding", + )); + } + Ok(source_column.clone()) + }) + .collect::<crate::Result<Vec<_>>>()?; + let key_converter = RowConverter::new( + keys.iter() + .map(|column| SortField::new(column.data_type().clone())) + .collect(), + ) + .map_err(|error| invalid(error.to_string()))?; + let rows = key_converter + .convert_columns(&keys) + .map_err(|error| invalid(error.to_string()))?; Review Comment: [P1] Normalize floating-point ON keys before row encoding Float32/Float64 are accepted here, but RowConverter preserves the distinction between -0.0 and +0.0 and between different NaN bit patterns. Paimon Spark SQL joins consider these equal. I reproduced this with a DOUBLE key: seed target (-0.0, 10), merge source (+0.0, 99) with an unconditional matched UPDATE, and merge_into returns no messages instead of updating the target. The same happens for f64::NAN versus f64::from_bits(0x7ff8_0000_0000_0001). A WHEN NOT MATCHED INSERT would consequently insert an extra row, and equivalent source keys can also escape cardinality validation. Please normalize both source and target matching keys before encoding (without changing the payload columns), and cover signed zero and distinct NaN representations in the MERGE tests. ########## crates/paimon/src/table/table_merge_into.rs: ########## @@ -0,0 +1,578 @@ +// 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. + +//! Data-evolution MERGE orchestration, independent of an expression engine. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use arrow_array::{Array, BooleanArray, Int64Array, RecordBatch, UInt64Array}; +use arrow_row::{RowConverter, SortField}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use arrow_select::{concat::concat_batches, take::take}; +use futures::{future::BoxFuture, TryStreamExt}; + +use super::data_evolution_writer::RowIdFileIndex; +use super::{CommitMessage, DataEvolutionWriter, Table, TableUpdateByRowId, UpdateAssignment}; + +const ROW_ID: &str = "_ROW_ID"; +type ConditionFunction = + dyn Fn(RecordBatch) -> BoxFuture<'static, crate::Result<BooleanArray>> + Send + Sync; + +/// An engine-compiled condition. NULL evaluates to false when choosing a clause. +/// Dependencies control target projection, including source fields for self MERGE. +/// Evaluation is asynchronous so expression engines can execute subquery plans. +#[derive(Clone)] +pub struct MergeCondition { + pub target_columns: Vec<String>, + pub source_columns: Vec<String>, + pub evaluate: Arc<ConditionFunction>, +} + +/// SET/INSERT values. Column references are evaluated by core, without callbacks. +#[derive(Clone)] +pub enum MergeAssignment { + SourceColumn(String), + TargetColumn(String), + Value(UpdateAssignment), +} + +/// Ordered matched clauses, corresponding to WHEN MATCHED UPDATE/DELETE. +#[derive(Clone)] +pub struct WhenMatched { + pub condition: Option<MergeCondition>, + pub delete: bool, + pub assignments: Vec<(String, MergeAssignment)>, +} + +/// Ordered WHEN NOT MATCHED INSERT clauses. +#[derive(Clone)] +pub struct WhenNotMatched { + pub condition: Option<MergeCondition>, + pub assignments: Vec<(String, MergeAssignment)>, +} + +/// Already materialized input, or the target itself joined on its row ID. +pub enum MergeSource { + Batches(Vec<RecordBatch>), + SelfTable, +} + +fn invalid(message: impl Into<String>) -> crate::Error { + crate::Error::DataInvalid { + message: message.into(), + source: None, + } +} + +fn gather(batch: &RecordBatch, rows: &[usize]) -> crate::Result<RecordBatch> { + let indices = UInt64Array::from(rows.iter().map(|&row| row as u64).collect::<Vec<_>>()); + let columns = batch + .columns() + .iter() + .map(|array| { + take(array.as_ref(), &indices, None).map_err(|error| invalid(error.to_string())) + }) + .collect::<crate::Result<_>>()?; + RecordBatch::try_new(batch.schema(), columns).map_err(|error| invalid(error.to_string())) +} + +fn aliases(target: Option<&RecordBatch>, source: &RecordBatch) -> crate::Result<RecordBatch> { + let mut fields = Vec::new(); + let mut columns = Vec::new(); + if let Some(target) = target { + fields.push(Arc::new(Field::new(ROW_ID, DataType::Int64, false))); + columns.push(target.column_by_name(ROW_ID).unwrap().clone()); + for (field, column) in target.schema().fields().iter().zip(target.columns()) { + fields.push(Arc::new( + field + .as_ref() + .clone() + .with_name(format!("t.{}", field.name())), + )); + columns.push(column.clone()); + } + } else { + fields.push(Arc::new(Field::new(ROW_ID, DataType::Int64, false))); + columns.push(Arc::new(Int64Array::from(vec![0; source.num_rows()]))); + } + for (field, column) in source.schema().fields().iter().zip(source.columns()) { + fields.push(Arc::new( + field + .as_ref() + .clone() + .with_name(format!("s.{}", field.name())), + )); + columns.push(column.clone()); + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns) + .map_err(|error| invalid(error.to_string())) +} + +async fn select_clause( + batch: RecordBatch, + condition: &Option<MergeCondition>, +) -> crate::Result<(RecordBatch, RecordBatch)> { + let mask = match condition { + Some(condition) => Some((condition.evaluate)(batch.clone()).await?), + None => None, + }; + if mask + .as_ref() + .is_some_and(|mask| mask.len() != batch.num_rows()) + { + return Err(invalid("MERGE condition length must match input row count")); + } + let mut selected = Vec::new(); + let mut remaining = Vec::new(); + for row in 0..batch.num_rows() { + if mask + .as_ref() + .is_none_or(|mask| !mask.is_null(row) && mask.value(row)) + { + selected.push(row); + } else { + remaining.push(row); + } + } + Ok((gather(&batch, &selected)?, gather(&batch, &remaining)?)) +} + +fn assigned( + batch: &RecordBatch, + assignments: &[(String, MergeAssignment)], + columns: &[String], + schema: SchemaRef, + insert: bool, +) -> crate::Result<RecordBatch> { + let mut values = Vec::new(); + for name in columns { + let value = assignments + .iter() + .find(|(column, _)| column == name) + .map(|(_, value)| value); + let assignment = match value { + Some(MergeAssignment::Value(value)) => value.clone(), + Some(MergeAssignment::SourceColumn(column)) + | Some(MergeAssignment::TargetColumn(column)) => { + let prefix = if matches!(value, Some(MergeAssignment::SourceColumn(_))) { + "s" + } else { + "t" + }; + let column = batch + .column_by_name(&format!("{prefix}.{column}")) + .ok_or_else(|| invalid(format!("Missing MERGE {prefix} column {column}")))?; + UpdateAssignment::Array(vec![column.clone()]) + } + None if insert => UpdateAssignment::Scalar(arrow_array::new_null_array( + schema.field_with_name(name).unwrap().data_type(), + 1, + )), + None => UpdateAssignment::Array(vec![batch + .column_by_name(&format!("t.{name}")) + .ok_or_else(|| invalid(format!("Missing target column {name}")))? + .clone()]), + }; + values.push((name.clone(), assignment)); + } + let batches = + super::update_assignment::assigned_batches(std::slice::from_ref(batch), values, schema)?; + let first = batches + .first() + .ok_or_else(|| invalid("MERGE assignment produced no rows"))?; + concat_batches(&first.schema(), &batches).map_err(|error| invalid(error.to_string())) +} + +pub(super) async fn merge_into( + table: &Table, + commit_user: &str, + source: MergeSource, + on: Vec<(String, String)>, + matched: Vec<WhenMatched>, + not_matched: Vec<WhenNotMatched>, +) -> crate::Result<Vec<CommitMessage>> { + if on.is_empty() || matched.is_empty() && not_matched.is_empty() { + return Err(invalid("MERGE requires ON keys and at least one action")); + } + let schema = crate::arrow::build_target_arrow_schema(table.schema().fields())?; + let mut update_columns = Vec::new(); + let mut projection: HashSet<String> = on.iter().map(|(target, _)| target.clone()).collect(); + projection.insert(ROW_ID.into()); + for clause in &matched { + if clause.delete && !clause.assignments.is_empty() { + return Err(invalid("MERGE DELETE cannot have assignments")); + } + for (name, value) in &clause.assignments { + schema + .field_with_name(name) + .map_err(|error| invalid(error.to_string()))?; + if table.schema().partition_keys().contains(name) { + return Err(invalid("MERGE does not support updating partition columns")); + } + if !update_columns.contains(name) { + update_columns.push(name.clone()); + } + if let MergeAssignment::TargetColumn(name) = value { + projection.insert(name.clone()); + } + if matches!(value, MergeAssignment::Value(UpdateAssignment::Function(_))) { + projection.extend(schema.fields().iter().map(|field| field.name().clone())); + } + } + if let Some(condition) = &clause.condition { + projection.extend(condition.target_columns.clone()); + } + } + projection.extend(update_columns.clone()); + for clause in ¬_matched { + if clause + .condition + .as_ref() + .is_some_and(|condition| !condition.target_columns.is_empty()) + || clause + .assignments + .iter() + .any(|(_, value)| matches!(value, MergeAssignment::TargetColumn(_))) + { + return Err(invalid("WHEN NOT MATCHED cannot reference target columns")); + } + } + for conditions in [ + matched + .iter() + .map(|clause| &clause.condition) + .collect::<Vec<_>>(), + not_matched.iter().map(|clause| &clause.condition).collect(), + ] { + if conditions + .iter() + .take(conditions.len().saturating_sub(1)) + .any(|condition| condition.is_none()) + { + return Err(invalid("Only the last MERGE clause may omit its condition")); + } + } + let _ = DataEvolutionWriter::new(table, update_columns.clone())?; + let snapshot = super::time_travel::resolve_snapshot(table).await?; + let scan_table = match &snapshot { + Some(snapshot) => table.copy_with_pinned_snapshot(snapshot), + None => table.clone(), + } + .copy_with_options(HashMap::from([ + ("scalar-index.search-mode".into(), "FULL".into()), + ("scan.mode".into(), "default".into()), + ])); + let self_merge = matches!(source, MergeSource::SelfTable); + if self_merge { + if on != [(ROW_ID.into(), ROW_ID.into())] || !not_matched.is_empty() { + return Err(invalid( + "Self MERGE requires ON _ROW_ID and no insert clauses", + )); + } + for clause in &matched { + if let Some(condition) = &clause.condition { + projection.extend(condition.source_columns.clone()); + } + for (_, assignment) in &clause.assignments { + if let MergeAssignment::SourceColumn(name) = assignment { + projection.insert(name.clone()); + } + } + } + } + let projection = projection.into_iter().collect::<Vec<_>>(); + let mut builder = scan_table.new_read_builder(); + builder.with_projection(&projection.iter().map(String::as_str).collect::<Vec<_>>())?; + let plan = if snapshot.is_some() { + builder.new_scan().with_scan_all_files().plan().await? + } else { + super::source::Plan::new(Vec::new()) + }; + let index = RowIdFileIndex::from_splits(scan_table.clone(), plan.splits())?; + let source = match source { + MergeSource::Batches(batches) => { + let first = batches + .first() + .ok_or_else(|| invalid("MERGE source needs an Arrow schema"))?; + if batches.iter().any(|batch| batch.schema() != first.schema()) { + return Err(invalid("MERGE source batches must share a schema")); + } + let names = first + .schema() + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::<HashSet<_>>(); + if names.len() != first.num_columns() { + return Err(invalid("MERGE source has duplicate column names")); + } + Some( + concat_batches(&first.schema(), &batches) + .map_err(|error| invalid(error.to_string()))?, + ) + } + MergeSource::SelfTable => None, + }; + let source_schema = source + .as_ref() + .map_or_else(|| schema.clone(), RecordBatch::schema); + for (assignments, condition, insert) in matched + .iter() + .map(|clause| (&clause.assignments, &clause.condition, false)) + .chain( + not_matched + .iter() + .map(|clause| (&clause.assignments, &clause.condition, true)), + ) + { + let mut seen = HashSet::new(); + for (name, value) in assignments { + if !seen.insert(name) { + return Err(invalid(format!("Duplicate MERGE assignment {name}"))); + } + schema + .field_with_name(name) + .map_err(|error| invalid(error.to_string()))?; + match value { + MergeAssignment::SourceColumn(name) => { + if !(self_merge && name == ROW_ID) { + source_schema + .field_with_name(name) + .map_err(|error| invalid(error.to_string()))?; + } + } + MergeAssignment::TargetColumn(name) => { + if insert { + return Err(invalid("INSERT cannot reference target columns")); + } + if name != ROW_ID { + schema + .field_with_name(name) + .map_err(|error| invalid(error.to_string()))?; + } + } + MergeAssignment::Value(_) => {} + } + } + if let Some(condition) = condition { + for name in &condition.target_columns { + if name != ROW_ID { + schema + .field_with_name(name) + .map_err(|error| invalid(error.to_string()))?; + } + } + } + } + if matched.iter().any(|clause| clause.delete) + && !crate::spec::CoreOptions::new(table.schema().options()).deletion_vectors_enabled() + { + return Err(invalid( + "MERGE DELETE requires deletion-vectors.enabled=true", + )); + } + let mut converter = None; + let mut source_keys = HashMap::<Vec<u8>, Vec<usize>>::new(); + let mut seen_source = vec![false; source.as_ref().map_or(0, RecordBatch::num_rows)]; + if let Some(source) = &source { + let keys = on + .iter() + .map(|(target, source_name)| { + let source_column = source + .column_by_name(source_name) + .ok_or_else(|| invalid(format!("Missing source ON key {source_name}")))?; + let target_type = if target == ROW_ID { + &DataType::Int64 + } else { + schema + .field_with_name(target) + .map_err(|error| invalid(error.to_string()))? + .data_type() + }; + if source_column.data_type() != target_type + || !super::upsert_key_matcher::supported_key_type(target_type) + { + return Err(invalid( + "MERGE ON key types must match and support row encoding", + )); + } + Ok(source_column.clone()) + }) + .collect::<crate::Result<Vec<_>>>()?; + let key_converter = RowConverter::new( + keys.iter() + .map(|column| SortField::new(column.data_type().clone())) + .collect(), + ) + .map_err(|error| invalid(error.to_string()))?; + let rows = key_converter + .convert_columns(&keys) + .map_err(|error| invalid(error.to_string()))?; + for row in 0..source.num_rows() { + if keys.iter().all(|column| !column.is_null(row)) { + source_keys + .entry(rows.row(row).as_ref().to_vec()) + .or_default() + .push(row); + } + } + converter = Some(key_converter); + } + // Spark validates cardinality before evaluating action predicates. A sole + // unconditional DELETE may match multiple source rows without ambiguity. + let check_cardinality = !matched.is_empty() + && !(matched.len() == 1 && matched[0].delete && matched[0].condition.is_none()); + let mut joined = Vec::new(); + let mut stream = builder.new_read()?.to_arrow(plan.splits())?; + while let Some(target) = stream.try_next().await? { + if let Some(source) = &source { + let keys = on + .iter() + .map(|(name, _)| target.column_by_name(name).unwrap().clone()) + .collect::<Vec<_>>(); + let rows = converter + .as_ref() + .unwrap() + .convert_columns(&keys) + .map_err(|error| invalid(error.to_string()))?; + let mut target_rows = Vec::new(); + let mut source_rows = Vec::new(); + for row in 0..target.num_rows() { + if keys.iter().any(|column| column.is_null(row)) { + continue; + } + if let Some(sources) = source_keys.get(rows.row(row).as_ref()) { + if check_cardinality && sources.len() > 1 { + return Err(invalid( + "MERGE matched multiple source rows to the same target _ROW_ID", + )); + } + for &source in sources { + seen_source[source] = true; + } + target_rows.push(row); + source_rows.push(sources[0]); + } + } + if !target_rows.is_empty() { + joined.push(aliases( + Some(&gather(&target, &target_rows)?), + &gather(source, &source_rows)?, + )?); + } + } else if target.num_rows() > 0 { + joined.push(aliases(Some(&target), &target)?); + } + } + let mut updates = Vec::new(); + let mut deletes = Vec::new(); + let mut pending = joined; + for clause in matched { + let mut remaining = Vec::new(); + for batch in pending { + let (selected, rest) = select_clause(batch, &clause.condition).await?; + if rest.num_rows() > 0 { + remaining.push(rest); + } + if selected.num_rows() == 0 { + continue; + } + if clause.delete { + let ids = selected + .column_by_name(ROW_ID) + .unwrap() + .as_any() + .downcast_ref::<Int64Array>() + .unwrap(); + deletes.extend(ids.values().iter().copied()); + } else if !update_columns.is_empty() { + updates.push(assigned( + &selected, + &clause.assignments, + &update_columns, + schema.clone(), + false, + )?); Review Comment: [P2] Apply array assignments across the complete clause selection The complete UpdateAssignment::Array is reused for each selected target RecordBatch. assigned() calls assigned_batches with only that one batch, so a valid array assignment starts failing when the target is split into multiple files/batches. Reproducer: seed (id=1,value=10) and (id=2,value=20) in separate commits, merge source ids [1,2], and assign value from two chunks [100] and [200]. This fails with "Assignment array length must match matched row count: 2 != 1"; the existing single-file version passes. Please evaluate/consume the assignment over all selected rows of the clause, rather than restarting it for each physical target batch, and add a multi-file/chunked regression. -- 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]
