diff --git a/bindings/python/project-description.md b/bindings/python/project-description.md index 86c408608..9738ab17b 100644 --- a/bindings/python/project-description.md +++ b/bindings/python/project-description.md @@ -85,6 +85,49 @@ updates, including incremental row-ID updaters, expose the same chainable setter `with_commit_user` is available on batch, stream, and postpone fixed bucket write builders, and the identity applies to both writers and committers from that builder. +### Typed MERGE Actions + +Batch and stream `merge_into` accept ordered `WhenMatched` and `WhenNotMatched` +objects. Assignments use `MergeAssignment.source_column`, `target_column`, +`literal`, or `function`. A function receives the selected clause's Arrow table +and returns an Array or ChunkedArray with the same number of rows. + +Execution engines that already joined target and source data can reuse these +types through the batch execution factories: + +```python +import pyarrow as pa +from pypaimon_rust.datafusion import MergeAssignment, MergeBatchExecution, WhenMatched + +joined = pa.table({ + "t._ROW_ID": [10, 11], + "t.value": pa.array([20, 30], type=pa.int32()), + "s.value": pa.array([40, 10], type=pa.int32()), +}) +matched = [ + WhenMatched.update( + {"value": MergeAssignment.source_column("value")}, + condition='"s.value" > "t.value"', + ), + WhenMatched.delete(), +] +update_schema = pa.schema([ + pa.field("_ROW_ID", pa.int64(), nullable=False), + pa.field("value", pa.int32()), +]) +updates = MergeBatchExecution.for_updates(joined.schema, update_schema, matched) +deletes = MergeBatchExecution.for_deletes(joined.schema, matched) +assert updates.execute(joined).to_pylist() == [{"_ROW_ID": 10, "value": 40}] +assert deletes.execute(joined).to_pylist() == [{"_ROW_ID": 11}] +``` + +Both matched factories consume the complete ordered branch list. The delete +factory skips UPDATE assignments and callbacks while preserving their conditions. +For unmatched source batches, use `for_inserts(source_schema, target_schema, +when_not_matched)` with `WhenNotMatched` objects. The factories compile conditions +once and return an executor with `execute(Table | RecordBatch) -> Table`. +The calling engine owns joins, authorization and table writes. + ### Multimodal SQL Helpers `SQLContext` also registers Python scalar UDFs for common BLOB media and vector diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index ae3937127..976c33090 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -18,7 +18,7 @@ from datetime import datetime from os import PathLike from typing import ( - Any, Callable, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, + Any, Callable, Dict, Iterable, List, Literal, Mapping, Optional, Sequence, Tuple, TypeAlias, Union, ) @@ -79,6 +79,40 @@ class RowIdPlan: ... def plan(self) -> Plan: ... +class MergeAssignment: + """An immutable source reference, target reference, literal or callback.""" + @staticmethod + def source_column(name: str) -> "MergeAssignment": ... + @staticmethod + def target_column(name: str) -> "MergeAssignment": ... + @staticmethod + def literal(value: Any) -> "MergeAssignment": + """Defer scalar conversion; Arrow arrays provide values for selected rows.""" + ... + @staticmethod + def function( + function: Callable[[pyarrow.Table], Union[pyarrow.Array, pyarrow.ChunkedArray]], + ) -> "MergeAssignment": ... + +class WhenMatched: + """An immutable, ordered UPDATE or DELETE clause with an optional SQL condition. + + Conditions reference quoted Arrow names, such as ``"s.value" > "t.value"``. + The first matching clause consumes a row; NULL conditions fall through. + """ + @staticmethod + def update( + assignments: Mapping[str, MergeAssignment], condition: Optional[str] = None, + ) -> "WhenMatched": ... + @staticmethod + def delete(condition: Optional[str] = None) -> "WhenMatched": ... + +class WhenNotMatched: + """An immutable, ordered INSERT clause.""" + def __init__( + self, assignments: Mapping[str, MergeAssignment], condition: Optional[str] = None, + ) -> None: ... + class MergeBatchExecution: """Evaluate ordered MERGE actions over already joined Arrow batches. @@ -89,10 +123,22 @@ class MergeBatchExecution: 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: ... + @staticmethod + def for_updates( + input_schema: pyarrow.Schema, update_schema: pyarrow.Schema, + when_matched: Sequence[WhenMatched], + ) -> "MergeBatchExecution": ... + @staticmethod + def for_deletes( + input_schema: pyarrow.Schema, when_matched: Sequence[WhenMatched], + ) -> "MergeBatchExecution": + """Keep all matched conditions; skip UPDATE assignment conversion and evaluation.""" + ... + @staticmethod + def for_inserts( + input_schema: pyarrow.Schema, target_schema: pyarrow.Schema, + when_not_matched: Sequence[WhenNotMatched], + ) -> "MergeBatchExecution": ... def execute(self, batch: Union[pyarrow.Table, pyarrow.RecordBatch]) -> pyarrow.Table: """Return selected actions; unselected assignments are never evaluated.""" ... @@ -542,15 +588,14 @@ class BatchTableUpdate: ) -> "BatchTableUpdate": ... def merge_into( self, source: Optional[Union[pyarrow.Table, Table]], *, on: Sequence[Tuple[str, str]], - when_matched: Sequence[Dict[str, Any]], when_not_matched: Sequence[Dict[str, Any]] + when_matched: Sequence[WhenMatched], when_not_matched: Sequence[WhenNotMatched] ) -> List[CommitMessage]: """Prepare MERGE against one target snapshot. None selects self-merge on _ROW_ID. A Table source reads its selected snapshot or numeric timestamp window in core, independently of the target. Arrow input preserves batches and source row order. - Clauses carry assignments as (column, source/target/literal/callable, value) - tuples, a delete flag, and optional SQL condition metadata. Callables + Typed clauses carry MergeAssignment values and optional SQL conditions. Callables consume the complete clause's qualified t./s. Arrow columns and return an Array or ChunkedArray with the same row count; empty clauses are lazy. Commit returned messages through the matching builder's committer. @@ -616,7 +661,7 @@ class StreamTableUpdate: ) -> "StreamTableUpdate": ... def merge_into( self, source: Optional[Union[pyarrow.Table, Table]], *, on: Sequence[Tuple[str, str]], - when_matched: Sequence[Dict[str, Any]], when_not_matched: Sequence[Dict[str, Any]], + when_matched: Sequence[WhenMatched], when_not_matched: Sequence[WhenNotMatched], commit_identifier: int ) -> List[CommitMessage]: """Prepare MERGE against one target snapshot. None selects self-merge on _ROW_ID. @@ -624,8 +669,7 @@ class StreamTableUpdate: A Table source reads its selected snapshot or numeric timestamp window in core, independently of the target. Arrow input preserves batches and source row order. - Clauses carry assignments as (column, source/target/literal/callable, value) - tuples, a delete flag, and optional SQL condition metadata. Callables + Typed clauses carry MergeAssignment values and optional SQL conditions. Callables consume the complete clause's qualified t./s. Arrow columns and return an Array or ChunkedArray with the same row count; empty clauses are lazy. Commit returned messages through the matching builder's committer. diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index 162c0444a..112a5014b 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -1039,6 +1039,9 @@ 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::()?; this.add_class::()?; this.add_class::()?; diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index e45485bf5..69f081bb0 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -31,6 +31,7 @@ mod hybrid_search; mod index_build; mod merge_execution; mod merge_into; +mod merge_types; mod oss_cpp; mod predicate; mod read; diff --git a/bindings/python/src/merge_execution.rs b/bindings/python/src/merge_execution.rs index 7310bfdc2..7d4f17db5 100644 --- a/bindings/python/src/merge_execution.rs +++ b/bindings/python/src/merge_execution.rs @@ -21,7 +21,7 @@ use std::collections::HashSet; use std::sync::{Arc, Mutex}; use arrow::array::Int64Array; -use arrow::datatypes::{DataType, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::pyarrow::{FromPyArrow, ToPyArrow}; use arrow::record_batch::RecordBatch; use paimon::table::{ @@ -30,7 +30,7 @@ use paimon::table::{ use paimon_datafusion::runtime::runtime; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList}; +use pyo3::types::PyList; use crate::error::to_py_err; @@ -79,28 +79,14 @@ fn row_ids(batch: RecordBatch) -> PyResult { RecordBatch::try_new(schema, columns).map_err(|error| PyValueError::new_err(error.to_string())) } -#[pymethods] impl PyMergeBatchExecution { - #[new] - fn new( + fn compile( py: Python<'_>, - input_schema: &Bound<'_, PyAny>, - output_schema: &Bound<'_, PyAny>, + input_schema: SchemaRef, + output_schema: SchemaRef, clauses: &Bound<'_, PyAny>, - mode: &str, + output: MergeActionOutput, ) -> 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") @@ -110,60 +96,46 @@ impl PyMergeBatchExecution { "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( + let target_schema = Arc::new(Schema::new_with_metadata( output_schema .fields() .iter() .filter(|field| output == MergeActionOutput::Inserts || field.name() != "_ROW_ID") .cloned() .collect::>(), + output_schema.metadata().clone(), )); 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( + let (mut matched, not_matched) = if output == MergeActionOutput::Inserts { + ( + Vec::new(), + crate::merge_into::not_matched( py, - clause, + clauses, + condition_schema, 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, - }); - } - } + )?, + ) + } else { + ( + crate::merge_into::matched( + py, + clauses, + condition_schema, + target_schema.clone(), + callback_error.clone(), + output != MergeActionOutput::Deletes, + )?, + Vec::new(), + ) + }; if output == MergeActionOutput::Updates { // The explicit output schema may request fields untouched by every // clause. Include them in core's shared update-column projection. @@ -206,6 +178,62 @@ impl PyMergeBatchExecution { execution: Mutex::new(()), }) } +} + +#[pymethods] +impl PyMergeBatchExecution { + #[staticmethod] + fn for_updates( + py: Python<'_>, + input_schema: &Bound<'_, PyAny>, + update_schema: &Bound<'_, PyAny>, + when_matched: &Bound<'_, PyAny>, + ) -> PyResult { + Self::compile( + py, + Arc::new(Schema::from_pyarrow_bound(input_schema)?), + Arc::new(Schema::from_pyarrow_bound(update_schema)?), + when_matched, + MergeActionOutput::Updates, + ) + } + + #[staticmethod] + fn for_deletes( + py: Python<'_>, + input_schema: &Bound<'_, PyAny>, + when_matched: &Bound<'_, PyAny>, + ) -> PyResult { + let input_schema = Arc::new(Schema::from_pyarrow_bound(input_schema)?); + let output_schema = Arc::new(Schema::new(vec![Field::new( + "_ROW_ID", + DataType::Int64, + false, + )])); + Self::compile( + py, + input_schema, + output_schema, + when_matched, + MergeActionOutput::Deletes, + ) + } + + #[staticmethod] + fn for_inserts( + py: Python<'_>, + input_schema: &Bound<'_, PyAny>, + target_schema: &Bound<'_, PyAny>, + when_not_matched: &Bound<'_, PyAny>, + ) -> PyResult { + Self::compile( + py, + Arc::new(Schema::from_pyarrow_bound(input_schema)?), + Arc::new(Schema::from_pyarrow_bound(target_schema)?), + when_not_matched, + MergeActionOutput::Inserts, + ) + } fn execute(&self, py: Python<'_>, batch: &Bound<'_, PyAny>) -> PyResult> { let schema = Schema::from_pyarrow_bound(&batch.getattr("schema")?)?; @@ -302,6 +330,7 @@ impl PyMergeBatchExecution { mod tests { use super::*; use pyo3::ffi::c_str; + use pyo3::types::PyDict; fn locals(py: Python<'_>) -> Bound<'_, PyDict> { let locals = PyDict::new(py); @@ -312,6 +341,24 @@ mod tests { ) .unwrap(); locals + .set_item( + "MergeAssignment", + py.get_type::(), + ) + .unwrap(); + locals + .set_item( + "WhenMatched", + py.get_type::(), + ) + .unwrap(); + locals + .set_item( + "WhenNotMatched", + py.get_type::(), + ) + .unwrap(); + locals } #[test] @@ -346,10 +393,11 @@ def update(rows): 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)]), + WhenMatched.delete(condition='"s.value" >= 90'), + WhenMatched.update({'value': MergeAssignment.function(update), + 'keep': MergeAssignment.target_column('keep')}, + condition='"s.value" > 10'), + WhenMatched.update({'value': MergeAssignment.literal(50)}), ] output_schema = pa.schema([ pa.field('_ROW_ID', pa.int64(), nullable=False), @@ -357,7 +405,7 @@ output_schema = pa.schema([ pa.field('unassigned', pa.int32()), pa.field('value', pa.int32()), ], metadata={'output': 'updates'}) -execution = MergeBatchExecution(input_schema, output_schema, clauses, 'update') +execution = MergeBatchExecution.for_updates(input_schema, output_schema, clauses) result = execution.execute(joined) assert result.schema.equals(output_schema, check_metadata=True) assert result.to_pydict() == { @@ -367,7 +415,7 @@ assert result.to_pydict() == { } 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) +deletes = MergeBatchExecution.for_deletes(input_schema, clauses).execute(joined) assert deletes.schema == delete_schema assert deletes.to_pydict() == {'_ROW_ID': [3]} assert calls == [[1, 4]] @@ -394,26 +442,29 @@ 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_())]) + pa.field('value', pa.int32()), pa.field('missing', pa.bool_())], + metadata={'output': 'inserts'}) 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)]), + WhenNotMatched({'id': MergeAssignment.source_column('id'), + 'value': MergeAssignment.function(insert), + 'label': MergeAssignment.literal('first')}, + condition='"s.value" > 10'), + WhenNotMatched({'id': MergeAssignment.source_column('id'), + 'value': MergeAssignment.literal(7), + 'label': MergeAssignment.literal('last'), + 'missing': MergeAssignment.literal(True)}), ] -execution = MergeBatchExecution(source.schema, target, clauses, 'insert') +execution = MergeBatchExecution.for_inserts(source.schema, target, clauses) result = execution.execute(source) -assert result.schema == target +assert result.schema.equals(target, check_metadata=True) 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 execution.execute(source.slice(0, 0)).schema.equals(target, check_metadata=True) assert calls == [[1]] "# ), @@ -443,9 +494,8 @@ def assignment(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') +clauses = [WhenMatched.update({'value': MergeAssignment.function(assignment)})] +execution = MergeBatchExecution.for_updates(joined.schema, output, clauses) try: execution.execute(joined) except LookupError as error: @@ -484,31 +534,171 @@ def rejected(function, text): 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') +clause = WhenMatched.update({'value': MergeAssignment.literal(1)}, + condition='"s.absent" > 0') +rejected(lambda: MergeBatchExecution.for_updates(joined.schema, output, [clause]), 'absent') +clause = WhenMatched.update({'value': MergeAssignment.function( + lambda rows: pa.array([1], pa.int32()))}) +execution = MergeBatchExecution.for_updates(joined.schema, output, [clause]) 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') +rejected(lambda: MergeBatchExecution.for_updates(joined.schema, + pa.schema([('value', pa.int32())]), [clause]), '_ROW_ID') +rejected(lambda: MergeBatchExecution.for_updates(joined.schema, + pa.schema([('_ROW_ID', pa.int32()), ('value', pa.int32())]), [clause]), 'Int64') +rejected(lambda: MergeBatchExecution.for_updates(joined.schema, output, [clause, clause]), '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), +rejected(lambda: MergeBatchExecution.for_updates(dual_ids.schema, output, [clause]), 'ambiguous') +rejected(lambda: MergeBatchExecution.for_deletes(dual_ids.schema, [WhenMatched.delete()]), 'ambiguous') +rejected(lambda: MergeBatchExecution.for_deletes(pa.schema([('s.value', pa.int32())]), + [WhenMatched.delete()]), '_ROW_ID') +wrong_ids = pa.table({'t._ROW_ID': pa.array([1], pa.int32())}) +rejected(lambda: MergeBatchExecution.for_deletes(wrong_ids.schema, + [WhenMatched.delete()]).execute(wrong_ids), 'Int64') +null_ids = pa.table({'_ROW_ID': pa.array([None], pa.int64())}) +rejected(lambda: MergeBatchExecution.for_deletes(null_ids.schema, + [WhenMatched.delete()]).execute(null_ids), 'NULL') +bad_type = WhenMatched.update({'value': MergeAssignment.function(lambda rows: [1, 2])}) +rejected(lambda: MergeBatchExecution.for_updates(joined.schema, output, [bad_type]).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 +clause = WhenMatched.update({'value': MergeAssignment.literal(object())}, + condition='"s.value" < 0') +assert MergeBatchExecution.for_updates(joined.schema, output, [clause]).execute(joined).num_rows == 0 +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn named_factories_reject_dictionary_protocol_and_wrong_clause_families() { + Python::attach(|py| { + let locals = locals(py); + py.run( + c_str!( + r#" +import pyarrow as pa +input_schema = pa.schema([('t._ROW_ID', pa.int64()), ('s.value', pa.int32())]) +update_schema = pa.schema([('_ROW_ID', pa.int64()), ('value', pa.int32())]) +target_schema = pa.schema([('value', pa.int32())]) +matched = WhenMatched.update({'value': MergeAssignment.source_column('value')}) +not_matched = WhenNotMatched({'value': MergeAssignment.source_column('value')}) +legacy = dict(condition=None, delete=False, + assignments=[('value', 'source', 'value')]) +def rejected(function): + try: + function() + except TypeError: + pass + else: + raise AssertionError('untyped MERGE protocol or wrong clause family was accepted') +for clause in [legacy, not_matched]: + rejected(lambda: MergeBatchExecution.for_updates(input_schema, update_schema, [clause])) + rejected(lambda: MergeBatchExecution.for_deletes(input_schema, [clause])) +for clause in [legacy, matched]: + rejected(lambda: MergeBatchExecution.for_inserts(input_schema, target_schema, [clause])) +rejected(lambda: MergeBatchExecution(input_schema, update_schema, [legacy], 'update')) +rejected(lambda: MergeBatchExecution()) +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn deletes_keep_update_conditions_without_converting_update_assignments() { + Python::attach(|py| { + let locals = locals(py); + py.run( + c_str!( + r#" +import pyarrow as pa +schema = pa.schema([ + pa.field('_ROW_ID', pa.int64(), nullable=True, metadata={'input': 'row-id'}), + pa.field('s.value', pa.int32()), +], metadata={'input': 'joined'}) +joined = pa.record_batch([pa.array([1, 2, 3], pa.int64()), + pa.array([10, 20, 30], pa.int32())], schema=schema) +calls = [] +def unused(rows): + calls.append(rows.num_rows) + raise AssertionError('DELETE evaluated an UPDATE callback') +clauses = [ + WhenMatched.update({'absent': MergeAssignment.literal(object()), + 'also_absent': MergeAssignment.source_column('missing')}, + condition='"_ROW_ID" = 1'), + WhenMatched.update({'unavailable': MergeAssignment.function(unused)}, + condition='"s.value" > 20'), + WhenMatched.delete(), +] +execution = MergeBatchExecution.for_deletes(input_schema=schema, when_matched=clauses) +result = execution.execute(joined) +expected = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False)]) +assert result.schema.equals(expected, check_metadata=True) +assert result.to_pydict() == {'_ROW_ID': [2]} +assert execution.execute(schema.empty_table()).schema.equals(expected, check_metadata=True) +assert calls == [] +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn cached_execution_keeps_concurrent_callback_errors_with_their_callers() { + Python::attach(|py| { + let locals = locals(py); + py.run( + c_str!( + r#" +import pyarrow as pa +import threading +from concurrent.futures import ThreadPoolExecutor +schema = pa.schema([('t._ROW_ID', pa.int64()), ('s.value', pa.int32())]) +output = pa.schema([('_ROW_ID', pa.int64()), ('value', pa.int32())]) +failures = {i: LookupError('assignment failed for ' + str(i)) for i in range(4)} +callback_entered = threading.Event() +release_callback = threading.Event() +def assignment(rows): + row_id = rows['t._ROW_ID'][0].as_py() + callback_entered.set() + assert release_callback.wait(timeout=10) + if row_id in failures: + raise failures[row_id] + return pa.array([9], pa.int32()) +execution = MergeBatchExecution.for_updates(schema, output, + [WhenMatched.update({'value': MergeAssignment.function(assignment)})]) +barrier = threading.Barrier(4) +def execute(row_id): + batch = pa.record_batch([pa.array([row_id], pa.int64()), pa.array([1], pa.int32())], + schema=schema) + barrier.wait() + try: + execution.execute(batch) + except LookupError as error: + return error + raise AssertionError('callback exception was lost') +with ThreadPoolExecutor(max_workers=4) as threads: + futures = [threads.submit(execute, row_id) for row_id in range(4)] + try: + assert callback_entered.wait(timeout=10) + finally: + release_callback.set() + results = [future.result(timeout=10) for future in futures] +assert all(error is failures[i] for i, error in enumerate(results)) +failures.clear() +batch = pa.record_batch([pa.array([7], pa.int64()), pa.array([1], pa.int32())], schema=schema) +assert execution.execute(batch).to_pydict() == {'_ROW_ID': [7], 'value': [9]} "# ), Some(&locals), diff --git a/bindings/python/src/merge_into.rs b/bindings/python/src/merge_into.rs index 3d6611aea..150bf36a1 100644 --- a/bindings/python/src/merge_into.rs +++ b/bindings/python/src/merge_into.rs @@ -15,79 +15,126 @@ // specific language governing permissions and limitations // under the License. -//! Language transport and DataFusion expression compilation for core MERGE. +//! Python MERGE descriptions and DataFusion expression compilation for core MERGE. use std::sync::{Arc, Mutex}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use paimon::table::{MergeAssignment, MergeCondition, WhenMatched, WhenNotMatched}; use paimon_datafusion::runtime::runtime; -use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use pyo3::types::PyDict; use crate::error::to_py_err; +use crate::merge_types::{AssignmentValue, Assignments, PyWhenMatched, PyWhenNotMatched}; type CallbackError = Arc>>; pub(crate) fn condition( py: Python<'_>, - clause: &Bound<'_, PyDict>, + sql: Option<&str>, schema: SchemaRef, ) -> PyResult> { - let Some(value) = clause - .get_item("condition")? - .filter(|value| !value.is_none()) - else { - return Ok(None); - }; - let value = value.cast::()?; - let sql: String = value - .get_item("sql")? - .ok_or_else(|| PyValueError::new_err("MERGE condition needs SQL"))? - .extract()?; - py.detach(|| runtime().block_on(paimon_datafusion::compile_merge_condition(sql, schema))) - .map(Some) + sql.map(|sql| { + py.detach(|| { + runtime().block_on(paimon_datafusion::compile_merge_condition( + sql.to_owned(), + schema, + )) + }) .map_err(to_py_err) + }) + .transpose() } pub(crate) fn assignments( py: Python<'_>, - clause: &Bound<'_, PyDict>, + values: &Assignments, schema: SchemaRef, errors: CallbackError, ) -> PyResult> { - let values = clause - .get_item("assignments")? - .ok_or_else(|| PyValueError::new_err("MERGE clause needs assignments"))?; values - .try_iter()? - .map(|item| { - let (name, kind, value): (String, String, Py) = item?.extract()?; - let assignment = match kind.as_str() { - "source" => MergeAssignment::SourceColumn(value.bind(py).extract()?), - "target" => MergeAssignment::TargetColumn(value.bind(py).extract()?), - "literal" | "callable" => { - if kind == "callable" && !value.bind(py).is_callable() { - return Err(PyValueError::new_err( - "MERGE callable assignment must be callable", - )); - } - let values = PyDict::new(py); - values.set_item(&name, value)?; - let assignment = crate::update_assignment::from_python_schema( + .iter() + .map(|(name, value)| { + let value = value.borrow(py); + let assignment = match &value.value { + AssignmentValue::SourceColumn(column) => { + MergeAssignment::SourceColumn(column.clone()) + } + AssignmentValue::TargetColumn(column) => { + MergeAssignment::TargetColumn(column.clone()) + } + AssignmentValue::Literal(value) => { + MergeAssignment::Value(crate::update_assignment::literal( py, - &values, + name.clone(), + value.clone_ref(py), schema.clone(), errors.clone(), - )? - .remove(0) - .1; - MergeAssignment::Value(assignment) + )?) } - _ => return Err(PyValueError::new_err("Unknown MERGE assignment kind")), + AssignmentValue::Function(value) => MergeAssignment::Value( + crate::update_assignment::function(value.clone_ref(py), errors.clone()), + ), }; - Ok((name, assignment)) + Ok((name.clone(), assignment)) + }) + .collect() +} + +pub(crate) fn matched( + py: Python<'_>, + clauses: &Bound<'_, PyAny>, + condition_schema: SchemaRef, + target_schema: SchemaRef, + errors: CallbackError, + include_assignments: bool, +) -> PyResult> { + clauses + .try_iter()? + .map(|item| { + let item = item?; + let clause = item.extract::>()?; + Ok(WhenMatched { + condition: condition(py, clause.condition.as_deref(), condition_schema.clone())?, + delete: clause.delete, + assignments: if include_assignments { + assignments( + py, + &clause.assignments, + target_schema.clone(), + errors.clone(), + )? + } else { + // DELETE output consumes preceding UPDATE conditions without + // inspecting values that this route will never materialize. + Vec::new() + }, + }) + }) + .collect() +} + +pub(crate) fn not_matched( + py: Python<'_>, + clauses: &Bound<'_, PyAny>, + condition_schema: SchemaRef, + target_schema: SchemaRef, + errors: CallbackError, +) -> PyResult> { + clauses + .try_iter()? + .map(|item| { + let item = item?; + let clause = item.extract::>()?; + Ok(WhenNotMatched { + condition: condition(py, clause.condition.as_deref(), condition_schema.clone())?, + assignments: assignments( + py, + &clause.assignments, + target_schema.clone(), + errors.clone(), + )?, + }) }) .collect() } @@ -96,8 +143,8 @@ pub(crate) fn from_python( py: Python<'_>, table: &paimon::table::Table, source_schema: SchemaRef, - matched: &Bound<'_, PyAny>, - not_matched: &Bound<'_, PyAny>, + when_matched: &Bound<'_, PyAny>, + when_not_matched: &Bound<'_, PyAny>, errors: CallbackError, ) -> PyResult<(Vec, Vec)> { let target_schema = @@ -115,22 +162,14 @@ pub(crate) fn from_python( .clone() .with_name(format!("s.{}", field.name())) })); - let schema = Arc::new(Schema::new(fields)); - let matched = matched - .try_iter()? - .map(|item| { - let item = item?; - let clause = item.cast::()?; - Ok(WhenMatched { - condition: condition(py, clause, schema.clone())?, - delete: clause - .get_item("delete")? - .ok_or_else(|| PyValueError::new_err("MERGE matched clause needs action"))? - .extract()?, - assignments: assignments(py, clause, target_schema.clone(), errors.clone())?, - }) - }) - .collect::>>()?; + let matched = matched( + py, + when_matched, + Arc::new(Schema::new(fields)), + target_schema.clone(), + errors.clone(), + true, + )?; let schema = Arc::new(Schema::new( source_schema .fields() @@ -143,16 +182,6 @@ pub(crate) fn from_python( }) .collect::>(), )); - let not_matched = not_matched - .try_iter()? - .map(|item| { - let item = item?; - let clause = item.cast::()?; - Ok(WhenNotMatched { - condition: condition(py, clause, schema.clone())?, - assignments: assignments(py, clause, target_schema.clone(), errors.clone())?, - }) - }) - .collect::>>()?; + let not_matched = not_matched(py, when_not_matched, schema, target_schema, errors)?; Ok((matched, not_matched)) } diff --git a/bindings/python/src/merge_types.rs b/bindings/python/src/merge_types.rs new file mode 100644 index 000000000..134a89665 --- /dev/null +++ b/bindings/python/src/merge_types.rs @@ -0,0 +1,228 @@ +// 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. + +//! Immutable Python descriptions of core MERGE clauses and assignments. + +use pyo3::exceptions::PyTypeError; +use pyo3::prelude::*; +use pyo3::types::PyMapping; +use pyo3::{PyTraverseError, PyVisit}; + +pub(crate) enum AssignmentValue { + SourceColumn(String), + TargetColumn(String), + Literal(Py), + Function(Py), +} + +/// A typed column reference, literal, or function evaluated for selected rows. +#[pyclass(name = "MergeAssignment", module = "pypaimon_rust.datafusion", frozen)] +pub struct PyMergeAssignment { + pub(crate) value: AssignmentValue, +} + +#[pymethods] +impl PyMergeAssignment { + #[staticmethod] + fn source_column(name: String) -> Self { + Self { + value: AssignmentValue::SourceColumn(name), + } + } + + #[staticmethod] + fn target_column(name: String) -> Self { + Self { + value: AssignmentValue::TargetColumn(name), + } + } + + #[staticmethod] + fn literal(value: Py) -> Self { + Self { + value: AssignmentValue::Literal(value), + } + } + + #[staticmethod] + fn function(function: Py, py: Python<'_>) -> PyResult { + if !function.bind(py).is_callable() { + return Err(PyTypeError::new_err( + "MergeAssignment.function requires a callable", + )); + } + Ok(Self { + value: AssignmentValue::Function(function), + }) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + match &self.value { + AssignmentValue::Literal(value) | AssignmentValue::Function(value) => visit.call(value), + AssignmentValue::SourceColumn(_) | AssignmentValue::TargetColumn(_) => Ok(()), + } + } +} + +pub(crate) type Assignments = Vec<(String, Py)>; + +fn assignments(values: &Bound<'_, PyMapping>) -> PyResult { + values.items()?.extract() +} + +/// An ordered WHEN MATCHED UPDATE or DELETE branch. +#[pyclass(name = "WhenMatched", module = "pypaimon_rust.datafusion", frozen)] +pub struct PyWhenMatched { + pub(crate) condition: Option, + pub(crate) delete: bool, + pub(crate) assignments: Assignments, +} + +#[pymethods] +impl PyWhenMatched { + #[staticmethod] + #[pyo3(signature = (assignments, condition=None))] + fn update(assignments: &Bound<'_, PyMapping>, condition: Option) -> PyResult { + Ok(Self { + condition, + delete: false, + assignments: self::assignments(assignments)?, + }) + } + + #[staticmethod] + #[pyo3(signature = (condition=None))] + fn delete(condition: Option) -> Self { + Self { + condition, + delete: true, + assignments: Vec::new(), + } + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + for (_, value) in &self.assignments { + visit.call(value)?; + } + Ok(()) + } +} + +/// An ordered WHEN NOT MATCHED INSERT branch. +#[pyclass(name = "WhenNotMatched", module = "pypaimon_rust.datafusion", frozen)] +pub struct PyWhenNotMatched { + pub(crate) condition: Option, + pub(crate) assignments: Assignments, +} + +#[pymethods] +impl PyWhenNotMatched { + #[new] + #[pyo3(signature = (assignments, condition=None))] + fn new(assignments: &Bound<'_, PyMapping>, condition: Option) -> PyResult { + Ok(Self { + condition, + assignments: self::assignments(assignments)?, + }) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + for (_, value) in &self.assignments { + visit.call(value)?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::ffi::c_str; + + #[test] + fn typed_factories_validate_descriptions_and_snapshot_assignments() { + Python::attach(|py| { + let locals = pyo3::types::PyDict::new(py); + locals + .set_item("MergeAssignment", py.get_type::()) + .unwrap(); + locals + .set_item("WhenMatched", py.get_type::()) + .unwrap(); + locals + .set_item("WhenNotMatched", py.get_type::()) + .unwrap(); + py.run( + c_str!( + r#" +from types import MappingProxyType +value = MergeAssignment.source_column('value') +mapping = {'value': value} +matched = WhenMatched.update(MappingProxyType(mapping), condition='"s.value" > 0') +unmatched = WhenNotMatched(mapping) +mapping.clear() +for make in [ + lambda: MergeAssignment(), + lambda: MergeAssignment.source_column(1), + lambda: MergeAssignment.target_column(None), + lambda: MergeAssignment.function(1), + lambda: WhenMatched(), + lambda: WhenMatched.update({'value': 1}), + lambda: WhenMatched.update({1: value}), + lambda: WhenMatched.update({'value': value}, condition={'sql': 'true'}), + lambda: WhenMatched.delete(condition={'sql': 'true'}), + lambda: WhenNotMatched({'value': 1}), + lambda: WhenNotMatched({'value': value}, condition={'sql': 'true'}), +]: + try: + make() + except TypeError: + pass + else: + raise AssertionError('untyped MERGE description was accepted') +for description in [value, matched, unmatched]: + try: + description.condition = None + except AttributeError: + pass + else: + raise AssertionError('MERGE description was mutable') +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let matched = locals + .get_item("matched") + .unwrap() + .unwrap() + .extract::>() + .unwrap(); + let unmatched = locals + .get_item("unmatched") + .unwrap() + .unwrap() + .extract::>() + .unwrap(); + assert_eq!(matched.assignments.len(), 1); + assert_eq!(unmatched.assignments.len(), 1); + assert_eq!(matched.assignments[0].0, "value"); + assert_eq!(unmatched.assignments[0].0, "value"); + }); + } +} diff --git a/bindings/python/src/update_assignment.rs b/bindings/python/src/update_assignment.rs index 3df5a3a5e..0779cbb7e 100644 --- a/bindings/python/src/update_assignment.rs +++ b/bindings/python/src/update_assignment.rs @@ -72,72 +72,85 @@ pub(crate) fn from_python_schema( schema: SchemaRef, error: CallbackError, ) -> PyResult> { - let pa = py.import("pyarrow")?; - let array_type = pa.getattr("Array")?; - let chunked_type = pa.getattr("ChunkedArray")?; assignments .iter() .map(|(name, value)| { let name: String = name.extract()?; let assignment = if value.is_callable() { - let value = value.unbind(); - let error = error.clone(); - UpdateAssignment::Function(Arc::new(move |batches| { - // A Python callback may invoke another native read or write. - tokio::task::block_in_place(|| { - Python::attach(|py| -> PyResult<_> { - let pa = py.import("pyarrow")?; - let batches = batches - .iter() - .map(|batch| batch.to_pyarrow(py)) - .collect::>>()?; - let matched = pa - .getattr("Table")? - .call_method1("from_batches", (PyList::new(py, batches)?,))?; - arrow_arrays(py, &value.bind(py).call1((matched,))?) - }) - }) - .map_err(|cause| bridge_error(cause, &error)) - })) - } else if value.is_instance(&array_type)? || value.is_instance(&chunked_type)? { - UpdateAssignment::Array(arrow_arrays(py, &value)?) + function(value.unbind(), error.clone()) } else { - // Literal conversion stays lazy: a no-match update must neither - // evaluate callbacks nor reject a scalar that is never assigned. - let value = value.unbind(); - let name = name.clone(); - let schema = schema.clone(); - let error = error.clone(); - UpdateAssignment::DeferredScalar(Arc::new(move || { - tokio::task::block_in_place(|| { - Python::attach(|py| -> PyResult<_> { - let pa = py.import("pyarrow")?; - let value = value.bind(py); - let scalar = if value.is_instance(&pa.getattr("Scalar")?)? { - value.call_method0("as_py")? - } else { - value.clone() - }; - let kwargs = PyDict::new(py); - kwargs.set_item( - "type", - schema - .to_pyarrow(py)? - .call_method1("field", (&name,))? - .getattr("type")?, - )?; - let one = pa.call_method( - "array", - (PyList::new(py, [scalar])?,), - Some(&kwargs), - )?; - crate::arrow_input::array_from_pyarrow(&one) - }) - }) - .map_err(|cause| bridge_error(cause, &error)) - })) + literal( + py, + name.clone(), + value.unbind(), + schema.clone(), + error.clone(), + )? }; Ok((name, assignment)) }) .collect() } + +pub(crate) fn function(value: Py, error: CallbackError) -> UpdateAssignment { + UpdateAssignment::Function(Arc::new(move |batches| { + // A Python callback may invoke another native read or write. + tokio::task::block_in_place(|| { + Python::attach(|py| -> PyResult<_> { + let pa = py.import("pyarrow")?; + let batches = batches + .iter() + .map(|batch| batch.to_pyarrow(py)) + .collect::>>()?; + let matched = pa + .getattr("Table")? + .call_method1("from_batches", (PyList::new(py, batches)?,))?; + arrow_arrays(py, &value.bind(py).call1((matched,))?) + }) + }) + .map_err(|cause| bridge_error(cause, &error)) + })) +} + +pub(crate) fn literal( + py: Python<'_>, + name: String, + value: Py, + schema: SchemaRef, + error: CallbackError, +) -> PyResult { + let pa = py.import("pyarrow")?; + let bound = value.bind(py); + if bound.is_instance(&pa.getattr("Array")?)? + || bound.is_instance(&pa.getattr("ChunkedArray")?)? + { + return Ok(UpdateAssignment::Array(arrow_arrays(py, bound)?)); + } + // Literal conversion stays lazy: a no-match update must neither evaluate + // callbacks nor reject a scalar that is never assigned. Explicit MERGE + // literals also stay literals when the Python value happens to be callable. + Ok(UpdateAssignment::DeferredScalar(Arc::new(move || { + tokio::task::block_in_place(|| { + Python::attach(|py| -> PyResult<_> { + let pa = py.import("pyarrow")?; + let value = value.bind(py); + let scalar = if value.is_instance(&pa.getattr("Scalar")?)? { + value.call_method0("as_py")? + } else { + value.clone() + }; + let kwargs = PyDict::new(py); + kwargs.set_item( + "type", + schema + .to_pyarrow(py)? + .call_method1("field", (&name,))? + .getattr("type")?, + )?; + let one = pa.call_method("array", (PyList::new(py, [scalar])?,), Some(&kwargs))?; + crate::arrow_input::array_from_pyarrow(&one) + }) + }) + .map_err(|cause| bridge_error(cause, &error)) + }))) +} diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs index e5688b9bd..af66d51ed 100644 --- a/bindings/python/src/write.rs +++ b/bindings/python/src/write.rs @@ -1847,6 +1847,29 @@ mod tests { use paimon::spec::{BigIntType, IntType, Schema, TableSchema, VarCharType}; use pyo3::ffi::c_str; + fn merge_locals(py: Python<'_>) -> Bound<'_, PyDict> { + let locals = PyDict::new(py); + locals + .set_item( + "MergeAssignment", + py.get_type::(), + ) + .unwrap(); + locals + .set_item( + "WhenMatched", + py.get_type::(), + ) + .unwrap(); + locals + .set_item( + "WhenNotMatched", + py.get_type::(), + ) + .unwrap(); + locals + } + #[test] fn captured_plan_rejects_other_operations_before_python_input_consumption() { Python::attach(|py| { @@ -2406,6 +2429,174 @@ first = install('first') }); } + #[test] + fn merge_into_rejects_legacy_dictionary_clauses() { + Python::attach(|py| { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .option("row-tracking.enabled", "true") + .option("data-evolution.enabled", "true") + .build() + .unwrap(); + let table = Arc::new(Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "typed_merge_clauses"), + "memory:/typed-merge-clauses".into(), + TableSchema::new(0, &schema), + None, + )); + let update = Py::new(py, PyBatchWriteBuilder::new(table)) + .unwrap() + .bind(py) + .call_method0("new_update") + .unwrap(); + let locals = PyDict::new(py); + locals.set_item("update", update).unwrap(); + py.run( + c_str!( + r#" +import pyarrow as pa +source = pa.table({'id': pa.array([], pa.int32())}) +for matched, unmatched in [ + ([dict(condition=None, delete=False, assignments=[('id', 'source', 'id')])], []), + ([], [dict(condition=None, assignments=[('id', 'source', 'id')])]), +]: + try: + update.merge_into(source, on=[('id', 'id')], + when_matched=matched, when_not_matched=unmatched) + except TypeError: + pass + else: + raise AssertionError('MERGE accepted a legacy dictionary clause') +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn typed_merge_into_preserves_literal_and_callback_evaluation() { + use arrow::array::Int32Array; + use futures::TryStreamExt; + + Python::attach(|py| { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .option("row-tracking.enabled", "true") + .option("data-evolution.enabled", "true") + .build() + .unwrap(); + let table = Arc::new(Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "typed_merge_lazy_values"), + "memory:/typed-merge-lazy-values".into(), + TableSchema::new(0, &schema), + None, + )); + let locals = merge_locals(py); + locals + .set_item( + "builder", + Py::new(py, PyBatchWriteBuilder::new(table.clone())).unwrap(), + ) + .unwrap(); + py.run( + c_str!( + r#" +import pyarrow as pa +source = pa.table({'id': pa.array([1, 2], pa.int32()), + 'value': pa.array([10, 20], pa.int32())}) +update = builder.new_update() +literal_calls = [] +def literal_value(rows): + literal_calls.append(rows.num_rows) + return pa.array([99, 99], pa.int32()) +unused = WhenNotMatched({'id': MergeAssignment.source_column('id'), + 'value': MergeAssignment.literal(literal_value)}, + condition='"s.id" < 0') +assert update.merge_into(source, on=[('id', 'id')], when_matched=[], + when_not_matched=[unused]) == [] +assert literal_calls == [] +selected = WhenNotMatched({'id': MergeAssignment.source_column('id'), + 'value': MergeAssignment.literal(literal_value)}) +try: + update.merge_into(source, on=[('id', 'id')], when_matched=[], + when_not_matched=[selected]) +except (TypeError, ValueError): + pass +else: + raise AssertionError('callable literal was evaluated as an assignment function') +assert literal_calls == [] + +failure = LookupError('typed MERGE callback failed') +callback_rows = [] +def assignment(rows): + callback_rows.append(rows['s.id'].to_pylist()) + if len(callback_rows) == 1: + raise failure + return pa.chunked_array([pa.array([21], pa.int32())]) +assignments = {'id': MergeAssignment.source_column('id'), + 'value': MergeAssignment.function(function=assignment)} +clause = WhenNotMatched(assignments, condition='"s.id" = 2') +assignments.clear() +for attempt in [1, 2]: + try: + messages = update.merge_into(source, on=[('id', 'id')], when_matched=[], + when_not_matched=[unused, clause]) + except LookupError as error: + assert attempt == 1 + assert error is failure + else: + assert attempt == 2 + builder.new_commit().commit(messages) +assert callback_rows == [[2], [2]] +assert literal_calls == [] +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let batches = py.detach(|| { + runtime().block_on(async { + let builder = table.new_read_builder(); + let plan = builder.new_scan().plan().await.unwrap(); + builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect::>() + .await + .unwrap() + }) + }); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2 + ); + assert_eq!( + batches[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 21 + ); + }); + } + #[test] fn merge_table_source_conditions_resolve_read_metadata() { Python::attach(|py| { @@ -2434,15 +2625,15 @@ first = install('first') )); let source = Py::new(py, crate::table::PyTable::new(source)).unwrap(); let context = UpdateContext::new(&WriteContext::new(table)).unwrap(); - let locals = PyDict::new(py); + let locals = merge_locals(py); locals.set_item("metadata", metadata).unwrap(); py.run( c_str!( r#" -condition = dict(sql='"s.' + metadata + '" >= 0') -matched = [dict(assignments=[('value', 'source', 'value')], condition=condition, delete=False)] -not_matched = [dict(assignments=[('id', 'source', metadata), ('value', 'source', 'value')], - condition=condition)] +condition = '"s.' + metadata + '" >= 0' +matched = [WhenMatched.update({'value': MergeAssignment.source_column('value')}, condition=condition)] +not_matched = [WhenNotMatched({'id': MergeAssignment.source_column(metadata), + 'value': MergeAssignment.source_column('value')}, condition=condition)] "# ), Some(&locals), @@ -2558,17 +2749,18 @@ not_matched = [dict(assignments=[('id', 'source', metadata), ('value', 'source', }); let source = Py::new(py, crate::table::PyTable::new(source)).unwrap(); let context = UpdateContext::new(&WriteContext::new(target.clone())).unwrap(); - let locals = PyDict::new(py); + let locals = merge_locals(py); locals.set_item("metadata", metadata).unwrap(); locals.set_item("first", first).unwrap(); locals.set_item("second", second).unwrap(); py.run( c_str!( r#" -matched = [dict(assignments=[('value', 'source', 'value')], - condition=dict(sql='"s.' + metadata + '" = ' + str(first)), delete=False)] -not_matched = [dict(assignments=[('id', 'source', metadata), ('value', 'source', 'value')], - condition=dict(sql='"s.' + metadata + '" = ' + str(second)))] +matched = [WhenMatched.update({'value': MergeAssignment.source_column('value')}, + condition='"s.' + metadata + '" = ' + str(first))] +not_matched = [WhenNotMatched({'id': MergeAssignment.source_column(metadata), + 'value': MergeAssignment.source_column('value')}, + condition='"s.' + metadata + '" = ' + str(second))] "# ), Some(&locals), @@ -2794,7 +2986,7 @@ invalid = source.slice(0, 1) }) }); let context = UpdateContext::new(&WriteContext::new(table)).unwrap(); - let locals = PyDict::new(py); + let locals = merge_locals(py); locals.set_item("reverse", reverse).unwrap(); py.run( c_str!( @@ -2817,7 +3009,7 @@ encoded = pa.chunked_array([pa.DictionaryArray.from_arrays( source = source.append_column(pa.field('encoded_unused', encoded.type, nullable=False), encoded) assert source.column('encoded_unused').null_count == 0 assert source.column('encoded_unused').to_pylist() == [None, None, None] -matched = [dict(assignments=[('value', 'source', 'value')], condition=None, delete=False)] +matched = [WhenMatched.update({'value': MergeAssignment.source_column('value')})] not_matched = [] "# ), diff --git a/bindings/python/tests/test_blob_uri_reader.py b/bindings/python/tests/test_blob_uri_reader.py index 54bbc54a2..aa457e225 100644 --- a/bindings/python/tests/test_blob_uri_reader.py +++ b/bindings/python/tests/test_blob_uri_reader.py @@ -22,7 +22,7 @@ import pytest from test_blob_consumer import _table -from pypaimon_rust.datafusion import PaimonCatalog, SQLContext +from pypaimon_rust.datafusion import MergeAssignment, PaimonCatalog, SQLContext, WhenMatched def descriptor(offset, length): @@ -217,9 +217,9 @@ def create(self, uri): extra = {'commit_identifier': 7} if stream else {} with pytest.raises(OSError) as caught: if operation == 'merge': - update.merge_into(source, on=[('id', 'id')], when_matched=[{ - 'delete': False, 'assignments': [('payload', 'source', 'payload')], - }], when_not_matched=[], **extra) + update.merge_into(source, on=[('id', 'id')], when_matched=[ + WhenMatched.update({'payload': MergeAssignment.source_column('payload')}), + ], when_not_matched=[], **extra) elif operation == 'row_id': update.update_by_arrow_with_row_id(data, **extra) elif operation == 'grouped_row_id': diff --git a/bindings/python/tests/test_update.py b/bindings/python/tests/test_update.py index 069a3a346..83dc9464c 100644 --- a/bindings/python/tests/test_update.py +++ b/bindings/python/tests/test_update.py @@ -19,7 +19,9 @@ import pytest import pypaimon_rust.datafusion as datafusion -from pypaimon_rust.datafusion import PaimonCatalog, SQLContext +from pypaimon_rust.datafusion import ( + MergeAssignment, PaimonCatalog, SQLContext, WhenMatched, WhenNotMatched, +) def test_upsert_batches_preserve_missing_fields_nulls_and_stream_updates(tmp_path): @@ -449,13 +451,14 @@ def test_prepared_merge_conditions_chunks_and_lazy_literals(tmp_path, stream): ('id', pa.int32()), ('value', pa.int32())])) update = builder.new_update() kwargs = dict(commit_identifier=42) if stream else {} - false_clause = dict(delete=False, condition=dict(sql='"s.value" < 0'), - assignments=[('value', 'literal', 'invalid int')]) + false_clause = WhenMatched.update( + {'value': MergeAssignment.literal('invalid int')}, condition='"s.value" < 0') assert update.merge_into(source, on=[('id', 'id')], when_matched=[false_clause], when_not_matched=[], **kwargs) == [] messages = update.merge_into(source, on=[('id', 'id')], when_matched=[ - dict(delete=False, condition=dict(sql='"s.value" IN (SELECT v FROM (VALUES (11), (22)) AS m(v))'), - assignments=[('value', 'literal', pa.chunked_array([[100], [200]]))])], + WhenMatched.update( + {'value': MergeAssignment.literal(pa.chunked_array([[100], [200]]))}, + condition='"s.value" IN (SELECT v FROM (VALUES (11), (22)) AS m(v))')], when_not_matched=[], **kwargs) # Preparation leaves the visible snapshot unchanged. @@ -495,12 +498,14 @@ def test_prepared_merge_accepts_native_table_sources(tmp_path, stream, historica builder = (target.new_stream_write_builder().with_commit_user('source-merge') if stream else target.new_batch_write_builder()) kwargs = dict(commit_identifier=42) if stream else {} - condition = dict(sql='"s.flag" = 1') + condition = '"s.flag" = 1' messages = builder.new_update().merge_into( source, on=[('id', 'source_id')], when_matched=[ - dict(delete=False, condition=condition, assignments=[('value', 'source', 'amount')])], - when_not_matched=[dict(condition=condition, assignments=[ - ('id', 'source', 'source_id'), ('value', 'source', 'amount')])], **kwargs) + WhenMatched.update({'value': MergeAssignment.source_column('amount')}, condition)], + when_not_matched=[WhenNotMatched({ + 'id': MergeAssignment.source_column('source_id'), + 'value': MergeAssignment.source_column('amount'), + }, condition)], **kwargs) assert target.latest_snapshot().id() == target_snapshot assert source.latest_snapshot().id() == source_snapshot commit = builder.new_commit() @@ -533,8 +538,9 @@ def values(batch): return pa.chunked_array([[100], [200]], type=pa.int32()) messages = update.merge_into(source, on=[('id', 'id')], when_matched=[ - dict(delete=False, assignments=[('value', 'callable', values)])], when_not_matched=[ - dict(assignments=[('id', 'source', 'id'), ('value', 'callable', lambda batch: batch['s.delta'])])], + WhenMatched.update({'value': MergeAssignment.function(values)})], when_not_matched=[ + WhenNotMatched({'id': MergeAssignment.source_column('id'), + 'value': MergeAssignment.function(lambda batch: batch['s.delta'])})], **kwargs) assert calls == [2] commit = builder.new_commit() @@ -547,7 +553,7 @@ def values(batch): def fail(batch): raise error - clause = [dict(delete=False, assignments=[('value', 'callable', fail)])] + clause = [WhenMatched.update({'value': MergeAssignment.function(fail)})] update = builder.new_update() assert update.merge_into(pa.table({'id': pa.array([99], pa.int32())}), on=[('id', 'id')], when_matched=clause, when_not_matched=[], **kwargs) == []