diff --git a/paimon-python/pypaimon/execution/native/write/merge_transform.py b/paimon-python/pypaimon/execution/native/write/merge_transform.py new file mode 100644 index 000000000000..f3a5d64178a7 --- /dev/null +++ b/paimon-python/pypaimon/execution/native/write/merge_transform.py @@ -0,0 +1,148 @@ +# 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. + +"""Transport joined Arrow batches to the Core MERGE action executor.""" + + +def _callable_assignment(function, on_pairs, callable_input_columns=None, hidden_columns=()): + from pypaimon.write.merge_transform import _merge_callable_input + + def evaluate(batch): + if hidden_columns: + batch = batch.drop(hidden_columns) + if callable_input_columns is not None: + import pyarrow as pa + view = pa.Table.from_arrays( + [batch.column('t.' + name) for name in callable_input_columns], + names=list(callable_input_columns), + ) + else: + view = _merge_callable_input(batch, on_pairs) + return function(view) + + return evaluate + + +def encode_merge_clause(clause, on_pairs=(), on_map=None, input_schema=None, + callable_input_columns=None, hidden_columns=()): + """Encode the existing clause contract without evaluating user callbacks.""" + from pypaimon.write.merge_types import SourceColumnRef, TargetColumnRef, LiteralValue + from pypaimon.write.merge_condition import remap_source_on_keys, rewrite_condition + + available = None if input_schema is None else set(input_schema.names) + values = [] + for name, value in clause.spec.items(): + if isinstance(value, SourceColumnRef): + kind, column = 'source', value.column + if available is not None and 's.' + column not in available: + target = next((target for source, target in on_pairs + if source == column and 't.' + target in available), None) + if target is None: + values.append((name, 'literal', None)) + continue + kind, column = 'target', target + values.append((name, kind, column)) + elif isinstance(value, TargetColumnRef): + if available is not None and 't.' + value.column not in available: + values.append((name, 'literal', None)) + else: + values.append((name, 'target', value.column)) + elif isinstance(value, LiteralValue): + values.append((name, 'literal', value.value)) + elif callable(value) and not isinstance(value, type): + values.append((name, 'callable', _callable_assignment( + value, on_pairs, callable_input_columns, hidden_columns))) + else: + return None + condition = None + if clause.condition is not None: + rewritten = rewrite_condition(clause.condition) + if on_map: + rewritten = remap_source_on_keys(rewritten, on_map) + condition = dict(sql=rewritten) + return dict(assignments=values, condition=condition, delete=clause.delete) + + +class _NativeMergeTransform: + def __init__(self, clauses, output_schema, mode, on_pairs=(), on_map=None, + callable_input_columns=None, row_id_name='_ROW_ID'): + # Require this release's protocol before consuming input. An installed + # incompatible SDK must not silently select another executor. + from pypaimon_rust.datafusion import MergeBatchExecution # noqa: F401 + + self.clauses = clauses + self.output_schema = output_schema + self.mode = mode + self.on_pairs = on_pairs + self.on_map = on_map + self.callable_input_columns = callable_input_columns + self.row_id_name = row_id_name + self._execution = None + + def __getstate__(self): + # Ray transports the request; each worker compiles its own operator. + state = self.__dict__.copy() + state['_execution'] = None + return state + + def __call__(self, batch): + if batch.num_rows == 0: + return self.output_schema.empty_table() + output_schema = self.output_schema + hidden_columns = () + if self.mode != 'insert' and self.row_id_name != '_ROW_ID': + batch = batch.append_column('t._ROW_ID', batch.column('t.' + self.row_id_name)) + hidden_columns = ('t._ROW_ID',) + index = output_schema.get_field_index(self.row_id_name) + output_schema = output_schema.set(index, output_schema.field(index).with_name('_ROW_ID')) + if self._execution is None: + from pypaimon_rust.datafusion import MergeBatchExecution + + clauses = [encode_merge_clause( + clause, self.on_pairs, self.on_map, batch.schema, + self.callable_input_columns, hidden_columns, + ) for clause in self.clauses] + if any(clause is None for clause in clauses): + raise TypeError('Unsupported normalized MERGE assignment') + self._execution = MergeBatchExecution( + batch.schema, output_schema, clauses, self.mode, + ) + result = self._execution.execute(batch) + if hidden_columns: + import pyarrow as pa + result = pa.Table.from_arrays(result.columns, schema=self.output_schema) + return result + + +def _build_matched_transform( + clauses, on_map, on_pairs, update_cols, row_id_name, + update_schema, callable_input_columns=None, +): + return _NativeMergeTransform( + clauses, update_schema, 'update', on_pairs, on_map, + callable_input_columns, row_id_name, + ) + + +def _build_matched_delete_transform(clauses, on_map, row_id_name, delete_schema): + return _NativeMergeTransform( + clauses, delete_schema, 'delete', on_map=on_map, row_id_name=row_id_name, + ) + + +def _build_insert_transform(clauses, target_field_names, target_pa_schema): + return _NativeMergeTransform(clauses, target_pa_schema, 'insert') diff --git a/paimon-python/pypaimon/execution/native/write/native_merge_into.py b/paimon-python/pypaimon/execution/native/write/native_merge_into.py index cd92aca1e3c7..c60540f8b49d 100644 --- a/paimon-python/pypaimon/execution/native/write/native_merge_into.py +++ b/paimon-python/pypaimon/execution/native/write/native_merge_into.py @@ -25,22 +25,12 @@ _dictionary_value_type, _native_row_id_table, _native_update_columns_supported, _native_update_formats_supported, _supported_upsert_key_type, ) - - -def _callable_assignment(function, on_pairs): - from pypaimon.write.merge_transform import _merge_callable_input - - def evaluate(batch): - return function(_merge_callable_input(batch, on_pairs)) - - return evaluate +from pypaimon.execution.native.write.merge_transform import encode_merge_clause def create_native_merge_into(table, source, on, matched, not_matched, commit_user, commit_identifier): from pypaimon.write.merge_inputs import _is_self_merge, _is_table_like, _prepare from pypaimon.write.merge_inputs import _normalize_on, _union_update_cols - from pypaimon.write.merge_types import SourceColumnRef, TargetColumnRef, LiteralValue - from pypaimon.write.merge_condition import rewrite_condition target_keys, source_keys = _normalize_on(on) self_merge = _is_self_merge(table, source, target_keys, source_keys) @@ -100,21 +90,7 @@ def create_native_merge_into(table, source, on, matched, not_matched, commit_use return None def encode(clause): - values = [] - for name, value in clause.spec.items(): - if isinstance(value, SourceColumnRef): - values.append((name, 'source', value.column)) - elif isinstance(value, TargetColumnRef): - values.append((name, 'target', value.column)) - elif isinstance(value, LiteralValue): - values.append((name, 'literal', value.value)) - elif callable(value) and not isinstance(value, type): - values.append((name, 'callable', _callable_assignment( - value, list(zip(source_keys, target_keys))))) - else: - return None - condition = None if clause.condition is None else dict(sql=rewrite_condition(clause.condition)) - return dict(assignments=values, condition=condition, delete=clause.delete) + return encode_merge_clause(clause, list(zip(source_keys, target_keys))) encoded_matched = [encode(clause) for clause in matched] encoded_not_matched = [encode(clause) for clause in not_matched] diff --git a/paimon-python/pypaimon/execution/python/write/merge_transform.py b/paimon-python/pypaimon/execution/python/write/merge_transform.py index 4418a8b914fa..d009eccc1b28 100644 --- a/paimon-python/pypaimon/execution/python/write/merge_transform.py +++ b/paimon-python/pypaimon/execution/python/write/merge_transform.py @@ -22,6 +22,7 @@ import pyarrow as pa +from pypaimon.write.merge_condition import _load_datafusion from pypaimon.write.merge_types import _NormalizedClause from pypaimon.write.merge_transform import ( cast_to_schema, @@ -54,6 +55,7 @@ def _build_matched_transform( _filter_batch = None if any(r is not None for _, r, _ in prepared_clauses): + _load_datafusion() from pypaimon.execution.python.write.merge_condition import filter_batch as _filter_batch def _transform(batch: pa.Table) -> pa.Table: @@ -120,6 +122,7 @@ def _build_matched_delete_transform( _filter_batch = None if any(r is not None for r, _ in prepared_clauses): + _load_datafusion() from pypaimon.execution.python.write.merge_condition import filter_batch as _filter_batch def _transform(batch: pa.Table) -> pa.Table: @@ -170,6 +173,7 @@ def _build_insert_transform( _filter_batch = None if any(r is not None for _, r in prepared_clauses): + _load_datafusion() from pypaimon.execution.python.write.merge_condition import filter_batch as _filter_batch def _transform(batch: pa.Table) -> pa.Table: diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_into.py b/paimon-python/pypaimon/ray/data_evolution_merge_into.py index 9209a3f5a830..77be7bd9d987 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_into.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_into.py @@ -237,10 +237,7 @@ def _prepare( for c in list(when_matched) + list(when_not_matched) ) if has_condition: - from pypaimon.write.merge_condition import ( - _load_datafusion, extract_target_columns, - ) - _load_datafusion() + from pypaimon.write.merge_condition import extract_target_columns for c in when_not_matched: if c.condition is not None: t_refs = extract_target_columns(c.condition) @@ -394,6 +391,7 @@ def _build_datasets( insert_ds = None update_cols_union: List[str] = [] target_empty = _is_target_empty(base_snapshot) + native_enabled = table.options.native_write_enabled() if ctx.is_self_merge: if matched_specs and not target_empty: @@ -409,6 +407,7 @@ def _build_datasets( snapshot_id=base_snapshot_id, scan_predicate=ctx.self_merge_scan_predicate, read_columns=ctx.read_columns, + native_enabled=native_enabled, ) if any(c.delete for c in matched_specs): delete_ds = build_self_merge_delete_ds( @@ -420,6 +419,7 @@ def _build_datasets( snapshot_id=base_snapshot_id, scan_predicate=ctx.self_merge_scan_predicate, ray_remote_args=ray_remote_args, + native_enabled=native_enabled, ) return update_ds, delete_ds, insert_ds, update_cols_union @@ -444,6 +444,7 @@ def _build_datasets( resolve_target_projection=_resolve_target_projection, snapshot_id=base_snapshot_id, ray_remote_args=ray_remote_args, + native_enabled=native_enabled, ) if any(c.delete for c in matched_specs): delete_ds = build_matched_delete_ds( @@ -459,6 +460,7 @@ def _build_datasets( resolve_target_projection=_resolve_target_projection, snapshot_id=base_snapshot_id, ray_remote_args=ray_remote_args, + native_enabled=native_enabled, ) if not_matched_specs: @@ -475,6 +477,7 @@ def _build_datasets( snapshot_id=base_snapshot_id, target_empty=target_empty, ray_remote_args=ray_remote_args, + native_enabled=native_enabled, ) return update_ds, delete_ds, insert_ds, update_cols_union diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index 6a9ecd1df493..edbb845d395b 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -82,6 +82,7 @@ class _SelfMergeUpdatePlan: snapshot_id: int callable_input_columns: Optional[List[str]] row_id_plan: Any + native_enabled: bool = False @dataclass(frozen=True) @@ -98,6 +99,7 @@ class _SelfMergeUpdateContext: row_id_name: str snapshot_id: int callable_input_columns: Optional[List[str]] + native_enabled: bool = False def _resolve_source_projection( @@ -132,9 +134,13 @@ def _build_matched_transform( row_id_name: str, update_schema: pa.Schema, callable_input_columns: Optional[Sequence[str]] = None, + native_enabled: bool = False, ): - from pypaimon.execution.python.write.merge_transform import _build_matched_transform as execute - return execute(clauses, on_map, on_pairs, update_cols, row_id_name, update_schema, callable_input_columns) + from pypaimon.write.backend_dispatch import create_merge_batch_transform + return create_merge_batch_transform( + 'update', clauses, on_map, on_pairs, update_cols, row_id_name, + update_schema, callable_input_columns, native_enabled=native_enabled, + ) def _build_matched_delete_transform( @@ -142,9 +148,13 @@ def _build_matched_delete_transform( on_map: Dict[str, str], row_id_name: str, delete_schema: pa.Schema, + native_enabled: bool = False, ): - from pypaimon.execution.python.write.merge_transform import _build_matched_delete_transform as execute - return execute(clauses, on_map, row_id_name, delete_schema) + from pypaimon.write.backend_dispatch import create_merge_batch_transform + return create_merge_batch_transform( + 'delete', clauses, on_map, row_id_name, delete_schema, + native_enabled=native_enabled, + ) def build_self_merge_update_plan( @@ -158,6 +168,7 @@ def build_self_merge_update_plan( snapshot_id: Optional[int] = None, scan_predicate=None, read_columns: Sequence[str] = (), + native_enabled: bool = False, ) -> _SelfMergeUpdatePlan: from pypaimon.common.options.core_options import ( CoreOptions, GlobalIndexSearchMode, @@ -229,6 +240,7 @@ def build_self_merge_update_plan( if read_columns else None ), row_id_plan=row_id_plan, + native_enabled=native_enabled, ) @@ -262,6 +274,7 @@ def _apply_self_merge_update_group(context, group_plan, collect_row_ids): row_id_name=context.row_id_name, update_schema=context.update_schema, callable_input_columns=context.callable_input_columns, + native_enabled=context.native_enabled, ) update_parts = [] batch_reader = table_read.to_arrow_batch_reader([file_group]) @@ -330,6 +343,7 @@ def distributed_self_merge_update_apply( row_id_name=plan.row_id_name, snapshot_id=plan.snapshot_id, callable_input_columns=plan.callable_input_columns, + native_enabled=plan.native_enabled, )) group_iter = iter(captured.groups) max_in_flight = max(1, min(num_partitions, len(captured.groups))) @@ -388,6 +402,7 @@ def build_self_merge_delete_ds( snapshot_id: Optional[int] = None, scan_predicate=None, ray_remote_args: Optional[Dict[str, Any]] = None, + native_enabled: bool = False, ) -> Tuple: from pypaimon.ray.ray_paimon import read_paimon from pypaimon.table.special_fields import SpecialFields @@ -453,6 +468,7 @@ def _add_source_aliases(batch: pa.Table) -> pa.Table: on_map={row_id_name: row_id_name}, row_id_name=row_id_name, delete_schema=delete_schema, + native_enabled=native_enabled, ) return aliased.map_batches(_transform, **_map_kwargs(ray_remote_args)) @@ -473,6 +489,7 @@ def build_matched_update_ds( estimated_size_bytes: Optional[int] = None, snapshot_id: Optional[int] = None, ray_remote_args: Optional[Dict[str, Any]] = None, + native_enabled: bool = False, ) -> Tuple: from pypaimon.ray.ray_paimon import read_paimon from pypaimon.table.special_fields import SpecialFields @@ -518,6 +535,7 @@ def build_matched_update_ds( update_cols=list(update_cols), row_id_name=row_id_name, update_schema=update_schema, + native_enabled=native_enabled, ) return joined.map_batches(_transform, **_map_kwargs(ray_remote_args)) @@ -536,6 +554,7 @@ def build_matched_delete_ds( estimated_size_bytes: Optional[int] = None, snapshot_id: Optional[int] = None, ray_remote_args: Optional[Dict[str, Any]] = None, + native_enabled: bool = False, ) -> Tuple: from pypaimon.ray.ray_paimon import read_paimon from pypaimon.table.special_fields import SpecialFields @@ -579,6 +598,7 @@ def build_matched_delete_ds( on_map=dict(zip(source_on, target_on)), row_id_name=row_id_name, delete_schema=delete_schema, + native_enabled=native_enabled, ) return joined.map_batches(_transform, **_map_kwargs(ray_remote_args)) @@ -1052,6 +1072,7 @@ def build_not_matched_insert_ds( target_empty: bool = False, snapshot_id: Optional[int] = None, ray_remote_args: Optional[Dict[str, Any]] = None, + native_enabled: bool = False, ): from pypaimon.ray.ray_paimon import read_paimon @@ -1084,9 +1105,12 @@ def build_not_matched_insert_ds( right_on=tuple(f"t.{c}" for c in target_on), ) - from pypaimon.execution.python.write.merge_transform import _build_insert_transform + from pypaimon.write.backend_dispatch import create_merge_batch_transform - _transform = _build_insert_transform(clauses, captured_field_names, out_schema) + _transform = create_merge_batch_transform( + 'insert', clauses, captured_field_names, out_schema, + native_enabled=native_enabled, + ) return unmatched.map_batches( _transform, **_map_kwargs(ray_remote_args) diff --git a/paimon-python/pypaimon/tests/merge_execution_protocol_test.py b/paimon-python/pypaimon/tests/merge_execution_protocol_test.py new file mode 100644 index 000000000000..083476ec4555 --- /dev/null +++ b/paimon-python/pypaimon/tests/merge_execution_protocol_test.py @@ -0,0 +1,446 @@ +# 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. + +"""Transport and backend selection for distributed MERGE action batches.""" + +import pickle +import subprocess +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock, patch + +import pyarrow as pa +import pytest + +from pypaimon.write.backend_dispatch import create_merge_batch_transform +from pypaimon.write.merge_types import ( + LiteralValue, SourceColumnRef, TargetColumnRef, _NormalizedClause, +) + + +@pytest.fixture +def native_protocol(monkeypatch): + calls = [] + execute = Mock() + + class MergeBatchExecution: + def __init__(self, input_schema, output_schema, clauses, mode): + calls.append((input_schema, output_schema, clauses, mode)) + + def execute(self, batch): + return execute(batch, calls[-1]) + + runtime = ModuleType('pypaimon_rust') + runtime.__path__ = [] + datafusion = ModuleType('pypaimon_rust.datafusion') + datafusion.MergeBatchExecution = MergeBatchExecution + monkeypatch.setitem(sys.modules, 'pypaimon_rust', runtime) + monkeypatch.setitem(sys.modules, 'pypaimon_rust.datafusion', datafusion) + return SimpleNamespace(calls=calls, execute=execute, module=datafusion) + + +def _joined_batch(): + return pa.table({ + 't._ROW_ID': pa.array([7, 8], pa.int64()), + 't.id': pa.array([1, 2], pa.int32()), + 't.value': pa.array([10, 20], pa.int32()), + 's.value': pa.array([100, 200], pa.int32()), + '_scratch': [0, 0], + }) + + +def _update_schema(): + return pa.schema([ + pa.field('_ROW_ID', pa.int64(), nullable=False), + pa.field('value', pa.int32()), + ]) + + +def _native_update(clauses, callable_input_columns=None): + return create_merge_batch_transform( + 'update', clauses, {'source_id': 'id'}, [('source_id', 'id')], + ['value'], '_ROW_ID', _update_schema(), callable_input_columns, + native_enabled=True, + ) + + +def test_native_update_restores_callable_aliases_and_preserves_result_chunks(native_protocol): + observed = [] + chunks = pa.chunked_array([ + pa.array([101], pa.int32()), pa.array([202], pa.int32()), + ]) + + def values(view): + observed.append(view) + return chunks + + transform = _native_update([_NormalizedClause( + {'value': values}, 's.source_id > 0', + )]) + batch = _joined_batch() + result = pa.Table.from_arrays([ + batch['t._ROW_ID'], chunks, + ], schema=_update_schema()) + + def execute(rows, compiled): + input_schema, output_schema, clauses, mode = compiled + assert input_schema == batch.schema + assert output_schema == _update_schema() + assert mode == 'update' + assert clauses[0]['condition'] == {'sql': '"t.id" > 0'} + name, kind, callback = clauses[0]['assignments'][0] + assert (name, kind) == ('value', 'callable') + assert callback(rows).equals(result['value']) + return result + + native_protocol.execute.side_effect = execute + assert not observed + actual = transform(batch) + assert actual is result + assert actual['value'].num_chunks == 2 + assert len(observed) == 1 + assert observed[0].column_names == ['s.source_id', 's.value', 't._ROW_ID', 't.id', 't.value'] + assert observed[0]['s.source_id'].equals(batch['t.id']) + + +def test_native_self_update_keeps_explicit_callable_read_order(native_protocol): + observed = [] + + def values(view): + observed.append(view) + return view['value'] + + transform = _native_update( + [_NormalizedClause({'value': values})], ['value', '_ROW_ID', 'id'], + ) + batch = _joined_batch() + + def execute(rows, compiled): + callback = compiled[2][0]['assignments'][0][2] + return pa.Table.from_arrays([rows['t._ROW_ID'], callback(rows)], schema=_update_schema()) + + native_protocol.execute.side_effect = execute + assert transform(batch).to_pylist() == [{'_ROW_ID': 7, 'value': 10}, {'_ROW_ID': 8, 'value': 20}] + assert observed[0].column_names == ['value', '_ROW_ID', 'id'] + + +@pytest.mark.parametrize('operation', ['update', 'delete']) +def test_native_custom_row_id_preserves_aliases_and_output_schema(native_protocol, operation): + observed = [] + + def values(view): + observed.append(view) + return view['t.value'] + + schema = pa.schema([pa.field('custom_id', pa.int64(), nullable=False)]) + if operation == 'update': + schema = schema.append(pa.field('value', pa.int32())) + transform = create_merge_batch_transform( + 'update', [_NormalizedClause({'value': values}, 't.custom_id > 0')], {}, [], + ['value'], 'custom_id', schema, native_enabled=True, + ) + else: + transform = create_merge_batch_transform( + 'delete', [_NormalizedClause({}, 't.custom_id > 0', delete=True)], {}, + 'custom_id', schema, native_enabled=True, + ) + batch = pa.table({'t.custom_id': pa.array([7, 8], pa.int64()), 't.value': pa.array([10, 20], pa.int32())}) + + def execute(rows, compiled): + assert rows['t._ROW_ID'].equals(batch['t.custom_id']) + assert compiled[2][0]['condition'] == {'sql': '"t.custom_id" > 0'} + assert compiled[1].names[0] == '_ROW_ID' + columns = [rows['t._ROW_ID']] + if operation == 'update': + columns.append(compiled[2][0]['assignments'][0][2](rows)) + return pa.Table.from_arrays(columns, schema=compiled[1]) + + native_protocol.execute.side_effect = execute + result = transform(batch) + assert result.schema == schema + assert result['custom_id'].to_pylist() == [7, 8] + if operation == 'update': + assert observed[0].column_names == ['t.custom_id', 't.value'] + + +def test_native_transports_coalesced_and_missing_column_assignments(native_protocol): + clauses = [_NormalizedClause({ + 'key': SourceColumnRef('source_id'), + 'source': SourceColumnRef('value'), + 'target': TargetColumnRef('value'), + 'missing_source': SourceColumnRef('absent'), + 'missing_target': TargetColumnRef('absent'), + 'constant': LiteralValue(9), + })] + transform = _native_update(clauses) + native_protocol.execute.return_value = _update_schema().empty_table() + transform(_joined_batch()) + assert native_protocol.calls[0][2][0]['assignments'] == [ + ('key', 'target', 'id'), ('source', 'source', 'value'), + ('target', 'target', 'value'), ('missing_source', 'literal', None), + ('missing_target', 'literal', None), ('constant', 'literal', 9), + ] + + +def test_native_insert_callback_retains_source_aliases_and_dictionary_types(native_protocol): + batch = pa.table({ + 's.source_id': pa.array([1, 2], pa.int32()).dictionary_encode(), + 's.value': pa.array([10, 20], pa.int32()), + '_scratch': [0, 0], + }) + observed = [] + + def values(view): + observed.append(view) + return view['s.value'] + + schema = pa.schema([('value', pa.int32())]) + transform = create_merge_batch_transform( + 'insert', [_NormalizedClause({'value': values})], ['value'], schema, + native_enabled=True, + ) + + def execute(rows, compiled): + assert compiled[3] == 'insert' + return pa.Table.from_arrays([compiled[2][0]['assignments'][0][2](rows)], schema=schema) + + native_protocol.execute.side_effect = execute + assert transform(batch)['value'].to_pylist() == [10, 20] + assert observed[0].column_names == ['s.source_id', 's.value'] + assert observed[0]['s.source_id'].type == batch['s.source_id'].type + + +@pytest.mark.parametrize('operation', ['update', 'delete', 'insert']) +def test_native_transform_keeps_empty_schema_and_compiled_state_local(native_protocol, operation): + clauses = [_NormalizedClause({'value': SourceColumnRef('value')})] + schema = _update_schema() + if operation == 'update': + transform = _native_update(clauses) + elif operation == 'delete': + schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False)]) + transform = create_merge_batch_transform( + 'delete', clauses, {'source_id': 'id'}, '_ROW_ID', schema, native_enabled=True, + ) + else: + schema = pa.schema([pa.field('value', pa.int32())]) + transform = create_merge_batch_transform( + 'insert', clauses, ['value'], schema, native_enabled=True, + ) + batch = _joined_batch() + empty = transform(batch.slice(0, 0)) + assert empty.schema == schema + assert empty.num_rows == 0 + assert native_protocol.calls == [] + native_protocol.execute.return_value = schema.empty_table() + transform(batch) + transform(batch) + assert len(native_protocol.calls) == 1 + assert native_protocol.calls[0][3] == operation + restored = pickle.loads(pickle.dumps(transform)) + restored(batch) + assert len(native_protocol.calls) == 2 + + +@pytest.mark.parametrize('phase', ['compile', 'execute', 'callable']) +def test_native_failure_preserves_exception_identity_without_python_replay(native_protocol, phase): + failure = ValueError('MERGE failure after selecting native execution') + + def callback(rows): + raise failure + + if phase == 'compile': + native_protocol.module.MergeBatchExecution = Mock(side_effect=failure) + elif phase == 'execute': + native_protocol.execute.side_effect = failure + else: + native_protocol.execute.side_effect = lambda rows, compiled: compiled[2][0]['assignments'][0][2](rows) + with patch('pypaimon.execution.python.write.merge_transform._build_matched_transform') as python: + transform = _native_update([_NormalizedClause({'value': callback})]) + with pytest.raises(ValueError) as caught: + transform(_joined_batch()) + assert caught.value is failure + python.assert_not_called() + + +@pytest.mark.parametrize('native_enabled', [False, True]) +@pytest.mark.parametrize('self_merge', [False, True]) +def test_ray_merge_seals_table_backend_for_every_action(native_enabled, self_merge): + from pypaimon.ray import data_evolution_merge_into as merge + + requested = Mock(return_value=native_enabled) + table = SimpleNamespace(options=SimpleNamespace(native_write_enabled=requested)) + ctx = SimpleNamespace( + is_self_merge=self_merge, full_target_field_names=['id', 'value'], + settable_field_names=['id', 'value'], target_on_cols=['id'], source_on_cols=['source_id'], + update_pa_schema=_update_schema(), full_pa_schema=_update_schema(), + catalog_options={'warehouse': '/tmp/merge'}, self_merge_scan_predicate=None, + read_columns=('value',), + ) + names = (['build_self_merge_update_plan', 'build_self_merge_delete_ds'] if self_merge else + ['build_matched_update_ds', 'build_matched_delete_ds', 'build_not_matched_insert_ds']) + builders = {name: Mock(return_value=object()) for name in names} + matched = [_NormalizedClause({'value': LiteralValue(1)}), _NormalizedClause({}, delete=True)] + with patch.multiple(merge, **builders): + merge._build_datasets( + table, 'default.t', object(), matched, [_NormalizedClause({'value': LiteralValue(2)})], + ctx, SimpleNamespace(id=3, total_record_count=1), 2, None, + ) + requested.assert_called_once_with() + for builder in builders.values(): + builder.assert_called_once() + assert builder.call_args.kwargs['native_enabled'] is native_enabled + + +def test_missing_optional_runtime_selects_python_before_consuming_input(): + with patch('pypaimon.execution.native.runtime.native_runtime_available', return_value=False) as available, \ + patch('pypaimon.execution.native.write.merge_transform._build_matched_transform') as native: + transform = _native_update([_NormalizedClause({'value': LiteralValue(5)})]) + result = transform(_joined_batch()) + assert result.to_pylist() == [{'_ROW_ID': 7, 'value': 5}, {'_ROW_ID': 8, 'value': 5}] + available.assert_called_once_with() + native.assert_not_called() + + +@pytest.mark.parametrize('operation', ['update', 'delete', 'insert', 'self-update', 'self-delete']) +def test_native_ray_preparation_does_not_require_python_datafusion( + tmp_path, monkeypatch, native_protocol, operation): + from pypaimon import CatalogFactory, Schema + from pypaimon.ray import data_evolution_merge_into as merge + from pypaimon.write.merge_types import WhenMatched, WhenNotMatched + + catalog_options = {'warehouse': str(tmp_path)} + catalog = CatalogFactory.create(catalog_options) + catalog.create_database('default', True) + schema = pa.schema([('id', pa.int32()), ('value', pa.int32())]) + catalog.create_table('default.t', Schema.from_pyarrow_schema(schema, options={ + 'data-evolution.enabled': 'true', 'row-tracking.enabled': 'true', + 'deletion-vectors.enabled': 'true', 'write.native.enabled': 'true', + }), False) + source = Mock() + source.schema.return_value = schema + monkeypatch.setitem(sys.modules, 'datafusion', None) + self_merge = operation.startswith('self-') + operation = operation[5:] if self_merge else operation + condition = 's.value > 0' + matched = ([WhenMatched.delete(condition)] if operation == 'delete' else + [WhenMatched.update({'value': 5}, condition)] if operation == 'update' else []) + not_matched = ([WhenNotMatched({'value': 5}, condition)] if operation == 'insert' else []) + with patch.object(merge, '_normalize_source', return_value=source): + table, _, matched, not_matched, ctx = merge._prepare( + 'default.t', 'default.t' if self_merge else object(), catalog_options, + matched, not_matched, ['_ROW_ID'] if self_merge else ['id'], + ) + assert table.options.native_write_enabled() + assert ctx.is_self_merge is self_merge + assert ctx.self_merge_scan_predicate is None + if operation == 'update': + output_schema = _update_schema() + transform = create_merge_batch_transform( + 'update', matched, dict(zip(ctx.source_on_cols, ctx.target_on_cols)), + list(zip(ctx.source_on_cols, ctx.target_on_cols)), ['value'], + '_ROW_ID', output_schema, native_enabled=True, + ) + elif operation == 'delete': + output_schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False)]) + transform = create_merge_batch_transform( + 'delete', matched, {}, '_ROW_ID', output_schema, native_enabled=True, + ) + else: + output_schema = schema + transform = create_merge_batch_transform( + 'insert', not_matched, schema.names, schema, native_enabled=True, + ) + native_protocol.execute.return_value = output_schema.empty_table() + batch = (pa.table({'s.id': pa.array([1], pa.int32()), 's.value': pa.array([10], pa.int32())}) + if operation == 'insert' else _joined_batch()) + assert transform(batch).schema == output_schema + assert native_protocol.calls[0][2][0]['condition'] == {'sql': '"s.value" > 0'} + + +@pytest.mark.parametrize('operation', ['update', 'delete', 'insert']) +@pytest.mark.parametrize('native_requested', [False, True]) +def test_python_conditional_merge_requires_datafusion_before_batch_execution( + monkeypatch, operation, native_requested): + monkeypatch.setitem(sys.modules, 'datafusion', None) + clauses = [_NormalizedClause( + {} if operation == 'delete' else {'value': LiteralValue(5)}, + 's.value > 0', delete=operation == 'delete', + )] + with patch('pypaimon.execution.native.runtime.native_runtime_available', return_value=False): + with pytest.raises(ImportError, match='require DataFusion'): + if operation == 'update': + create_merge_batch_transform( + 'update', clauses, {}, [], ['value'], '_ROW_ID', _update_schema(), + native_enabled=native_requested, + ) + elif operation == 'delete': + create_merge_batch_transform( + 'delete', clauses, {}, '_ROW_ID', + pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False)]), + native_enabled=native_requested, + ) + else: + create_merge_batch_transform( + 'insert', clauses, ['value'], pa.schema([('value', pa.int32())]), + native_enabled=native_requested, + ) + + +def test_incompatible_installed_runtime_is_not_replayed_in_python(native_protocol): + del native_protocol.module.MergeBatchExecution + with patch('pypaimon.execution.python.write.merge_transform._build_matched_transform') as python: + with pytest.raises(ImportError): + _native_update([_NormalizedClause({'value': LiteralValue(5)})]) + python.assert_not_called() + + +@pytest.mark.parametrize('native_enabled', [False, True]) +def test_merge_transform_imports_only_its_selected_backend(native_enabled): + source = r''' +import importlib.abc +import sys +from types import ModuleType +import pyarrow as pa + +native_enabled = NATIVE_ENABLED +blocked = 'pypaimon.execution.python' if native_enabled else 'pypaimon.execution.native' +class RejectOtherBackend(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.startswith(blocked) or (not native_enabled and fullname.startswith('pypaimon_rust')): + raise AssertionError('Loaded disabled backend: ' + fullname) +sys.meta_path.insert(0, RejectOtherBackend()) +if native_enabled: + class MergeBatchExecution: + def __init__(self, input_schema, output_schema, clauses, mode): + self.output_schema = output_schema + def execute(self, batch): + return self.output_schema.empty_table() + module = ModuleType('pypaimon_rust.datafusion') + module.MergeBatchExecution = MergeBatchExecution + sys.modules['pypaimon_rust.datafusion'] = module +from pypaimon.ray.data_evolution_merge_join import _build_matched_transform +from pypaimon.write.merge_types import LiteralValue, _NormalizedClause +schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False), pa.field('value', pa.int32())]) +transform = _build_matched_transform([_NormalizedClause({'value': LiteralValue(5)})], {}, [], + ['value'], '_ROW_ID', schema, native_enabled=native_enabled) +result = transform(pa.table({'t._ROW_ID': pa.array([7], pa.int64()), 't.value': pa.array([10], pa.int32())})) +assert result.schema == schema +assert result.num_rows == (0 if native_enabled else 1) +'''.replace('NATIVE_ENABLED', repr(native_enabled)) + result = subprocess.run([sys.executable, '-c', source], stdout=subprocess.PIPE, stderr=subprocess.PIPE, + universal_newlines=True, timeout=30) + assert result.returncode == 0, result.stderr diff --git a/paimon-python/pypaimon/tests/native_merge_batch_execution_test.py b/paimon-python/pypaimon/tests/native_merge_batch_execution_test.py new file mode 100644 index 000000000000..7711be866713 --- /dev/null +++ b/paimon-python/pypaimon/tests/native_merge_batch_execution_test.py @@ -0,0 +1,120 @@ +# 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. + +"""Real SDK coverage of the joined-batch MERGE boundary used by Ray workers.""" + +import pyarrow as pa +import pytest + +from pypaimon.write.backend_dispatch import create_merge_batch_transform +from pypaimon.write.merge_types import LiteralValue, SourceColumnRef, _NormalizedClause + +pytestmark = pytest.mark.native_plan + + +def test_native_matched_actions_preserve_clause_order_and_callable_view(): + batch = pa.table({ + 't._ROW_ID': pa.chunked_array([[0, 1], [2, 3]], type=pa.int64()), + 't.id': [10, 11, 12, 13], + 't.value': pa.array([90, 91, 92, 93], type=pa.int32()), + 't.other': ['a', 'b', 'c', 'd'], + 's.value': pa.array([None, 5, 15, 25], type=pa.int32()), + }) + calls = [] + + def fallback(rows): + calls.append(rows.to_pydict()) + assert rows.column_names == ['s.key', 's.value', 't._ROW_ID', 't.id', 't.other', 't.value'] + return pa.chunked_array([[str(value)] for value in rows.column('s.key').to_pylist()]) + + clauses = [ + _NormalizedClause({}, condition='s.value > 20', delete=True), + _NormalizedClause({'value': SourceColumnRef('value')}, condition='s.value > 10'), + _NormalizedClause({'other': fallback}), + ] + # Output order can differ from the first appearance of fields in clauses. + schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False), + ('other', pa.string()), ('value', pa.int32())]) + transform = create_merge_batch_transform( + 'update', clauses, {'key': 'id'}, [('key', 'id')], ['other', 'value'], + '_ROW_ID', schema, native_enabled=True) + actual = transform(batch) + assert actual.schema == schema + assert actual.to_pylist() == [ + {'_ROW_ID': 2, 'other': 'c', 'value': 15}, + {'_ROW_ID': 0, 'other': '10', 'value': 90}, + {'_ROW_ID': 1, 'other': '11', 'value': 91}, + ] + assert len(calls) == 1 + assert calls[0]['t._ROW_ID'] == [0, 1] + delete_schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False)]) + deletes = create_merge_batch_transform( + 'delete', clauses, {'key': 'id'}, '_ROW_ID', delete_schema, native_enabled=True) + assert deletes(batch).to_pylist() == [{'_ROW_ID': 3}] + assert len(calls) == 1 # UPDATE assignments are irrelevant to the delete route. + + +def test_native_insert_actions_fill_typed_nulls_and_keep_chunk_order(): + schema = pa.schema([('value', pa.int32()), ('optional', pa.string())]) + clauses = [ + _NormalizedClause({'value': LiteralValue(20)}, condition='s.value >= (SELECT 20)'), + _NormalizedClause({'value': LiteralValue(10)}, condition='s.value >= 10'), + _NormalizedClause({'value': SourceColumnRef('value')}), + ] + batch = pa.table({'s.value': pa.chunked_array([ + pa.array([None, 5], type=pa.int32()), pa.array([15, 25], type=pa.int32()), + ])}) + transform = create_merge_batch_transform( + 'insert', clauses, schema.names, schema, native_enabled=True) + actual = transform(batch) + assert actual.schema == schema + assert actual.to_pylist() == [ + {'value': 20, 'optional': None}, {'value': 10, 'optional': None}, + {'value': None, 'optional': None}, {'value': 5, 'optional': None}, + ] + + +def test_native_worker_transform_is_serializable_after_compilation(): + import cloudpickle + + schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False), ('value', pa.int32())]) + transform = create_merge_batch_transform( + 'update', [_NormalizedClause({'value': SourceColumnRef('key')})], + {'key': 'id'}, [('key', 'id')], ['value'], '_ROW_ID', schema, native_enabled=True) + batch = pa.table({'t._ROW_ID': [7], 't.id': pa.array([42], type=pa.int32())}) + assert transform(batch).to_pylist() == [{'_ROW_ID': 7, 'value': 42}] + restored = cloudpickle.loads(cloudpickle.dumps(transform)) + assert restored(batch).to_pylist() == [{'_ROW_ID': 7, 'value': 42}] + + +def test_native_self_update_callable_keeps_explicit_unqualified_projection(): + calls = [] + + def assignment(rows): + calls.append(rows.column_names) + assert rows.schema == pa.schema([('value', pa.int32()), ('id', pa.int64())]) + return rows.column('value') + + schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False), ('value', pa.int32())]) + transform = create_merge_batch_transform( + 'update', [_NormalizedClause({'value': assignment}, condition='s.id > 1')], + {'id': 'id'}, [('id', 'id')], ['value'], '_ROW_ID', schema, + callable_input_columns=['value', 'id'], native_enabled=True) + batch = pa.table({'t._ROW_ID': [0, 1], 't.id': [1, 2], + 't.value': pa.array([3, 4], type=pa.int32())}) + assert transform(batch).to_pylist() == [{'_ROW_ID': 1, 'value': 4}] + assert calls == [['value', 'id']] diff --git a/paimon-python/pypaimon/write/backend_dispatch.py b/paimon-python/pypaimon/write/backend_dispatch.py index 4460950cf208..5961d946a2fd 100644 --- a/paimon-python/pypaimon/write/backend_dispatch.py +++ b/paimon-python/pypaimon/write/backend_dispatch.py @@ -274,6 +274,23 @@ def prepare_native_merge_into(*args, **kwargs): return create_native_merge_into(*args, **kwargs) +def create_merge_batch_transform(operation, *args, native_enabled=False, **kwargs): + """Select the joined-batch executor before a Ray task consumes input.""" + if native_enabled: + from pypaimon.execution.native.runtime import native_runtime_available + native_enabled = native_runtime_available() + if native_enabled: + from pypaimon.execution.native.write import merge_transform + else: + from pypaimon.execution.python.write import merge_transform + builders = { + 'update': merge_transform._build_matched_transform, + 'delete': merge_transform._build_matched_delete_transform, + 'insert': merge_transform._build_insert_transform, + } + return builders[operation](*args, **kwargs) + + def create_python_row_id_update(table, commit_user, commit_identifier, precomputed=None): """Create an updater after Python execution has already been selected.""" from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId