[Python] Execute Ray MERGE actions through native core

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')

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]

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:

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