From b0e5cdba62cc958d2272d177eb03a6090fa48839 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Fri, 23 Oct 2026 09:05:42 +0800 Subject: [PATCH] [core] Share MERGE action evaluation for distributed execution --- .../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 +- .../src/table/table_merge_actions_tests.rs | 403 ++++++++++++++ crates/paimon/src/table/table_merge_into.rs | 313 ++++++++--- 9 files changed, 1202 insertions(+), 89 deletions(-) create mode 100644 bindings/python/src/merge_execution.rs create mode 100644 crates/paimon/src/table/table_merge_actions_tests.rs diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index f50293ef0..ae3937127 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.`` and ``s.`` 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 33fdd2659..162c0444a 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::()?; this.add_class::()?; this.add_class::()?; + this.add_class::()?; this.add_class::()?; this.add_class::()?; this.add_class::()?; diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 49b694209..e45485bf5 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 000000000..7310bfdc2 --- /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, + not_matched: Vec, + output: MergeActionOutput, + callback_error: Arc>>, + execution: Mutex<()>, +} + +fn matched_schema(schema: SchemaRef) -> PyResult { + 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 { + 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 { + 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::>(), + )); + 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::()?; + 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::>(); + 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> { + 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::>>()? + }; + 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::>()?, + 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::>()?; + RecordBatch::try_new(self.output_schema.clone(), columns) + .map_err(|error| PyValueError::new_err(error.to_string()))? + .to_pyarrow(py) + }) + .collect::>>()?; + 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::(), + ) + .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 4d237d2fe..3d6611aea 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>>; -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> { 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::>>()?; @@ -146,7 +150,7 @@ pub(crate) fn from_python( let clause = item.cast::()?; 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::>>()?; diff --git a/bindings/python/src/update_assignment.rs b/bindings/python/src/update_assignment.rs index 290d95607..3df5a3a5e 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> { 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> { 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 1f95c5af6..dc463d171 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 000000000..b35d2c7b2 --- /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, source: Vec>) -> RecordBatch { + let values = ids.iter().map(|id| *id as i32 * 10).collect::>(); + let other = ids.iter().map(|id| *id as i32 * 100).collect::>(); + 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>) -> 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::() + .unwrap(); + Ok(BooleanArray::from( + values + .iter() + .map(|value| value.map(|value| value > bound)) + .collect::>(), + )) + }) + }), + } +} + +fn int_values(batches: &[RecordBatch], name: &str) -> Vec> { + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>() + }) + .collect() +} + +fn row_ids(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("_ROW_ID") + .unwrap() + .as_any() + .downcast_ref::() + .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![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::() + .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 4abdff749..c330c097e 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, + pub deletes: Vec, + pub inserts: Vec, + pub update_columns: Vec, +} + +/// 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, + unmatched_batches: Vec, + matched: &[WhenMatched], + not_matched: &[WhenNotMatched], + output: MergeActionOutput, +) -> crate::Result { + 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::>(), + 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::>(); + for batch in &pending { + let ids = batch + .column_by_name(ROW_ID) + .and_then(|array| array.as_any().downcast_ref::()) + .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::() + .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::>>()?; + let columns = schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + 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::>()) + .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 { + 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), @@ -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::() - .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::>(); - 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::>(); - 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::>()) - .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;