This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 3def4702 [core] Share MERGE action evaluation for distributed
execution (#1203)
3def4702 is described below
commit 3def4702fc0ed39026b6bba6c299d09c76558e36
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Oct 23 09:26:27 2026 +0800
[core] Share MERGE action evaluation for distributed execution (#1203)
---
.../python/python/pypaimon_rust/datafusion.pyi | 18 +
bindings/python/src/context.rs | 1 +
bindings/python/src/lib.rs | 1 +
bindings/python/src/merge_execution.rs | 520 +++++++++++++++++++++
bindings/python/src/merge_into.rs | 22 +-
bindings/python/src/update_assignment.rs | 10 +
crates/paimon/src/table/mod.rs | 3 +-
.../paimon/src/table/table_merge_actions_tests.rs | 403 ++++++++++++++++
crates/paimon/src/table/table_merge_into.rs | 313 +++++++++----
9 files changed, 1202 insertions(+), 89 deletions(-)
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index f50293ef..ae393712 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -79,6 +79,24 @@ class RowIdPlan:
...
def plan(self) -> Plan: ...
+class MergeBatchExecution:
+ """Evaluate ordered MERGE actions over already joined Arrow batches.
+
+ Input fields use the ``t.<column>`` and ``s.<column>`` aliases. This
+ operator evaluates expressions and assignments without reading or writing
+ a table; callers retain responsibility for joins and authorization.
+ Matched input carries ``t._ROW_ID`` or ``_ROW_ID``, never both. Update
+ output contains ``_ROW_ID`` and the requested target columns; delete output
+ contains only ``_ROW_ID``. Insert output uses the target schema.
+ """
+ def __init__(
+ self, input_schema: pyarrow.Schema, output_schema: pyarrow.Schema,
+ clauses: Sequence[Dict[str, Any]], mode: Literal["update", "delete",
"insert"],
+ ) -> None: ...
+ def execute(self, batch: Union[pyarrow.Table, pyarrow.RecordBatch]) ->
pyarrow.Table:
+ """Return selected actions; unselected assignments are never
evaluated."""
+ ...
+
class ScanStats:
"""Read-only counters from normal metadata planning."""
@property
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index 33fdd265..162c0444 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -1038,6 +1038,7 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_,
PyModule>) -> PyResult<()>
this.add_class::<crate::write::PyBatchTableWrite>()?;
this.add_class::<crate::write::PyBatchTableUpdate>()?;
this.add_class::<crate::write::PyTableUpdateByRowId>()?;
+ this.add_class::<crate::merge_execution::PyMergeBatchExecution>()?;
this.add_class::<crate::shard_update::PyShardTableUpdator>()?;
this.add_class::<crate::write::PyStreamTableUpdate>()?;
this.add_class::<crate::write::PyStreamTableWrite>()?;
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
index 49b69420..e45485bf 100644
--- a/bindings/python/src/lib.rs
+++ b/bindings/python/src/lib.rs
@@ -29,6 +29,7 @@ mod full_text_search;
mod global_index_result;
mod hybrid_search;
mod index_build;
+mod merge_execution;
mod merge_into;
mod oss_cpp;
mod predicate;
diff --git a/bindings/python/src/merge_execution.rs
b/bindings/python/src/merge_execution.rs
new file mode 100644
index 00000000..7310bfdc
--- /dev/null
+++ b/bindings/python/src/merge_execution.rs
@@ -0,0 +1,520 @@
+// 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.
+
+//! Arrow and Python transport for the core prejoined MERGE evaluator.
+
+use std::collections::HashSet;
+use std::sync::{Arc, Mutex};
+
+use arrow::array::Int64Array;
+use arrow::datatypes::{DataType, Schema, SchemaRef};
+use arrow::pyarrow::{FromPyArrow, ToPyArrow};
+use arrow::record_batch::RecordBatch;
+use paimon::table::{
+ evaluate_merge_actions, MergeActionOutput, MergeAssignment, WhenMatched,
WhenNotMatched,
+};
+use paimon_datafusion::runtime::runtime;
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyList};
+
+use crate::error::to_py_err;
+
+#[pyclass(name = "MergeBatchExecution")]
+pub struct PyMergeBatchExecution {
+ input_schema: SchemaRef,
+ output_schema: SchemaRef,
+ target_schema: SchemaRef,
+ matched: Vec<WhenMatched>,
+ not_matched: Vec<WhenNotMatched>,
+ output: MergeActionOutput,
+ callback_error: Arc<Mutex<Option<PyErr>>>,
+ execution: Mutex<()>,
+}
+
+fn matched_schema(schema: SchemaRef) -> PyResult<SchemaRef> {
+ if schema.column_with_name("_ROW_ID").is_some() {
+ if schema.column_with_name("t._ROW_ID").is_some() {
+ return Err(PyValueError::new_err(
+ "Matched MERGE input has ambiguous _ROW_ID and t._ROW_ID
fields",
+ ));
+ }
+ return Ok(schema);
+ }
+ let index = schema.index_of("t._ROW_ID").map_err(|error| {
+ PyValueError::new_err(format!("Matched MERGE batch needs t._ROW_ID:
{error}"))
+ })?;
+ let mut fields = schema.fields().to_vec();
+ fields.push(Arc::new(
+ fields[index].as_ref().clone().with_name("_ROW_ID"),
+ ));
+ Ok(Arc::new(Schema::new_with_metadata(
+ fields,
+ schema.metadata().clone(),
+ )))
+}
+
+fn row_ids(batch: RecordBatch) -> PyResult<RecordBatch> {
+ let schema = matched_schema(batch.schema())?;
+ if schema.fields().len() == batch.num_columns() {
+ return Ok(batch);
+ }
+ let index = batch.schema().index_of("t._ROW_ID").unwrap();
+ let mut columns = batch.columns().to_vec();
+ columns.push(batch.column(index).clone());
+ RecordBatch::try_new(schema, columns).map_err(|error|
PyValueError::new_err(error.to_string()))
+}
+
+#[pymethods]
+impl PyMergeBatchExecution {
+ #[new]
+ fn new(
+ py: Python<'_>,
+ input_schema: &Bound<'_, PyAny>,
+ output_schema: &Bound<'_, PyAny>,
+ clauses: &Bound<'_, PyAny>,
+ mode: &str,
+ ) -> PyResult<Self> {
+ let output = match mode {
+ "update" => MergeActionOutput::Updates,
+ "delete" => MergeActionOutput::Deletes,
+ "insert" => MergeActionOutput::Inserts,
+ _ => {
+ return Err(PyValueError::new_err(
+ "MERGE execution mode must be update, delete or insert",
+ ));
+ }
+ };
+ let input_schema = Arc::new(Schema::from_pyarrow_bound(input_schema)?);
+ let output_schema =
Arc::new(Schema::from_pyarrow_bound(output_schema)?);
+ if output != MergeActionOutput::Inserts {
+ let row_id = output_schema.field_with_name("_ROW_ID").map_err(|_| {
+ PyValueError::new_err("Matched MERGE output needs an Int64
_ROW_ID field")
+ })?;
+ if row_id.data_type() != &DataType::Int64 {
+ return Err(PyValueError::new_err(
+ "Matched MERGE output needs an Int64 _ROW_ID field",
+ ));
+ }
+ if output == MergeActionOutput::Deletes &&
output_schema.fields().len() != 1 {
+ return Err(PyValueError::new_err(
+ "MERGE DELETE output must contain only _ROW_ID",
+ ));
+ }
+ }
+ let condition_schema = if output == MergeActionOutput::Inserts {
+ input_schema.clone()
+ } else {
+ matched_schema(input_schema.clone())?
+ };
+ let target_schema = Arc::new(Schema::new(
+ output_schema
+ .fields()
+ .iter()
+ .filter(|field| output == MergeActionOutput::Inserts ||
field.name() != "_ROW_ID")
+ .cloned()
+ .collect::<Vec<_>>(),
+ ));
+ let callback_error = Arc::new(Mutex::new(None));
+ let mut matched = Vec::new();
+ let mut not_matched = Vec::new();
+ for clause in clauses.try_iter()? {
+ let clause = clause?;
+ let clause = clause.cast::<PyDict>()?;
+ let condition = crate::merge_into::condition(py, clause,
condition_schema.clone())?;
+ let assignments = if output == MergeActionOutput::Deletes {
+ // A DELETE route still consumes preceding UPDATE conditions,
+ // but must not inspect or evaluate those assignments.
+ Vec::new()
+ } else {
+ crate::merge_into::assignments(
+ py,
+ clause,
+ target_schema.clone(),
+ callback_error.clone(),
+ )?
+ };
+ if output == MergeActionOutput::Inserts {
+ not_matched.push(WhenNotMatched {
+ condition,
+ assignments,
+ });
+ } else {
+ matched.push(WhenMatched {
+ condition,
+ delete: clause
+ .get_item("delete")?
+ .ok_or_else(|| PyValueError::new_err("MERGE matched
clause needs action"))?
+ .extract()?,
+ assignments,
+ });
+ }
+ }
+ if output == MergeActionOutput::Updates {
+ // The explicit output schema may request fields untouched by every
+ // clause. Include them in core's shared update-column projection.
+ let assigned = matched
+ .iter()
+ .flat_map(|clause| clause.assignments.iter().map(|(name, _)|
name.clone()))
+ .collect::<HashSet<_>>();
+ if let Some(clause) = matched.iter_mut().find(|clause|
!clause.delete) {
+ for field in target_schema.fields() {
+ if !assigned.contains(field.name()) {
+ clause.assignments.push((
+ field.name().clone(),
+
MergeAssignment::TargetColumn(field.name().clone()),
+ ));
+ }
+ }
+ }
+ }
+ // Reuse core's clause and assignment validation before a worker starts
+ // consuming data. Empty input cannot evaluate lazy assignment values.
+ py.detach(|| {
+ runtime().block_on(evaluate_merge_actions(
+ target_schema.clone(),
+ Vec::new(),
+ Vec::new(),
+ &matched,
+ ¬_matched,
+ output,
+ ))
+ })
+ .map_err(to_py_err)?;
+ Ok(Self {
+ input_schema,
+ output_schema,
+ target_schema,
+ matched,
+ not_matched,
+ output,
+ callback_error,
+ execution: Mutex::new(()),
+ })
+ }
+
+ fn execute(&self, py: Python<'_>, batch: &Bound<'_, PyAny>) ->
PyResult<Py<PyAny>> {
+ let schema = Schema::from_pyarrow_bound(&batch.getattr("schema")?)?;
+ if &schema != self.input_schema.as_ref() {
+ return Err(PyValueError::new_err(
+ "MERGE batch schema differs from the compiled input schema",
+ ));
+ }
+ let pa = py.import("pyarrow")?;
+ let mut batches = if batch.is_instance(&pa.getattr("RecordBatch")?)? {
+ vec![crate::arrow_input::batch_from_pyarrow(batch)?]
+ } else {
+ batch
+ .call_method0("to_batches")?
+ .try_iter()?
+ .map(|batch| crate::arrow_input::batch_from_pyarrow(&batch?))
+ .collect::<PyResult<Vec<_>>>()?
+ };
+ if batches.is_empty() {
+ batches.push(RecordBatch::new_empty(self.input_schema.clone()));
+ }
+ let (matched, unmatched) = if self.output ==
MergeActionOutput::Inserts {
+ (Vec::new(), batches)
+ } else {
+ (
+ batches.into_iter().map(row_ids).collect::<PyResult<_>>()?,
+ Vec::new(),
+ )
+ };
+ let (result, callback_error) = py.detach(|| {
+ // Cached operators may be called by multiple Python threads. Take
+ // this lock without the GIL so callbacks retain their error owner.
+ let _execution = self.execution.lock().unwrap();
+ let result = runtime().block_on(evaluate_merge_actions(
+ self.target_schema.clone(),
+ matched,
+ unmatched,
+ &self.matched,
+ &self.not_matched,
+ self.output,
+ ));
+ (result, self.callback_error.lock().unwrap().take())
+ });
+ if let Some(error) = callback_error {
+ return Err(error);
+ }
+ let actions = result.map_err(to_py_err)?;
+ let batches = match self.output {
+ MergeActionOutput::Updates => actions.updates,
+ MergeActionOutput::Inserts => actions.inserts,
+ MergeActionOutput::Deletes => vec![RecordBatch::try_new(
+ self.output_schema.clone(),
+ vec![Arc::new(Int64Array::from(actions.deletes))],
+ )
+ .map_err(|error| PyValueError::new_err(error.to_string()))?],
+ MergeActionOutput::All => unreachable!(),
+ };
+ let batches = batches
+ .iter()
+ .map(|batch| {
+ let columns = self
+ .output_schema
+ .fields()
+ .iter()
+ .map(|field| {
+
batch.column_by_name(field.name()).cloned().ok_or_else(|| {
+ PyValueError::new_err(format!(
+ "Missing MERGE output column {}",
+ field.name()
+ ))
+ })
+ })
+ .collect::<PyResult<_>>()?;
+ RecordBatch::try_new(self.output_schema.clone(), columns)
+ .map_err(|error| PyValueError::new_err(error.to_string()))?
+ .to_pyarrow(py)
+ })
+ .collect::<PyResult<Vec<_>>>()?;
+ Ok(py
+ .import("pyarrow")?
+ .getattr("Table")?
+ .call_method1(
+ "from_batches",
+ (
+ PyList::new(py, batches)?,
+ self.output_schema.to_pyarrow(py)?,
+ ),
+ )?
+ .unbind())
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use pyo3::ffi::c_str;
+
+ fn locals(py: Python<'_>) -> Bound<'_, PyDict> {
+ let locals = PyDict::new(py);
+ locals
+ .set_item(
+ "MergeBatchExecution",
+ py.get_type::<PyMergeBatchExecution>(),
+ )
+ .unwrap();
+ locals
+ }
+
+ #[test]
+ fn ordered_matched_actions_keep_schema_and_callback_grouping() {
+ Python::attach(|py| {
+ let locals = locals(py);
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+input_schema = pa.schema([
+ pa.field('t._ROW_ID', pa.int64(), nullable=False),
+ pa.field('t.value', pa.int32()),
+ pa.field('s.value', pa.int32()),
+ pa.field('t.keep', pa.string()),
+ pa.field('t.unassigned', pa.int32()),
+], metadata={'input': 'joined'})
+first = pa.record_batch([
+ pa.array([1, 2, 3], pa.int64()), pa.array([1, 2, 3], pa.int32()),
+ pa.array([15, None, 99], pa.int32()), pa.array(['one', 'two', 'three']),
+ pa.array([101, 102, 103], pa.int32()),
+], schema=input_schema)
+second = pa.record_batch([
+ pa.array([4, 5], pa.int64()), pa.array([4, 5], pa.int32()),
+ pa.array([20, 5], pa.int32()), pa.array(['four', 'five']),
+ pa.array([104, 105], pa.int32()),
+], schema=input_schema)
+joined = pa.Table.from_batches([first, second])
+calls = []
+def update(rows):
+ calls.append(rows['t._ROW_ID'].to_pylist())
+ assert rows['s.value'].to_pylist() == [15, 20]
+ return pa.chunked_array([pa.array([30]), pa.array([40])])
+clauses = [
+ dict(condition={'sql': '"s.value" >= 90'}, delete=True, assignments=[]),
+ dict(condition={'sql': '"s.value" > 10'}, delete=False,
+ assignments=[('value', 'callable', update), ('keep', 'target',
'keep')]),
+ dict(condition=None, delete=False, assignments=[('value', 'literal', 50)]),
+]
+output_schema = pa.schema([
+ pa.field('_ROW_ID', pa.int64(), nullable=False),
+ pa.field('keep', pa.string(), metadata={'field': 'kept'}),
+ pa.field('unassigned', pa.int32()),
+ pa.field('value', pa.int32()),
+], metadata={'output': 'updates'})
+execution = MergeBatchExecution(input_schema, output_schema, clauses, 'update')
+result = execution.execute(joined)
+assert result.schema.equals(output_schema, check_metadata=True)
+assert result.to_pydict() == {
+ '_ROW_ID': [1, 4, 2, 5], 'keep': ['one', 'four', 'two', 'five'],
+ 'unassigned': [101, 104, 102, 105],
+ 'value': [30, 40, 50, 50],
+}
+assert calls == [[1, 4]]
+delete_schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False)])
+deletes = MergeBatchExecution(input_schema, delete_schema, clauses,
'delete').execute(joined)
+assert deletes.schema == delete_schema
+assert deletes.to_pydict() == {'_ROW_ID': [3]}
+assert calls == [[1, 4]]
+assert
execution.execute(input_schema.empty_table()).schema.equals(output_schema,
check_metadata=True)
+assert calls == [[1, 4]]
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ });
+ }
+
+ #[test]
+ fn ordered_inserts_preserve_typed_nulls_and_evaluate_selected_rows_once() {
+ Python::attach(|py| {
+ let locals = locals(py);
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+source = pa.table({'s.value': pa.array([None, 20, 5], pa.int32()),
+ 's.id': pa.array([3, 1, 2], pa.int32())})
+source = pa.Table.from_batches(source.to_batches(max_chunksize=1))
+target = pa.schema([pa.field('label', pa.string()), pa.field('id', pa.int32()),
+ pa.field('value', pa.int32()), pa.field('missing',
pa.bool_())])
+calls = []
+def insert(rows):
+ calls.append(rows['s.id'].to_pylist())
+ return pa.array([21], pa.int32())
+clauses = [
+ dict(condition={'sql': '"s.value" > 10'},
+ assignments=[('id', 'source', 'id'), ('value', 'callable', insert),
+ ('label', 'literal', 'first')]),
+ dict(condition=None,
+ assignments=[('id', 'source', 'id'), ('value', 'literal', 7),
+ ('label', 'literal', 'last'), ('missing', 'literal',
True)]),
+]
+execution = MergeBatchExecution(source.schema, target, clauses, 'insert')
+result = execution.execute(source)
+assert result.schema == target
+assert result.to_pydict() == {'label': ['first', 'last', 'last'], 'id': [1, 3,
2],
+ 'value': [21, 7, 7], 'missing': [None, True, True]}
+assert calls == [[1]]
+assert execution.execute(source.slice(0, 0)).schema == target
+assert calls == [[1]]
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ });
+ }
+
+ #[test]
+ fn callback_errors_keep_identity_and_do_not_poison_cached_execution() {
+ Python::attach(|py| {
+ let locals = locals(py);
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+joined = pa.table({'t._ROW_ID': pa.array([7], pa.int64()),
+ 's.value': pa.array([1], pa.int32())})
+output = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False),
+ pa.field('value', pa.int32())])
+failure = LookupError('assignment failed')
+attempts = []
+def assignment(rows):
+ attempts.append(rows.num_rows)
+ if len(attempts) == 1:
+ raise failure
+ return pa.array([9], pa.int32())
+clauses = [dict(condition=None, delete=False,
+ assignments=[('value', 'callable', assignment)])]
+execution = MergeBatchExecution(joined.schema, output, clauses, 'update')
+try:
+ execution.execute(joined)
+except LookupError as error:
+ assert error is failure
+else:
+ raise AssertionError('callback did not raise')
+assert execution.execute(joined).to_pydict() == {'_ROW_ID': [7], 'value': [9]}
+assert attempts == [1, 1]
+assert execution.execute(joined.to_batches()[0]).to_pydict() == {'_ROW_ID':
[7], 'value': [9]}
+assert attempts == [1, 1, 1]
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ });
+ }
+
+ #[test]
+ fn invalid_conditions_lengths_and_schema_changes_fail_before_output() {
+ Python::attach(|py| {
+ let locals = locals(py);
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+joined = pa.table({'t._ROW_ID': pa.array([1, 2], pa.int64()),
+ 's.value': pa.array([2, 3], pa.int32())})
+output = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False),
+ pa.field('value', pa.int32())])
+def rejected(function, text):
+ try:
+ function()
+ except ValueError as error:
+ assert text in str(error), str(error)
+ else:
+ raise AssertionError('invalid MERGE operation succeeded')
+clause = dict(condition={'sql': '"s.absent" > 0'}, delete=False,
+ assignments=[('value', 'literal', 1)])
+rejected(lambda: MergeBatchExecution(joined.schema, output, [clause],
'update'), 'absent')
+clause = dict(condition=None, delete=False,
+ assignments=[('value', 'callable', lambda rows: pa.array([1],
pa.int32()))])
+execution = MergeBatchExecution(joined.schema, output, [clause], 'update')
+rejected(lambda: execution.execute(joined), 'length')
+rejected(lambda: execution.execute(joined.rename_columns(['t._ROW_ID',
's.other'])), 'schema')
+rejected(lambda: MergeBatchExecution(joined.schema, output, [clause],
'unknown'), 'mode')
+rejected(lambda: MergeBatchExecution(joined.schema, pa.schema([('value',
pa.int32())]),
+ [clause], 'update'), '_ROW_ID')
+rejected(lambda: MergeBatchExecution(joined.schema, pa.schema([('_ROW_ID',
pa.int32())]),
+ [clause], 'delete'), 'Int64')
+rejected(lambda: MergeBatchExecution(joined.schema, output, [clause],
'delete'), 'only')
+rejected(lambda: MergeBatchExecution(joined.schema, output, [clause, clause],
'update'), 'last')
+dual_ids = joined.append_column('_ROW_ID', pa.array([99, 98], pa.int64()))
+rejected(lambda: MergeBatchExecution(dual_ids.schema, output, [clause],
'update'), 'ambiguous')
+bad_type = dict(condition=None, delete=False,
+ assignments=[('value', 'callable', lambda rows: [1, 2])])
+rejected(lambda: MergeBatchExecution(joined.schema, output, [bad_type],
'update').execute(joined),
+ 'pyarrow.Array')
+# A no-match clause must not convert an unused literal or invoke a callback.
+clause = dict(condition={'sql': '"s.value" < 0'}, delete=False,
+ assignments=[('value', 'literal', object())])
+assert MergeBatchExecution(joined.schema, output, [clause],
'update').execute(joined).num_rows == 0
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ });
+ }
+}
diff --git a/bindings/python/src/merge_into.rs
b/bindings/python/src/merge_into.rs
index 4d237d2f..3d6611ae 100644
--- a/bindings/python/src/merge_into.rs
+++ b/bindings/python/src/merge_into.rs
@@ -30,7 +30,7 @@ use crate::error::to_py_err;
type CallbackError = Arc<Mutex<Option<PyErr>>>;
-fn condition(
+pub(crate) fn condition(
py: Python<'_>,
clause: &Bound<'_, PyDict>,
schema: SchemaRef,
@@ -51,10 +51,10 @@ fn condition(
.map_err(to_py_err)
}
-fn assignments(
+pub(crate) fn assignments(
py: Python<'_>,
clause: &Bound<'_, PyDict>,
- fields: &[paimon::spec::DataField],
+ schema: SchemaRef,
errors: CallbackError,
) -> PyResult<Vec<(String, MergeAssignment)>> {
let values = clause
@@ -75,10 +75,14 @@ fn assignments(
}
let values = PyDict::new(py);
values.set_item(&name, value)?;
- let assignment =
- crate::update_assignment::from_python(py, &values,
fields, errors.clone())?
- .remove(0)
- .1;
+ let assignment =
crate::update_assignment::from_python_schema(
+ py,
+ &values,
+ schema.clone(),
+ errors.clone(),
+ )?
+ .remove(0)
+ .1;
MergeAssignment::Value(assignment)
}
_ => return Err(PyValueError::new_err("Unknown MERGE
assignment kind")),
@@ -123,7 +127,7 @@ pub(crate) fn from_python(
.get_item("delete")?
.ok_or_else(|| PyValueError::new_err("MERGE matched clause
needs action"))?
.extract()?,
- assignments: assignments(py, clause, table.schema().fields(),
errors.clone())?,
+ assignments: assignments(py, clause, target_schema.clone(),
errors.clone())?,
})
})
.collect::<PyResult<Vec<_>>>()?;
@@ -146,7 +150,7 @@ pub(crate) fn from_python(
let clause = item.cast::<PyDict>()?;
Ok(WhenNotMatched {
condition: condition(py, clause, schema.clone())?,
- assignments: assignments(py, clause, table.schema().fields(),
errors.clone())?,
+ assignments: assignments(py, clause, target_schema.clone(),
errors.clone())?,
})
})
.collect::<PyResult<Vec<_>>>()?;
diff --git a/bindings/python/src/update_assignment.rs
b/bindings/python/src/update_assignment.rs
index 290d9560..3df5a3a5 100644
--- a/bindings/python/src/update_assignment.rs
+++ b/bindings/python/src/update_assignment.rs
@@ -18,6 +18,7 @@
//! Python object and callback conversion for core update assignments.
use arrow::array::ArrayRef;
+use arrow::datatypes::SchemaRef;
use arrow::pyarrow::ToPyArrow;
use paimon::spec::DataField;
use paimon::table::UpdateAssignment;
@@ -62,6 +63,15 @@ pub(crate) fn from_python(
error: CallbackError,
) -> PyResult<Vec<(String, UpdateAssignment)>> {
let schema =
paimon::arrow::build_target_arrow_schema(fields).map_err(to_py_err)?;
+ from_python_schema(py, assignments, schema, error)
+}
+
+pub(crate) fn from_python_schema(
+ py: Python<'_>,
+ assignments: &Bound<'_, PyDict>,
+ schema: SchemaRef,
+ error: CallbackError,
+) -> PyResult<Vec<(String, UpdateAssignment)>> {
let pa = py.import("pyarrow")?;
let array_type = pa.getattr("Array")?;
let chunked_type = pa.getattr("ChunkedArray")?;
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index 1f95c5af..dc463d17 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -260,7 +260,8 @@ pub use source::{
pub use stream_table_scan::StreamTableScan;
pub use table_commit::TableCommit;
pub use table_merge_into::{
- MergeAssignment, MergeCondition, MergeSource, WhenMatched, WhenNotMatched,
+ evaluate_merge_actions, MergeActionOutput, MergeActions, MergeAssignment,
MergeCondition,
+ MergeSource, WhenMatched, WhenNotMatched,
};
pub use table_read::{AuditLogRead, TableRead};
pub use table_scan::{AuditLogScan, FileScanPlan, TableScan};
diff --git a/crates/paimon/src/table/table_merge_actions_tests.rs
b/crates/paimon/src/table/table_merge_actions_tests.rs
new file mode 100644
index 00000000..b35d2c7b
--- /dev/null
+++ b/crates/paimon/src/table/table_merge_actions_tests.rs
@@ -0,0 +1,403 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use super::{
+ evaluate_merge_actions, MergeActionOutput, MergeAssignment,
MergeCondition, UpdateAssignment,
+ WhenMatched, WhenNotMatched,
+};
+use arrow_array::{
+ Array, ArrayRef, BooleanArray, Int32Array, Int64Array, RecordBatch,
StringArray,
+};
+use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+
+fn target_schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![
+ Field::new("value", DataType::Int32, true),
+ Field::new("other", DataType::Int32, true),
+ ]))
+}
+
+fn joined(ids: Vec<i64>, source: Vec<Option<i64>>) -> RecordBatch {
+ let values = ids.iter().map(|id| *id as i32 * 10).collect::<Vec<_>>();
+ let other = ids.iter().map(|id| *id as i32 * 100).collect::<Vec<_>>();
+ RecordBatch::try_from_iter([
+ ("_ROW_ID", Arc::new(Int64Array::from(ids)) as ArrayRef),
+ ("t.value", Arc::new(Int32Array::from(values)) as ArrayRef),
+ ("t.other", Arc::new(Int32Array::from(other)) as ArrayRef),
+ ("s.value", Arc::new(Int64Array::from(source)) as ArrayRef),
+ ])
+ .unwrap()
+}
+
+fn source(values: Vec<Option<i64>>) -> RecordBatch {
+ RecordBatch::try_from_iter([("s.value", Arc::new(Int64Array::from(values))
as ArrayRef)])
+ .unwrap()
+}
+
+fn greater_than(bound: i64) -> MergeCondition {
+ MergeCondition {
+ target_columns: Vec::new(),
+ source_columns: vec!["value".into()],
+ evaluate: Arc::new(move |batch| {
+ Box::pin(async move {
+ let values = batch
+ .column_by_name("s.value")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ Ok(BooleanArray::from(
+ values
+ .iter()
+ .map(|value| value.map(|value| value > bound))
+ .collect::<Vec<_>>(),
+ ))
+ })
+ }),
+ }
+}
+
+fn int_values(batches: &[RecordBatch], name: &str) -> Vec<Option<i32>> {
+ batches
+ .iter()
+ .flat_map(|batch| {
+ batch
+ .column_by_name(name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .iter()
+ .collect::<Vec<_>>()
+ })
+ .collect()
+}
+
+fn row_ids(batches: &[RecordBatch]) -> Vec<i64> {
+ batches
+ .iter()
+ .flat_map(|batch| {
+ batch
+ .column_by_name("_ROW_ID")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect()
+}
+
+#[tokio::test]
+async fn
prejoined_matched_actions_preserve_clause_order_null_fallthrough_and_target_values()
{
+ let calls = Arc::new(AtomicUsize::new(0));
+ let observed = calls.clone();
+ let assignment = UpdateAssignment::Function(Arc::new(move |batches| {
+ observed.fetch_add(1, Ordering::Relaxed);
+ assert_eq!(row_ids(batches), vec![1, 2]);
+ // The function returns different chunk boundaries from the selected
input.
+ Ok(vec![Arc::new(Int64Array::from(vec![-1, -2])) as ArrayRef])
+ }));
+ let actions = evaluate_merge_actions(
+ target_schema(),
+ vec![
+ joined(vec![1, 3], vec![None, Some(15)]),
+ joined(vec![2, 4, 5], vec![Some(5), Some(25), Some(35)]),
+ ],
+ Vec::new(),
+ &[
+ WhenMatched {
+ condition: Some(greater_than(20)),
+ delete: true,
+ assignments: Vec::new(),
+ },
+ WhenMatched {
+ condition: Some(greater_than(10)),
+ delete: false,
+ assignments: vec![(
+ "value".into(),
+ MergeAssignment::SourceColumn("value".into()),
+ )],
+ },
+ WhenMatched {
+ condition: None,
+ delete: false,
+ assignments: vec![("other".into(),
MergeAssignment::Value(assignment))],
+ },
+ ],
+ &[],
+ MergeActionOutput::All,
+ )
+ .await
+ .unwrap();
+ assert_eq!(actions.update_columns, vec!["value", "other"]);
+ assert_eq!(row_ids(&actions.updates), vec![3, 1, 2]);
+ assert_eq!(
+ int_values(&actions.updates, "value"),
+ vec![Some(15), Some(10), Some(20)]
+ );
+ assert_eq!(
+ int_values(&actions.updates, "other"),
+ vec![Some(300), Some(-1), Some(-2)]
+ );
+ assert_eq!(actions.deletes, vec![4, 5]);
+ assert!(actions.inserts.is_empty());
+ assert_eq!(calls.load(Ordering::Relaxed), 1);
+}
+
+#[tokio::test]
+async fn prejoined_insert_actions_fill_typed_nulls_and_select_first_match() {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("value", DataType::Int32, true),
+ Field::new("label", DataType::Utf8, true),
+ ]));
+ let literal = |value| {
+ vec![(
+ "value".into(),
+
MergeAssignment::Value(UpdateAssignment::Scalar(Arc::new(Int64Array::from(vec![
+ value,
+ ])))),
+ )]
+ };
+ let calls = Arc::new(AtomicUsize::new(0));
+ let observed = calls.clone();
+ let fallback = UpdateAssignment::Function(Arc::new(move |batches| {
+ observed.fetch_add(1, Ordering::Relaxed);
+ assert_eq!(
+ batches
+ .iter()
+ .map(RecordBatch::num_rows)
+ .collect::<Vec<_>>(),
+ vec![1, 1]
+ );
+ Ok(batches
+ .iter()
+ .map(|batch| batch.column_by_name("s.value").unwrap().clone())
+ .collect())
+ }));
+ let actions = evaluate_merge_actions(
+ schema.clone(),
+ Vec::new(),
+ vec![
+ source(vec![None, Some(15)]),
+ source(vec![Some(5), Some(25)]),
+ ],
+ &[],
+ &[
+ WhenNotMatched {
+ condition: Some(greater_than(20)),
+ assignments: literal(20),
+ },
+ WhenNotMatched {
+ condition: Some(greater_than(10)),
+ assignments: literal(10),
+ },
+ WhenNotMatched {
+ condition: None,
+ assignments: vec![("value".into(),
MergeAssignment::Value(fallback))],
+ },
+ ],
+ MergeActionOutput::Inserts,
+ )
+ .await
+ .unwrap();
+ assert_eq!(
+ int_values(&actions.inserts, "value"),
+ vec![Some(20), Some(10), None, Some(5)]
+ );
+ for batch in &actions.inserts {
+ assert_eq!(batch.schema(), schema);
+ let labels = batch
+ .column_by_name("label")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ assert_eq!(labels.null_count(), batch.num_rows());
+ }
+ assert!(actions.updates.is_empty());
+ assert!(actions.deletes.is_empty());
+ assert_eq!(calls.load(Ordering::Relaxed), 1);
+}
+
+#[tokio::test]
+async fn prejoined_delete_output_consumes_updates_without_evaluating_values() {
+ let clauses = [
+ WhenMatched {
+ condition: Some(greater_than(10)),
+ delete: false,
+ assignments: vec![(
+ "value".into(),
+ MergeAssignment::Value(UpdateAssignment::Function(Arc::new(|_|
{
+ panic!("delete output evaluated UPDATE")
+ }))),
+ )],
+ },
+ WhenMatched {
+ condition: None,
+ delete: true,
+ assignments: Vec::new(),
+ },
+ ];
+ // A delete-only caller needs no target output fields.
+ let actions = evaluate_merge_actions(
+ Arc::new(Schema::empty()),
+ vec![joined(vec![1, 2], vec![None, Some(15)])],
+ Vec::new(),
+ &clauses,
+ &[],
+ MergeActionOutput::Deletes,
+ )
+ .await
+ .unwrap();
+ assert_eq!(actions.deletes, vec![1]);
+ assert!(actions.updates.is_empty());
+}
+
+#[tokio::test]
+async fn prejoined_updates_consume_deletes_and_keep_empty_selections_lazy() {
+ let calls = Arc::new(AtomicUsize::new(0));
+ let observed = calls.clone();
+ let clauses = [
+ WhenMatched {
+ condition: Some(greater_than(10)),
+ delete: true,
+ assignments: Vec::new(),
+ },
+ WhenMatched {
+ condition: None,
+ delete: false,
+ assignments: vec![(
+ "value".into(),
+
MergeAssignment::Value(UpdateAssignment::Function(Arc::new(move |batches| {
+ observed.fetch_add(1, Ordering::Relaxed);
+ assert_eq!(row_ids(batches), vec![1]);
+ Ok(vec![Arc::new(Int32Array::from(vec![7])) as ArrayRef])
+ }))),
+ )],
+ },
+ ];
+ let actions = evaluate_merge_actions(
+ target_schema(),
+ vec![joined(vec![1, 2], vec![None, Some(15)])],
+ Vec::new(),
+ &clauses,
+ &[],
+ MergeActionOutput::Updates,
+ )
+ .await
+ .unwrap();
+ assert_eq!(row_ids(&actions.updates), vec![1]);
+ assert_eq!(int_values(&actions.updates, "value"), vec![Some(7)]);
+ assert!(actions.deletes.is_empty());
+ assert_eq!(calls.load(Ordering::Relaxed), 1);
+
+ evaluate_merge_actions(
+ target_schema(),
+ vec![joined(Vec::new(), Vec::new())],
+ Vec::new(),
+ &clauses,
+ &[],
+ MergeActionOutput::Updates,
+ )
+ .await
+ .unwrap();
+ assert_eq!(calls.load(Ordering::Relaxed), 1);
+}
+
+#[tokio::test]
+async fn
prejoined_assignment_cast_errors_and_bad_condition_lengths_are_rejected() {
+ let error = evaluate_merge_actions(
+ target_schema(),
+ vec![joined(vec![1], vec![Some(15)])],
+ Vec::new(),
+ &[WhenMatched {
+ condition: None,
+ delete: false,
+ assignments: vec![(
+ "value".into(),
+
MergeAssignment::Value(UpdateAssignment::Scalar(Arc::new(StringArray::from(
+ vec!["invalid int"],
+ )))),
+ )],
+ }],
+ &[],
+ MergeActionOutput::Updates,
+ )
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("invalid int"), "{error}");
+
+ let condition = MergeCondition {
+ target_columns: Vec::new(),
+ source_columns: Vec::new(),
+ evaluate: Arc::new(|_| Box::pin(async {
Ok(BooleanArray::from(vec![true, false])) })),
+ };
+ let error = evaluate_merge_actions(
+ target_schema(),
+ vec![joined(vec![1], vec![Some(15)])],
+ Vec::new(),
+ &[WhenMatched {
+ condition: Some(condition),
+ delete: true,
+ assignments: Vec::new(),
+ }],
+ &[],
+ MergeActionOutput::Deletes,
+ )
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("condition length"), "{error}");
+}
+
+#[tokio::test]
+async fn prejoined_matched_row_ids_are_validated_before_clause_callbacks() {
+ let missing = joined(vec![1], vec![Some(15)]).project(&[1, 2, 3]).unwrap();
+ let wrong_type =
+ RecordBatch::try_from_iter([("_ROW_ID",
Arc::new(Int32Array::from(vec![1])) as ArrayRef)])
+ .unwrap();
+ let null_id = RecordBatch::try_from_iter([(
+ "_ROW_ID",
+ Arc::new(Int64Array::from(vec![None])) as ArrayRef,
+ )])
+ .unwrap();
+ let clauses = [WhenMatched {
+ condition: Some(MergeCondition {
+ target_columns: Vec::new(),
+ source_columns: Vec::new(),
+ evaluate: Arc::new(|_| panic!("invalid row IDs reached clause
evaluation")),
+ }),
+ delete: true,
+ assignments: Vec::new(),
+ }];
+ for batch in [missing, wrong_type, null_id] {
+ let error = evaluate_merge_actions(
+ target_schema(),
+ vec![batch],
+ Vec::new(),
+ &clauses,
+ &[],
+ MergeActionOutput::Deletes,
+ )
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("_ROW_ID"), "{error}");
+ }
+}
diff --git a/crates/paimon/src/table/table_merge_into.rs
b/crates/paimon/src/table/table_merge_into.rs
index 4abdff74..c330c097 100644
--- a/crates/paimon/src/table/table_merge_into.rs
+++ b/crates/paimon/src/table/table_merge_into.rs
@@ -69,6 +69,218 @@ pub struct WhenNotMatched {
pub assignments: Vec<(String, MergeAssignment)>,
}
+/// Which MERGE outputs to materialize from already joined input.
+/// Other actions still consume their rows, without evaluating their
assignments.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum MergeActionOutput {
+ All,
+ Updates,
+ Deletes,
+ Inserts,
+}
+
+/// Materialized MERGE actions, grouped by clause and then input batch order.
+#[derive(Debug, Default)]
+pub struct MergeActions {
+ pub updates: Vec<RecordBatch>,
+ pub deletes: Vec<i64>,
+ pub inserts: Vec<RecordBatch>,
+ pub update_columns: Vec<String>,
+}
+
+/// Apply ordered clauses to prejoined matched and unmatched Arrow batches.
+/// Matched input has `_ROW_ID` and qualified `t.`/`s.` columns; unmatched
input
+/// has qualified `s.` columns. Assignment callbacks see the complete selection
+/// of one clause, even when that selection spans several input batches.
+/// Updates contain `_ROW_ID` and the union of assigned columns in clause
order,
+/// preserving target values for assignments omitted by a selected clause.
+/// Inserts contain every schema field, with typed NULLs for omitted
assignments.
+pub async fn evaluate_merge_actions(
+ schema: SchemaRef,
+ matched_batches: Vec<RecordBatch>,
+ unmatched_batches: Vec<RecordBatch>,
+ matched: &[WhenMatched],
+ not_matched: &[WhenNotMatched],
+ output: MergeActionOutput,
+) -> crate::Result<MergeActions> {
+ let materialize_updates = matches!(output, MergeActionOutput::All |
MergeActionOutput::Updates);
+ let materialize_deletes = matches!(output, MergeActionOutput::All |
MergeActionOutput::Deletes);
+ let materialize_inserts = matches!(output, MergeActionOutput::All |
MergeActionOutput::Inserts);
+ let mut actions = MergeActions::default();
+ for clause in matched {
+ if clause.delete && !clause.assignments.is_empty() {
+ return Err(invalid("MERGE DELETE cannot have assignments"));
+ }
+ if materialize_updates {
+ validate_assignments(&schema, &clause.assignments, false)?;
+ for (name, _) in &clause.assignments {
+ if !actions.update_columns.contains(name) {
+ actions.update_columns.push(name.clone());
+ }
+ }
+ }
+ }
+ if materialize_inserts {
+ for clause in not_matched {
+ validate_assignments(&schema, &clause.assignments, true)?;
+ if clause
+ .condition
+ .as_ref()
+ .is_some_and(|condition| !condition.target_columns.is_empty())
+ {
+ 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"));
+ }
+ }
+ if materialize_updates || materialize_deletes {
+ let mut pending = matched_batches
+ .into_iter()
+ .filter(|batch| batch.num_rows() > 0)
+ .collect::<Vec<_>>();
+ for batch in &pending {
+ let ids = batch
+ .column_by_name(ROW_ID)
+ .and_then(|array| array.as_any().downcast_ref::<Int64Array>())
+ .ok_or_else(|| {
+ invalid("Matched MERGE input must contain an Int64 _ROW_ID
column")
+ })?;
+ if ids.null_count() != 0 {
+ return Err(invalid("Matched MERGE _ROW_ID must not contain
NULL"));
+ }
+ }
+ for clause in matched {
+ let mut remaining = Vec::new();
+ let mut clause_rows = 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 {
+ if materialize_deletes {
+ let ids = selected
+ .column_by_name(ROW_ID)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ actions.deletes.extend(ids.values().iter().copied());
+ }
+ } else if materialize_updates &&
!actions.update_columns.is_empty() {
+ clause_rows.push(selected);
+ }
+ }
+ // Assignment cursors and callbacks cover the complete clause
selection,
+ // independent of physical input and result chunk boundaries.
+ if !clause_rows.is_empty() {
+ actions.updates.extend(assigned(
+ &clause_rows,
+ &clause.assignments,
+ &actions.update_columns,
+ schema.clone(),
+ false,
+ )?);
+ }
+ pending = remaining;
+ }
+ }
+ if materialize_inserts {
+ let mut pending = unmatched_batches
+ .into_iter()
+ .filter(|batch| batch.num_rows() > 0)
+ .map(with_insert_row_ids)
+ .collect::<crate::Result<Vec<_>>>()?;
+ let columns = schema
+ .fields()
+ .iter()
+ .map(|field| field.name().clone())
+ .collect::<Vec<_>>();
+ for clause in not_matched {
+ let mut remaining = Vec::new();
+ let mut clause_rows = 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 {
+ clause_rows.push(selected);
+ }
+ }
+ if !clause_rows.is_empty() {
+ for batch in assigned(
+ &clause_rows,
+ &clause.assignments,
+ &columns,
+ schema.clone(),
+ true,
+ )? {
+ actions.inserts.push(
+ batch
+
.project(&(1..batch.num_columns()).collect::<Vec<_>>())
+ .map_err(|error| invalid(error.to_string()))?,
+ );
+ }
+ }
+ pending = remaining;
+ }
+ }
+ Ok(actions)
+}
+
+fn validate_assignments(
+ schema: &SchemaRef,
+ assignments: &[(String, MergeAssignment)],
+ insert: bool,
+) -> crate::Result<()> {
+ let mut names = HashSet::new();
+ for (name, value) in assignments {
+ if name == ROW_ID || !names.insert(name) {
+ return Err(invalid(format!(
+ "Invalid or duplicate MERGE assignment {name}"
+ )));
+ }
+ schema
+ .field_with_name(name)
+ .map_err(|error| invalid(error.to_string()))?;
+ if insert && matches!(value, MergeAssignment::TargetColumn(_)) {
+ return Err(invalid("WHEN NOT MATCHED cannot reference target
columns"));
+ }
+ }
+ Ok(())
+}
+
+fn with_insert_row_ids(batch: RecordBatch) -> crate::Result<RecordBatch> {
+ if batch.column_by_name(ROW_ID).is_some() {
+ return Ok(batch);
+ }
+ let mut fields = vec![Arc::new(Field::new(ROW_ID, DataType::Int64,
false))];
+ fields.extend(batch.schema().fields().iter().cloned());
+ let mut columns =
+ vec![Arc::new(Int64Array::from(vec![0; batch.num_rows()])) as
arrow_array::ArrayRef];
+ columns.extend(batch.columns().iter().cloned());
+ RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
+ .map_err(|error| invalid(error.to_string()))
+}
+
/// Materialized input, a Paimon source table, or the target joined on its row
ID.
pub enum MergeSource {
Batches(Vec<RecordBatch>),
@@ -516,93 +728,32 @@ pub(super) async fn merge_into(
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();
- let mut clause_rows = 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() {
- clause_rows.push(selected);
- }
- }
- // Array/function assignments cover the complete clause selection, not
- // each physical target batch. Shared assignment cursors consume chunks
- // across these batches while preserving row IDs and scan order.
- if !clause_rows.is_empty() {
- updates.extend(assigned(
- &clause_rows,
- &clause.assignments,
- &update_columns,
- schema.clone(),
- false,
- )?);
- }
- pending = remaining;
- }
- let mut inserts = Vec::new();
+ let mut unmatched = Vec::new();
if let Some(source) = source {
- let mut pending = Vec::new();
for (chunk, batch) in source.batches.iter().enumerate() {
let offset = source.row_offset(chunk);
- let unmatched = (0..batch.num_rows())
+ let rows = (0..batch.num_rows())
.filter(|row| !seen_source[offset + row])
.collect::<Vec<_>>();
- if !unmatched.is_empty() {
- pending.push(aliases(None, &gather(batch, &unmatched)?)?);
+ if !rows.is_empty() {
+ unmatched.push(aliases(None, &gather(batch, &rows)?)?);
}
}
- let columns = schema
- .fields()
- .iter()
- .map(|field| field.name().clone())
- .collect::<Vec<_>>();
- for clause in not_matched {
- let mut remaining = Vec::new();
- let mut clause_rows = 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 {
- clause_rows.push(selected);
- }
- }
- if !clause_rows.is_empty() {
- for batch in assigned(
- &clause_rows,
- &clause.assignments,
- &columns,
- schema.clone(),
- true,
- )? {
- inserts.push(
- batch
-
.project(&(1..batch.num_columns()).collect::<Vec<_>>())
- .map_err(|error| invalid(error.to_string()))?,
- );
- }
- }
- pending = remaining;
- }
}
+ let MergeActions {
+ updates,
+ deletes,
+ inserts,
+ update_columns,
+ } = evaluate_merge_actions(
+ schema,
+ joined,
+ unmatched,
+ &matched,
+ ¬_matched,
+ MergeActionOutput::All,
+ )
+ .await?;
// Materialize/validate every action before staging. Once prepared,
messages
// may already be referenced by publication; a later error must retain
their files.
let mut messages = Vec::new();
@@ -634,3 +785,7 @@ pub(super) async fn merge_into(
}
Ok(messages)
}
+
+#[cfg(test)]
+#[path = "table_merge_actions_tests.rs"]
+mod tests;
