From 9b34827f9175fe088fe7e84827caebcefe7e6cb9 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 22 Oct 2026 23:16:33 +0800 Subject: [PATCH 1/2] [Python] Execute captured Ray mutation and index plans through native core --- .../pypaimon/execution/native/write/commit.py | 24 +- .../execution/native/write/native_commit.py | 34 ++ .../native/write/native_index_build.py | 17 +- .../execution/native/write/native_update.py | 28 +- .../execution/native/write/row_id_plan.py | 66 ++++ .../python/write/commit/conflict_detection.py | 28 +- .../python/write/commit/row_id_conflict.py | 55 +++ .../execution/python/write/row_id_plan.py | 42 +++ .../pypaimon/globalindex/index_dispatch.py | 6 +- .../pypaimon/ray/data_evolution_merge_join.py | 69 ++-- .../pypaimon/ray/row_id_conflict_rewriter.py | 19 +- .../pypaimon/ray/vector_index_build.py | 44 ++- .../tests/native_ray_index_build_test.py | 202 +++++++++++ .../tests/native_ray_row_id_plan_test.py | 112 ++++++ .../ray_data_evolution_merge_into_test.py | 1 + .../ray_row_id_conflict_rewriter_test.py | 2 +- .../write/row_id_conflict_protocol_test.py | 326 ++++++++++++++++++ .../write/row_id_execution_boundary_test.py | 6 +- .../pypaimon/tests/write/row_id_plan_test.py | 260 ++++++++++++++ .../pypaimon/write/backend_dispatch.py | 31 +- .../pypaimon/write/row_id_conflict.py | 132 +++++++ paimon-python/pypaimon/write/row_id_plan.py | 112 ++++++ .../pypaimon/write/update_dispatch.py | 12 +- 23 files changed, 1520 insertions(+), 108 deletions(-) create mode 100644 paimon-python/pypaimon/execution/native/write/row_id_plan.py create mode 100644 paimon-python/pypaimon/execution/python/write/commit/row_id_conflict.py create mode 100644 paimon-python/pypaimon/execution/python/write/row_id_plan.py create mode 100644 paimon-python/pypaimon/tests/native_ray_index_build_test.py create mode 100644 paimon-python/pypaimon/tests/native_ray_row_id_plan_test.py create mode 100644 paimon-python/pypaimon/tests/write/row_id_conflict_protocol_test.py create mode 100644 paimon-python/pypaimon/tests/write/row_id_plan_test.py create mode 100644 paimon-python/pypaimon/write/row_id_conflict.py create mode 100644 paimon-python/pypaimon/write/row_id_plan.py diff --git a/paimon-python/pypaimon/execution/native/write/commit.py b/paimon-python/pypaimon/execution/native/write/commit.py index 6f6f4ed7ac1a..71629f75fd90 100644 --- a/paimon-python/pypaimon/execution/native/write/commit.py +++ b/paimon-python/pypaimon/execution/native/write/commit.py @@ -26,9 +26,27 @@ def __init__(self, native_commit, is_batch=False): def commit(self, messages, commit_identifier, snapshot_properties=None): kwargs = ({} if snapshot_properties is None else {'snapshot_properties': snapshot_properties}) - if self._is_batch: - return self._native_commit.commit(messages, **kwargs) - return self._native_commit.commit(commit_identifier, messages, **kwargs) + try: + if self._is_batch: + return self._native_commit.commit(messages, **kwargs) + return self._native_commit.commit(commit_identifier, messages, **kwargs) + except Exception as error: + self._raise_commit_error(error) + + def check_row_id_conflicts(self, messages, latest_snapshot_id, check_from_snapshot): + try: + return self._native_commit.check_row_id_conflicts(messages, latest_snapshot_id, check_from_snapshot) + except Exception as error: + self._raise_commit_error(error) + + @staticmethod + def _raise_commit_error(error): + from pypaimon.execution.native.write.native_commit import native_commit_error + + converted = native_commit_error(error) + if converted is error: + raise error + raise converted from error def abort(self, messages): return self._native_commit.abort(messages) diff --git a/paimon-python/pypaimon/execution/native/write/native_commit.py b/paimon-python/pypaimon/execution/native/write/native_commit.py index 5e4e77d52937..adc819e1aa43 100644 --- a/paimon-python/pypaimon/execution/native/write/native_commit.py +++ b/paimon-python/pypaimon/execution/native/write/native_commit.py @@ -202,6 +202,40 @@ def from_native_commit_messages(table, messages): return decoded +def native_commit_error(error): + """Keep retry classification structured; never infer it from error text.""" + from pypaimon_rust import RowIdExistenceConflict as NativeRowIdConflict + from pypaimon.write.row_id_conflict import RowIdExistenceConflict + + if not isinstance(error, NativeRowIdConflict): + return error + return RowIdExistenceConflict._from_fields( + file_name=error.file_name, first_row_id=error.first_row_id, + row_count=error.row_count, bucket=error.bucket, + current_first_row_id=error.current_first_row_id, + current_row_count=error.current_row_count) + + +def prepare_native_row_id_conflict_check(table, messages): + """Prepare a check without creating Python publication or conflict state.""" + from pypaimon.execution.native.write.commit import NativeTableCommit + + if not native_messages_supported(table, messages): + return None + native = create_native_commit(table, 'row-id-conflict-check', batch=True) + if native is None: + return None + try: + converted = to_native_commit_messages(table, messages) + except BaseException: + try: + native.close() + except BaseException: + pass + raise + return NativeTableCommit(native, is_batch=True), converted + + def native_commit_callback(table, callbacks): """Translate the native notification into the existing Python context.""" from pypaimon.manifest.schema.manifest_entry import ManifestEntry diff --git a/paimon-python/pypaimon/execution/native/write/native_index_build.py b/paimon-python/pypaimon/execution/native/write/native_index_build.py index 4eeeb63c2ecc..c25491a81267 100644 --- a/paimon-python/pypaimon/execution/native/write/native_index_build.py +++ b/paimon-python/pypaimon/execution/native/write/native_index_build.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -"""Local global-index builds delegated to Rust, retaining Python's commit API.""" +"""Global-index builds delegated to Rust, retaining Python's commit API.""" from pypaimon.execution.native.table import (_native_table, _option_value_to_string) from pypaimon.execution.native.read.builder import (_predicate_to_native) @@ -23,7 +23,7 @@ from pypaimon.execution.native.write.native_commit import _rest_catalog_supported, from_native_commit_messages -def build_native_global_index(builder, partition_filter): +def build_native_global_index(builder, partition_filter, *, execution="local", concurrency=None, ray_remote_args=None): """Return unpublished messages, or None before building for an ineligible table. Rust owns snapshot selection, range planning and index file generation. @@ -34,7 +34,6 @@ def build_native_global_index(builder, partition_filter): if (builder._index_type not in ('btree', 'bitmap', 'full-text') + tuple(VINDEX_IDENTIFIERS) or not table.options.native_write_enabled() or (table.options.deletion_vectors_enabled() and builder._index_type == 'full-text') - or table.options.query_auth_enabled or table.is_primary_key_table or table.options.file_format() != 'parquet' or not native_runtime_available() @@ -77,4 +76,16 @@ def build_native_global_index(builder, partition_filter): native.with_options(options) if partition_filter is not None: native.with_partition_filter(_predicate_to_native(partition_filter)) + if execution == "ray": + from pypaimon.ray.vector_index_build import build_native_vector_index + # Planning may authorize and capture inputs; its errors must not select + # another engine or restart a build that may already have side effects. + plan = native.plan() + return build_native_vector_index(table, plan, concurrency, ray_remote_args) return from_native_commit_messages(table, native.build()) + + +def build_native_index_shard(table, shard): + """Execute captured metadata in Core; workers never replan or publish.""" + native = _native_table(table).new_global_index_build_builder() + return from_native_commit_messages(table, native.build_shard(shard)) diff --git a/paimon-python/pypaimon/execution/native/write/native_update.py b/paimon-python/pypaimon/execution/native/write/native_update.py index f7e58eab7e46..7844079a1f41 100644 --- a/paimon-python/pypaimon/execution/native/write/native_update.py +++ b/paimon-python/pypaimon/execution/native/write/native_update.py @@ -84,21 +84,25 @@ def create_native_update(table, commit_user, columns): return NativeBatchTableUpdate(table, writer) -def create_native_update_by_row_id(table, commit_user, commit_identifier): +def create_native_update_by_row_id(table, commit_user, commit_identifier, plan=None): """Create a core updater sharing one snapshot across incremental calls.""" # Columns are selected per call; one core updater pins the snapshot and # tracks overlaps across Parquet and dedicated Blob updates. if not _native_update_columns_supported(table, None): return None - native_table = _native_row_id_table(table) + native_table = _native_row_id_table( + table._copy_with_snapshot(plan.table._read_snapshot) if plan is not None else table) if native_table is None: return None if commit_identifier == BATCH_COMMIT_IDENTIFIER: - writer = (native_table.new_batch_write_builder().with_commit_user(commit_user) - .new_update().new_update_by_row_id()) + update = native_table.new_batch_write_builder().with_commit_user(commit_user).new_update() else: - writer = (native_table.new_stream_write_builder().with_commit_user(commit_user) - .new_update().new_update_by_row_id(commit_identifier)) + update = native_table.new_stream_write_builder().with_commit_user(commit_user).new_update() + if plan is not None: + from pypaimon.execution.native.write.row_id_plan import native_plan + update.with_plan(native_plan(plan.table, plan.plan)) + writer = (update.new_update_by_row_id() if commit_identifier == BATCH_COMMIT_IDENTIFIER + else update.new_update_by_row_id(commit_identifier)) return NativeTableUpdateByRowId( table, commit_user, commit_identifier, writer) @@ -438,22 +442,26 @@ def canonical_paths(node): return NativePredicateTableUpdate(table, writer, native_predicate) -def create_native_delete(table, commit_user): +def create_native_delete(table, commit_user, plan=None): """Select Rust's deletion-vector writer for supported batch deletes.""" if not table.options.deletion_vectors_enabled(False): return None - native_table = _native_row_id_table(table) + native_table = _native_row_id_table( + table._copy_with_snapshot(plan.table._read_snapshot) if plan is not None else table) if native_table is None: return None writer = (native_table.new_batch_write_builder() .with_commit_user(commit_user) .new_update()) + if plan is not None: + from pypaimon.execution.native.write.row_id_plan import native_plan + writer.with_plan(native_plan(plan.table, plan.plan)) return NativeBatchTableUpdate(table, writer) -def create_native_row_id_delete(table): +def create_native_row_id_delete(table, plan=None): """Adapt the direct deletion API to the existing core table updater.""" - operation = create_native_delete(table, table.new_batch_write_builder().commit_user) + operation = create_native_delete(table, table.new_batch_write_builder().commit_user, plan) return None if operation is None else NativeTableDeleteByRowId(operation) diff --git a/paimon-python/pypaimon/execution/native/write/row_id_plan.py b/paimon-python/pypaimon/execution/native/write/row_id_plan.py new file mode 100644 index 000000000000..d040b5b2059d --- /dev/null +++ b/paimon-python/pypaimon/execution/native/write/row_id_plan.py @@ -0,0 +1,66 @@ +# 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 split transport for Core's captured row-ID plans.""" + +from pypaimon.execution.native.runtime import native_runtime_available + + +def _with_bucket_path(table, split): + from copy import copy + from pypaimon.globalindex.indexed_split import IndexedSplit + from pypaimon.read.query_auth_split import QueryAuthSplit + from pypaimon.read.split import DataSplit + + if isinstance(split, QueryAuthSplit): + inner = _with_bucket_path(table, split.split) + return split if inner is split.split else QueryAuthSplit(inner, split.auth_result) + if isinstance(split, IndexedSplit): + inner = _with_bucket_path(table, split.data_split()) + return split if inner is split.data_split() else IndexedSplit(inner, split.row_ranges(), split.scores()) + if not isinstance(split, DataSplit) or split.bucket_path: + return split + # Java keeps the table's canonical bucket path even when individual files + # live elsewhere. Serialization retains each exact external/file path. + captured = copy(split) + captured.bucket_path = table.path_factory().bucket_path( + tuple(split.partition.values), split.bucket, canonical_partition=True) + return captured + + +def native_plan(table, plan): + from pypaimon_rust.datafusion import Plan, Split + from pypaimon.read.split_serializer import serialize_split_v1 + return Plan(plan.snapshot_id, [Split.deserialize(serialize_split_v1(_with_bucket_path(table, split))) + for split in plan.splits()]) + + +def create_native_row_id_plan(table, plan): + if not (table.options.native_write_enabled() or table.options.native_read_enabled() + or table.options.native_plan_enabled()) or not native_runtime_available(): + return None + from pypaimon_rust.datafusion import RowIdPlan + return RowIdPlan(native_plan(table, plan)) + + +def python_groups(table, plan): + from pypaimon.execution.native.table import _partition_fields + from pypaimon.read.split_serializer import deserialize_split_v1 + return [(start, end, deserialize_split_v1(split.serialize(), _partition_fields(table), + table.trimmed_primary_keys_fields)) + for start, end, split in plan.groups()] diff --git a/paimon-python/pypaimon/execution/python/write/commit/conflict_detection.py b/paimon-python/pypaimon/execution/python/write/commit/conflict_detection.py index 4c18e168686f..3f6c3e5553b5 100644 --- a/paimon-python/pypaimon/execution/python/write/commit/conflict_detection.py +++ b/paimon-python/pypaimon/execution/python/write/commit/conflict_detection.py @@ -28,6 +28,7 @@ from pypaimon.table.special_fields import SpecialFields from pypaimon.utils.range import Range from pypaimon.utils.range_helper import RangeHelper +from pypaimon.write.row_id_conflict import RowIdExistenceConflict from pypaimon.execution.python.write.commit.commit_scanner import CommitScanner @@ -165,23 +166,6 @@ def __init__(self, range_, field_ids): self.field_ids = field_ids -class RowIdExistenceConflict(RuntimeError): - """A staged row-id file no longer matches the current base-file layout.""" - - def __init__(self, entry): - self.entry = entry - super().__init__( - "Row ID existence conflict: file '{}' references " - "firstRowId={}, rowCount={} in bucket {}, " - "but no matching file exists in the current snapshot. " - "The referenced file may have been rewritten by a " - "concurrent compaction or removed by an overwrite.".format( - entry.file.file_name, - entry.file.first_row_id, - entry.file.row_count, - entry.bucket)) - - class ConflictDetection: """Detects conflicts between base and delta files during commit.""" @@ -489,6 +473,16 @@ def check_row_id_existence(self, base_entries, delta_entries, next_row_id=None): if not self.data_evolution_enabled: return None + # File identity excludes row IDs. A DELETE must still describe the + # current ADD before identifier-based manifest merging cancels it. + deleted = {entry.identifier(): entry for entry in delta_entries if entry.kind == 1} + for current in base_entries: + stale = deleted.get(current.identifier()) if current.kind == 0 else None + if stale is not None and ( + current.file.first_row_id != stale.file.first_row_id + or current.file.row_count != stale.file.row_count): + return RowIdExistenceConflict._from_delete(stale, current) + if next_row_id is None: return None diff --git a/paimon-python/pypaimon/execution/python/write/commit/row_id_conflict.py b/paimon-python/pypaimon/execution/python/write/commit/row_id_conflict.py new file mode 100644 index 000000000000..773fd2235467 --- /dev/null +++ b/paimon-python/pypaimon/execution/python/write/commit/row_id_conflict.py @@ -0,0 +1,55 @@ +# 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. + +"""Read-only logical conflict checks for the selected Python backend.""" + +from pypaimon.manifest.manifest_list_manager import ManifestListManager +from pypaimon.manifest.schema.manifest_entry import ManifestEntry +from pypaimon.execution.python.write.commit.commit_scanner import CommitScanner +from pypaimon.execution.python.write.commit.conflict_detection import ConflictDetection +from pypaimon.table.row.generic_row import GenericRow + + +def check_row_id_conflicts(table, messages, latest_snapshot_id, check_from_snapshot): + if not table.options.data_evolution_enabled(): + raise NotImplementedError('Row-ID conflict checks require a data-evolution table') + if any(message.compact_before or message.compact_after or message.compact_changelog_files + or message.compact_index_adds or message.compact_index_deletes for message in messages): + raise NotImplementedError('Row-ID conflict checks do not accept compact increments') + manager = table.snapshot_manager() + latest = manager.get_snapshot_by_id(latest_snapshot_id) + baseline = manager.get_snapshot_by_id(check_from_snapshot) + if latest is None or baseline is None: + raise ValueError('Row-ID conflict check snapshots must exist in the current table') + current_schema = table.schema_manager.latest() + if (baseline.schema_id != latest.schema_id + or current_schema is not None and current_schema.id != table.table_schema.id): + raise ValueError('Row-ID conflict checks require an unchanged, current table schema') + entries = [] + for message in messages: + partition = GenericRow(list(message.partition), table.partition_keys_fields) + for kind, files in ((0, message.new_files), (1, message.deleted_files)): + entries.extend(ManifestEntry( + kind, partition, message.bucket, + table.options.bucket() if message.total_buckets is None else message.total_buckets, file) + for file in files) + manifests = ManifestListManager(table) + detection = ConflictDetection(True, manager, manifests, table, CommitScanner(table, manifests)) + detection._row_id_check_from_snapshot = check_from_snapshot + conflict = detection.check_row_id_from_snapshot(latest, entries, check_compaction=False) + if conflict is not None: + raise conflict diff --git a/paimon-python/pypaimon/execution/python/write/row_id_plan.py b/paimon-python/pypaimon/execution/python/write/row_id_plan.py new file mode 100644 index 000000000000..72f486afafdc --- /dev/null +++ b/paimon-python/pypaimon/execution/python/write/row_id_plan.py @@ -0,0 +1,42 @@ +# 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. + + +"""Python grouping for distributed row-ID plans, selected before execution.""" + +from pypaimon.execution.python.write.table_update import PythonTableUpdate +from pypaimon.manifest.schema.data_file_meta import DataFileMeta +from pypaimon.read.query_auth_split import QueryAuthSplit + + +def _file_groups(split): + if isinstance(split, QueryAuthSplit): + for group in _file_groups(split.split): + yield QueryAuthSplit(group, split.auth_result) + else: + yield from PythonTableUpdate._predicate_update_file_groups([split]) + + +def row_id_groups(plan): + groups = [] + for planned in plan.splits(): + for split in _file_groups(planned): + ranges = [file.row_id_range() for file in split.files + if file.first_row_id is not None and not DataFileMeta.is_blob_file(file.file_name)] + if ranges: + groups.append((min(r.from_ for r in ranges), max(r.to for r in ranges), split)) + return sorted(groups, key=lambda group: group[0]) diff --git a/paimon-python/pypaimon/globalindex/index_dispatch.py b/paimon-python/pypaimon/globalindex/index_dispatch.py index 97042990963a..39c18df03b96 100644 --- a/paimon-python/pypaimon/globalindex/index_dispatch.py +++ b/paimon-python/pypaimon/globalindex/index_dispatch.py @@ -83,9 +83,11 @@ def close(self): def build_global_index(request, partition_filter, *, execution="local", concurrency=None, ray_remote_args=None): - if execution == "local" and request._table.options.native_write_enabled(): + if request._table.options.native_write_enabled(): from pypaimon.execution.native.write.native_index_build import build_native_global_index - messages = build_native_global_index(request, partition_filter) + messages = build_native_global_index( + request, partition_filter, execution=execution, concurrency=concurrency, + ray_remote_args=ray_remote_args) if messages is not None: return messages from pypaimon.execution.python.globalindex.create_global_index import PythonGlobalIndexBuilder diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index 433f35a19f78..1692a34c564f 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -81,6 +81,7 @@ class _SelfMergeUpdatePlan: row_id_name: str snapshot_id: int callable_input_columns: Optional[List[str]] + row_id_plan: Any @dataclass(frozen=True) @@ -97,6 +98,7 @@ class _SelfMergeUpdateContext: row_id_name: str snapshot_id: int callable_input_columns: Optional[List[str]] + row_id_plan: Any def _resolve_source_projection( @@ -162,7 +164,7 @@ def build_self_merge_update_plan( CoreOptions, GlobalIndexSearchMode, ) from pypaimon.table.special_fields import SpecialFields - from pypaimon.execution.python.write.table_update import PythonTableUpdate + from pypaimon.write.row_id_plan import capture_row_id_plan row_id_name = SpecialFields.ROW_ID.name needed_cols = set(resolve_target_projection( @@ -197,6 +199,8 @@ def build_self_merge_update_plan( table.copy_without_time_travel(dynamic_options) if dynamic_options else table ) + from pypaimon.snapshot.time_travel_util import TimeTravelUtil + scan_table = scan_table._copy_with_snapshot(TimeTravelUtil.resolve_snapshot(scan_table)) read_builder = scan_table.new_read_builder().with_projection(projection) if scan_predicate is not None: read_builder.with_filter(scan_predicate) @@ -207,9 +211,8 @@ def build_self_merge_update_plan( ) # A packed scan split may contain multiple logical row-id groups. Keep # each group intact, but do not materialize unrelated groups together. - file_groups = list(PythonTableUpdate._predicate_update_file_groups( - scan_plan.splits() - )) + row_id_plan = capture_row_id_plan(scan_table, scan_plan) + file_groups = [split for _, _, split in row_id_plan.groups] update_schema = build_update_schema(target_pa_schema, update_cols, row_id_name) return _SelfMergeUpdatePlan( table=table, @@ -226,6 +229,7 @@ def build_self_merge_update_plan( list(read_columns) + [row_id_name] if read_columns else None ), + row_id_plan=row_id_plan, ) @@ -244,7 +248,6 @@ def _self_merge_aliases(batch: pa.Table, row_id_name: str) -> pa.Table: def _apply_self_merge_update_group(context, file_group, collect_row_ids): """Read, transform, and stage one complete first-row-id file group.""" from pypaimon.read.table_read import TableRead - from pypaimon.execution.python.write.row_id_file_index import RowIdFileIndex table_read = TableRead( context.scan_table, @@ -284,9 +287,7 @@ def _apply_self_merge_update_group(context, file_group, collect_row_ids): if collect_row_ids else [] ) - files_info = RowIdFileIndex.from_splits( - context.snapshot_id, [file_group], - ) + files_info = context.row_id_plan.for_split(file_group) updater = context.table.new_batch_write_builder().new_update().new_update_by_row_id( _precomputed_files_info=files_info) try: @@ -327,6 +328,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, + row_id_plan=plan.row_id_plan, )) group_iter = iter(plan.file_groups) max_in_flight = max(1, min(num_partitions, len(plan.file_groups))) @@ -612,11 +614,11 @@ def distributed_update_apply( # the delete path). from pypaimon.common.options.core_options import CoreOptions scan_table = ( - table.copy({CoreOptions.SCAN_SNAPSHOT_ID.key(): str(base_snapshot_id)}) + table.copy_without_time_travel({CoreOptions.SCAN_SNAPSHOT_ID.key(): str(base_snapshot_id)}) if base_snapshot_id is not None else table ) - from pypaimon.execution.python.write.row_id_file_index import RowIdFileIndex - planner = RowIdFileIndex.from_table(scan_table) + from pypaimon.write.row_id_plan import capture_row_id_plan + planner = capture_row_id_plan(scan_table) sorted_first_row_ids = list(planner.first_row_ids) if not sorted_first_row_ids: return [], 0, [] @@ -632,23 +634,9 @@ def distributed_update_apply( # Keep the matched join and file-routing shuffle in separate Ray jobs. update_ds = update_ds.materialize() - # Pin commit-time conflict check to the snapshot the join was built on, - # so concurrent commits between read and planner are detected. - check_from_snapshot = ( - base_snapshot_id if base_snapshot_id is not None - else planner.snapshot_id - ) - - # Put file metadata into Ray's object store and pass a single ref to - # workers. Avoids per-task manifest re-scans (Jingsong review #6) and - # avoids serializing the metadata into every task's closure. Override - # snapshot_id with the join's base snapshot so commit-time conflict - # detection covers the read→planner window. - from dataclasses import replace - files_info = replace( - planner, - snapshot_id=check_from_snapshot, - ) + # Capture files and the resolved snapshot together. Workers must not plan + # again after later append/compaction commits. + files_info = planner precomputed_info_ref = ray.put(files_info) frid_col = "_FIRST_ROW_ID" @@ -708,7 +696,7 @@ def _assign_frid(batch: pa.Table) -> pa.Table: ) with_frid = with_frid.union(ray.data.from_arrow(sentinel)) - captured_table = table + captured_table = planner.table captured_cols = cols def _apply_group(group: pa.Table) -> pa.Table: @@ -731,15 +719,16 @@ def _apply_group(group: pa.Table) -> pa.Table: for_update = group.drop_columns([frid_col]) info = ray.get(precomputed_info_ref) - _validate_live_row_ids( - captured_table, info, int(group[frid_col][0].as_py()), for_update[row_id_name].to_pylist()) row_ids = ( for_update.column(row_id_name).to_pylist() if collect_row_ids else [] ) update = captured_table.new_batch_write_builder().new_update() - worker = update.new_update_by_row_id( - _precomputed_files_info=ray.get(precomputed_info_ref)) + worker = update.new_update_by_row_id(_precomputed_files_info=info) + from pypaimon.write.backend_dispatch import native_row_id_update_selected + if not native_row_id_update_selected(worker): + _validate_live_row_ids( + captured_table, info, int(group[frid_col][0].as_py()), for_update[row_id_name].to_pylist()) msgs = worker.update_columns(for_update, list(captured_cols)) return pa.Table.from_pydict({ "msgs_blob": [pickle.dumps(msgs)], @@ -817,8 +806,8 @@ def distributed_read_by_row_id( empty_out = _read_output_schema(table, read_cols).empty_table() # The caller pinned the resolved snapshot, including any retained tag metadata. - from pypaimon.execution.python.write.row_id_file_index import RowIdFileIndex - planner = RowIdFileIndex.from_table(table) + from pypaimon.write.row_id_plan import capture_row_id_plan + planner = capture_row_id_plan(table) sorted_first_row_ids = list(planner.first_row_ids) if not sorted_first_row_ids: return None @@ -869,7 +858,7 @@ def _assign_frid(batch: pa.Table) -> pa.Table: frid_col, pa.array(sorted_arr[idx], type=pa.int64()) ) - captured_table = table # read at the same pinned snapshot the planner routed on + captured_table = planner.table # read the captured view used for routing captured_read_cols = read_cols captured_empty = empty_out @@ -920,12 +909,12 @@ def distributed_delete_apply( row_id_name = SpecialFields.ROW_ID.name scan_table = ( - table.copy({CoreOptions.SCAN_SNAPSHOT_ID.key(): str(base_snapshot_id)}) + table.copy_without_time_travel({CoreOptions.SCAN_SNAPSHOT_ID.key(): str(base_snapshot_id)}) if base_snapshot_id is not None else table ) - planner = TableDeleteByRowId._from_python(scan_table) - anchor_info = planner._snapshot_anchor_ranges() + from pypaimon.write.row_id_plan import capture_row_id_plan + anchor_info = capture_row_id_plan(scan_table) if not anchor_info.anchors: return [], 0, [] @@ -990,7 +979,7 @@ def _assign_group(batch: pa.Table) -> pa.Table: map_kwargs = _map_kwargs(ray_remote_args) with_group = delete_ds.map_batches(_assign_group, **map_kwargs) - captured_table = scan_table + captured_table = anchor_info.table def _apply_group(group: pa.Table) -> pa.Table: if group.num_rows == 0: diff --git a/paimon-python/pypaimon/ray/row_id_conflict_rewriter.py b/paimon-python/pypaimon/ray/row_id_conflict_rewriter.py index d289ad439ea4..4995bcaca7bd 100644 --- a/paimon-python/pypaimon/ray/row_id_conflict_rewriter.py +++ b/paimon-python/pypaimon/ray/row_id_conflict_rewriter.py @@ -31,7 +31,7 @@ from pypaimon.table.row.generic_row import GenericRow from pypaimon.table.special_fields import SpecialFields from pypaimon.utils.range import Range -from pypaimon.execution.python.write.commit.conflict_detection import RowIdExistenceConflict +from pypaimon.write.row_id_conflict import RowIdExistenceConflict, check_row_id_conflicts from pypaimon.write.commit_message import CommitMessage logger = logging.getLogger(__name__) @@ -367,22 +367,7 @@ def _validate_no_logical_conflict( base_snapshot_id: int, ) -> None: """Reject non-compaction changes before advancing the rewrite baseline.""" - commit = table.new_batch_write_builder().new_commit() - try: - file_store_commit = commit.file_store_commit - file_store_commit.conflict_detection._row_id_check_from_snapshot = ( - base_snapshot_id - ) - entries = file_store_commit._collect_manifest_entries(update_messages) - conflict = file_store_commit.conflict_detection.check_row_id_from_snapshot( - latest_snapshot, - entries, - check_compaction=False, - ) - if conflict is not None: - raise conflict - finally: - commit.close() + check_row_id_conflicts(table, update_messages, latest_snapshot.id, base_snapshot_id) def _find_row_id_conflict(error) -> Optional[RowIdExistenceConflict]: diff --git a/paimon-python/pypaimon/ray/vector_index_build.py b/paimon-python/pypaimon/ray/vector_index_build.py index 31bf1a73a5a3..fa5f722578eb 100644 --- a/paimon-python/pypaimon/ray/vector_index_build.py +++ b/paimon-python/pypaimon/ray/vector_index_build.py @@ -17,8 +17,6 @@ """Ray shard construction with one driver-side global-index commit.""" -from pypaimon.execution.python.globalindex.build_plan import split_by_global_index_shard -from pypaimon.execution.python.globalindex.index_file_utils import new_global_index_file_name from pypaimon.globalindex.vindex_options import VINDEX_IDENTIFIERS from pypaimon.ray.search_execution import execution_options @@ -40,6 +38,8 @@ def validate_build_options(builder, concurrency, ray_remote_args): def build_vector_index(builder, splits, unindexed_ranges, index_field, table_read, index_path, concurrency, remote_args): import ray + from pypaimon.execution.python.globalindex.build_plan import split_by_global_index_shard + from pypaimon.execution.python.globalindex.index_file_utils import new_global_index_file_name shards = split_by_global_index_shard( splits, builder._core_options.global_index_row_count_per_shard(), unindexed_ranges) @@ -79,3 +79,43 @@ def _build_shard(context, shard, file_name): index_split, index_range = shard return builder._build_generic_shard( index_split, index_range, index_field, table_read, index_path, file_name=file_name) + + +def build_native_vector_index(table, plan, concurrency, remote_args): + """Schedule immutable Core tasks; all publication remains with the caller.""" + import ray + + shards = plan.shards() + if not shards: + return [] + context = ray.put(table) + remote = ray.remote(_build_native_shard).options(**remote_args) + pending, results = {}, {} + next_shard = 0 + try: + while pending or next_shard < len(shards): + while len(pending) < concurrency and next_shard < len(shards): + ref = remote.remote(context, shards[next_shard]) + pending[ref] = next_shard + next_shard += 1 + ready, _ = ray.wait(list(pending), num_returns=1) + ref = ready[0] + ordinal = pending[ref] + results[ordinal] = ray.get(ref) + del pending[ref] + return [message for ordinal in range(len(shards)) for message in results[ordinal]] + except BaseException: + # A worker may have completed even when result delivery failed. Wait + # for dispatched tasks, but retain every output: its commit ownership + # is unknown and another attempt may already hold the messages. + for ref in pending: + try: + ray.get(ref) + except BaseException: + pass + raise + + +def _build_native_shard(table, shard): + from pypaimon.execution.native.write.native_index_build import build_native_index_shard + return build_native_index_shard(table, shard) diff --git a/paimon-python/pypaimon/tests/native_ray_index_build_test.py b/paimon-python/pypaimon/tests/native_ray_index_build_test.py new file mode 100644 index 000000000000..57bf5630472c --- /dev/null +++ b/paimon-python/pypaimon/tests/native_ray_index_build_test.py @@ -0,0 +1,202 @@ +# 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. + +"""Native Ray tasks build captured shards without Python materialization or replay.""" + +import sys +import subprocess +import textwrap +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock, patch + +import pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.globalindex.create_global_index import GlobalIndexBuilder + +pytestmark = [pytest.mark.python_plan, pytest.mark.python_read, + pytest.mark.python_write, pytest.mark.python_commit] + + +class _Reference: + def __init__(self, function, args): + self.function, self.args, self.done = function, args, False + + def get(self): + if not self.done: + self.done = True + try: + self.value, self.error = self.function(*self.args), None + except BaseException as error: + self.error = error + if self.error is not None: + raise self.error + return self.value + + +class _Ray(ModuleType): + def __init__(self): + super().__init__('ray') + self.submitted, self.remote_options, self.delivery_failure = [], None, None + + def put(self, value): + return value + + def remote(self, function): + owner = self + + class Remote: + def options(self, **options): + owner.remote_options = options + return self + + def remote(self, *args): + ref = _Reference(function, args) + owner.submitted.append(ref) + return ref + + return Remote() + + def wait(self, refs, num_returns): + return [refs[-1]], refs[:-1] + + def get(self, ref): + value = ref.get() + if self.delivery_failure is not None and ref is self.submitted[0]: + raise self.delivery_failure + return value + + +def _request(tmp_path): + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', True) + schema = pa.schema([('id', pa.int32()), ('embedding', pa.list_(pa.float32()))]) + catalog.create_table('default.vectors', Schema.from_pyarrow_schema(schema, options={ + 'bucket': '-1', 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', + 'write.native.enabled': 'true', 'read.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', 'commit.native.enabled': 'false', + }), False) + return GlobalIndexBuilder(catalog.get_table('default.vectors'), 'embedding', 'ivf-flat', options={ + 'ivf-flat.dimension': '2', 'ivf-flat.nlist': '1', 'global-index.row-count-per-shard': '2'}) + + +@pytest.fixture +def native_protocol(): + module = 'pypaimon.execution.native.write.native_index_build' + core = Mock() + shards = [SimpleNamespace(ordinal=ordinal) for ordinal in range(3)] + core.plan.return_value = SimpleNamespace(shards=lambda: shards) + core.build_shard.side_effect = lambda shard: ['message-' + str(shard.ordinal)] + native_table = SimpleNamespace(new_global_index_build_builder=lambda: core) + sdk = ModuleType('pypaimon_rust.datafusion') + sdk.GlobalIndexBuildBuilder = SimpleNamespace(supports_index_type=lambda kind: True) + scheduler = _Ray() + with patch.dict(sys.modules, {'pypaimon_rust.datafusion': sdk, 'ray': scheduler}), \ + patch(module + '.native_runtime_available', return_value=True), \ + patch(module + '._rest_catalog_supported', return_value=True), \ + patch(module + '._native_table', return_value=native_table), \ + patch(module + '.from_native_commit_messages', side_effect=lambda table, messages: messages), \ + patch('pypaimon.execution.python.globalindex.create_global_index.PythonGlobalIndexBuilder', + side_effect=AssertionError('Python index execution')): + yield core, shards, scheduler + + +def test_native_ray_build_plans_once_and_builds_only_selected_shards(tmp_path, native_protocol): + core, _, scheduler = native_protocol + assert _request(tmp_path).build(execution='ray', concurrency=2) == ['message-0', 'message-1', 'message-2'] + core.plan.assert_called_once_with() + core.build.assert_not_called() + assert len(scheduler.submitted) == 3 + assert sorted(call[0][0].ordinal for call in core.build_shard.call_args_list) == [0, 1, 2] + assert scheduler.remote_options['max_retries'] == 0 + assert scheduler.remote_options['retry_exceptions'] is False + + +def test_native_ray_delivery_failure_preserves_outputs_and_never_replays(tmp_path, native_protocol): + core, _, scheduler = native_protocol + request = _request(tmp_path) + failure = RuntimeError('result delivery failed after the worker wrote its file') + scheduler.delivery_failure = failure + with patch.object(request._table.file_io, 'delete_quietly') as delete: + with pytest.raises(RuntimeError) as raised: + request.build(execution='ray', concurrency=2) + assert raised.value is failure + delete.assert_not_called() + core.build.assert_not_called() + assert len(core.build_shard.call_args_list) == 3 + assert all(ref.done for ref in scheduler.submitted) + + +def test_native_ray_plan_permission_failure_does_not_select_python(tmp_path, native_protocol): + core, _, scheduler = native_protocol + failure = PermissionError('index build denied') + core.plan.side_effect = failure + with pytest.raises(PermissionError) as raised: + _request(tmp_path).build(execution='ray', concurrency=2) + assert raised.value is failure + core.build_shard.assert_not_called() + assert not scheduler.submitted + + +def test_native_ray_empty_plan_does_not_submit_or_select_python(tmp_path, native_protocol): + core, _, scheduler = native_protocol + core.plan.return_value = SimpleNamespace(shards=lambda: []) + assert _request(tmp_path).build(execution='ray', concurrency=2) == [] + core.plan.assert_called_once_with() + core.build.assert_not_called() + core.build_shard.assert_not_called() + assert not scheduler.submitted + + +def test_native_ray_worker_failure_stops_admission_and_retains_outputs(tmp_path, native_protocol): + core, _, scheduler = native_protocol + failure = PermissionError('worker authorization changed') + + def build(shard): + if shard.ordinal == 1: + raise failure + return ['message-' + str(shard.ordinal)] + + core.build_shard.side_effect = build + request = _request(tmp_path) + with patch.object(request._table.file_io, 'delete_quietly') as delete: + with pytest.raises(PermissionError) as raised: + request.build(execution='ray', concurrency=2) + assert raised.value is failure + delete.assert_not_called() + core.build.assert_not_called() + assert sorted(call[0][0].ordinal for call in core.build_shard.call_args_list) == [0, 1] + assert len(scheduler.submitted) == 2 + assert all(ref.done for ref in scheduler.submitted) + + +def test_native_ray_scheduler_import_does_not_load_python_execution(): + script = """ + import importlib.abc + import sys + + class RejectPythonExecution(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.startswith('pypaimon.execution.python'): + raise AssertionError('Python execution imported: ' + fullname) + + sys.meta_path.insert(0, RejectPythonExecution()) + from pypaimon.ray.vector_index_build import build_native_vector_index + assert callable(build_native_vector_index) + """ + subprocess.run([sys.executable, '-c', textwrap.dedent(script)], check=True) diff --git a/paimon-python/pypaimon/tests/native_ray_row_id_plan_test.py b/paimon-python/pypaimon/tests/native_ray_row_id_plan_test.py new file mode 100644 index 000000000000..1ff0a6c40022 --- /dev/null +++ b/paimon-python/pypaimon/tests/native_ray_row_id_plan_test.py @@ -0,0 +1,112 @@ +# 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. + +"""Exercise the actual SDK plan bridge without rebuilding the extension.""" + +from pathlib import Path +from unittest.mock import patch + +import pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.execution.native.runtime import native_runtime_available +from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER +from pypaimon.write.backend_dispatch import native_row_id_update_selected +from pypaimon.write.row_id_plan import capture_row_id_plan +from pypaimon.write.table_delete import TableDeleteByRowId +from pypaimon.write.table_update_by_row_id import TableUpdateByRowId + +pytestmark = [pytest.mark.python_plan, pytest.mark.python_read, + pytest.mark.python_write, pytest.mark.python_commit] +_ENGINES = [False, pytest.param(True, marks=pytest.mark.skipif( + not native_runtime_available(), reason='pypaimon-rust not installed'))] +_SCHEMA = pa.schema([('id', pa.int32()), ('value', pa.int32())]) + + +def _table(tmp_path): + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', True) + catalog.create_table('default.t', Schema.from_pyarrow_schema(_SCHEMA, options={ + 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', + 'deletion-vectors.enabled': 'true', 'write.native.enabled': 'false', + 'commit.native.enabled': 'false', 'scan.native-plan.enabled': 'false', + 'read.native.enabled': 'false'}), False) + return catalog.get_table('default.t') + + +def _append(table, ids): + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(pa.table({'id': ids, 'value': [0] * len(ids)}, schema=_SCHEMA)) + builder.new_commit().commit(writer.prepare_commit()) + writer.close() + + +@pytest.mark.parametrize('native', _ENGINES) +def test_planned_worker_preserves_snapshot_and_cross_backend_commit(tmp_path, native): + table = _table(tmp_path) + _append(table, [1, 2]) + worker_table = table.copy({'write.native.enabled': str(native).lower()}) + captured = capture_row_id_plan(worker_table) + _append(table, [3]) + with patch('pypaimon.read.table_scan.TableScan.plan_all', side_effect=AssertionError('worker replanned')): + worker = TableUpdateByRowId(worker_table, 'planned', BATCH_COMMIT_IDENTIFIER, captured) + assert native_row_id_update_selected(worker) == native + updates = worker.update_columns(pa.table({'_ROW_ID': [1], 'value': [9]}), ['value']) + deletes = TableDeleteByRowId(worker_table, captured).delete([0]) + assert all(message.check_from_snapshot == captured.snapshot_id for message in updates + deletes) + table.new_batch_write_builder().new_commit().commit(updates + deletes) + builder = table.new_read_builder() + data = builder.new_read().to_arrow(builder.new_scan().plan().splits()).sort_by('id') + assert data.to_pydict() == {'id': [2, 3], 'value': [9, 0]} + + +@pytest.mark.skipif(not native_runtime_available(), reason='pypaimon-rust not installed') +def test_native_planned_deleted_row_rejected_before_output(tmp_path): + table = _table(tmp_path) + _append(table, [1, 2]) + table.new_batch_write_builder().new_commit().commit(TableDeleteByRowId(table).delete([0])) + table = table.copy({'write.native.enabled': 'true'}) + captured = capture_row_id_plan(table) + worker = TableUpdateByRowId(table, 'deleted', BATCH_COMMIT_IDENTIFIER, captured) + assert native_row_id_update_selected(worker) + before = set(Path(tmp_path).rglob('*.parquet')) + with pytest.raises(ValueError, match='deleted.*ROW_ID|deleted.*row ID'): + worker.update_columns(pa.table({'_ROW_ID': [1, 0], 'value': [9, 9]}), ['value']) + assert set(Path(tmp_path).rglob('*.parquet')) == before + + +@pytest.mark.parametrize('native', _ENGINES) +def test_planned_update_keeps_resolved_write_schema_with_older_files(tmp_path, native): + from pypaimon.schema.data_types import AtomicType + from pypaimon.schema.schema_change import SchemaChange + table = _table(tmp_path) + _append(table, [1, 2]) + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.alter_table(table.identifier, [SchemaChange.add_column('extra', AtomicType('INT'))]) + table = catalog.get_table(table.identifier).copy({'write.native.enabled': str(native).lower()}) + captured = capture_row_id_plan(table.copy_without_time_travel({'scan.snapshot-id': '1'})) + assert captured.table.table_schema.id == table.table_schema.id == 1 + assert captured.table._read_snapshot.schema_id == 0 + worker = TableUpdateByRowId(table, 'new-field', BATCH_COMMIT_IDENTIFIER, captured) + assert native_row_id_update_selected(worker) == native + messages = worker.update_columns(pa.table({'_ROW_ID': [0], 'extra': [7]}), ['extra']) + table.new_batch_write_builder().new_commit().commit(messages) + read = table.new_read_builder() + data = read.new_read().to_arrow(read.new_scan().plan().splits()).sort_by('id') + assert data.to_pydict() == {'id': [1, 2], 'value': [0, 0], 'extra': [7, None]} diff --git a/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py b/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py index fd1a63d2f3d3..453faf1455ad 100644 --- a/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py +++ b/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py @@ -3289,6 +3289,7 @@ def test_self_merge_filters_file_group_in_batches(self): row_id_name=plan.row_id_name, snapshot_id=plan.snapshot_id, callable_input_columns=plan.callable_input_columns, + row_id_plan=plan.row_id_plan, ) batch_sizes = [] real_build_transform = join_module._build_matched_transform diff --git a/paimon-python/pypaimon/tests/ray_row_id_conflict_rewriter_test.py b/paimon-python/pypaimon/tests/ray_row_id_conflict_rewriter_test.py index 6bbd778faa0c..eda5c8792862 100644 --- a/paimon-python/pypaimon/tests/ray_row_id_conflict_rewriter_test.py +++ b/paimon-python/pypaimon/tests/ray_row_id_conflict_rewriter_test.py @@ -26,7 +26,7 @@ ) from pypaimon.schema.data_types import AtomicType, DataField from pypaimon.schema.table_schema import TableSchema -from pypaimon.execution.python.write.commit.conflict_detection import RowIdExistenceConflict +from pypaimon.write.row_id_conflict import RowIdExistenceConflict class RayRowIdConflictRewriterTest(unittest.TestCase): diff --git a/paimon-python/pypaimon/tests/write/row_id_conflict_protocol_test.py b/paimon-python/pypaimon/tests/write/row_id_conflict_protocol_test.py new file mode 100644 index 000000000000..091cc57d2057 --- /dev/null +++ b/paimon-python/pypaimon/tests/write/row_id_conflict_protocol_test.py @@ -0,0 +1,326 @@ +# 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. + +from dataclasses import replace +import pickle +import sys +from types import ModuleType +from unittest.mock import Mock, patch + +import pytest + +from pypaimon.tests.write.conflict_detection_test import _make_entry +from pypaimon.execution.python.write.commit.conflict_detection import ConflictDetection +from pypaimon.write.commit_message import CommitMessage +from pypaimon.write.row_id_conflict import RowIdExistenceConflict, check_row_id_conflicts + + +def _native_error_module(): + module = ModuleType('pypaimon_rust') + + class NativeRowIdConflict(RuntimeError): + pass + + module.RowIdExistenceConflict = NativeRowIdConflict + return module + + +def _native_conflict(module): + error = module.RowIdExistenceConflict('an opaque transport message') + error.file_name = "opaque '(4): file.parquet" + error.first_row_id = 4 + error.row_count = 8 + error.bucket = 3 + error.current_first_row_id = None + error.current_row_count = 9 + return error + + +@pytest.mark.parametrize('from_entry', [False, True]) +def test_structured_conflict_survives_worker_serialization(from_entry): + if from_entry: + error = RowIdExistenceConflict(_make_entry('staged.parquet', first_row_id=7)) + else: + error = RowIdExistenceConflict._from_fields( + file_name='staged.parquet', first_row_id=7, + row_count=100, bucket=0, current_row_count=101) + restored = pickle.loads(pickle.dumps(error)) + assert str(restored) == str(error) + assert restored.file_name == error.file_name + assert restored.first_row_id == error.first_row_id + assert restored.current_row_count == error.current_row_count + assert (restored.entry is None) == (error.entry is None) + + +@pytest.mark.parametrize('batch', [False, True]) +def test_native_commit_translates_only_typed_row_id_conflicts(batch): + from pypaimon.execution.native.write.commit import NativeTableCommit + from pypaimon.ray.row_id_conflict_rewriter import _find_row_id_conflict + + module = _native_error_module() + native = _native_conflict(module) + core = Mock() + core.commit.side_effect = native + commit = NativeTableCommit(core, is_batch=batch) + with patch.dict(sys.modules, {'pypaimon_rust': module}): + with pytest.raises(RowIdExistenceConflict) as raised: + commit.commit([], 7) + error = raised.value + assert error.__cause__ is native + assert error.entry is None + assert (error.file_name, error.first_row_id, error.row_count, error.bucket) == ( + "opaque '(4): file.parquet", 4, 8, 3) + assert error.current_first_row_id is None + assert error.current_row_count == 9 + assert _find_row_id_conflict(error) is error + core.abort.assert_not_called() + + +@pytest.mark.parametrize('opaque_message', ['unknown commit result', 'Row ID existence conflict: not a typed error']) +def test_native_unknown_outcome_is_not_a_rewrite_signal(opaque_message): + from pypaimon.execution.native.write.commit import NativeTableCommit + from pypaimon.ray.row_id_conflict_rewriter import _find_row_id_conflict + + module = _native_error_module() + failure = ValueError(opaque_message) + failure.__cause__ = _native_conflict(module) + core = Mock() + core.commit.side_effect = failure + with patch.dict(sys.modules, {'pypaimon_rust': module}): + with pytest.raises(ValueError) as raised: + NativeTableCommit(core, is_batch=True).commit([], 7) + assert raised.value is failure + assert _find_row_id_conflict(failure) is None + core.abort.assert_not_called() + + +@pytest.mark.parametrize('refuse', [False, True]) +def test_native_logical_check_never_constructs_python_commit_or_replays(refuse): + message = CommitMessage((), 0, [_make_entry('staged.parquet', first_row_id=0).file], check_from_snapshot=1) + table = Mock() + table.options.native_commit_enabled.return_value = True + core = Mock() + failure = RuntimeError('logical overlap') + if refuse: + core.check_row_id_conflicts.side_effect = failure + core.close.side_effect = RuntimeError('closing failed') + with patch('pypaimon.write.update_validation._authorize_update') as authorize, patch( + 'pypaimon.execution.native.write.native_commit.prepare_native_row_id_conflict_check', + return_value=(core, ['wire-message'])) as prepare, patch( + 'pypaimon.execution.python.write.commit.row_id_conflict.check_row_id_conflicts') as pure, patch( + 'pypaimon.execution.python.write.file_store_commit.FileStoreCommit') as publisher: + if refuse: + with pytest.raises(RuntimeError) as raised: + check_row_id_conflicts(table, [message], 2, 1) + assert raised.value is failure + else: + check_row_id_conflicts(table, [message], 2, 1) + authorize.assert_called_once_with(table) + prepare.assert_called_once_with(table, [message]) + core.check_row_id_conflicts.assert_called_once_with(['wire-message'], 2, 1) + core.close.assert_called_once_with() + core.abort.assert_not_called() + pure.assert_not_called() + publisher.assert_not_called() + + +@pytest.mark.parametrize('failure', [PermissionError('grant refused'), ValueError('unknown transport failure')]) +def test_native_check_preparation_errors_are_never_replayed(failure): + table = Mock() + table.options.native_commit_enabled.return_value = True + message = CommitMessage((), 0, [], check_from_snapshot=1) + with patch('pypaimon.write.update_validation._authorize_update'), patch( + 'pypaimon.execution.native.write.native_commit.prepare_native_row_id_conflict_check', + side_effect=failure), patch( + 'pypaimon.execution.python.write.commit.row_id_conflict.check_row_id_conflicts') as pure: + with pytest.raises(type(failure)) as raised: + check_row_id_conflicts(table, [message], 2, 1) + assert raised.value is failure + pure.assert_not_called() + + +@pytest.mark.parametrize('latest,baseline', [ + (2, 0), (1, 2), (0, 1), (2, -1), (2, 2), (2.5, 1), (True, 1), (2, False), (1 << 63, 1)]) +def test_check_rejects_invalid_baseline_before_authorization_or_io(latest, baseline): + table = Mock() + message = CommitMessage((), 0, [], check_from_snapshot=1) + with patch('pypaimon.write.update_validation._authorize_update') as authorize: + with pytest.raises(ValueError): + check_row_id_conflicts(table, [message], latest, baseline) + authorize.assert_not_called() + table.snapshot_manager.assert_not_called() + + +@pytest.mark.parametrize('baseline', [None, True, 1.0, '1', -1, 1 << 63]) +def test_check_rejects_invalid_message_baseline_before_io(baseline): + table = Mock() + message = CommitMessage((), 0, [], check_from_snapshot=baseline) + with patch('pypaimon.write.update_validation._authorize_update') as authorize: + with pytest.raises(ValueError): + check_row_id_conflicts(table, [message], 2, 1) + authorize.assert_not_called() + table.snapshot_manager.assert_not_called() + + +def test_pure_logical_check_uses_history_without_a_publisher(tmp_path): + import pyarrow as pa + from pypaimon import CatalogFactory, Schema + + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', False) + catalog.create_table('default.t', Schema.from_pyarrow_schema( + pa.schema([('id', pa.int32()), ('name', pa.string())]), options={ + 'data-evolution.enabled': 'true', 'row-tracking.enabled': 'true', + 'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false', + 'write.native.enabled': 'false', 'commit.native.enabled': 'false'}), False) + table = catalog.get_table('default.t') + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(pa.table({'id': pa.array([1, 2], pa.int32()), 'name': ['a', 'b']})) + initial = writer.prepare_commit() + writer.close() + commit = builder.new_commit() + commit.commit(initial) + commit.close() + update = table.new_batch_write_builder().new_update().with_update_type(['name']) + changes = update.update_by_arrow_with_row_id(pa.table({'_ROW_ID': [0, 1], 'name': ['x', 'y']})) + commit = table.new_batch_write_builder().new_commit() + commit.commit(changes) + commit.close() + staged = replace(changes[0], new_files=[replace(changes[0].new_files[0], file_name='still-prepared.parquet')]) + before = sorted(path.name for path in tmp_path.rglob('*') if path.is_file()) + with patch('pypaimon.execution.python.write.file_store_commit.FileStoreCommit', + side_effect=AssertionError('a publisher was constructed')), patch.object( + table, 'new_snapshot_commit', side_effect=AssertionError('publication route requested')): + with pytest.raises(RuntimeError, match='MERGE INTO') as raised: + check_row_id_conflicts(table, [staged], 2, 1) + assert not isinstance(raised.value, RowIdExistenceConflict) + disjoint = replace(staged, new_files=[replace(staged.new_files[0], write_cols=['id'])]) + check_row_id_conflicts(table, [disjoint], 2, 1) + assert before == sorted(path.name for path in tmp_path.rglob('*') if path.is_file()) + + +@pytest.mark.parametrize('first_row_id,row_count', [(10, 100), (None, 100), (0, 200)]) +def test_delete_rejects_reassigned_current_file(first_row_id, row_count): + detection = ConflictDetection(True, None, None, None, None) + current = _make_entry('same-file.parquet', first_row_id=first_row_id, row_count=row_count) + stale = _make_entry('same-file.parquet', kind=1, first_row_id=0) + conflict = detection.check_row_id_existence([current], [stale], None) + assert conflict is not None + assert conflict.entry is stale + assert conflict.file_name == 'same-file.parquet' + assert conflict.first_row_id == 0 + assert conflict.row_count == 100 + assert conflict.current_first_row_id == first_row_id + assert conflict.current_row_count == row_count + + +def test_conflict_type_remains_available_from_previous_module(): + from pypaimon.write.row_id_conflict import RowIdExistenceConflict + from pypaimon.execution.python.write.commit.conflict_detection import RowIdExistenceConflict as Previous + + assert Previous is RowIdExistenceConflict + entry = _make_entry('stale.parquet', first_row_id=7) + error = RowIdExistenceConflict(entry) + assert error.entry is entry + assert (error.file_name, error.first_row_id, error.row_count, error.bucket) == ('stale.parquet', 7, 100, 0) + assert error.current_row_count is None + + +def test_ray_conflict_module_does_not_import_python_execution(): + # A fresh process catches imports hidden by the module cache of this test run. + import subprocess + import sys + + program = ''' +import importlib.abc +import sys +class Guard(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == 'pypaimon.execution.python' or fullname.startswith('pypaimon.execution.python.'): + raise AssertionError('Python execution imported: ' + fullname) +sys.meta_path.insert(0, Guard()) +import pypaimon.ray.row_id_conflict_rewriter +''' + subprocess.run([sys.executable, '-c', program], check=True) + + +def test_native_logical_check_does_not_import_python_execution(): + import subprocess + + program = ''' +import importlib.abc +import sys +from tempfile import TemporaryDirectory +from types import ModuleType +from unittest.mock import patch +class Guard(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == 'pypaimon.execution.python' or fullname.startswith('pypaimon.execution.python.'): + raise AssertionError('Python execution imported: ' + fullname) +sys.meta_path.insert(0, Guard()) +class NativeMessage: + @staticmethod + def deserialize(payload, version): + assert isinstance(payload, bytes) and payload and version == 14 + return NativeMessage() +runtime = ModuleType('pypaimon_rust') +runtime.datafusion = ModuleType('pypaimon_rust.datafusion') +runtime.datafusion.CommitMessage = NativeMessage +sys.modules['pypaimon_rust'] = runtime +sys.modules['pypaimon_rust.datafusion'] = runtime.datafusion +import pyarrow as pa +from pypaimon import CatalogFactory, Schema +from pypaimon.execution.native.write import native_commit +from pypaimon.write.commit_message import CommitMessage +from pypaimon.write.row_id_conflict import check_row_id_conflicts +class Core: + checked = False + closed = False + def check_row_id_conflicts(self, messages, latest, baseline): + from pypaimon_rust.datafusion import CommitMessage as NativeMessage + assert len(messages) == 1 and isinstance(messages[0], NativeMessage) + assert (latest, baseline) == (2, 1) + self.checked = True + def close(self): + self.closed = True +with TemporaryDirectory() as path: + catalog = CatalogFactory.create({'warehouse': path}) + catalog.create_database('default', False) + catalog.create_table('default.t', Schema.from_pyarrow_schema( + pa.schema([('id', pa.int32())]), options={ + 'data-evolution.enabled': 'true', 'row-tracking.enabled': 'true', + 'commit.native.enabled': 'true'}), False) + table = catalog.get_table('default.t') + core = Core() + with patch.object(native_commit, 'create_native_commit', return_value=core), patch.object( + table, 'new_snapshot_commit', side_effect=AssertionError('publisher requested')): + check_row_id_conflicts(table, [CommitMessage((), 0, [], check_from_snapshot=1)], 2, 1) + assert core.checked and core.closed +''' + subprocess.run([sys.executable, '-c', program], check=True) + + +def test_ordinary_delete_and_empty_messages_are_not_existence_conflicts(): + current = _make_entry('ordinary.parquet') + deleted = _make_entry('ordinary.parquet', kind=1) + detection = ConflictDetection(True, None, None, None, None) + assert detection.check_row_id_existence([current], [deleted], None) is None + assert detection.check_row_id_existence([], [], None) is None + disabled = ConflictDetection(False, None, None, None, None) + stale = _make_entry('ordinary.parquet', kind=1, first_row_id=7) + assert disabled.check_row_id_existence([current], [stale], 100) is None diff --git a/paimon-python/pypaimon/tests/write/row_id_execution_boundary_test.py b/paimon-python/pypaimon/tests/write/row_id_execution_boundary_test.py index b0a4d9d3faa6..bb17d33888d0 100644 --- a/paimon-python/pypaimon/tests/write/row_id_execution_boundary_test.py +++ b/paimon-python/pypaimon/tests/write/row_id_execution_boundary_test.py @@ -106,10 +106,8 @@ def test_builder_row_id_fallback_prepares_native_only_once(): table = Mock() index = RowIdFileIndex(snapshot_id=7, first_row_ids=[]) - with patch('pypaimon.write.update_dispatch._authorize_update'), \ - patch('pypaimon.write.update_dispatch.prepare_native_write_operation', return_value=None) as prepare, \ - patch('pypaimon.write.backend_dispatch.prepare_native_write_operation', - side_effect=AssertionError('Native prepared twice')), \ + with patch('pypaimon.write.update_validation._authorize_update'), \ + patch('pypaimon.write.backend_dispatch.prepare_native_write_operation', return_value=None) as prepare, \ patch.object(RowIdFileIndex, 'from_table', return_value=index): update = BatchTableUpdate(table, 'user').new_update_by_row_id() prepare.assert_called_once_with('update_by_row_id', table, 'user', BATCH_COMMIT_IDENTIFIER) diff --git a/paimon-python/pypaimon/tests/write/row_id_plan_test.py b/paimon-python/pypaimon/tests/write/row_id_plan_test.py new file mode 100644 index 000000000000..6e8d7db08198 --- /dev/null +++ b/paimon-python/pypaimon/tests/write/row_id_plan_test.py @@ -0,0 +1,260 @@ +# 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. + + +import pickle +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock, patch + +import pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.write.row_id_plan import capture_row_id_plan + +pytestmark = [pytest.mark.python_plan, pytest.mark.python_read, + pytest.mark.python_write, pytest.mark.python_commit] + + +@pytest.fixture +def table(tmp_path): + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', True) + catalog.create_table('default.t', Schema.from_pyarrow_schema( + pa.schema([('id', pa.int32()), ('v', pa.int32())]), options={ + 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', + 'deletion-vectors.enabled': 'true'}), False) + return catalog.get_table('default.t') + + +def _append(table, ids): + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(pa.table({'id': ids, 'v': [0] * len(ids)}, + schema=pa.schema([('id', pa.int32()), ('v', pa.int32())]))) + builder.new_commit().commit(writer.prepare_commit()) + writer.close() + + +def test_captured_update_and_delete_keep_snapshot_without_manifest_rescan(table): + _append(table, [1, 2]) + captured = pickle.loads(pickle.dumps(capture_row_id_plan(table))) + _append(table, [3]) + assert captured.snapshot_id == 1 + assert captured.table._read_snapshot.id == 1 + assert [(r.from_, r.to) for r in captured.valid_row_id_ranges] == [(0, 1)] + assert captured.anchors[0].row_range.to == 1 + + from pypaimon.write.table_update_by_row_id import TableUpdateByRowId + from pypaimon.write.table_delete import TableDeleteByRowId + with patch('pypaimon.read.table_scan.TableScan.plan_all', side_effect=AssertionError('worker replanned')): + updater = TableUpdateByRowId(table, 'planned', 1, captured) + updates = updater.update_columns(pa.table({'_ROW_ID': [1], 'v': [9]}), ['v']) + deletes = TableDeleteByRowId(table, captured).delete([0]) + assert updates and deletes + assert all(message.check_from_snapshot == 1 for message in updates + deletes) + with pytest.raises(ValueError, match='valid range'): + TableUpdateByRowId(table, 'foreign', 2, captured).update_columns( + pa.table({'_ROW_ID': [2], 'v': [9]}), ['v']) + + +def test_native_captured_plan_does_not_construct_python_index(table): + _append(table, [1, 2]) + plan = table.new_read_builder().new_scan().plan_all() + split = plan.splits()[0] + native_plan = Mock() + native_plan.groups.return_value = [(0, 1, split)] + with patch('pypaimon.execution.native.write.row_id_plan.create_native_row_id_plan', return_value=native_plan), \ + patch('pypaimon.execution.native.write.row_id_plan.python_groups', + return_value=[(0, 1, split)]) as core_groups, \ + patch('pypaimon.execution.python.write.row_id_file_index.RowIdFileIndex.from_splits', + side_effect=AssertionError('Python index loaded')): + captured = capture_row_id_plan(table, plan) + core_groups.assert_called_once_with(captured.table, native_plan) + assert captured.first_row_id_index[0][1] == split.files + assert captured.for_split(split).plan.splits() == [split] + assert pickle.loads(pickle.dumps(captured)).snapshot_id == plan.snapshot_id + + +def test_planned_native_factory_applies_plan_before_creating_writer(table): + from pypaimon.execution.native.write.native_update import create_native_update_by_row_id + _append(table, [1]) + captured = capture_row_id_plan(table) + core = Mock() + update = core.new_batch_write_builder.return_value.with_commit_user.return_value.new_update.return_value + wire = object() + with patch('pypaimon.execution.native.write.native_update._native_row_id_table', return_value=core), \ + patch('pypaimon.execution.native.write.row_id_plan.native_plan', return_value=wire): + operation = create_native_update_by_row_id(table, 'user', (1 << 63) - 1, captured) + update.with_plan.assert_called_once_with(wire) + update.new_update_by_row_id.assert_called_once_with() + assert operation.writer is update.new_update_by_row_id.return_value + + +def test_planned_native_authorization_error_is_not_replayed(table): + from pypaimon.write.backend_dispatch import create_row_id_update_execution + _append(table, [1]) + captured = capture_row_id_plan(table) + failure = PermissionError('worker grant denied') + with patch.object(table.options, 'native_write_enabled', return_value=True), \ + patch('pypaimon.execution.native.write.native_update.create_native_update_by_row_id', + side_effect=failure), \ + patch('pypaimon.write.backend_dispatch.create_python_row_id_update', + side_effect=AssertionError('Python replay')): + with pytest.raises(PermissionError) as raised: + create_row_id_update_execution(table, 'user', 1, captured) + assert raised.value is failure + + +@pytest.mark.parametrize('operation', ['update', 'delete']) +def test_foreign_captured_table_is_rejected_before_backend_preparation(table, operation): + from dataclasses import replace + from pypaimon.write.backend_dispatch import create_row_id_update_execution, create_row_id_delete_execution + _append(table, [1]) + captured = capture_row_id_plan(table) + foreign = table.copy_without_time_travel({}) + foreign.table_path = table.table_path + '-foreign' + captured = replace(captured, table=foreign) + with patch('pypaimon.execution.native.write.native_update.create_native_update_by_row_id', + side_effect=AssertionError('worker created')): + with pytest.raises(ValueError, match='different table'): + if operation == 'update': + create_row_id_update_execution(table, 'user', 1, captured) + else: + create_row_id_delete_execution(table, captured) + + +def test_captured_split_requires_the_complete_group(table): + _append(table, [1, 2]) + captured = capture_row_id_plan(table) + from dataclasses import replace + unrelated = replace(captured.plan.splits()[0].files[0], file_name='not-captured.parquet') + from pypaimon.read.split import DataSplit + split = DataSplit(files=[unrelated], partition=captured.plan.splits()[0].partition, bucket=0) + with pytest.raises(ValueError, match='complete captured'): + captured.for_split(split) + + +def test_capture_normalizes_index_selection_before_core_mutation_routing(table): + from pypaimon.globalindex.indexed_split import IndexedSplit + from pypaimon.read.plan import Plan + from pypaimon.utils.range import Range + + _append(table, [1, 2]) + plan = table.new_read_builder().new_scan().plan_all() + original = plan.splits()[0] + indexed = Plan([IndexedSplit(original, [Range(1, 1)])], plan.snapshot_id) + native = object() + + def prepare(table, routing): + assert routing.snapshot_id == plan.snapshot_id + assert routing.splits() == [original], 'Core mutation routing requires the complete physical group' + return native + + with patch('pypaimon.execution.native.write.row_id_plan.create_native_row_id_plan', side_effect=prepare), \ + patch('pypaimon.execution.native.write.row_id_plan.python_groups', + return_value=[(0, 1, original)]) as core_groups: + captured = capture_row_id_plan(table, indexed) + core_groups.assert_called_once_with(captured.table, native) + assert captured.plan.splits() == [original] + # Normalization does not mutate the read plan used to refine the predicate. + assert indexed.splits()[0].row_ranges() == [Range(1, 1)] + + +@pytest.mark.parametrize('native', [False, True]) +def test_capture_keeps_query_auth_rules_when_normalizing_and_grouping(table, native): + from pypaimon.catalog.table_query_auth import TableQueryAuthResult + from pypaimon.globalindex.indexed_split import IndexedSplit + from pypaimon.read.plan import Plan + from pypaimon.read.query_auth_split import QueryAuthSplit + from pypaimon.utils.range import Range + + _append(table, [1, 2]) + plan = table.new_read_builder().new_scan().plan_all() + original = plan.splits()[0] + grant = TableQueryAuthResult(['restricted-row-filter'], {}) + indexed = Plan([QueryAuthSplit(IndexedSplit(original, [Range(1, 1)]), grant)], plan.snapshot_id) + with patch('pypaimon.execution.native.write.row_id_plan.create_native_row_id_plan', + return_value=object() if native else None) as prepare, \ + patch('pypaimon.execution.native.write.row_id_plan.python_groups', + return_value=[(0, 1, QueryAuthSplit(original, grant))]) as native_groups: + captured = capture_row_id_plan(table, indexed) + routed = prepare.call_args[0][1].splits()[0] + assert isinstance(routed, QueryAuthSplit) + assert routed.auth_result is grant + assert routed.split is original + assert native_groups.call_count == int(native) + complete = captured.plan.splits()[0] + assert isinstance(complete, QueryAuthSplit) + assert complete.auth_result is grant + assert complete.split is original + group = captured.groups[0][2] + assert isinstance(group, QueryAuthSplit) + assert group.auth_result is grant + # A same-name caller split cannot replace the captured grant or snapshot. + retained = captured.for_split(original).plan.splits()[0] + assert retained is group + assert retained.auth_result is grant + + +@pytest.mark.parametrize('layout', ['external', 'dedicated']) +@pytest.mark.parametrize('authorized', [False, True]) +@pytest.mark.parametrize('bucket_path', [None, '', 'memory:/foreign/bucket-0']) +def test_native_mutation_transport_retains_canonical_bucket_and_external_layout( + table, layout, authorized, bucket_path): + from dataclasses import replace + from pypaimon.catalog.table_query_auth import TableQueryAuthResult + from pypaimon.execution.native.write.row_id_plan import native_plan + from pypaimon.read.plan import Plan + from pypaimon.read.query_auth_split import QueryAuthSplit + from pypaimon.read.split import DataSplit + from pypaimon.read.split_serializer import deserialize_split_v1 + from pypaimon.table.source.deletion_file import DeletionFile + + _append(table, [1, 2]) + read_plan = table.new_read_builder().new_scan().plan_all() + source = read_plan.splits()[0] + files = list(source.files) + if layout == 'external': + path = str(table.table_path) + '-external/' + files[0].file_name + files[0] = replace(files[0], external_path=path, file_path=path) + else: + files.append(replace(files[0], file_name='payload.blob', write_cols=['payload'], + file_path=str(table.table_path) + '-blobs/payload.blob')) + deletion = DeletionFile(str(table.table_path) + '/index/dv', 4, 12, 1) + split = DataSplit(files, source.partition, source.bucket, data_deletion_files=[deletion] * len(files), + snapshot_id=read_plan.snapshot_id, bucket_path=bucket_path) + grant = TableQueryAuthResult(['restricted-row-filter'], {}) + selected = QueryAuthSplit(split, grant) if authorized else split + module = ModuleType('pypaimon_rust.datafusion') + module.Plan = lambda snapshot_id, splits: SimpleNamespace(snapshot_id=snapshot_id, splits=lambda: splits) + module.Split = SimpleNamespace(deserialize=lambda value: deserialize_split_v1( + value, table.partition_keys_fields, table.trimmed_primary_keys_fields)) + with patch.dict(sys.modules, {'pypaimon_rust.datafusion': module}): + transported = native_plan(table, Plan([selected], read_plan.snapshot_id)).splits()[0] + if authorized: + assert isinstance(transported, QueryAuthSplit) + assert transported.auth_result.filter == grant.filter + transported = transported.split + expected = bucket_path or table.path_factory().bucket_path( + tuple(source.partition.values), source.bucket, canonical_partition=True) + assert transported.bucket_path == expected + assert transported.snapshot_id == read_plan.snapshot_id + assert transported.data_deletion_files == [deletion] * len(files) + assert [file.file_path for file in transported.files] == [file.file_path for file in files] + assert split.bucket_path == bucket_path, 'transport must not mutate the caller\'s read metadata' diff --git a/paimon-python/pypaimon/write/backend_dispatch.py b/paimon-python/pypaimon/write/backend_dispatch.py index f30ccd7c1a37..4460950cf208 100644 --- a/paimon-python/pypaimon/write/backend_dispatch.py +++ b/paimon-python/pypaimon/write/backend_dispatch.py @@ -335,6 +335,14 @@ def create_row_id_update_execution(table, commit_user, commit_identifier, precom """Select one incremental updater before capturing target files.""" from pypaimon.write.update_validation import _authorize_update _authorize_update(table) + from pypaimon.write.row_id_plan import CapturedRowIdPlan + if isinstance(precomputed, CapturedRowIdPlan): + _validate_captured_row_id_plan(table, precomputed) + if table.options.native_write_enabled(): + from pypaimon.execution.native.write.native_update import create_native_update_by_row_id + native = create_native_update_by_row_id(table, commit_user, commit_identifier, precomputed) + if native is not None: + return native if precomputed is None: native = prepare_native_write_operation('update_by_row_id', table, commit_user, commit_identifier) if native is not None: @@ -343,9 +351,17 @@ def create_row_id_update_execution(table, commit_user, commit_identifier, precom def create_row_id_delete_execution(table, precomputed=None): - """Keep precomputed broadcast metadata within its selected Python executor.""" + """Select one deleter, reusing captured metadata when supplied.""" from pypaimon.write.update_validation import _authorize_update _authorize_update(table) + from pypaimon.write.row_id_plan import CapturedRowIdPlan + if isinstance(precomputed, CapturedRowIdPlan): + _validate_captured_row_id_plan(table, precomputed) + if table.options.native_write_enabled(): + from pypaimon.execution.native.write.native_update import create_native_row_id_delete + native = create_native_row_id_delete(table, precomputed) + if native is not None: + return native if precomputed is None and table.options.native_write_enabled(): def prepare(): from pypaimon.execution.native.write.native_update import create_native_row_id_delete @@ -356,6 +372,19 @@ def prepare(): return create_python_row_id_delete(table, precomputed) +def _validate_captured_row_id_plan(table, captured): + if (table.table_path != captured.table.table_path or table.current_branch() != captured.table.current_branch() + or table.table_schema.id != captured.table.table_schema.id + or table.fields != captured.table.fields): + raise ValueError('Captured row-ID plan belongs to a different table, branch or schema') + + +def native_row_id_update_selected(worker): + from pypaimon.execution.native.write.native_update import NativeTableUpdateByRowId + execution = vars(worker).get('_execution', worker) + return isinstance(execution, NativeTableUpdateByRowId) + + def python_commit(owner): """Create Python commit resources only after that backend is selected.""" if owner._python_commit is None: diff --git a/paimon-python/pypaimon/write/row_id_conflict.py b/paimon-python/pypaimon/write/row_id_conflict.py new file mode 100644 index 000000000000..61cda36f30a3 --- /dev/null +++ b/paimon-python/pypaimon/write/row_id_conflict.py @@ -0,0 +1,132 @@ +# 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. + +"""Structured row-ID conflicts shared by both execution backends.""" + + +class RowIdExistenceConflict(RuntimeError): + """A staged row-id file no longer matches the current base-file layout.""" + + def __init__(self, entry): + self.entry = entry + self._initialize(entry.file.file_name, entry.file.first_row_id, + entry.file.row_count, entry.bucket) + + @classmethod + def _from_fields(cls, file_name, first_row_id, row_count, bucket, + current_first_row_id=None, current_row_count=None): + error = cls.__new__(cls) + error.entry = None + error._initialize(file_name, first_row_id, row_count, bucket, + current_first_row_id, current_row_count) + return error + + @classmethod + def _from_delete(cls, deleted, current): + error = cls._from_fields( + deleted.file.file_name, deleted.file.first_row_id, + deleted.file.row_count, deleted.bucket, + current.file.first_row_id, current.file.row_count) + error.entry = deleted + return error + + def _initialize(self, file_name, first_row_id, row_count, bucket, + current_first_row_id=None, current_row_count=None): + self.file_name = file_name + self.first_row_id = first_row_id + self.row_count = row_count + self.bucket = bucket + self.current_first_row_id = current_first_row_id + self.current_row_count = current_row_count + first_id_text = 'null' if first_row_id is None else first_row_id + if current_row_count is None: + message = ( + "Row ID existence conflict: file '{}' references " + "firstRowId={}, rowCount={} in bucket {}, " + "but no matching file exists in the current snapshot. " + "The referenced file may have been rewritten by a " + "concurrent compaction or removed by an overwrite.".format( + file_name, first_id_text, row_count, bucket)) + else: + message = ( + "Row ID existence conflict: DELETE for file '{}' references " + "firstRowId={}, rowCount={} in bucket {}, but the current file " + "has firstRowId={}, rowCount={}. The file may have been " + "reassigned by a concurrent commit.".format( + file_name, first_id_text, row_count, bucket, + 'null' if current_first_row_id is None else current_first_row_id, current_row_count)) + super().__init__(message) + + def __reduce__(self): + # Ray transports errors between worker processes. RuntimeError's + # default reduction would pass the formatted message as a file entry. + return (_restore_row_id_conflict, ( + self.entry, self.file_name, self.first_row_id, self.row_count, + self.bucket, self.current_first_row_id, self.current_row_count)) + + +def _restore_row_id_conflict(entry, file_name, first_row_id, row_count, bucket, + current_first_row_id, current_row_count): + error = RowIdExistenceConflict._from_fields( + file_name, first_row_id, row_count, bucket, + current_first_row_id, current_row_count) + error.entry = entry + return error + + +def check_row_id_conflicts(table, messages, latest_snapshot_id, check_from_snapshot): + """Check logical writes before a rebase, without opening a publisher.""" + from pypaimon.write.update_validation import _authorize_update + + if (any(not isinstance(value, int) or isinstance(value, bool) + or value <= 0 or value > (1 << 63) - 1 + for value in (latest_snapshot_id, check_from_snapshot)) + or check_from_snapshot > latest_snapshot_id): + raise ValueError('Row-ID conflict checks require 0 < check-from snapshot <= latest snapshot') + baselines = [message.check_from_snapshot for message in messages if message.check_from_snapshot is not None] + invalid_baseline = any( + not isinstance(value, int) or isinstance(value, bool) or value < 0 or value > (1 << 63) - 1 + for value in baselines) + missing_baseline = any( + message.check_from_snapshot is None and any( + file.first_row_id is not None for file in message.new_files + message.deleted_files) + for message in messages) + if (not baselines or invalid_baseline or min(baselines) != check_from_snapshot or missing_baseline): + raise ValueError('Row-ID conflict check baseline must match the commit messages') + _authorize_update(table) + if table.options.native_commit_enabled(): + from pypaimon.execution.native.write.native_commit import prepare_native_row_id_conflict_check + # Unsupported configurations return None explicitly. Permission or + # transport failures must never restart the check in another backend. + native = prepare_native_row_id_conflict_check(table, messages) + if native is not None: + commit, native_messages = native + failure = None + try: + commit.check_row_id_conflicts(native_messages, latest_snapshot_id, check_from_snapshot) + except BaseException as error: + failure = error + raise + finally: + try: + commit.close() + except BaseException: + if failure is None: + raise + return + from pypaimon.execution.python.write.commit.row_id_conflict import check_row_id_conflicts as check + check(table, messages, latest_snapshot_id, check_from_snapshot) diff --git a/paimon-python/pypaimon/write/row_id_plan.py b/paimon-python/pypaimon/write/row_id_plan.py new file mode 100644 index 000000000000..3b859b54fef4 --- /dev/null +++ b/paimon-python/pypaimon/write/row_id_plan.py @@ -0,0 +1,112 @@ +# 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. + + +"""Serializable routing metadata for a captured row-ID mutation plan. + +Ray transports Python metadata; file grouping belongs to the selected backend. +Workers reuse complete logical groups and the original snapshot. +""" + +from dataclasses import dataclass, replace + +from pypaimon.read.plan import Plan +from pypaimon.utils.data_evolution_utils import retrieve_anchor_file +from pypaimon.utils.range import Range + + +@dataclass(frozen=True) +class _RowIdAnchor: + row_range: Range + partition: object + bucket: int + anchor_file: object + + +@dataclass(frozen=True) +class CapturedRowIdPlan: + table: object + plan: Plan + groups: list + + @property + def snapshot_id(self): + return self.plan.snapshot_id if self.plan.snapshot_id is not None else -1 + + @property + def first_row_ids(self): + return [start for start, _, _ in self.groups] + + @property + def first_row_id_index(self): + return {start: (split, split.files) for start, _, split in self.groups} + + @property + def valid_row_id_ranges(self): + return Range.sort_and_merge_overlap( + [Range(start, end) for start, end, _ in self.groups], True, True) + + @property + def anchors(self): + return [_RowIdAnchor(anchor.row_id_range(), split.partition, split.bucket, anchor) + for _, _, split in self.groups + for anchor in [retrieve_anchor_file(split.files)]] + + def for_split(self, split): + """Keep one complete logical group and its captured snapshot metadata.""" + names = {file.file_name for file in split.files} + groups = [group for group in self.groups + if group[2].bucket == split.bucket and group[2].partition == split.partition + and {file.file_name for file in group[2].files} == names] + if len(groups) != 1: + raise ValueError('Split is not a complete captured row-ID group') + return CapturedRowIdPlan(self.table, Plan([groups[0][2]], self.plan.snapshot_id), groups) + + +def _mutation_routing_split(split): + from pypaimon.globalindex.indexed_split import IndexedSplit + from pypaimon.read.query_auth_split import QueryAuthSplit + + if isinstance(split, QueryAuthSplit): + complete = _mutation_routing_split(split.split) + return split if complete is split.split else QueryAuthSplit(complete, split.auth_result) + if isinstance(split, IndexedSplit): + # A trusted scan's index selection limits reads, not the physical file + # group rewritten by an update. Keep every file and deletion vector; + # the caller's read predicate still selects the rows to mutate. + return _mutation_routing_split(split.data_split()) + return split + + +def capture_row_id_plan(table, plan=None): + """Resolve one read view before planning, then retain it across Ray tasks.""" + from pypaimon.snapshot.time_travel_util import TimeTravelUtil + snapshot = TimeTravelUtil.resolve_snapshot(table) + if plan is not None and (snapshot.id if snapshot is not None else None) != plan.snapshot_id: + raise ValueError('Row-ID plan does not belong to the captured snapshot') + table = table._copy_with_snapshot(snapshot) + if plan is None: + plan = table.new_read_builder().new_scan().plan_all() + plan = replace(plan, _splits=[_mutation_routing_split(split) for split in plan.splits()]) + from pypaimon.execution.native.write.row_id_plan import create_native_row_id_plan, python_groups + native = create_native_row_id_plan(table, plan) + if native is not None: + groups = python_groups(table, native) + else: + from pypaimon.execution.python.write.row_id_plan import row_id_groups + groups = row_id_groups(plan) + return CapturedRowIdPlan(table, plan, groups) diff --git a/paimon-python/pypaimon/write/update_dispatch.py b/paimon-python/pypaimon/write/update_dispatch.py index 0b170d9855c1..383192225ce3 100644 --- a/paimon-python/pypaimon/write/update_dispatch.py +++ b/paimon-python/pypaimon/write/update_dispatch.py @@ -42,14 +42,10 @@ def __init__(self, context): def _new_row_id_updater(self, commit_identifier: int, _precomputed_files_info=None) -> TableUpdateByRowId: context = self._context - _authorize_update(context.table) - if _precomputed_files_info is None: - native = prepare_native_write_operation( - 'update_by_row_id', context.table, context.commit_user, commit_identifier) - if native is not None: - return native - return TableUpdateByRowId._from_python(context.table, context.commit_user, commit_identifier, - _precomputed_files_info=_precomputed_files_info) + # The public constructor and builders share one selection boundary, + # including complete plans transported to distributed workers. + return TableUpdateByRowId(context.table, context.commit_user, commit_identifier, + _precomputed_files_info=_precomputed_files_info) def new_shard_updator(self, shard_num: int, total_shard_count: int): context = self._context From 62f7efe49caca54e6d7b2f344a4d603a96d11c08 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Fri, 23 Oct 2026 08:44:07 +0800 Subject: [PATCH 2/2] [Python] Preserve authorized index dispatch and bound self-merge task plans --- .../native/write/native_index_build.py | 1 + .../pypaimon/ray/data_evolution_merge_join.py | 25 ++++++----- .../tests/native_ray_index_build_test.py | 18 ++++++++ .../tests/native_sorted_index_build_test.py | 41 +++++++++++++++++- .../ray_data_evolution_merge_into_test.py | 3 +- .../pypaimon/tests/write/row_id_plan_test.py | 42 +++++++++++++++++++ 6 files changed, 117 insertions(+), 13 deletions(-) diff --git a/paimon-python/pypaimon/execution/native/write/native_index_build.py b/paimon-python/pypaimon/execution/native/write/native_index_build.py index c25491a81267..7c0735b0bca1 100644 --- a/paimon-python/pypaimon/execution/native/write/native_index_build.py +++ b/paimon-python/pypaimon/execution/native/write/native_index_build.py @@ -34,6 +34,7 @@ def build_native_global_index(builder, partition_filter, *, execution="local", c if (builder._index_type not in ('btree', 'bitmap', 'full-text') + tuple(VINDEX_IDENTIFIERS) or not table.options.native_write_enabled() or (table.options.deletion_vectors_enabled() and builder._index_type == 'full-text') + or table.options.query_auth_enabled or table.is_primary_key_table or table.options.file_format() != 'parquet' or not native_runtime_available() diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index 1692a34c564f..6a9ecd1df493 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -98,7 +98,6 @@ class _SelfMergeUpdateContext: row_id_name: str snapshot_id: int callable_input_columns: Optional[List[str]] - row_id_plan: Any def _resolve_source_projection( @@ -245,10 +244,11 @@ def _self_merge_aliases(batch: pa.Table, row_id_name: str) -> pa.Table: return pa.table(columns, names=names) -def _apply_self_merge_update_group(context, file_group, collect_row_ids): +def _apply_self_merge_update_group(context, group_plan, collect_row_ids): """Read, transform, and stage one complete first-row-id file group.""" from pypaimon.read.table_read import TableRead + file_group = group_plan.plan.splits()[0] table_read = TableRead( context.scan_table, context.predicate, @@ -287,9 +287,8 @@ def _apply_self_merge_update_group(context, file_group, collect_row_ids): if collect_row_ids else [] ) - files_info = context.row_id_plan.for_split(file_group) updater = context.table.new_batch_write_builder().new_update().new_update_by_row_id( - _precomputed_files_info=files_info) + _precomputed_files_info=group_plan) try: messages = updater.update_columns(updates, context.update_cols) except Exception: @@ -308,8 +307,11 @@ def distributed_self_merge_update_apply( ) -> Tuple[list, int, list]: """Stage self-merge updates in scan tasks without a routing shuffle.""" import ray + from pypaimon.read.plan import Plan + from pypaimon.write.row_id_plan import CapturedRowIdPlan - if not plan.file_groups: + captured = plan.row_id_plan + if not captured.groups: return [], 0, [] remote_args = dict(ray_remote_args or {}) @@ -328,10 +330,9 @@ 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, - row_id_plan=plan.row_id_plan, )) - group_iter = iter(plan.file_groups) - max_in_flight = max(1, min(num_partitions, len(plan.file_groups))) + group_iter = iter(captured.groups) + max_in_flight = max(1, min(num_partitions, len(captured.groups))) pending = set() messages = [] num_updated = 0 @@ -340,11 +341,15 @@ def distributed_self_merge_update_apply( def submit_next(): try: - file_group = next(group_iter) + group = next(group_iter) except StopIteration: return False + # Transport only this canonical group. Ray deserializes shared Python + # objects per task, so broadcasting the full plan repeats all metadata. + group_plan = CapturedRowIdPlan( + captured.table, Plan([group[2]], captured.plan.snapshot_id), [group]) pending.add(apply_remote.remote( - context, file_group, collect_row_ids + context, group_plan, collect_row_ids )) return True diff --git a/paimon-python/pypaimon/tests/native_ray_index_build_test.py b/paimon-python/pypaimon/tests/native_ray_index_build_test.py index 57bf5630472c..5f6fd1b6261d 100644 --- a/paimon-python/pypaimon/tests/native_ray_index_build_test.py +++ b/paimon-python/pypaimon/tests/native_ray_index_build_test.py @@ -163,6 +163,24 @@ def test_native_ray_empty_plan_does_not_submit_or_select_python(tmp_path, native assert not scheduler.submitted +@pytest.mark.parametrize('execution', ['local', 'ray']) +def test_query_auth_selects_python_before_native_index_execution(tmp_path, native_protocol, execution): + core, _, scheduler = native_protocol + request = _request(tmp_path) + request._table = request._table.copy({'query-auth.enabled': 'true'}) + with patch('pypaimon.execution.python.globalindex.create_global_index.PythonGlobalIndexBuilder') as python: + python.return_value.build.return_value = ['authorized-python-message'] + kwargs = {'execution': execution} + if execution == 'ray': + kwargs['concurrency'] = 2 + assert request.build(**kwargs) == ['authorized-python-message'] + assert python.return_value.build.call_args[1]['execution'] == execution + core.plan.assert_not_called() + core.build.assert_not_called() + core.build_shard.assert_not_called() + assert not scheduler.submitted + + def test_native_ray_worker_failure_stops_admission_and_retains_outputs(tmp_path, native_protocol): core, _, scheduler = native_protocol failure = PermissionError('worker authorization changed') diff --git a/paimon-python/pypaimon/tests/native_sorted_index_build_test.py b/paimon-python/pypaimon/tests/native_sorted_index_build_test.py index 204a95c56359..13fce8767c26 100644 --- a/paimon-python/pypaimon/tests/native_sorted_index_build_test.py +++ b/paimon-python/pypaimon/tests/native_sorted_index_build_test.py @@ -26,7 +26,8 @@ import pytest from pypaimon import Schema -from pypaimon.api.api_response import ErrorResponse, GetTableSnapshotResponse +from pypaimon.api.api_response import AuthTableQueryResponse, ErrorResponse, GetTableSnapshotResponse +from pypaimon.catalog.catalog_exception import TableNoPermissionException from pypaimon.globalindex.create_global_index import GlobalIndexBuilder from pypaimon.execution.python.globalindex.create_global_index import PythonGlobalIndexBuilder from pypaimon.execution.python.globalindex.data_evolution_global_index_scanner import DataEvolutionGlobalIndexScanner @@ -35,6 +36,7 @@ from pypaimon.execution.native.read.builder import (_predicate_to_native) from pypaimon.execution.native.runtime import (native_runtime_available) from pypaimon.snapshot.table_snapshot import TableSnapshot +from pypaimon.schema.schema_change import SchemaChange from pypaimon.tests import native_plan_rest_test from pypaimon.execution.native.write.native_commit import from_native_commit_messages from pypaimon.write.table_commit import BatchTableCommit @@ -118,6 +120,43 @@ def _delete(table, ids): _commit(table, builder.new_update().delete_by_row_id(ids)) +@pytest.mark.python_plan +@pytest.mark.python_read +@pytest.mark.python_write +@pytest.mark.python_commit +@pytest.mark.parametrize('policy', ['unrestricted', 'filtered', 'masked', 'denied']) +def test_query_auth_keeps_python_index_build_authorization(rest_catalog, policy): + catalog, server = rest_catalog + table = _create(catalog, options={'scan.native-plan.enabled': 'false'}) + _append(table, ROWS) + catalog.alter_table(table.identifier, [SchemaChange.set_option('query-auth.enabled', 'true')]) + table = catalog.get_table(table.identifier) + if policy == 'denied': + reply = server._mock_response(ErrorResponse('TABLE', 'indexes', 'index build denied', 403), 403) + else: + response = AuthTableQueryResponse( + filter=['{"kind":"LEAF","transform":{"name":"FIELD_REF",' + '"fieldRef":{"name":"id","index":0,"type":"INT"}},' + '"function":"GREATER_THAN","literals":[1]}'] if policy == 'filtered' else None, + column_masking={'name': '{"name":"NULL"}'} if policy == 'masked' else None) + reply = server._mock_response(response, 200) + native_module = 'pypaimon.execution.native.write.native_index_build' + with patch.object(server, '_table_query_auth_handle', return_value=reply) as authorize, \ + patch(native_module + '._native_table', side_effect=AssertionError('Native index build selected')): + for native in [False, True]: + request = GlobalIndexBuilder(table.copy({'write.native.enabled': str(native).lower()}), 'name') + if policy == 'unrestricted': + files = _message_files(request.build()) + assert sum(file.row_count for file in files) == len(ROWS) + assert all(table.file_io.exists(_path(table, file)) for file in files) + else: + with pytest.raises(TableNoPermissionException): + request.build() + assert authorize.call_count >= 2 + assert _files(table) == [] + assert table.snapshot_manager().get_latest_snapshot().id == 1 + + @pytest.mark.parametrize('kind', ['btree', 'bitmap']) @pytest.mark.parametrize('deleted', [[0, 2], [0, 1, 2, 3]]) @pytest.mark.parametrize('native_commit', [False, True]) diff --git a/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py b/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py index 453faf1455ad..9f2440d49214 100644 --- a/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py +++ b/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py @@ -3289,7 +3289,6 @@ def test_self_merge_filters_file_group_in_batches(self): row_id_name=plan.row_id_name, snapshot_id=plan.snapshot_id, callable_input_columns=plan.callable_input_columns, - row_id_plan=plan.row_id_plan, ) batch_sizes = [] real_build_transform = join_module._build_matched_transform @@ -3312,7 +3311,7 @@ def tracked_transform(batch): ): messages, count, row_ids = ( join_module._apply_self_merge_update_group( - context, plan.file_groups[0], True, + context, plan.row_id_plan.for_split(plan.file_groups[0]), True, ) ) diff --git a/paimon-python/pypaimon/tests/write/row_id_plan_test.py b/paimon-python/pypaimon/tests/write/row_id_plan_test.py index 6e8d7db08198..f2c970974e0d 100644 --- a/paimon-python/pypaimon/tests/write/row_id_plan_test.py +++ b/paimon-python/pypaimon/tests/write/row_id_plan_test.py @@ -73,6 +73,48 @@ def test_captured_update_and_delete_keep_snapshot_without_manifest_rescan(table) pa.table({'_ROW_ID': [2], 'v': [9]}), ['v']) +def test_self_merge_tasks_transport_only_their_complete_captured_group(table): + from pypaimon.ray.data_evolution_merge_join import ( + build_self_merge_update_plan, distributed_self_merge_update_apply) + from pypaimon.write.row_id_plan import CapturedRowIdPlan + from pypaimon.write.table_update_by_row_id import TableUpdateByRowId + + _append(table, [1, 2]) + _append(table, [3, 4]) + messages = TableUpdateByRowId(table, 'partial', 1).update_columns( + pa.table({'_ROW_ID': [0], 'v': [9]}), ['v']) + table.new_batch_write_builder().new_commit().commit(messages) + plan = build_self_merge_update_plan( + table=table, clauses=[], target_field_names=table.field_names, + target_pa_schema=pa.schema([('id', pa.int32()), ('v', pa.int32())]), + update_cols=['v'], resolve_target_projection=lambda *args, **kwargs: ['v']) + assert len(plan.file_groups) == 2 + expected = {start: {file.file_name for file in split.files} + for start, _, split in plan.row_id_plan.groups} + assert max(len(names) for names in expected.values()) == 2 + _append(table, [5]) + submitted = [] + + def submit(context, captured, collect_row_ids): + assert not hasattr(context, 'row_id_plan'), 'shared context includes whole-table metadata' + assert isinstance(captured, CapturedRowIdPlan) + captured = pickle.loads(pickle.dumps(captured)) + assert len(captured.groups) == len(captured.plan.splits()) == 1 + assert captured.table._read_snapshot.id == captured.snapshot_id == plan.snapshot_id + start, _, split = captured.groups[0] + assert {file.file_name for file in split.files} == expected[start] + assert captured.plan.splits() == [split] + submitted.append(start) + return start + + scheduler = SimpleNamespace( + put=lambda value: value, remote=lambda function: SimpleNamespace(remote=submit), + wait=lambda pending, **kwargs: (pending[:1], pending[1:]), get=lambda ref: ([], 0, [])) + with patch.dict(sys.modules, {'ray': scheduler}): + assert distributed_self_merge_update_apply(plan, num_partitions=2) == ([], 0, []) + assert submitted == plan.row_id_plan.first_row_ids + + def test_native_captured_plan_does_not_construct_python_index(table): _append(table, [1, 2]) plan = table.new_read_builder().new_scan().plan_all()