diff --git a/paimon-python/pypaimon/execution/native/write/merge_transform.py b/paimon-python/pypaimon/execution/native/write/merge_transform.py index f3a5d64178a7..8a6ff9e08c05 100644 --- a/paimon-python/pypaimon/execution/native/write/merge_transform.py +++ b/paimon-python/pypaimon/execution/native/write/merge_transform.py @@ -38,13 +38,14 @@ def evaluate(batch): 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.""" + callable_input_columns=None, hidden_columns=(), not_matched=False): + """Adapt a normalized PyPaimon clause to the typed SDK without running callbacks.""" + from pypaimon_rust.datafusion import MergeAssignment, WhenMatched, WhenNotMatched 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 = [] + values = {} for name, value in clause.spec.items(): if isinstance(value, SourceColumnRef): kind, column = 'source', value.column @@ -52,20 +53,21 @@ def encode_merge_clause(clause, on_pairs=(), on_map=None, input_schema=None, 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)) + values[name] = MergeAssignment.literal(None) continue kind, column = 'target', target - values.append((name, kind, column)) + values[name] = (MergeAssignment.source_column(column) if kind == 'source' else + MergeAssignment.target_column(column)) elif isinstance(value, TargetColumnRef): if available is not None and 't.' + value.column not in available: - values.append((name, 'literal', None)) + values[name] = MergeAssignment.literal(None) else: - values.append((name, 'target', value.column)) + values[name] = MergeAssignment.target_column(value.column) elif isinstance(value, LiteralValue): - values.append((name, 'literal', value.value)) + values[name] = MergeAssignment.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))) + values[name] = MergeAssignment.function(_callable_assignment( + value, on_pairs, callable_input_columns, hidden_columns)) else: return None condition = None @@ -73,8 +75,12 @@ def encode_merge_clause(clause, on_pairs=(), on_map=None, input_schema=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) + condition = rewritten + if not_matched: + return WhenNotMatched(values, condition) + if clause.delete: + return WhenMatched.delete(condition) + return WhenMatched.update(values, condition) class _NativeMergeTransform: @@ -82,7 +88,9 @@ 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 + from pypaimon_rust.datafusion import ( # noqa: F401 + MergeAssignment, MergeBatchExecution, WhenMatched, WhenNotMatched, + ) self.clauses = clauses self.output_schema = output_schema @@ -114,15 +122,19 @@ def __call__(self, batch): clauses = [encode_merge_clause( clause, self.on_pairs, self.on_map, batch.schema, - self.callable_input_columns, hidden_columns, + self.callable_input_columns, hidden_columns, not_matched=self.mode == 'insert', ) 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, - ) + if self.mode == 'update': + self._execution = MergeBatchExecution.for_updates(batch.schema, output_schema, clauses) + elif self.mode == 'delete': + self._execution = MergeBatchExecution.for_deletes(batch.schema, clauses) + else: + self._execution = MergeBatchExecution.for_inserts(batch.schema, output_schema, clauses) result = self._execution.execute(batch) - if hidden_columns: + if hidden_columns or (self.mode == 'delete' and not result.schema.equals( + self.output_schema, check_metadata=True)): import pyarrow as pa result = pa.Table.from_arrays(result.columns, schema=self.output_schema) return result 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 c60540f8b49d..8aa072b54ff9 100644 --- a/paimon-python/pypaimon/execution/native/write/native_merge_into.py +++ b/paimon-python/pypaimon/execution/native/write/native_merge_into.py @@ -89,11 +89,9 @@ def create_native_merge_into(table, source, on, matched, not_matched, commit_use if source_type != target_type or not _supported_upsert_key_type(target_type): return None - def encode(clause): - 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] + on_pairs = list(zip(source_keys, target_keys)) + encoded_matched = [encode_merge_clause(clause, on_pairs) for clause in matched] + encoded_not_matched = [encode_merge_clause(clause, on_pairs, not_matched=True) for clause in not_matched] if any(clause is None for clause in encoded_matched + encoded_not_matched): return None builder = (native_table.new_batch_write_builder().with_commit_user(commit_user) diff --git a/paimon-python/pypaimon/tests/merge_execution_boundary_test.py b/paimon-python/pypaimon/tests/merge_execution_boundary_test.py index 92a524e3ab36..c89ce69cc324 100644 --- a/paimon-python/pypaimon/tests/merge_execution_boundary_test.py +++ b/paimon-python/pypaimon/tests/merge_execution_boundary_test.py @@ -77,6 +77,8 @@ def find_spec(self, fullname, path=None, target=None): assert read.new_read().to_arrow(read.new_scan().plan().splits()).sort_by('id').to_pylist() == [ {'id': 1, 'value': 9}, {'id': 2, 'value': 7}] else: + from pypaimon.tests.merge_sdk_test_support import merge_sdk_module + sys.modules['pypaimon_rust.datafusion'] = merge_sdk_module() core = Mock() builder = core.new_batch_write_builder.return_value builder.with_commit_user.return_value = builder @@ -91,10 +93,9 @@ def assignment(batch): side_effect=lambda table, messages: messages): assert merge_into(table, data, on=['id'], when_matched=matched, commit_user='test') == [] writer.merge_into.assert_called_once() - callback = writer.merge_into.call_args[1]['when_matched'][0]['assignments'][0][2] + callback = writer.merge_into.call_args[1]['when_matched'][0].assignments['value'].value assert callback(pa.table({'t.id': [1], 't.value': [3], 's.value': [9]})).to_pylist() == [9] - assert writer.merge_into.call_args[1]['when_matched'][0]['condition'] == { - 'sql': '\"s.value\" > \"t.value\"'} + assert writer.merge_into.call_args[1]['when_matched'][0].condition == '\"s.value\" > \"t.value\"' assert not any(name.startswith('pypaimon.ray') or name == blocked or name.startswith(blocked + '.') for name in sys.modules) """ diff --git a/paimon-python/pypaimon/tests/merge_execution_protocol_test.py b/paimon-python/pypaimon/tests/merge_execution_protocol_test.py index 083476ec4555..61a694080298 100644 --- a/paimon-python/pypaimon/tests/merge_execution_protocol_test.py +++ b/paimon-python/pypaimon/tests/merge_execution_protocol_test.py @@ -26,6 +26,7 @@ import pyarrow as pa import pytest +from pypaimon.tests.merge_sdk_test_support import MergeAssignment, WhenMatched, WhenNotMatched, merge_sdk_module from pypaimon.write.backend_dispatch import create_merge_batch_transform from pypaimon.write.merge_types import ( LiteralValue, SourceColumnRef, TargetColumnRef, _NormalizedClause, @@ -38,15 +39,32 @@ def native_protocol(monkeypatch): execute = Mock() class MergeBatchExecution: - def __init__(self, input_schema, output_schema, clauses, mode): - calls.append((input_schema, output_schema, clauses, mode)) + @classmethod + def for_updates(cls, input_schema, update_schema, when_matched): + assert all(isinstance(clause, WhenMatched) for clause in when_matched) + return cls((input_schema, update_schema, when_matched, 'update')) + + @classmethod + def for_deletes(cls, input_schema, when_matched): + assert all(isinstance(clause, WhenMatched) for clause in when_matched) + schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False)]) + return cls((input_schema, schema, when_matched, 'delete')) + + @classmethod + def for_inserts(cls, input_schema, target_schema, when_not_matched): + assert all(isinstance(clause, WhenNotMatched) for clause in when_not_matched) + return cls((input_schema, target_schema, when_not_matched, 'insert')) + + def __init__(self, compiled): + self.compiled = compiled + calls.append(compiled) def execute(self, batch): - return execute(batch, calls[-1]) + return execute(batch, self.compiled) runtime = ModuleType('pypaimon_rust') runtime.__path__ = [] - datafusion = ModuleType('pypaimon_rust.datafusion') + datafusion = merge_sdk_module() datafusion.MergeBatchExecution = MergeBatchExecution monkeypatch.setitem(sys.modules, 'pypaimon_rust', runtime) monkeypatch.setitem(sys.modules, 'pypaimon_rust.datafusion', datafusion) @@ -101,9 +119,10 @@ def execute(rows, 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 clauses[0].condition == '"t.id" > 0' + assignment = clauses[0].assignments['value'] + assert assignment.kind == 'function' + callback = assignment.value assert callback(rows).equals(result['value']) return result @@ -130,7 +149,7 @@ def values(view): batch = _joined_batch() def execute(rows, compiled): - callback = compiled[2][0]['assignments'][0][2] + callback = compiled[2][0].assignments['value'].value return pa.Table.from_arrays([rows['t._ROW_ID'], callback(rows)], schema=_update_schema()) native_protocol.execute.side_effect = execute @@ -162,11 +181,11 @@ def values(view): 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[2][0].condition == '"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)) + columns.append(compiled[2][0].assignments['value'].value(rows)) return pa.Table.from_arrays(columns, schema=compiled[1]) native_protocol.execute.side_effect = execute @@ -189,11 +208,87 @@ def test_native_transports_coalesced_and_missing_column_assignments(native_proto 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), - ] + assert native_protocol.calls[0][2][0].assignments == { + 'key': MergeAssignment.target_column('id'), + 'source': MergeAssignment.source_column('value'), + 'target': MergeAssignment.target_column('value'), + 'missing_source': MergeAssignment.literal(None), + 'missing_target': MergeAssignment.literal(None), + 'constant': MergeAssignment.literal(9), + } + + +@pytest.mark.parametrize('stream', [False, True], ids=['batch', 'stream']) +def test_native_full_merge_uses_typed_branches_without_evaluating_callbacks(tmp_path, native_protocol, stream): + from pypaimon import CatalogFactory, Schema + from pypaimon.write.merge_types import WhenMatched as PublicWhenMatched, WhenNotMatched as PublicWhenNotMatched + from pypaimon.write.merge_types import lit, source_col, target_col + + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', True) + schema = pa.schema([('id', pa.int32()), ('value', pa.int32()), ('other', pa.int32())]) + catalog.create_table('default.t', Schema.from_pyarrow_schema(schema, options={ + 'data-evolution.enabled': 'true', 'row-tracking.enabled': 'true', 'write.native.enabled': 'true', + 'deletion-vectors.enabled': 'true', + }), False) + table = catalog.get_table('default.t') + native_table = Mock() + builder = (native_table.new_stream_write_builder.return_value if stream else + native_table.new_batch_write_builder.return_value) + builder.with_commit_user.return_value = builder + writer = builder.new_update.return_value + writer.merge_into.return_value = [] + callback = Mock(return_value=pa.array([11], pa.int32())) + source = pa.table({'source_id': pa.array([1], pa.int32()), 'value': pa.array([11], pa.int32())}) + public_builder = table.new_stream_write_builder() if stream else table.new_batch_write_builder() + with patch('pypaimon.execution.native.write.native_merge_into._native_row_id_table', return_value=native_table), \ + patch('pypaimon.execution.native.write.native_merge_into.from_native_commit_messages', + side_effect=lambda table, messages: messages): + assert public_builder.new_update().merge_into( + source, on={'id': 'source_id'}, when_matched=[ + PublicWhenMatched.delete('s.value < 0'), + PublicWhenMatched.update({'value': callback}, 's.source_id > 0'), + PublicWhenMatched.update({'value': source_col('value'), 'other': target_col('value')}), + ], when_not_matched=[PublicWhenNotMatched( + {'id': source_col('source_id'), 'value': lit(5)}, 's.value > 0')], + **({'commit_identifier': 31} if stream else {})) == [] + callback.assert_not_called() + args, kwargs = writer.merge_into.call_args + assert args[0] is source + assert kwargs['on'] == [('id', 'source_id')] + if stream: + assert kwargs['commit_identifier'] == 31 + else: + assert 'commit_identifier' not in kwargs + matched = kwargs['when_matched'] + assert all(isinstance(clause, WhenMatched) for clause in matched) + assert [clause.action for clause in matched] == ['delete', 'update', 'update'] + assert [clause.condition for clause in matched] == ['"s.value" < 0', '"s.source_id" > 0', None] + assert matched[2].assignments == { + 'value': MergeAssignment.source_column('value'), 'other': MergeAssignment.target_column('value'), + } + inserted = kwargs['when_not_matched'] + assert all(isinstance(clause, WhenNotMatched) for clause in inserted) + assert inserted[0].condition == '"s.value" > 0' + assert inserted[0].assignments == { + 'id': MergeAssignment.source_column('source_id'), 'value': MergeAssignment.literal(5), + } + joined = pa.table({'t.id': [1], 't.value': [3], 't.other': [4], 's.value': [11]}) + assert matched[1].assignments['value'].value(joined).to_pylist() == [11] + assert callback.call_args[0][0].column_names == ['s.source_id', 's.value', 't.id', 't.other', 't.value'] + + +def test_native_delete_restores_requested_field_and_schema_metadata(native_protocol): + schema = pa.schema([pa.field('_ROW_ID', pa.int64(), nullable=False, metadata={b'field': b'value'})], + metadata={b'schema': b'value'}) + transform = create_merge_batch_transform( + 'delete', [_NormalizedClause({}, delete=True)], {}, '_ROW_ID', schema, native_enabled=True, + ) + native_protocol.execute.side_effect = lambda batch, compiled: pa.Table.from_arrays( + [batch['t._ROW_ID']], schema=compiled[1]) + result = transform(_joined_batch()) + assert result.schema.equals(schema, check_metadata=True) + assert result['_ROW_ID'].to_pylist() == [7, 8] def test_native_insert_callback_retains_source_aliases_and_dictionary_types(native_protocol): @@ -216,7 +311,7 @@ def values(view): def execute(rows, compiled): assert compiled[3] == 'insert' - return pa.Table.from_arrays([compiled[2][0]['assignments'][0][2](rows)], schema=schema) + return pa.Table.from_arrays([compiled[2][0].assignments['value'].value(rows)], schema=schema) native_protocol.execute.side_effect = execute assert transform(batch)['value'].to_pylist() == [10, 20] @@ -263,11 +358,11 @@ def callback(rows): raise failure if phase == 'compile': - native_protocol.module.MergeBatchExecution = Mock(side_effect=failure) + native_protocol.module.MergeBatchExecution.for_updates = 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) + native_protocol.execute.side_effect = lambda rows, compiled: compiled[2][0].assignments['value'].value(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: @@ -368,7 +463,7 @@ def test_native_ray_preparation_does_not_require_python_datafusion( 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'} + assert native_protocol.calls[0][2][0].condition == '"s.value" > 0' @pytest.mark.parametrize('operation', ['update', 'delete', 'insert']) @@ -425,11 +520,15 @@ def find_spec(self, fullname, path=None, target=None): sys.meta_path.insert(0, RejectOtherBackend()) if native_enabled: class MergeBatchExecution: - def __init__(self, input_schema, output_schema, clauses, mode): + @classmethod + def for_updates(cls, input_schema, update_schema, when_matched): + return cls(update_schema) + def __init__(self, output_schema): self.output_schema = output_schema def execute(self, batch): return self.output_schema.empty_table() - module = ModuleType('pypaimon_rust.datafusion') + from pypaimon.tests.merge_sdk_test_support import merge_sdk_module + module = merge_sdk_module() module.MergeBatchExecution = MergeBatchExecution sys.modules['pypaimon_rust.datafusion'] = module from pypaimon.ray.data_evolution_merge_join import _build_matched_transform diff --git a/paimon-python/pypaimon/tests/merge_sdk_test_support.py b/paimon-python/pypaimon/tests/merge_sdk_test_support.py new file mode 100644 index 000000000000..db2cef51e260 --- /dev/null +++ b/paimon-python/pypaimon/tests/merge_sdk_test_support.py @@ -0,0 +1,76 @@ +# 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. + +"""Typed companion SDK doubles for transport tests that do not load Rust.""" + +from dataclasses import dataclass +from types import ModuleType + + +@dataclass(frozen=True) +class MergeAssignment: + kind: str + value: object + + @classmethod + def source_column(cls, name): + return cls('source', name) + + @classmethod + def target_column(cls, name): + return cls('target', name) + + @classmethod + def literal(cls, value): + return cls('literal', value) + + @classmethod + def function(cls, function): + return cls('function', function) + + +@dataclass(frozen=True) +class WhenMatched: + assignments: dict + condition: object + action: str + + @classmethod + def update(cls, assignments, condition=None): + assert all(isinstance(value, MergeAssignment) for value in assignments.values()) + return cls(dict(assignments), condition, 'update') + + @classmethod + def delete(cls, condition=None): + return cls({}, condition, 'delete') + + +@dataclass(frozen=True) +class WhenNotMatched: + assignments: dict + condition: object = None + + def __post_init__(self): + assert all(isinstance(value, MergeAssignment) for value in self.assignments.values()) + + +def merge_sdk_module(): + module = ModuleType('pypaimon_rust.datafusion') + module.MergeAssignment = MergeAssignment + module.WhenMatched = WhenMatched + module.WhenNotMatched = WhenNotMatched + return module diff --git a/paimon-python/pypaimon/tests/native_merge_source_metadata_test.py b/paimon-python/pypaimon/tests/native_merge_source_metadata_test.py index 9a5251026db1..03c7b29b7134 100644 --- a/paimon-python/pypaimon/tests/native_merge_source_metadata_test.py +++ b/paimon-python/pypaimon/tests/native_merge_source_metadata_test.py @@ -19,6 +19,7 @@ from contextlib import ExitStack import importlib.util +import sys from unittest.mock import Mock, patch import pyarrow as pa @@ -27,6 +28,7 @@ from pypaimon import Schema from pypaimon.execution.native.runtime import (native_runtime_available) +from pypaimon.tests.merge_sdk_test_support import merge_sdk_module from pypaimon.table.data_evolution_merge_into import WhenMatched, WhenNotMatched, source_col @@ -128,7 +130,8 @@ def assignment(batch): @pytest.mark.parametrize('table_source', [False, True], ids=['arrow-source', 'table-source']) -def test_native_condition_transport_does_not_load_python_datafusion(native_rest_catalog, table_source): +def test_native_condition_transport_does_not_load_python_datafusion(native_rest_catalog, monkeypatch, table_source): + monkeypatch.setitem(sys.modules, 'pypaimon_rust.datafusion', merge_sdk_module()) target = _new_table(native_rest_catalog, 'default.no_python_engine').copy({'write.native.enabled': 'true'}) source = (_new_table(native_rest_catalog, 'default.metadata_source') if table_source else pa.table({'key': [0], 'value': [91]}, schema=_SCHEMA)) @@ -150,7 +153,7 @@ def test_native_condition_transport_does_not_load_python_datafusion(native_rest_ core.merge_into.assert_called_once() args, kwargs = core.merge_into.call_args assert args[0] is (native_source if table_source else source) - assert kwargs['when_matched'][0]['condition'] == {'sql': '"s.{}" = 0'.format(source_key)} + assert kwargs['when_matched'][0].condition == '"s.{}" = 0'.format(source_key) def test_python_merge_conditions_still_require_the_python_execution_engine(native_rest_catalog):