This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new a65238d017 [Python] Execute Ray MERGE actions through native core
(#10761)
a65238d017 is described below
commit a65238d017ee11cadf81340c0fef63b094ca2deb
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Oct 23 09:26:35 2026 +0800
[Python] Execute Ray MERGE actions through native core (#10761)
---
.../execution/native/write/merge_transform.py | 148 +++++++
.../execution/native/write/native_merge_into.py | 28 +-
.../execution/python/write/merge_transform.py | 4 +
.../pypaimon/ray/data_evolution_merge_into.py | 11 +-
.../pypaimon/ray/data_evolution_merge_join.py | 36 +-
.../tests/merge_execution_protocol_test.py | 446 +++++++++++++++++++++
.../tests/native_merge_batch_execution_test.py | 120 ++++++
paimon-python/pypaimon/write/backend_dispatch.py | 17 +
8 files changed, 774 insertions(+), 36 deletions(-)
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 0000000000..f3a5d64178
--- /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 cd92aca1e3..c60540f8b4 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 @@ from pypaimon.execution.native.write.native_update import (
_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 4418a8b914..d009eccc1b 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 @@ from typing import Dict, List, Optional, Sequence, Tuple
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 9209a3f5a8..77be7bd9d9 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 6a9ecd1df4..edbb845d39 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 build_self_merge_delete_ds(
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 0000000000..083476ec45
--- /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,
+)
+
+
[email protected]
+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']
+
+
[email protected]('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
+
+
[email protected]('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
+
+
[email protected]('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()
+
+
[email protected]('native_enabled', [False, True])
[email protected]('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()
+
+
[email protected]('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'}
+
+
[email protected]('operation', ['update', 'delete', 'insert'])
[email protected]('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()
+
+
[email protected]('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 0000000000..7711be8667
--- /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 4460950cf2..5961d946a2 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
