From 883da038d44c5f6f7f94be444fbb68be457d5561 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 28 Jul 2026 20:31:39 -0700 Subject: [PATCH 01/11] [python][ray] Incrementally commit update_by_row_id file groups --- docs/docs/pypaimon/ray-data.md | 8 + .../pypaimon/ray/data_evolution_merge_join.py | 121 +++++--- .../pypaimon/ray/update_by_row_id.py | 116 +++++++- .../tests/ray_update_by_row_id_test.py | 269 +++++++++++++++++- 4 files changed, 469 insertions(+), 45 deletions(-) diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md index 88c0a7143ac4..7c90c5718e2c 100644 --- a/docs/docs/pypaimon/ray-data.md +++ b/docs/docs/pypaimon/ray-data.md @@ -524,6 +524,7 @@ metrics = update_by_row_id( source=ray_dataset, # ray.data.Dataset / pa.Table / pandas, carrying _ROW_ID catalog_options={"warehouse": "/path/to/warehouse"}, update_cols=["feature"], # non-blob columns to overwrite + max_groups_per_commit=64, # optional incremental commit window ) print(metrics) # {"num_updated": 50} ``` @@ -537,6 +538,8 @@ print(metrics) # {"num_updated": 50} - `num_partitions`: parallelism for grouping the update rows by target file; defaults to `max(1, cluster_cpus * 2)`. - `ray_remote_args`: Ray remote options applied to the update tasks. +- `max_groups_per_commit`: optional positive number of completed target file groups + per commit. By default, all groups are committed together. **Returns:** `{"num_updated": }`. @@ -548,6 +551,11 @@ print(metrics) # {"num_updated": 50} - Partition columns cannot be updated (in-place rewrite can't move a row across partitions). - Deletion-vectors-enabled tables are not supported yet: a DV-deleted row still lives in its data file, so it can't be told apart from a live row without reading the target. +- Incremental commits are not atomic across the whole operation. If a group fails, + completed groups are flushed before the error is raised. +- Incremental mode reports group exceptions as results so other groups can finish. + Ray task retry options in `ray_remote_args` do not retry these exceptions; a + transient group failure can therefore leave a partial update. ## Read By Row Id diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index 359af8e6efa2..c189961cb153 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -16,7 +16,7 @@ # limitations under the License. ################################################################################ -from typing import Any, Dict, List, Optional, Sequence, Tuple +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple import pyarrow as pa @@ -32,6 +32,21 @@ ) +class GroupApplyError(RuntimeError): + """A file group failed after other groups may have completed.""" + + +def _group_error_text(error: BaseException) -> str: + import traceback + + return "update_by_row_id file group failed: {}: {}\n{}".format( + type(error).__name__, + str(error), + "".join(traceback.format_exception( + type(error), error, error.__traceback__)), + ) + + def _map_kwargs( ray_remote_args: Optional[Dict[str, Any]], ) -> Dict[str, Any]: @@ -445,6 +460,7 @@ def distributed_update_apply( ray_remote_args: Optional[Dict[str, Any]] = None, base_snapshot_id: Optional[int] = None, collect_row_ids: bool = False, + on_group_result: Optional[Callable[[list, int, list], None]] = None, ) -> Tuple[list, int, list]: import numpy as np import pickle @@ -546,6 +562,15 @@ def _assign_frid(batch: pa.Table) -> pa.Table: captured_table = table captured_cols = cols + capture_group_errors = on_group_result is not None + + def _group_result(msgs_blob, n_updated, row_ids_blob, error): + return pa.Table.from_pydict({ + "msgs_blob": pa.array([msgs_blob], type=pa.binary()), + "n_updated": pa.array([n_updated], type=pa.int64()), + "row_ids_blob": pa.array([row_ids_blob], type=pa.binary()), + "error": pa.array([error], type=pa.string()), + }) def _apply_group(group: pa.Table) -> pa.Table: if group.num_rows == 0: @@ -553,39 +578,43 @@ def _apply_group(group: pa.Table) -> pa.Table: "msgs_blob": pa.array([], type=pa.binary()), "n_updated": pa.array([], type=pa.int64()), "row_ids_blob": pa.array([], type=pa.binary()), + "error": pa.array([], type=pa.string()), }) - if ( - pc.count_distinct(group.column(row_id_name)).as_py() - != group.num_rows - ): - raise ValueError( - "MERGE matched multiple source rows to the same " - "target _ROW_ID. Deduplicate the source before " - "merging." - ) + try: + if ( + pc.count_distinct(group.column(row_id_name)).as_py() + != group.num_rows + ): + raise ValueError( + "MERGE matched multiple source rows to the same " + "target _ROW_ID. Deduplicate the source before " + "merging." + ) - for_update = group.drop_columns([frid_col]) - row_ids = ( - for_update.column(row_id_name).to_pylist() - if collect_row_ids else [] - ) - worker = TableUpdateByRowId( - captured_table, - "_merge_into_shard_" + uuid.uuid4().hex[:8], - BATCH_COMMIT_IDENTIFIER, - _precomputed_files_info=ray.get(precomputed_info_ref), + for_update = group.drop_columns([frid_col]) + row_ids = ( + for_update.column(row_id_name).to_pylist() + if collect_row_ids else [] + ) + worker = TableUpdateByRowId( + captured_table, + "_merge_into_shard_" + uuid.uuid4().hex[:8], + BATCH_COMMIT_IDENTIFIER, + _precomputed_files_info=ray.get(precomputed_info_ref), + ) + msgs = worker.update_columns(for_update, list(captured_cols)) + except Exception as error: + if not capture_group_errors: + raise + return _group_result(b"", 0, b"", _group_error_text(error)) + + return _group_result( + pickle.dumps(msgs), + for_update.num_rows, + pickle.dumps(row_ids), + None, ) - msgs = worker.update_columns(for_update, list(captured_cols)) - return pa.Table.from_pydict({ - "msgs_blob": [pickle.dumps(msgs)], - "n_updated": pa.array( - [for_update.num_rows], type=pa.int64() - ), - "row_ids_blob": pa.array( - [pickle.dumps(row_ids)], type=pa.binary() - ), - }) # One group per target data file; bounded by file count and num_partitions. group_partitions = max( @@ -598,14 +627,34 @@ def _apply_group(group: pa.Table) -> pa.Table: all_msgs: list = [] num_updated = 0 action_row_ids = [] + group_error = None for batch in msgs_ds.iter_batches(batch_format="pyarrow"): - for blob in batch.column("msgs_blob").to_pylist(): - all_msgs.extend(pickle.loads(blob)) - for n in batch.column("n_updated").to_pylist(): + message_blobs = batch.column("msgs_blob").to_pylist() + updated_counts = batch.column("n_updated").to_pylist() + errors = batch.column("error").to_pylist() + row_id_blobs = ( + batch.column("row_ids_blob").to_pylist() + if collect_row_ids else [None] * len(message_blobs) + ) + for blob, n, row_ids_blob, error in zip( + message_blobs, updated_counts, row_id_blobs, errors): + if error is not None: + if group_error is None: + group_error = error + continue + group_msgs = pickle.loads(blob) + group_row_ids = ( + pickle.loads(row_ids_blob) if collect_row_ids else [] + ) + if on_group_result is None: + all_msgs.extend(group_msgs) + else: + on_group_result(group_msgs, n, group_row_ids) num_updated += n - if collect_row_ids: - for blob in batch.column("row_ids_blob").to_pylist(): - action_row_ids.extend(pickle.loads(blob)) + if collect_row_ids: + action_row_ids.extend(group_row_ids) + if group_error is not None: + raise GroupApplyError(group_error) return all_msgs, num_updated, action_row_ids diff --git a/paimon-python/pypaimon/ray/update_by_row_id.py b/paimon-python/pypaimon/ray/update_by_row_id.py index 41b968c4fdc8..4350bc60e122 100644 --- a/paimon-python/pypaimon/ray/update_by_row_id.py +++ b/paimon-python/pypaimon/ray/update_by_row_id.py @@ -34,7 +34,10 @@ _require_ray_join, _resolve_num_partitions, ) -from pypaimon.ray.data_evolution_merge_join import distributed_update_apply +from pypaimon.ray.data_evolution_merge_join import ( + GroupApplyError, + distributed_update_apply, +) from pypaimon.ray.data_evolution_merge_transform import build_update_schema from pypaimon.schema.data_types import is_blob_file_field @@ -56,6 +59,7 @@ def update_by_row_id( update_cols: List[str], num_partitions: Optional[int] = None, ray_remote_args: Optional[Dict[str, Any]] = None, + max_groups_per_commit: Optional[int] = None, ) -> Dict[str, int]: """Update ``update_cols`` of a data-evolution table by ``_ROW_ID``. @@ -65,6 +69,12 @@ def update_by_row_id( the target is never fully read and there is no join against it. Requires ``ray >= 2.50`` and a target with ``data-evolution.enabled`` + ``row-tracking.enabled``. + By default, all file groups are committed atomically. Set + ``max_groups_per_commit`` to commit completed groups in smaller windows. + If a group fails, completed groups are flushed before the error is raised. + Group exceptions are returned as results in incremental mode, so Ray task + retry options in ``ray_remote_args`` do not retry them. + Returns ``{"num_updated": }``. """ from pypaimon.catalog.catalog_factory import CatalogFactory @@ -74,6 +84,11 @@ def update_by_row_id( _require_ray_join() if not update_cols: raise ValueError("update_cols must be non-empty.") + if max_groups_per_commit is not None: + if (isinstance(max_groups_per_commit, bool) + or not isinstance(max_groups_per_commit, int) + or max_groups_per_commit <= 0): + raise ValueError("max_groups_per_commit must be a positive integer.") update_cols = list(dict.fromkeys(update_cols)) # de-dup, keep order num_partitions = _resolve_num_partitions(num_partitions) @@ -139,22 +154,111 @@ def _project_cast(batch: pa.Table) -> pa.Table: raise ValueError( f"target '{target}' has no rows; every _ROW_ID in the source is foreign.") return {"num_updated": 0} + incremental_committer = ( + _IncrementalUpdateCommitter(table, max_groups_per_commit) + if max_groups_per_commit is not None else None + ) try: + apply_kwargs = { + "num_partitions": num_partitions, + "ray_remote_args": ray_remote_args, + "base_snapshot_id": base.id, + } + if incremental_committer is not None: + apply_kwargs["on_group_result"] = incremental_committer.add_group msgs, num_updated, _ = distributed_update_apply( - update_ds, table, update_cols, - num_partitions=num_partitions, - ray_remote_args=ray_remote_args, - base_snapshot_id=base.id, + update_ds, table, update_cols, **apply_kwargs ) + if incremental_committer is not None: + incremental_committer.finish() + except GroupApplyError: + if incremental_committer is not None: + try: + incremental_committer.finish() + except Exception: + incremental_committer.abort_pending() + raise + raise except Exception as e: + if incremental_committer is not None: + incremental_committer.abort_pending() + if incremental_committer._commit_failed: + raise _reraise_inner(e) raise # _reraise_inner always raises; keeps msgs/num_updated defined for linters + finally: + if incremental_committer is not None: + incremental_committer.close() - if msgs: + if incremental_committer is None and msgs: _commit_update_messages(table, msgs) return {"num_updated": num_updated} +class _IncrementalUpdateCommitter: + + def __init__(self, table, max_groups_per_commit: int): + self._table = table + self._max_groups_per_commit = max_groups_per_commit + self._pending_messages: list = [] + self._pending_groups = 0 + self._table_commit = None + self._next_commit_identifier = 1 + self._commit_failed = False + + def add_group(self, commit_messages, _num_updated, _row_ids) -> None: + self._pending_messages.extend(commit_messages) + self._pending_groups += 1 + if self._pending_groups >= self._max_groups_per_commit: + self._commit_pending() + + def finish(self) -> None: + self._commit_pending() + + def _commit_pending(self) -> None: + if self._pending_groups == 0: + return + if not self._pending_messages: + self._pending_groups = 0 + return + + try: + if self._table_commit is None: + self._table_commit = ( + self._table.new_stream_write_builder().new_commit() + ) + + messages = self._pending_messages + self._pending_messages = [] + self._pending_groups = 0 + commit_identifier = self._next_commit_identifier + self._table_commit.commit(messages, commit_identifier) + except Exception: + self._commit_failed = True + raise + self._next_commit_identifier += 1 + + def abort_pending(self) -> None: + if not self._pending_messages: + return + messages = self._pending_messages + self._pending_messages = [] + self._pending_groups = 0 + _abort_pending_update_messages(self._table, messages) + + def close(self) -> None: + if self._table_commit is None: + return + try: + self._table_commit.close() + except Exception as close_error: + logger.warning( + "Failed to close incremental update_by_row_id commit: %s", + close_error, + exc_info=close_error, + ) + + def _commit_update_messages(table, commit_messages) -> None: pending_msgs: list = list(commit_messages) commit_started = False diff --git a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py index 613fa50d0319..2d2c51aee916 100644 --- a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py +++ b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py @@ -130,6 +130,222 @@ def test_updates_correct_row_across_files(self): self.assertEqual(got[21], 999) self.assertTrue(all(v == 0 for k, v in got.items() if k != 21)) + def test_incrementally_commits_file_group_windows(self): + from pypaimon.write.table_commit import StreamTableCommit + + target = self._create() + chunks = [[start, start + 1] for start in range(10, 60, 10)] + for chunk in chunks: + self._write(target, pa.Table.from_pydict( + {"id": chunk, "name": ["x", "x"], "age": [0, 0]}, + schema=self.pa_schema, + )) + + table = self.catalog.get_table(target) + base_snapshot_id = table.snapshot_manager().get_latest_snapshot().id + row_ids = self._rowid_by_id(target) + updated_ids = [chunk[0] for chunk in chunks] + source = pa.table( + { + "_ROW_ID": [row_ids[row_id] for row_id in updated_ids], + "age": updated_ids, + }, + schema=pa.schema([("_ROW_ID", pa.int64()), ("age", pa.int32())]), + ) + commits = [] + original_commit = StreamTableCommit.commit + + def record_commit(stream_commit, messages, commit_identifier): + commits.append((len(messages), commit_identifier)) + return original_commit( + stream_commit, messages, commit_identifier) + + with mock.patch.object(StreamTableCommit, "commit", record_commit): + stats = update_by_row_id( + target, + ray.data.from_arrow(source).repartition(4), + self.catalog_options, + update_cols=["age"], + num_partitions=4, + max_groups_per_commit=2, + ) + + self.assertEqual({"num_updated": 5}, stats) + self.assertEqual([(2, 1), (2, 2), (1, 3)], commits) + self.assertEqual( + base_snapshot_id + 3, + table.snapshot_manager().get_latest_snapshot().id, + ) + result = self._read(target).sort_by("id").to_pydict() + ages = dict(zip(result["id"], result["age"])) + self.assertEqual( + {row_id: row_id for row_id in updated_ids}, + {row_id: ages[row_id] for row_id in updated_ids}, + ) + + def test_group_failure_preserves_completed_groups(self): + for max_groups, expected_snapshots in [(1, 2), (5, 1)]: + with self.subTest(max_groups_per_commit=max_groups): + target = self._create() + for row_id in range(1, 4): + self._write(target, pa.Table.from_pydict( + {"id": [row_id], "name": ["a"], "age": [0]}, + schema=self.pa_schema, + )) + + table = self.catalog.get_table(target) + base_snapshot_id = ( + table.snapshot_manager().get_latest_snapshot().id + ) + row_ids = self._rowid_by_id(target) + source = pa.table( + { + "_ROW_ID": [ + row_ids[1], row_ids[2], row_ids[3], row_ids[3], + ], + "age": [100, 200, 300, 301], + }, + schema=pa.schema([ + ("_ROW_ID", pa.int64()), + ("age", pa.int32()), + ]), + ) + + with self.assertRaisesRegex(RuntimeError, "Deduplicate"): + update_by_row_id( + target, + ray.data.from_arrow(source), + self.catalog_options, + update_cols=["age"], + num_partitions=1, + max_groups_per_commit=max_groups, + ) + + self.assertEqual( + base_snapshot_id + expected_snapshots, + table.snapshot_manager().get_latest_snapshot().id, + ) + self.assertEqual( + [100, 200, 0], + self._read(target).sort_by("id")["age"].to_pylist(), + ) + + def test_atomic_group_failure_commits_nothing(self): + target = self._create() + for row_id in range(1, 4): + self._write(target, pa.Table.from_pydict( + {"id": [row_id], "name": ["a"], "age": [0]}, + schema=self.pa_schema, + )) + + table = self.catalog.get_table(target) + base_snapshot_id = table.snapshot_manager().get_latest_snapshot().id + row_ids = self._rowid_by_id(target) + source = pa.table( + { + "_ROW_ID": [ + row_ids[1], row_ids[2], row_ids[3], row_ids[3], + ], + "age": [100, 200, 300, 301], + }, + schema=pa.schema([ + ("_ROW_ID", pa.int64()), + ("age", pa.int32()), + ]), + ) + + with self.assertRaisesRegex(ValueError, "Deduplicate"): + update_by_row_id( + target, + ray.data.from_arrow(source), + self.catalog_options, + update_cols=["age"], + num_partitions=1, + ) + + self.assertEqual( + base_snapshot_id, + table.snapshot_manager().get_latest_snapshot().id, + ) + self.assertEqual( + [0, 0, 0], + self._read(target).sort_by("id")["age"].to_pylist(), + ) + + def test_incremental_committer_batches_and_aborts_pending_groups(self): + import importlib + + module = importlib.import_module("pypaimon.ray.update_by_row_id") + commits = [] + aborted = [] + close_calls = [] + + class FakeCommit: + def commit(self, messages, commit_identifier): + commits.append((list(messages), commit_identifier)) + + def close(self): + close_calls.append(True) + + class FakeBuilder: + def new_commit(self): + return FakeCommit() + + class FakeTable: + def new_stream_write_builder(self): + return FakeBuilder() + + committer = module._IncrementalUpdateCommitter(FakeTable(), 2) + with mock.patch.object( + module, + "_abort_pending_update_messages", + side_effect=lambda table, messages: aborted.append(list(messages))): + committer.add_group(["group-1"], 1, []) + committer.add_group(["group-2"], 1, []) + committer.add_group(["group-3"], 1, []) + committer.abort_pending() + committer.close() + + self.assertEqual( + [(["group-1", "group-2"], 1)], + commits, + ) + self.assertEqual([["group-3"]], aborted) + self.assertEqual([True], close_calls) + + def test_incremental_commit_failure_does_not_abort_unknown_outcome(self): + import importlib + + module = importlib.import_module("pypaimon.ray.update_by_row_id") + aborted = [] + + class FakeCommit: + def commit(self, messages, commit_identifier): + raise RuntimeError("commit failed") + + def close(self): + pass + + class FakeBuilder: + def new_commit(self): + return FakeCommit() + + class FakeTable: + def new_stream_write_builder(self): + return FakeBuilder() + + committer = module._IncrementalUpdateCommitter(FakeTable(), 1) + with mock.patch.object( + module, + "_abort_pending_update_messages", + side_effect=lambda table, messages: aborted.append(list(messages))): + with self.assertRaisesRegex(RuntimeError, "commit failed"): + committer.add_group(["group-1"], 1, []) + committer.abort_pending() + committer.close() + + self.assertEqual([], aborted) + def test_pins_base_snapshot_for_conflict_detection(self): # The update pins its base snapshot and threads it to distributed_update_apply, # which uses it for commit-time conflict detection against concurrent writers. @@ -183,6 +399,20 @@ def test_commit_failure_does_not_abort_after_commit_started(self): self.assertEqual(recorder["abort_calls"], 0) self.assertEqual(recorder["close_calls"], 1) + def test_incremental_driver_commit_failure_is_not_unwrapped(self): + retry_error = ValueError("retry failed") + commit_error = RuntimeError("commit failed") + commit_error.__cause__ = retry_error + + with self.assertRaises(RuntimeError) as raised: + self._run_with_fake_commit( + commit_error=commit_error, + incremental=True, + ) + + self.assertIs(commit_error, raised.exception) + self.assertIs(retry_error, raised.exception.__cause__) + def test_close_failure_after_success_warns_and_returns_stats(self): close_error = RuntimeError("close failed") @@ -324,8 +554,30 @@ def test_rejects_unknown_and_empty_update_cols(self): with self.assertRaises(ValueError): update_by_row_id(target, src, self.catalog_options, update_cols=[]) + def test_rejects_invalid_max_groups_per_commit(self): + target = self._create() + self._write(target, pa.Table.from_pydict( + {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) + source = pa.table( + {"_ROW_ID": [0], "age": [9]}, + schema=pa.schema([("_ROW_ID", pa.int64()), ("age", pa.int32())]), + ) + + for value in [0, -1, True, 1.5, "2"]: + with self.subTest(value=value): + with self.assertRaisesRegex( + ValueError, "must be a positive integer"): + update_by_row_id( + target, + source, + self.catalog_options, + update_cols=["age"], + max_groups_per_commit=value, + ) + def _run_with_fake_commit(self, *, recorder=None, new_commit_errors=None, - commit_error=None, close_error=None): + commit_error=None, close_error=None, + incremental=False): import importlib m = importlib.import_module("pypaimon.ray.update_by_row_id") @@ -354,7 +606,7 @@ def deletion_vectors_enabled(self): return False class FakeCommit: - def commit(self, msgs): + def commit(self, msgs, *args): recorder["commit_calls"] += 1 recorder["commit_msgs"] = list(msgs) if recorder["commit_error"] is not None: @@ -396,6 +648,9 @@ def snapshot_manager(self): def new_batch_write_builder(self): return FakeWriteBuilder() + def new_stream_write_builder(self): + return FakeWriteBuilder() + class FakeCatalog: def get_table(self, target): return FakeTable() @@ -410,6 +665,13 @@ def map_batches(self, fn, batch_format=None): parser_path = ( "pypaimon.schema.data_types.PyarrowFieldParser.from_paimon_schema" ) + + def fake_apply(*args, **kwargs): + if incremental: + kwargs["on_group_result"](recorder["msgs"], 3, []) + return [], 3, [] + return recorder["msgs"], 3, [] + with mock.patch( "pypaimon.catalog.catalog_factory.CatalogFactory.create", return_value=FakeCatalog()), \ @@ -425,12 +687,13 @@ def map_batches(self, fn, batch_format=None): ("age", pa.int32()), ])), \ mock.patch.object(m, "distributed_update_apply", - return_value=(recorder["msgs"], 3, [])): + side_effect=fake_apply): recorder["result"] = m.update_by_row_id( "default.fake", FakeSource(), self.catalog_options, update_cols=["age"], + max_groups_per_commit=1 if incremental else None, ) return recorder From 8c2369b07bd65bb667cf4c399a6376109ca4ecec Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 29 Jul 2026 06:57:29 -0700 Subject: [PATCH 02/11] [python][ray] Abort buffered groups after commit failure --- docs/docs/pypaimon/ray-data.md | 10 +-- .../pypaimon/ray/data_evolution_merge_join.py | 2 + .../pypaimon/ray/update_by_row_id.py | 21 ++++-- .../tests/ray_update_by_row_id_test.py | 71 ++++++++++++++++++- 4 files changed, 92 insertions(+), 12 deletions(-) diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md index 7c90c5718e2c..13becf26491c 100644 --- a/docs/docs/pypaimon/ray-data.md +++ b/docs/docs/pypaimon/ray-data.md @@ -551,11 +551,13 @@ print(metrics) # {"num_updated": 50} - Partition columns cannot be updated (in-place rewrite can't move a row across partitions). - Deletion-vectors-enabled tables are not supported yet: a DV-deleted row still lives in its data file, so it can't be told apart from a live row without reading the target. -- Incremental commits are not atomic across the whole operation. If a group fails, +- Incremental commits are not atomic across the whole operation. Ordinary exceptions + raised while applying a group are reported as results so other groups can finish; completed groups are flushed before the error is raised. -- Incremental mode reports group exceptions as results so other groups can finish. - Ray task retry options in `ray_remote_args` do not retry these exceptions; a - transient group failure can therefore leave a partial update. +- Ray task retry options in `ray_remote_args` do not retry group exceptions reported + as results, so a transient group failure can leave a partial update. Worker loss and + other task-level failures which bypass Python exception handling may still lose + results buffered by that task. ## Read By Row Id diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index c189961cb153..4f1a64064e8e 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -639,6 +639,8 @@ def _apply_group(group: pa.Table) -> pa.Table: for blob, n, row_ids_blob, error in zip( message_blobs, updated_counts, row_id_blobs, errors): if error is not None: + # Keep the first failure as the primary error, but continue + # draining results so successful groups can still be committed. if group_error is None: group_error = error continue diff --git a/paimon-python/pypaimon/ray/update_by_row_id.py b/paimon-python/pypaimon/ray/update_by_row_id.py index 4350bc60e122..2335d3a67526 100644 --- a/paimon-python/pypaimon/ray/update_by_row_id.py +++ b/paimon-python/pypaimon/ray/update_by_row_id.py @@ -71,9 +71,11 @@ def update_by_row_id( By default, all file groups are committed atomically. Set ``max_groups_per_commit`` to commit completed groups in smaller windows. - If a group fails, completed groups are flushed before the error is raised. - Group exceptions are returned as results in incremental mode, so Ray task - retry options in ``ray_remote_args`` do not retry them. + Ordinary exceptions raised while applying a group are returned as results + in incremental mode, so completed groups are flushed before the error is + raised and Ray task retry options in ``ray_remote_args`` do not retry the + group. Worker loss and other task-level failures which bypass Python + exception handling may still lose results buffered by that task. Returns ``{"num_updated": }``. """ @@ -205,14 +207,23 @@ def __init__(self, table, max_groups_per_commit: int): self._table_commit = None self._next_commit_identifier = 1 self._commit_failed = False + self._deferred_commit_error = None def add_group(self, commit_messages, _num_updated, _row_ids) -> None: self._pending_messages.extend(commit_messages) self._pending_groups += 1 - if self._pending_groups >= self._max_groups_per_commit: - self._commit_pending() + if (self._deferred_commit_error is None + and self._pending_groups >= self._max_groups_per_commit): + try: + self._commit_pending() + except Exception as error: + # Keep draining materialized Ray results so their staged files + # are known to the driver and can be aborted by the caller. + self._deferred_commit_error = error def finish(self) -> None: + if self._deferred_commit_error is not None: + raise self._deferred_commit_error self._commit_pending() def _commit_pending(self) -> None: diff --git a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py index 2d2c51aee916..45a412afe001 100644 --- a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py +++ b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py @@ -88,6 +88,16 @@ def _rowid_by_id(self, target): tab = self._read(target, ["_ROW_ID", "id"]) return dict(zip(tab.column("id").to_pylist(), tab.column("_ROW_ID").to_pylist())) + @staticmethod + def _data_files_under(table): + table_path = table.file_io.to_filesystem_path(table.table_path) + return { + os.path.join(root, file_name) + for root, _, files in os.walk(table_path) + for file_name in files + if file_name.endswith(".parquet") + } + def test_update_by_row_id_basic(self): target = self._create() self._write(target, pa.Table.from_pydict( @@ -313,14 +323,16 @@ def new_stream_write_builder(self): self.assertEqual([["group-3"]], aborted) self.assertEqual([True], close_calls) - def test_incremental_commit_failure_does_not_abort_unknown_outcome(self): + def test_incremental_commit_failure_aborts_later_groups(self): import importlib module = importlib.import_module("pypaimon.ray.update_by_row_id") aborted = [] + commit_calls = [] class FakeCommit: def commit(self, messages, commit_identifier): + commit_calls.append((list(messages), commit_identifier)) raise RuntimeError("commit failed") def close(self): @@ -339,12 +351,65 @@ def new_stream_write_builder(self): module, "_abort_pending_update_messages", side_effect=lambda table, messages: aborted.append(list(messages))): + committer.add_group(["group-1"], 1, []) + committer.add_group(["group-2"], 1, []) with self.assertRaisesRegex(RuntimeError, "commit failed"): - committer.add_group(["group-1"], 1, []) + committer.finish() committer.abort_pending() committer.close() - self.assertEqual([], aborted) + self.assertEqual([(["group-1"], 1)], commit_calls) + self.assertEqual([["group-2"]], aborted) + + def test_incremental_commit_conflict_aborts_buffered_group_files(self): + from pypaimon.write.commit.conflict_detection import CommitConflictError + from pypaimon.write.file_store_commit import FileStoreCommit + + target = self._create() + for row_id in range(1, 4): + self._write(target, pa.Table.from_pydict( + {"id": [row_id], "name": ["a"], "age": [0]}, + schema=self.pa_schema, + )) + + table = self.catalog.get_table(target) + base_snapshot_id = table.snapshot_manager().get_latest_snapshot().id + files_before = self._data_files_under(table) + row_ids = self._rowid_by_id(target) + source = pa.table( + { + "_ROW_ID": [row_ids[1], row_ids[2], row_ids[3]], + "age": [100, 200, 300], + }, + schema=pa.schema([ + ("_ROW_ID", pa.int64()), + ("age", pa.int32()), + ]), + ) + + with mock.patch.object( + FileStoreCommit, + "commit", + side_effect=CommitConflictError("forced conflict")): + with self.assertRaisesRegex(CommitConflictError, "forced conflict"): + update_by_row_id( + target, + ray.data.from_arrow(source), + self.catalog_options, + update_cols=["age"], + num_partitions=1, + max_groups_per_commit=1, + ) + + self.assertEqual(files_before, self._data_files_under(table)) + self.assertEqual( + base_snapshot_id, + table.snapshot_manager().get_latest_snapshot().id, + ) + self.assertEqual( + [0, 0, 0], + self._read(target).sort_by("id")["age"].to_pylist(), + ) def test_pins_base_snapshot_for_conflict_detection(self): # The update pins its base snapshot and threads it to distributed_update_apply, From a053ddd69b53619e7e13e3ec93e0cf23fbcfc2f4 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 29 Jul 2026 07:48:08 -0700 Subject: [PATCH 03/11] [python][ray] Make incremental commit mode explicit --- docs/docs/pypaimon/ray-data.md | 37 ++++++++++---- .../pypaimon/ray/data_evolution_merge_join.py | 50 +++++++++---------- .../pypaimon/ray/update_by_row_id.py | 23 ++++++--- .../tests/ray_update_by_row_id_test.py | 42 ++++++++++++++++ 4 files changed, 109 insertions(+), 43 deletions(-) diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md index 13becf26491c..1a98b9444ba5 100644 --- a/docs/docs/pypaimon/ray-data.md +++ b/docs/docs/pypaimon/ray-data.md @@ -524,11 +524,24 @@ metrics = update_by_row_id( source=ray_dataset, # ray.data.Dataset / pa.Table / pandas, carrying _ROW_ID catalog_options={"warehouse": "/path/to/warehouse"}, update_cols=["feature"], # non-blob columns to overwrite - max_groups_per_commit=64, # optional incremental commit window ) print(metrics) # {"num_updated": 50} ``` +The default is one atomic commit. To explicitly opt into non-atomic, +incremental commits: + +```python +metrics = update_by_row_id( + target="database_name.table_name", + source=ray_dataset, + catalog_options={"warehouse": "/path/to/warehouse"}, + update_cols=["feature"], + commit_mode="incremental", + max_groups_per_commit=64, +) +``` + **Parameters:** - `source`: a `ray.data.Dataset`, `pyarrow.Table`, or `pandas.DataFrame` carrying the target `_ROW_ID` and every column in `update_cols`; extra columns are ignored, and @@ -538,8 +551,10 @@ print(metrics) # {"num_updated": 50} - `num_partitions`: parallelism for grouping the update rows by target file; defaults to `max(1, cluster_cpus * 2)`. - `ray_remote_args`: Ray remote options applied to the update tasks. -- `max_groups_per_commit`: optional positive number of completed target file groups - per commit. By default, all groups are committed together. +- `commit_mode`: `"atomic"` (default) commits all groups together. + `"incremental"` explicitly opts into non-atomic, windowed commits. +- `max_groups_per_commit`: required in incremental mode; positive number of + completed target file groups per commit. It is rejected in atomic mode. **Returns:** `{"num_updated": }`. @@ -551,13 +566,15 @@ print(metrics) # {"num_updated": 50} - Partition columns cannot be updated (in-place rewrite can't move a row across partitions). - Deletion-vectors-enabled tables are not supported yet: a DV-deleted row still lives in its data file, so it can't be told apart from a live row without reading the target. -- Incremental commits are not atomic across the whole operation. Ordinary exceptions - raised while applying a group are reported as results so other groups can finish; - completed groups are flushed before the error is raised. -- Ray task retry options in `ray_remote_args` do not retry group exceptions reported - as results, so a transient group failure can leave a partial update. Worker loss and - other task-level failures which bypass Python exception handling may still lose - results buffered by that task. +- Incremental commits are not atomic across the whole operation. Commits completed + before a later failure remain visible. +- A duplicate `_ROW_ID` within a target file group is a deterministic validation + error. In incremental mode it is reported as a result so successful groups can + finish and be committed before the error is raised. +- Other application and I/O exceptions propagate to Ray and remain subject to the + retry options in `ray_remote_args`. If retries are exhausted, earlier incremental + commits remain visible. Worker loss may also discard successful results still + buffered in the failed Ray task. ## Read By Row Id diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index 4f1a64064e8e..d19b942dce27 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -562,7 +562,7 @@ def _assign_frid(batch: pa.Table) -> pa.Table: captured_table = table captured_cols = cols - capture_group_errors = on_group_result is not None + capture_group_validation_errors = on_group_result is not None def _group_result(msgs_blob, n_updated, row_ids_blob, error): return pa.Table.from_pydict({ @@ -581,34 +581,32 @@ def _apply_group(group: pa.Table) -> pa.Table: "error": pa.array([], type=pa.string()), }) - try: - if ( - pc.count_distinct(group.column(row_id_name)).as_py() - != group.num_rows - ): - raise ValueError( - "MERGE matched multiple source rows to the same " - "target _ROW_ID. Deduplicate the source before " - "merging." - ) - - for_update = group.drop_columns([frid_col]) - row_ids = ( - for_update.column(row_id_name).to_pylist() - if collect_row_ids else [] - ) - worker = TableUpdateByRowId( - captured_table, - "_merge_into_shard_" + uuid.uuid4().hex[:8], - BATCH_COMMIT_IDENTIFIER, - _precomputed_files_info=ray.get(precomputed_info_ref), + if ( + pc.count_distinct(group.column(row_id_name)).as_py() + != group.num_rows + ): + error = ValueError( + "MERGE matched multiple source rows to the same " + "target _ROW_ID. Deduplicate the source before " + "merging." ) - msgs = worker.update_columns(for_update, list(captured_cols)) - except Exception as error: - if not capture_group_errors: - raise + if not capture_group_validation_errors: + raise error return _group_result(b"", 0, b"", _group_error_text(error)) + for_update = group.drop_columns([frid_col]) + row_ids = ( + for_update.column(row_id_name).to_pylist() + if collect_row_ids else [] + ) + worker = TableUpdateByRowId( + captured_table, + "_merge_into_shard_" + uuid.uuid4().hex[:8], + BATCH_COMMIT_IDENTIFIER, + _precomputed_files_info=ray.get(precomputed_info_ref), + ) + msgs = worker.update_columns(for_update, list(captured_cols)) + return _group_result( pickle.dumps(msgs), for_update.num_rows, diff --git a/paimon-python/pypaimon/ray/update_by_row_id.py b/paimon-python/pypaimon/ray/update_by_row_id.py index 2335d3a67526..acafd27de4e2 100644 --- a/paimon-python/pypaimon/ray/update_by_row_id.py +++ b/paimon-python/pypaimon/ray/update_by_row_id.py @@ -59,6 +59,7 @@ def update_by_row_id( update_cols: List[str], num_partitions: Optional[int] = None, ray_remote_args: Optional[Dict[str, Any]] = None, + commit_mode: str = "atomic", max_groups_per_commit: Optional[int] = None, ) -> Dict[str, int]: """Update ``update_cols`` of a data-evolution table by ``_ROW_ID``. @@ -70,12 +71,12 @@ def update_by_row_id( ``ray >= 2.50`` and a target with ``data-evolution.enabled`` + ``row-tracking.enabled``. By default, all file groups are committed atomically. Set - ``max_groups_per_commit`` to commit completed groups in smaller windows. - Ordinary exceptions raised while applying a group are returned as results - in incremental mode, so completed groups are flushed before the error is - raised and Ray task retry options in ``ray_remote_args`` do not retry the - group. Worker loss and other task-level failures which bypass Python - exception handling may still lose results buffered by that task. + ``commit_mode="incremental"`` together with ``max_groups_per_commit`` to + commit completed groups in smaller windows. Incremental mode is not atomic: + commits made before a later failure remain visible. Deterministic duplicate + ``_ROW_ID`` errors are returned as group results so completed groups can be + flushed first. Other application errors propagate to Ray and remain subject + to the retry options in ``ray_remote_args``. Returns ``{"num_updated": }``. """ @@ -86,6 +87,14 @@ def update_by_row_id( _require_ray_join() if not update_cols: raise ValueError("update_cols must be non-empty.") + if commit_mode not in ("atomic", "incremental"): + raise ValueError("commit_mode must be 'atomic' or 'incremental'.") + if commit_mode == "atomic" and max_groups_per_commit is not None: + raise ValueError( + "max_groups_per_commit requires commit_mode='incremental'.") + if commit_mode == "incremental" and max_groups_per_commit is None: + raise ValueError( + "commit_mode='incremental' requires max_groups_per_commit.") if max_groups_per_commit is not None: if (isinstance(max_groups_per_commit, bool) or not isinstance(max_groups_per_commit, int) @@ -158,7 +167,7 @@ def _project_cast(batch: pa.Table) -> pa.Table: return {"num_updated": 0} incremental_committer = ( _IncrementalUpdateCommitter(table, max_groups_per_commit) - if max_groups_per_commit is not None else None + if commit_mode == "incremental" else None ) try: apply_kwargs = { diff --git a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py index 45a412afe001..b8851a485994 100644 --- a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py +++ b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py @@ -177,6 +177,7 @@ def record_commit(stream_commit, messages, commit_identifier): self.catalog_options, update_cols=["age"], num_partitions=4, + commit_mode="incremental", max_groups_per_commit=2, ) @@ -228,6 +229,7 @@ def test_group_failure_preserves_completed_groups(self): self.catalog_options, update_cols=["age"], num_partitions=1, + commit_mode="incremental", max_groups_per_commit=max_groups, ) @@ -398,6 +400,7 @@ def test_incremental_commit_conflict_aborts_buffered_group_files(self): self.catalog_options, update_cols=["age"], num_partitions=1, + commit_mode="incremental", max_groups_per_commit=1, ) @@ -637,9 +640,47 @@ def test_rejects_invalid_max_groups_per_commit(self): source, self.catalog_options, update_cols=["age"], + commit_mode="incremental", max_groups_per_commit=value, ) + def test_requires_explicit_incremental_commit_mode(self): + target = self._create() + self._write(target, pa.Table.from_pydict( + {"id": [1], "name": ["a"], "age": [1]}, schema=self.pa_schema)) + source = pa.table( + {"_ROW_ID": [0], "age": [9]}, + schema=pa.schema([("_ROW_ID", pa.int64()), ("age", pa.int32())]), + ) + + with self.assertRaisesRegex( + ValueError, "requires commit_mode='incremental'"): + update_by_row_id( + target, + source, + self.catalog_options, + update_cols=["age"], + max_groups_per_commit=1, + ) + with self.assertRaisesRegex( + ValueError, "requires max_groups_per_commit"): + update_by_row_id( + target, + source, + self.catalog_options, + update_cols=["age"], + commit_mode="incremental", + ) + with self.assertRaisesRegex( + ValueError, "must be 'atomic' or 'incremental'"): + update_by_row_id( + target, + source, + self.catalog_options, + update_cols=["age"], + commit_mode="unknown", + ) + def _run_with_fake_commit(self, *, recorder=None, new_commit_errors=None, commit_error=None, close_error=None, incremental=False): @@ -758,6 +799,7 @@ def fake_apply(*args, **kwargs): FakeSource(), self.catalog_options, update_cols=["age"], + commit_mode="incremental" if incremental else "atomic", max_groups_per_commit=1 if incremental else None, ) return recorder From 9a3138a03b2b1ac6f537165eb2ab2fd23399139f Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 30 Jul 2026 05:01:12 -0700 Subject: [PATCH 04/11] [python][ray] Resume incremental row-id updates from checkpoints --- docs/docs/pypaimon/ray-data.md | 41 ++ paimon-python/pypaimon/ray/__init__.py | 6 +- .../pypaimon/ray/data_evolution_merge_join.py | 60 +- .../pypaimon/ray/update_by_row_id.py | 603 +++++++++++++++++- .../pypaimon/tests/file_store_commit_test.py | 7 + .../tests/ray_update_by_row_id_test.py | 391 +++++++++++- .../pypaimon/write/commit/commit_scanner.py | 71 ++- .../write/commit/conflict_detection.py | 97 +++ .../pypaimon/write/file_store_commit.py | 36 ++ paimon-python/pypaimon/write/table_commit.py | 23 + 10 files changed, 1290 insertions(+), 45 deletions(-) diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md index 1a98b9444ba5..ddcd79cd9c1f 100644 --- a/docs/docs/pypaimon/ray-data.md +++ b/docs/docs/pypaimon/ray-data.md @@ -539,6 +539,23 @@ metrics = update_by_row_id( update_cols=["feature"], commit_mode="incremental", max_groups_per_commit=64, + operation_id="feature-backfill-2026-07", +) +``` + +If the call fails, fix the failing input or infrastructure issue and invoke it +again with the same `operation_id`. The checkpoint is stored with the table, so +the new driver skips target file groups committed by the previous invocation: + +```python +metrics = update_by_row_id( + target="database_name.table_name", + source=ray_dataset, + catalog_options={"warehouse": "/path/to/warehouse"}, + update_cols=["feature"], + commit_mode="incremental", + max_groups_per_commit=64, + operation_id="feature-backfill-2026-07", # resume ) ``` @@ -555,6 +572,10 @@ metrics = update_by_row_id( `"incremental"` explicitly opts into non-atomic, windowed commits. - `max_groups_per_commit`: required in incremental mode; positive number of completed target file groups per commit. It is rejected in atomic mode. +- `operation_id`: required in incremental mode. It identifies one immutable + update input and its durable progress across driver restarts. Do not run two + callers concurrently with the same id or reuse an id for different values or + `update_cols`. **Returns:** `{"num_updated": }`. @@ -575,6 +596,26 @@ metrics = update_by_row_id( retry options in `ray_remote_args`. If retries are exhausted, earlier incremental commits remain visible. Worker loss may also discard successful results still buffered in the failed Ray task. +- Each incremental snapshot stores the cumulative completed row-id ranges, and + internal Paimon tags retain the latest checkpoint across normal snapshot + expiration. A retry with the same `operation_id` filters those ranges before + target file-group rewrite and returns the cumulative `num_updated`. +- The checkpoint avoids repeating Paimon's target file read/rewrite work. It + cannot checkpoint arbitrary transformations upstream of `update_by_row_id`. + Persist a multi-day or nondeterministic source first (for example to Parquet or + a staging Paimon table), then pass the persisted `(_ROW_ID, values)` dataset to + every attempt. +- Keep the checkpoint until a successful result has been acknowledged. Afterwards + it can be removed explicitly: + ```python + from pypaimon.ray import delete_update_by_row_id_checkpoint + + delete_update_by_row_id_checkpoint( + "database_name.table_name", + {"warehouse": "/path/to/warehouse"}, + "feature-backfill-2026-07", + ) + ``` ## Read By Row Id diff --git a/paimon-python/pypaimon/ray/__init__.py b/paimon-python/pypaimon/ray/__init__.py index bc91c8da4570..84e662158a8f 100644 --- a/paimon-python/pypaimon/ray/__init__.py +++ b/paimon-python/pypaimon/ray/__init__.py @@ -27,7 +27,10 @@ target_col, lit, ) -from pypaimon.ray.update_by_row_id import update_by_row_id +from pypaimon.ray.update_by_row_id import ( + delete_update_by_row_id_checkpoint, + update_by_row_id, +) from pypaimon.ray.read_by_row_id import read_by_row_id __all__ = [ @@ -37,6 +40,7 @@ "bucket_join", "merge_into", "update_by_row_id", + "delete_update_by_row_id_checkpoint", "read_by_row_id", "WhenMatched", "WhenNotMatched", diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index d19b942dce27..6db1bb20184e 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -460,7 +460,8 @@ def distributed_update_apply( ray_remote_args: Optional[Dict[str, Any]] = None, base_snapshot_id: Optional[int] = None, collect_row_ids: bool = False, - on_group_result: Optional[Callable[[list, int, list], None]] = None, + on_group_result: Optional[Callable[[list, int, list, list], None]] = None, + skip_row_id_ranges: Optional[Sequence[Tuple[int, int]]] = None, ) -> Tuple[list, int, list]: import numpy as np import pickle @@ -471,6 +472,9 @@ def distributed_update_apply( from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER from pypaimon.table.special_fields import SpecialFields + from pypaimon.write.commit.conflict_detection import ( + row_id_file_signature, + ) from pypaimon.write.table_update_by_row_id import TableUpdateByRowId row_id_name = SpecialFields.ROW_ID.name @@ -524,6 +528,7 @@ def distributed_update_apply( valid_ranges = planner.valid_row_id_ranges range_starts = np.asarray([r.from_ for r in valid_ranges], dtype=np.int64) range_ends = np.asarray([r.to for r in valid_ranges], dtype=np.int64) + completed_ranges = tuple(skip_row_id_ranges or ()) def _assign_frid(batch: pa.Table) -> pa.Table: if batch.num_rows == 0: @@ -553,9 +558,17 @@ def _assign_frid(batch: pa.Table) -> pa.Table: captured_sorted_arr, rids, side="right" ) - 1 frids = captured_sorted_arr[idx] - return batch.append_column( + with_first_row_id = batch.append_column( frid_col, pa.array(frids, type=pa.int64()) ) + if not completed_ranges: + return with_first_row_id + completed = np.zeros(len(rids), dtype=bool) + for start, end in completed_ranges: + completed |= (rids >= start) & (rids <= end) + if not completed.any(): + return with_first_row_id + return with_first_row_id.filter(pa.array(~completed)) map_kwargs = _map_kwargs(ray_remote_args) with_frid = update_ds.map_batches(_assign_frid, **map_kwargs) @@ -564,11 +577,18 @@ def _assign_frid(batch: pa.Table) -> pa.Table: captured_cols = cols capture_group_validation_errors = on_group_result is not None - def _group_result(msgs_blob, n_updated, row_ids_blob, error): + def _group_result( + msgs_blob, + n_updated, + row_ids_blob, + planned_files_blob, + error): return pa.Table.from_pydict({ "msgs_blob": pa.array([msgs_blob], type=pa.binary()), "n_updated": pa.array([n_updated], type=pa.int64()), "row_ids_blob": pa.array([row_ids_blob], type=pa.binary()), + "planned_files_blob": pa.array( + [planned_files_blob], type=pa.binary()), "error": pa.array([error], type=pa.string()), }) @@ -578,6 +598,7 @@ def _apply_group(group: pa.Table) -> pa.Table: "msgs_blob": pa.array([], type=pa.binary()), "n_updated": pa.array([], type=pa.int64()), "row_ids_blob": pa.array([], type=pa.binary()), + "planned_files_blob": pa.array([], type=pa.binary()), "error": pa.array([], type=pa.string()), }) @@ -592,18 +613,27 @@ def _apply_group(group: pa.Table) -> pa.Table: ) if not capture_group_validation_errors: raise error - return _group_result(b"", 0, b"", _group_error_text(error)) + return _group_result( + b"", 0, b"", b"", _group_error_text(error)) for_update = group.drop_columns([frid_col]) row_ids = ( for_update.column(row_id_name).to_pylist() if collect_row_ids else [] ) + files_info = ray.get(precomputed_info_ref) + first_row_id = group.column(frid_col)[0].as_py() + split, target_files = files_info.first_row_id_index[first_row_id] + planned_file_signatures = [ + row_id_file_signature( + split.partition, split.bucket, data_file) + for data_file in target_files + ] worker = TableUpdateByRowId( captured_table, "_merge_into_shard_" + uuid.uuid4().hex[:8], BATCH_COMMIT_IDENTIFIER, - _precomputed_files_info=ray.get(precomputed_info_ref), + _precomputed_files_info=files_info, ) msgs = worker.update_columns(for_update, list(captured_cols)) @@ -611,6 +641,7 @@ def _apply_group(group: pa.Table) -> pa.Table: pickle.dumps(msgs), for_update.num_rows, pickle.dumps(row_ids), + pickle.dumps(planned_file_signatures), None, ) @@ -630,12 +661,18 @@ def _apply_group(group: pa.Table) -> pa.Table: message_blobs = batch.column("msgs_blob").to_pylist() updated_counts = batch.column("n_updated").to_pylist() errors = batch.column("error").to_pylist() + planned_file_blobs = ( + batch.column("planned_files_blob").to_pylist()) row_id_blobs = ( batch.column("row_ids_blob").to_pylist() if collect_row_ids else [None] * len(message_blobs) ) - for blob, n, row_ids_blob, error in zip( - message_blobs, updated_counts, row_id_blobs, errors): + for blob, n, row_ids_blob, planned_files_blob, error in zip( + message_blobs, + updated_counts, + row_id_blobs, + planned_file_blobs, + errors): if error is not None: # Keep the first failure as the primary error, but continue # draining results so successful groups can still be committed. @@ -646,10 +683,17 @@ def _apply_group(group: pa.Table) -> pa.Table: group_row_ids = ( pickle.loads(row_ids_blob) if collect_row_ids else [] ) + planned_file_signatures = pickle.loads( + planned_files_blob) if on_group_result is None: all_msgs.extend(group_msgs) else: - on_group_result(group_msgs, n, group_row_ids) + on_group_result( + group_msgs, + n, + group_row_ids, + planned_file_signatures, + ) num_updated += n if collect_row_ids: action_row_ids.extend(group_row_ids) diff --git a/paimon-python/pypaimon/ray/update_by_row_id.py b/paimon-python/pypaimon/ray/update_by_row_id.py index acafd27de4e2..6caa64adb6a1 100644 --- a/paimon-python/pypaimon/ray/update_by_row_id.py +++ b/paimon-python/pypaimon/ray/update_by_row_id.py @@ -23,6 +23,8 @@ target). Pairs with ``bucket_join``, which produces the row ids. """ +import hashlib +import json import logging from typing import Any, Dict, List, Optional @@ -40,11 +42,19 @@ ) from pypaimon.ray.data_evolution_merge_transform import build_update_schema from pypaimon.schema.data_types import is_blob_file_field +from pypaimon.utils.range import Range -__all__ = ["update_by_row_id"] +__all__ = [ + "update_by_row_id", + "delete_update_by_row_id_checkpoint", +] logger = logging.getLogger(__name__) +_CHECKPOINT_PROPERTY = "pypaimon.ray.update-by-row-id.checkpoint" +_CHECKPOINT_TAG_PREFIX = "_pypaimon_ray_update_" +_CHECKPOINT_VERSION = 1 + def _blob_col_names(table: "FileStoreTable") -> set: return {f.name for f in table.table_schema.fields @@ -61,6 +71,7 @@ def update_by_row_id( ray_remote_args: Optional[Dict[str, Any]] = None, commit_mode: str = "atomic", max_groups_per_commit: Optional[int] = None, + operation_id: Optional[str] = None, ) -> Dict[str, int]: """Update ``update_cols`` of a data-evolution table by ``_ROW_ID``. @@ -76,7 +87,9 @@ def update_by_row_id( commits made before a later failure remain visible. Deterministic duplicate ``_ROW_ID`` errors are returned as group results so completed groups can be flushed first. Other application errors propagate to Ray and remain subject - to the retry options in ``ray_remote_args``. + to the retry options in ``ray_remote_args``. Incremental mode requires a + stable ``operation_id``. A later invocation with the same id reads its + durable checkpoint and skips file groups which are already committed. Returns ``{"num_updated": }``. """ @@ -92,9 +105,20 @@ def update_by_row_id( if commit_mode == "atomic" and max_groups_per_commit is not None: raise ValueError( "max_groups_per_commit requires commit_mode='incremental'.") + if commit_mode == "atomic" and operation_id is not None: + raise ValueError( + "operation_id requires commit_mode='incremental'.") if commit_mode == "incremental" and max_groups_per_commit is None: raise ValueError( "commit_mode='incremental' requires max_groups_per_commit.") + if commit_mode == "incremental" and operation_id is None: + raise ValueError( + "commit_mode='incremental' requires operation_id.") + if operation_id is not None: + if not isinstance(operation_id, str) or not operation_id.strip(): + raise ValueError("operation_id must be a non-empty string.") + if len(operation_id) > 256: + raise ValueError("operation_id must contain at most 256 characters.") if max_groups_per_commit is not None: if (isinstance(max_groups_per_commit, bool) or not isinstance(max_groups_per_commit, int) @@ -103,7 +127,8 @@ def update_by_row_id( update_cols = list(dict.fromkeys(update_cols)) # de-dup, keep order num_partitions = _resolve_num_partitions(num_partitions) - table = CatalogFactory.create(catalog_options).get_table(target) + catalog = CatalogFactory.create(catalog_options) + table = catalog.get_table(target) if not table.options.data_evolution_enabled(): raise ValueError( f"update_by_row_id requires 'data-evolution.enabled'='true' on '{target}'.") @@ -157,18 +182,35 @@ def _project_cast(batch: pa.Table) -> pa.Table: update_ds = source_ds.map_batches(_project_cast, batch_format="pyarrow") base = table.snapshot_manager().get_latest_snapshot() + incremental_committer = ( + _IncrementalUpdateCommitter( + table, + max_groups_per_commit, + operation_id=operation_id, + update_cols=update_cols, + initial_snapshot=base, + catalog=catalog, + target=target, + ) + if commit_mode == "incremental" else None + ) # Without deletion vectors (rejected above), total_record_count is the live row # count, so 0 means the target is empty (never written, or emptied by overwrite). if base is None or base.total_record_count == 0: # Every source row id is foreign; don't silently no-op non-empty input. - if update_ds.limit(1).count() > 0: - raise ValueError( - f"target '{target}' has no rows; every _ROW_ID in the source is foreign.") - return {"num_updated": 0} - incremental_committer = ( - _IncrementalUpdateCommitter(table, max_groups_per_commit) - if commit_mode == "incremental" else None - ) + try: + if update_ds.limit(1).count() > 0: + raise ValueError( + f"target '{target}' has no rows; every _ROW_ID in the " + "source is foreign.") + return { + "num_updated": ( + incremental_committer.num_updated + if incremental_committer is not None else 0) + } + finally: + if incremental_committer is not None: + incremental_committer.close() try: apply_kwargs = { "num_partitions": num_partitions, @@ -177,6 +219,8 @@ def _project_cast(batch: pa.Table) -> pa.Table: } if incremental_committer is not None: apply_kwargs["on_group_result"] = incremental_committer.add_group + apply_kwargs["skip_row_id_ranges"] = ( + incremental_committer.completed_row_id_ranges) msgs, num_updated, _ = distributed_update_apply( update_ds, table, update_cols, **apply_kwargs ) @@ -203,24 +247,118 @@ def _project_cast(batch: pa.Table) -> pa.Table: if incremental_committer is None and msgs: _commit_update_messages(table, msgs) + if incremental_committer is not None: + num_updated = incremental_committer.num_updated return {"num_updated": num_updated} class _IncrementalUpdateCommitter: - def __init__(self, table, max_groups_per_commit: int): + def __init__( + self, + table, + max_groups_per_commit: int, + operation_id: Optional[str] = None, + update_cols: Optional[List[str]] = None, + initial_snapshot=None, + catalog=None, + target=None): self._table = table + self._catalog = catalog + self._target = target self._max_groups_per_commit = max_groups_per_commit + self._operation_id = operation_id + self._update_cols = list(update_cols or []) self._pending_messages: list = [] self._pending_groups = 0 + self._pending_num_updated = 0 + self._pending_planned_file_signatures = set() self._table_commit = None self._next_commit_identifier = 1 self._commit_failed = False self._deferred_commit_error = None + self._checkpoint_snapshot = None + self._completed_row_id_ranges = [] + self._completed_num_updated = 0 + self._commit_user = None + self._checkpoint_tags = None + self._planner_tag = None + + if self._operation_id is not None: + if self._catalog is None or self._target is None: + raise ValueError( + "catalog and target are required for resumable updates.") + self._commit_user = _operation_commit_user(self._operation_id) + self._checkpoint_tags = _operation_checkpoint_tags( + self._operation_id) + self._planner_tag = _operation_planner_tag( + self._operation_id) + checkpoint = _load_operation_checkpoint( + self._catalog, + self._target, + self._table, + self._operation_id, + self._update_cols, + self._commit_user, + self._checkpoint_tags, + ) + if checkpoint is None and initial_snapshot is not None: + checkpoint = _initialize_operation_checkpoint( + self._catalog, + self._target, + self._table, + initial_snapshot, + self._operation_id, + self._update_cols, + self._checkpoint_tags, + ) + if checkpoint is not None: + snapshot, state = checkpoint + self._checkpoint_snapshot = snapshot + self._completed_row_id_ranges = _decode_ranges( + state["completed_ranges"]) + self._completed_num_updated = int( + state["num_updated"]) + if _checkpoint_state(snapshot) is not None: + self._next_commit_identifier = ( + snapshot.commit_identifier + 1) + if initial_snapshot is not None: + _store_operation_planner( + self._catalog, + self._target, + self._planner_tag, + initial_snapshot, + ) + try: + self._ensure_table_commit() + self._configure_protected_scope() + self._table_commit.validate_protected_row_id_ranges() + except Exception: + self.close() + raise - def add_group(self, commit_messages, _num_updated, _row_ids) -> None: + @property + def completed_row_id_ranges(self): + return [ + (row_range.from_, row_range.to) + for row_range in self._completed_row_id_ranges + ] + + @property + def num_updated(self): + return self._completed_num_updated + + def add_group( + self, + commit_messages, + num_updated, + _row_ids, + planned_file_signatures=()) -> None: self._pending_messages.extend(commit_messages) self._pending_groups += 1 + self._pending_num_updated += num_updated + self._pending_planned_file_signatures.update( + planned_file_signatures) if (self._deferred_commit_error is None and self._pending_groups >= self._max_groups_per_commit): try: @@ -234,29 +372,105 @@ def finish(self) -> None: if self._deferred_commit_error is not None: raise self._deferred_commit_error self._commit_pending() + if (self._operation_id is not None + and self._table_commit is not None): + self._table_commit.validate_protected_row_id_ranges() def _commit_pending(self) -> None: if self._pending_groups == 0: return if not self._pending_messages: self._pending_groups = 0 + self._pending_num_updated = 0 + self._pending_planned_file_signatures.clear() return + messages = self._pending_messages + group_count = self._pending_groups + pending_num_updated = self._pending_num_updated + pending_ranges = [] + candidate_ranges = self._completed_row_id_ranges + if self._operation_id is not None: + pending_ranges = _row_id_ranges_from_messages(messages) + candidate_ranges = _merge_row_id_ranges( + self._completed_row_id_ranges, + pending_ranges, + ) + candidate_num_updated = ( + self._completed_num_updated + pending_num_updated) + commit_identifier = self._next_commit_identifier try: - if self._table_commit is None: - self._table_commit = ( - self._table.new_stream_write_builder().new_commit() + self._ensure_table_commit() + self._configure_protected_scope() + if self._operation_id is not None: + latest = self._table.snapshot_manager().get_latest_snapshot() + if latest is None: + raise RuntimeError( + "Target has no snapshot before incremental commit.") + self._table_commit.protect_planned_row_id_files( + pending_ranges, + self._pending_planned_file_signatures, + ) + self._table_commit.with_snapshot_properties( + _checkpoint_properties( + self._operation_id, + self._table.table_schema.id, + self._update_cols, + candidate_ranges, + candidate_num_updated, + ) ) - - messages = self._pending_messages self._pending_messages = [] self._pending_groups = 0 - commit_identifier = self._next_commit_identifier + self._pending_num_updated = 0 + self._pending_planned_file_signatures = set() self._table_commit.commit(messages, commit_identifier) + if self._operation_id is not None: + snapshot = _find_operation_snapshot( + self._table, + self._commit_user, + commit_identifier, + ) + if snapshot is None: + raise RuntimeError( + "Committed resumable update snapshot cannot be found.") + _store_operation_checkpoint( + self._catalog, + self._target, + self._checkpoint_tags, + snapshot, + ) + self._checkpoint_snapshot = snapshot + self._completed_row_id_ranges = candidate_ranges + self._completed_num_updated = candidate_num_updated + self._configure_protected_scope() except Exception: self._commit_failed = True raise self._next_commit_identifier += 1 + logger.info( + "Incrementally committed %d update_by_row_id file groups " + "for operation %s.", + group_count, + self._operation_id, + ) + + def _ensure_table_commit(self): + if self._table_commit is not None: + return + builder = self._table.new_stream_write_builder() + if self._commit_user is not None: + builder.commit_user = self._commit_user + self._table_commit = builder.new_commit() + + def _configure_protected_scope(self): + if (self._operation_id is None + or self._table_commit is None): + return + self._table_commit.protect_row_id_ranges( + self._checkpoint_snapshot, + self._completed_row_id_ranges, + ) def abort_pending(self) -> None: if not self._pending_messages: @@ -264,19 +478,350 @@ def abort_pending(self) -> None: messages = self._pending_messages self._pending_messages = [] self._pending_groups = 0 + self._pending_num_updated = 0 + self._pending_planned_file_signatures.clear() _abort_pending_update_messages(self._table, messages) def close(self) -> None: - if self._table_commit is None: - return - try: - self._table_commit.close() - except Exception as close_error: - logger.warning( - "Failed to close incremental update_by_row_id commit: %s", - close_error, - exc_info=close_error, - ) + if self._table_commit is not None: + try: + self._table_commit.close() + except Exception as close_error: + logger.warning( + "Failed to close incremental update_by_row_id commit: %s", + close_error, + exc_info=close_error, + ) + if self._planner_tag is not None: + try: + _delete_checkpoint_tag( + self._catalog, + self._target, + self._planner_tag, + ignore_missing=True, + ) + except Exception as cleanup_error: + logger.warning( + "Failed to delete update_by_row_id planner tag: %s", + cleanup_error, + exc_info=cleanup_error, + ) + + +def delete_update_by_row_id_checkpoint( + target: str, + catalog_options: Dict[str, str], + operation_id: str) -> bool: + """Delete a completed incremental update's durable resume checkpoint.""" + from pypaimon.catalog.catalog_factory import CatalogFactory + + if not isinstance(operation_id, str) or not operation_id.strip(): + raise ValueError("operation_id must be a non-empty string.") + catalog = CatalogFactory.create(catalog_options) + deleted = False + operation_tags = ( + _operation_checkpoint_tags(operation_id) + + (_operation_planner_tag(operation_id),) + ) + for checkpoint_tag in operation_tags: + deleted = ( + _delete_checkpoint_tag( + catalog, target, checkpoint_tag, ignore_missing=True) + or deleted + ) + return deleted + + +def _operation_digest(operation_id): + return hashlib.sha256(operation_id.encode("utf-8")).hexdigest()[:32] + + +def _operation_commit_user(operation_id): + return "pypaimon-ray-update-" + _operation_digest(operation_id) + + +def _operation_checkpoint_tags(operation_id): + base = _CHECKPOINT_TAG_PREFIX + _operation_digest(operation_id) + return base + "_0", base + "_1" + + +def _operation_planner_tag(operation_id): + return _CHECKPOINT_TAG_PREFIX + _operation_digest(operation_id) + "_plan" + + +def _checkpoint_properties( + operation_id, + schema_id, + update_cols, + completed_ranges, + num_updated): + state = { + "version": _CHECKPOINT_VERSION, + "operation_id": operation_id, + "schema_id": schema_id, + "update_cols": list(update_cols), + "completed_ranges": [ + [row_range.from_, row_range.to] + for row_range in completed_ranges + ], + "num_updated": num_updated, + } + return { + _CHECKPOINT_PROPERTY: json.dumps( + state, sort_keys=True, separators=(",", ":")) + } + + +def _checkpoint_state(snapshot): + properties = snapshot.properties or {} + encoded = properties.get(_CHECKPOINT_PROPERTY) + if encoded is None: + return None + try: + state = json.loads(encoded) + except Exception as error: + raise RuntimeError( + "Invalid resumable update checkpoint JSON.") from error + if state.get("version") != _CHECKPOINT_VERSION: + raise RuntimeError( + "Unsupported resumable update checkpoint version: {}.".format( + state.get("version"))) + required = { + "operation_id", + "schema_id", + "update_cols", + "completed_ranges", + "num_updated", + } + if not required.issubset(state): + raise RuntimeError( + "Incomplete resumable update checkpoint.") + _decode_ranges(state["completed_ranges"]) + if (isinstance(state["num_updated"], bool) + or not isinstance(state["num_updated"], int) + or state["num_updated"] < 0): + raise RuntimeError( + "Invalid resumable update checkpoint row count.") + return state + + +def _validate_checkpoint_state( + snapshot, + operation_id, + update_cols, + schema_id): + state = _checkpoint_state(snapshot) + if state is None: + raise RuntimeError( + "Resumable update checkpoint snapshot has no checkpoint state.") + if state["operation_id"] != operation_id: + raise RuntimeError( + "Resumable update operation id hash collision.") + if state["schema_id"] != schema_id: + raise RuntimeError( + "Target schema changed since resumable update operation " + f"{operation_id!r} started.") + if state["update_cols"] != list(update_cols): + raise ValueError( + "operation_id {!r} was already used with update_cols {}; " + "got {}.".format( + operation_id, state["update_cols"], list(update_cols))) + return state + + +def _load_operation_checkpoint( + catalog, + target, + table, + operation_id, + update_cols, + commit_user, + checkpoint_tags): + snapshot_manager = table.snapshot_manager() + latest = snapshot_manager.get_latest_snapshot() + tagged_snapshots = [] + for checkpoint_tag in checkpoint_tags: + tagged = _get_checkpoint_tag( + catalog, target, checkpoint_tag) + if tagged is not None: + tagged_snapshots.append(tagged.snapshot) + + checkpoint_snapshot = ( + max(tagged_snapshots, key=lambda snapshot: snapshot.id) + if tagged_snapshots else None + ) + if latest is not None: + if checkpoint_snapshot is not None: + lower_bound = checkpoint_snapshot.id + 1 + else: + earliest = snapshot_manager.try_get_earliest_snapshot(latest.id) + lower_bound = earliest.id + for snapshot_id in range(latest.id, lower_bound - 1, -1): + snapshot = snapshot_manager.get_snapshot_by_id(snapshot_id) + if (snapshot is not None + and snapshot.commit_user == commit_user + and _checkpoint_state(snapshot) is not None): + checkpoint_snapshot = snapshot + break + + if checkpoint_snapshot is None: + return None + + state = _checkpoint_state(checkpoint_snapshot) + if state is None: + state = _initial_checkpoint_state( + operation_id, table.table_schema.id, update_cols) + return checkpoint_snapshot, state + if checkpoint_snapshot.commit_user != commit_user: + raise RuntimeError( + "Resumable update checkpoint tag belongs to another writer.") + + state = _validate_checkpoint_state( + checkpoint_snapshot, + operation_id, + update_cols, + table.table_schema.id, + ) + _store_operation_checkpoint( + catalog, target, checkpoint_tags, checkpoint_snapshot) + return checkpoint_snapshot, state + + +def _initialize_operation_checkpoint( + catalog, + target, + table, + snapshot, + operation_id, + update_cols, + checkpoint_tags): + catalog.create_tag( + target, + checkpoint_tags[0], + snapshot_id=snapshot.id, + ) + return snapshot, _initial_checkpoint_state( + operation_id, table.table_schema.id, update_cols) + + +def _initial_checkpoint_state(operation_id, schema_id, update_cols): + return { + "version": _CHECKPOINT_VERSION, + "operation_id": operation_id, + "schema_id": schema_id, + "update_cols": list(update_cols), + "completed_ranges": [], + "num_updated": 0, + } + + +def _find_operation_snapshot(table, commit_user, commit_identifier): + snapshot_manager = table.snapshot_manager() + latest = snapshot_manager.get_latest_snapshot() + if latest is None: + return None + for snapshot_id in range(latest.id, 0, -1): + snapshot = snapshot_manager.get_snapshot_by_id(snapshot_id) + if (snapshot is not None + and snapshot.commit_user == commit_user + and snapshot.commit_identifier == commit_identifier + and _checkpoint_state(snapshot) is not None): + return snapshot + return None + + +def _store_operation_checkpoint( + catalog, target, checkpoint_tags, snapshot): + slot = snapshot.commit_identifier % 2 + checkpoint_tag = checkpoint_tags[slot] + previous = _get_checkpoint_tag(catalog, target, checkpoint_tag) + if previous is None or previous.snapshot.id != snapshot.id: + if previous is not None: + _delete_checkpoint_tag(catalog, target, checkpoint_tag) + catalog.create_tag( + target, checkpoint_tag, snapshot_id=snapshot.id) + _delete_checkpoint_tag( + catalog, target, checkpoint_tags[1 - slot], ignore_missing=True) + + +def _store_operation_planner(catalog, target, planner_tag, snapshot): + previous = _get_checkpoint_tag(catalog, target, planner_tag) + if previous is not None and previous.snapshot.id == snapshot.id: + return + if previous is not None: + _delete_checkpoint_tag(catalog, target, planner_tag) + catalog.create_tag(target, planner_tag, snapshot_id=snapshot.id) + + +def _get_checkpoint_tag(catalog, target, checkpoint_tag): + from pypaimon.catalog.catalog_exception import TagNotExistException + + try: + return catalog.get_tag(target, checkpoint_tag) + except TagNotExistException: + return None + + +def _delete_checkpoint_tag( + catalog, target, checkpoint_tag, ignore_missing=False): + from pypaimon.catalog.catalog_exception import TagNotExistException + + try: + catalog.delete_tag(target, checkpoint_tag) + return True + except TagNotExistException: + if not ignore_missing: + raise + return False + + +def _decode_ranges(encoded_ranges): + ranges = [] + if not isinstance(encoded_ranges, list): + raise RuntimeError( + "Invalid resumable update checkpoint ranges.") + for encoded in encoded_ranges: + if (not isinstance(encoded, list) + or len(encoded) != 2 + or isinstance(encoded[0], bool) + or isinstance(encoded[1], bool) + or not isinstance(encoded[0], int) + or not isinstance(encoded[1], int) + or encoded[0] > encoded[1]): + raise RuntimeError( + "Invalid resumable update checkpoint range.") + ranges.append(Range(encoded[0], encoded[1])) + merged = Range.sort_and_merge_overlap(ranges, True, True) + if len(merged) != len(ranges): + raise RuntimeError( + "Resumable update checkpoint ranges are not canonical.") + return merged + + +def _merge_row_id_ranges(*range_groups): + return Range.sort_and_merge_overlap( + [ + row_range + for ranges in range_groups + for row_range in ranges + ], + True, + True, + ) + + +def _row_id_ranges_from_messages(messages): + ranges = [] + for message in messages: + for data_file in message.new_files: + row_range = data_file.row_id_range() + if row_range is not None: + ranges.append(row_range) + merged = _merge_row_id_ranges(ranges) + if messages and not merged: + raise RuntimeError( + "Cannot checkpoint update messages without row-id ranges.") + return merged def _commit_update_messages(table, commit_messages) -> None: diff --git a/paimon-python/pypaimon/tests/file_store_commit_test.py b/paimon-python/pypaimon/tests/file_store_commit_test.py index 88350bfd5f01..772b5f460d4b 100644 --- a/paimon-python/pypaimon/tests/file_store_commit_test.py +++ b/paimon-python/pypaimon/tests/file_store_commit_test.py @@ -24,6 +24,7 @@ from pypaimon.manifest.schema.manifest_entry import ManifestEntry from pypaimon.snapshot.snapshot_commit import PartitionStatistics from pypaimon.table.row.generic_row import GenericRow +from pypaimon.utils.range import Range from pypaimon.write.commit.row_id_conflict_rewriter import RowIdRewriteResult from pypaimon.write.commit_message import CommitMessage from pypaimon.write.file_store_commit import ( @@ -522,6 +523,8 @@ def test_row_id_rewrite_respects_commit_retry_limit( commit_entry = Mock() rewrite = RewriteResult(RowIdRewriteResult([commit_entry], 1)) + file_store_commit.conflict_detection.protect_planned_row_id_files( + [Range(0, 1)], [("planned",)]) file_store_commit._try_commit_once = Mock(side_effect=[ rewrite, rewrite, @@ -538,6 +541,10 @@ def test_row_id_rewrite_respects_commit_retry_limit( self.assertIn("with 1 retries", str(ctx.exception)) self.assertEqual(2, file_store_commit._try_commit_once.call_count) file_store_commit._commit_retry_wait.assert_called_once_with(0) + self.assertEqual( + [], + file_store_commit.conflict_detection._planned_row_id_ranges, + ) @staticmethod def _to_entries(commit_messages): diff --git a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py index b8851a485994..43c370fb21ae 100644 --- a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py +++ b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py @@ -30,7 +30,10 @@ ray = pytest.importorskip("ray") from pypaimon import CatalogFactory, Schema -from pypaimon.ray import update_by_row_id +from pypaimon.ray import ( + delete_update_by_row_id_checkpoint, + update_by_row_id, +) class RayUpdateByRowIdTest(unittest.TestCase): @@ -166,7 +169,11 @@ def test_incrementally_commits_file_group_windows(self): original_commit = StreamTableCommit.commit def record_commit(stream_commit, messages, commit_identifier): - commits.append((len(messages), commit_identifier)) + commits.append(( + len(messages), + commit_identifier, + {message.check_from_snapshot for message in messages}, + )) return original_commit( stream_commit, messages, commit_identifier) @@ -179,10 +186,18 @@ def record_commit(stream_commit, messages, commit_identifier): num_partitions=4, commit_mode="incremental", max_groups_per_commit=2, + operation_id="windowed-update", ) self.assertEqual({"num_updated": 5}, stats) - self.assertEqual([(2, 1), (2, 2), (1, 3)], commits) + self.assertEqual( + [ + (2, 1, {base_snapshot_id}), + (2, 2, {base_snapshot_id}), + (1, 3, {base_snapshot_id}), + ], + commits, + ) self.assertEqual( base_snapshot_id + 3, table.snapshot_manager().get_latest_snapshot().id, @@ -231,6 +246,7 @@ def test_group_failure_preserves_completed_groups(self): num_partitions=1, commit_mode="incremental", max_groups_per_commit=max_groups, + operation_id="group-failure", ) self.assertEqual( @@ -242,6 +258,317 @@ def test_group_failure_preserves_completed_groups(self): self._read(target).sort_by("id")["age"].to_pylist(), ) + def test_resumes_after_partial_failure_without_rewriting_completed_groups(self): + target = self._create() + for row_id in range(1, 4): + self._write(target, pa.Table.from_pydict( + {"id": [row_id], "name": ["a"], "age": [0]}, + schema=self.pa_schema, + )) + + table = self.catalog.get_table(target) + row_ids = self._rowid_by_id(target) + operation_id = "feature-backfill-" + uuid.uuid4().hex + failing_source = pa.table( + { + "_ROW_ID": [ + row_ids[1], row_ids[2], row_ids[3], row_ids[3], + ], + "age": [100, 200, 300, 301], + }, + schema=pa.schema([ + ("_ROW_ID", pa.int64()), + ("age", pa.int32()), + ]), + ) + + with self.assertRaisesRegex(RuntimeError, "Deduplicate"): + update_by_row_id( + target, + ray.data.from_arrow(failing_source), + self.catalog_options, + update_cols=["age"], + num_partitions=1, + commit_mode="incremental", + max_groups_per_commit=1, + operation_id=operation_id, + ) + + partial_snapshot_id = ( + table.snapshot_manager().get_latest_snapshot().id) + self.assertEqual( + [100, 200, 0], + self._read(target).sort_by("id")["age"].to_pylist(), + ) + with self.assertRaisesRegex(ValueError, "already used"): + update_by_row_id( + target, + pa.table({ + "_ROW_ID": [row_ids[3]], + "name": ["b"], + }), + self.catalog_options, + update_cols=["name"], + num_partitions=1, + commit_mode="incremental", + max_groups_per_commit=1, + operation_id=operation_id, + ) + corrected_source = pa.table( + { + "_ROW_ID": [row_ids[1], row_ids[2], row_ids[3]], + "age": [100, 200, 300], + }, + schema=failing_source.schema, + ) + stats = update_by_row_id( + target, + ray.data.from_arrow(corrected_source), + self.catalog_options, + update_cols=["age"], + num_partitions=1, + commit_mode="incremental", + max_groups_per_commit=1, + operation_id=operation_id, + ) + + self.assertEqual({"num_updated": 3}, stats) + self.assertEqual( + partial_snapshot_id + 1, + table.snapshot_manager().get_latest_snapshot().id, + ) + self.assertEqual( + [100, 200, 300], + self._read(target).sort_by("id")["age"].to_pylist(), + ) + completed_snapshot_id = ( + table.snapshot_manager().get_latest_snapshot().id) + self.assertEqual( + {"num_updated": 3}, + update_by_row_id( + target, + corrected_source, + self.catalog_options, + update_cols=["age"], + num_partitions=1, + commit_mode="incremental", + max_groups_per_commit=1, + operation_id=operation_id, + ), + ) + self.assertEqual( + completed_snapshot_id, + table.snapshot_manager().get_latest_snapshot().id, + ) + self.assertTrue(delete_update_by_row_id_checkpoint( + target, self.catalog_options, operation_id)) + self.assertFalse(delete_update_by_row_id_checkpoint( + target, self.catalog_options, operation_id)) + + def test_resume_rejects_external_change_to_completed_group(self): + target = self._create() + for row_id in range(1, 4): + self._write(target, pa.Table.from_pydict( + {"id": [row_id], "name": ["a"], "age": [0]}, + schema=self.pa_schema, + )) + + row_ids = self._rowid_by_id(target) + operation_id = "protected-backfill-" + uuid.uuid4().hex + failing_source = pa.table( + { + "_ROW_ID": [ + row_ids[1], row_ids[2], row_ids[3], row_ids[3], + ], + "age": [100, 200, 300, 301], + }, + schema=pa.schema([ + ("_ROW_ID", pa.int64()), + ("age", pa.int32()), + ]), + ) + with self.assertRaisesRegex(RuntimeError, "Deduplicate"): + update_by_row_id( + target, + failing_source, + self.catalog_options, + update_cols=["age"], + num_partitions=1, + commit_mode="incremental", + max_groups_per_commit=1, + operation_id=operation_id, + ) + + update_by_row_id( + target, + pa.table( + {"_ROW_ID": [row_ids[1]], "age": [777]}, + schema=failing_source.schema, + ), + self.catalog_options, + update_cols=["age"], + ) + corrected_source = pa.table( + { + "_ROW_ID": [row_ids[1], row_ids[2], row_ids[3]], + "age": [100, 200, 300], + }, + schema=failing_source.schema, + ) + with self.assertRaisesRegex( + RuntimeError, "already completed"): + update_by_row_id( + target, + corrected_source, + self.catalog_options, + update_cols=["age"], + num_partitions=1, + commit_mode="incremental", + max_groups_per_commit=1, + operation_id=operation_id, + ) + + self.assertEqual( + [777, 200, 0], + self._read(target).sort_by("id")["age"].to_pylist(), + ) + + def test_commit_rejects_target_group_changed_after_planning(self): + from pypaimon.write.table_commit import StreamTableCommit + + target = self._create() + for row_id in range(1, 3): + self._write(target, pa.Table.from_pydict( + {"id": [row_id], "name": ["a"], "age": [0]}, + schema=self.pa_schema, + )) + + row_ids = self._rowid_by_id(target) + schema = pa.schema([ + ("_ROW_ID", pa.int64()), + ("age", pa.int32()), + ]) + source = pa.table( + { + "_ROW_ID": [row_ids[1], row_ids[2]], + "age": [100, 200], + }, + schema=schema, + ) + original_commit = StreamTableCommit.commit + external_committed = False + + def commit_after_external_update( + stream_commit, messages, commit_identifier): + nonlocal external_committed + if not external_committed: + external_committed = True + update_by_row_id( + target, + pa.table( + {"_ROW_ID": [row_ids[1]], "age": [777]}, + schema=schema, + ), + self.catalog_options, + update_cols=["age"], + ) + return original_commit( + stream_commit, messages, commit_identifier) + + with mock.patch.object( + StreamTableCommit, + "commit", + commit_after_external_update): + with self.assertRaises(RuntimeError): + update_by_row_id( + target, + source, + self.catalog_options, + update_cols=["age"], + num_partitions=1, + commit_mode="incremental", + max_groups_per_commit=1, + operation_id="planner-conflict-" + uuid.uuid4().hex, + ) + + self.assertEqual( + [777, 0], + self._read(target).sort_by("id")["age"].to_pylist(), + ) + + def test_resume_recovers_commit_before_checkpoint_tag_update(self): + target = self._create() + for row_id in range(1, 3): + self._write(target, pa.Table.from_pydict( + {"id": [row_id], "name": ["a"], "age": [0]}, + schema=self.pa_schema, + )) + + row_ids = self._rowid_by_id(target) + source = pa.table( + { + "_ROW_ID": [row_ids[1], row_ids[2]], + "age": [100, 200], + }, + schema=pa.schema([ + ("_ROW_ID", pa.int64()), + ("age", pa.int32()), + ]), + ) + operation_id = "tag-gap-" + uuid.uuid4().hex + catalog_type = type(self.catalog) + create_tag = catalog_type.create_tag + create_calls = 0 + + def fail_progress_tag(catalog, *args, **kwargs): + nonlocal create_calls + create_calls += 1 + if create_calls == 3: + raise RuntimeError("tag update failed") + return create_tag(catalog, *args, **kwargs) + + with mock.patch.object( + catalog_type, + "create_tag", + new=fail_progress_tag): + with self.assertRaisesRegex( + RuntimeError, "tag update failed"): + update_by_row_id( + target, + source, + self.catalog_options, + update_cols=["age"], + num_partitions=1, + commit_mode="incremental", + max_groups_per_commit=1, + operation_id=operation_id, + ) + + table = self.catalog.get_table(target) + partial_snapshot_id = ( + table.snapshot_manager().get_latest_snapshot().id) + self.assertEqual( + [100, 0], + self._read(target).sort_by("id")["age"].to_pylist(), + ) + self.assertEqual( + {"num_updated": 2}, + update_by_row_id( + target, + source, + self.catalog_options, + update_cols=["age"], + num_partitions=1, + commit_mode="incremental", + max_groups_per_commit=1, + operation_id=operation_id, + ), + ) + self.assertEqual( + partial_snapshot_id + 1, + table.snapshot_manager().get_latest_snapshot().id, + ) + def test_atomic_group_failure_commits_nothing(self): target = self._create() for row_id in range(1, 4): @@ -402,6 +729,7 @@ def test_incremental_commit_conflict_aborts_buffered_group_files(self): num_partitions=1, commit_mode="incremental", max_groups_per_commit=1, + operation_id="forced-conflict", ) self.assertEqual(files_before, self._data_files_under(table)) @@ -642,6 +970,7 @@ def test_rejects_invalid_max_groups_per_commit(self): update_cols=["age"], commit_mode="incremental", max_groups_per_commit=value, + operation_id="invalid-window", ) def test_requires_explicit_incremental_commit_mode(self): @@ -670,6 +999,26 @@ def test_requires_explicit_incremental_commit_mode(self): self.catalog_options, update_cols=["age"], commit_mode="incremental", + operation_id="missing-window", + ) + with self.assertRaisesRegex( + ValueError, "requires operation_id"): + update_by_row_id( + target, + source, + self.catalog_options, + update_cols=["age"], + commit_mode="incremental", + max_groups_per_commit=1, + ) + with self.assertRaisesRegex( + ValueError, "requires commit_mode='incremental'"): + update_by_row_id( + target, + source, + self.catalog_options, + update_cols=["age"], + operation_id="atomic-operation", ) with self.assertRaisesRegex( ValueError, "must be 'atomic' or 'incremental'"): @@ -691,7 +1040,7 @@ def _run_with_fake_commit(self, *, recorder=None, new_commit_errors=None, if recorder is None: recorder = {} recorder.update({ - "msgs": [object()], + "msgs": [types.SimpleNamespace(check_from_snapshot=None)], "new_commit_errors": list(new_commit_errors or []), "commit_error": commit_error, "close_error": close_error, @@ -712,6 +1061,18 @@ def deletion_vectors_enabled(self): return False class FakeCommit: + def with_snapshot_properties(self, properties): + return self + + def protect_row_id_ranges(self, checkpoint, ranges): + return self + + def protect_planned_row_id_files(self, ranges, signatures): + return self + + def validate_protected_row_id_ranges(self): + pass + def commit(self, msgs, *args): recorder["commit_calls"] += 1 recorder["commit_msgs"] = list(msgs) @@ -743,7 +1104,7 @@ class FakeTable: name="age", type=types.SimpleNamespace(type="INT"), ) - ]) + ], id=1) def snapshot_manager(self): return types.SimpleNamespace(get_latest_snapshot=lambda: types.SimpleNamespace( @@ -793,7 +1154,24 @@ def fake_apply(*args, **kwargs): ("age", pa.int32()), ])), \ mock.patch.object(m, "distributed_update_apply", - side_effect=fake_apply): + side_effect=fake_apply), \ + mock.patch.object(m, "_load_operation_checkpoint", + return_value=None), \ + mock.patch.object(m, "_initialize_operation_checkpoint", + return_value=None), \ + mock.patch.object(m, "_store_operation_planner"), \ + mock.patch.object( + m, + "_row_id_ranges_from_messages", + return_value=[m.Range(0, 0)]), \ + mock.patch.object( + m, + "_find_operation_snapshot", + return_value=types.SimpleNamespace( + commit_identifier=1)), \ + mock.patch.object(m, "_store_operation_checkpoint"), \ + mock.patch.object( + m, "_delete_checkpoint_tag", return_value=False): recorder["result"] = m.update_by_row_id( "default.fake", FakeSource(), @@ -801,6 +1179,7 @@ def fake_apply(*args, **kwargs): update_cols=["age"], commit_mode="incremental" if incremental else "atomic", max_groups_per_commit=1 if incremental else None, + operation_id="fake-incremental" if incremental else None, ) return recorder diff --git a/paimon-python/pypaimon/write/commit/commit_scanner.py b/paimon-python/pypaimon/write/commit/commit_scanner.py index 4612054e6c5c..aecdd2a0ced7 100644 --- a/paimon-python/pypaimon/write/commit/commit_scanner.py +++ b/paimon-python/pypaimon/write/commit/commit_scanner.py @@ -18,6 +18,8 @@ """ Manifest entries scanner for commit operations. """ +import bisect +import os from typing import Optional, List from pypaimon.common.predicate_builder import PredicateBuilder @@ -25,8 +27,51 @@ from pypaimon.manifest.manifest_file_manager import ManifestFileManager from pypaimon.manifest.manifest_list_manager import ManifestListManager from pypaimon.manifest.schema.manifest_entry import ManifestEntry -from pypaimon.read.scanner.file_scanner import FileScanner +from pypaimon.read.scanner.file_scanner import ( + FileScanner, + _filter_manifest_files_by_row_ranges, +) from pypaimon.snapshot.snapshot import Snapshot +from pypaimon.utils.range import Range + + +class _RowIdRangeScope: + """Compact global row-id ranges used by resumable update checkpoints.""" + + def __init__(self, row_id_ranges): + self.ranges = Range.sort_and_merge_overlap( + list(row_id_ranges or []), True, True) + self._range_ends = [row_range.to for row_range in self.ranges] + + def is_empty(self): + return not self.ranges + + def matches_record(self, record): + file_dict = record.get("_FILE") + if file_dict is None: + return False + first_row_id = file_dict.get("_FIRST_ROW_ID") + row_count = file_dict.get("_ROW_COUNT") + if first_row_id is None or row_count is None: + return False + return self.matches_values( + int(first_row_id), + int(first_row_id) + int(row_count) - 1, + ) + + def matches_entry(self, entry): + row_range = entry.file.row_id_range() + return ( + row_range is not None + and self.matches_values(row_range.from_, row_range.to) + ) + + def matches_values(self, from_, to): + index = bisect.bisect_left(self._range_ends, from_) + return ( + index < len(self.ranges) + and self.ranges[index].from_ <= to + ) class CommitScanner: @@ -74,6 +119,30 @@ def read_all_entries_from_changed_partitions(self, self.table, lambda: ([], None), partition_predicate=partition_filter ).read_manifest_entries(all_manifests) + def read_entries_for_row_id_ranges( + self, + snapshot: Optional[Snapshot], + row_id_ranges): + """Read live entries intersecting table-global row-id ranges.""" + if snapshot is None: + return [] + + scope = _RowIdRangeScope(row_id_ranges) + if scope.is_empty(): + return [] + + manifest_files = self.manifest_list_manager.read_all(snapshot) + manifest_files = _filter_manifest_files_by_row_ranges( + manifest_files, scope.ranges) + max_workers = self.table.options.scan_manifest_parallelism( + os.cpu_count() or 8) + return ManifestFileManager(self.table).read_entries_parallel( + manifest_files, + manifest_entry_filter=scope.matches_entry, + max_workers=max_workers, + early_record_filter=scope.matches_record, + ) + def read_incremental_entries_from_changed_partitions(self, snapshot: Snapshot, commit_entries: List[ManifestEntry], diff --git a/paimon-python/pypaimon/write/commit/conflict_detection.py b/paimon-python/pypaimon/write/commit/conflict_detection.py index 37992239593d..03b1b5c37a6c 100644 --- a/paimon-python/pypaimon/write/commit/conflict_detection.py +++ b/paimon-python/pypaimon/write/commit/conflict_detection.py @@ -170,6 +170,27 @@ def __init__(self, entry): entry.bucket)) +def row_id_file_signature(partition, bucket, data_file): + values = ( + tuple(partition.values) + if hasattr(partition, "values") else tuple(partition) + ) + return ( + values, + bucket, + data_file.level, + data_file.file_name, + tuple(data_file.extra_files) if data_file.extra_files else (), + data_file.embedded_index, + data_file.external_path, + data_file.first_row_id, + data_file.row_count, + data_file.schema_id, + tuple(data_file.write_cols) + if data_file.write_cols is not None else None, + ) + + class ConflictDetection: """Detects conflicts between base and delta files during commit.""" @@ -180,8 +201,84 @@ def __init__(self, data_evolution_enabled, snapshot_manager, self.manifest_list_manager = manifest_list_manager self.table = table self._row_id_check_from_snapshot = None + self._protected_row_id_checkpoint = None + self._protected_row_id_ranges = [] + self._planned_row_id_ranges = [] + self._planned_row_id_signatures = set() self.commit_scanner = commit_scanner + def protect_row_id_ranges(self, checkpoint_snapshot, row_id_ranges): + """Protect already-committed row-id groups during a later commit.""" + self._protected_row_id_checkpoint = checkpoint_snapshot + self._protected_row_id_ranges = Range.sort_and_merge_overlap( + list(row_id_ranges or []), True, True) + + def check_protected_row_id_ranges(self, latest_snapshot): + """Reject a commit if a completed resumable-update group changed.""" + checkpoint = self._protected_row_id_checkpoint + if checkpoint is None or not self._protected_row_id_ranges: + return None + if latest_snapshot is None or latest_snapshot.id < checkpoint.id: + return RuntimeError( + "Resumable update checkpoint is no longer in the current " + "snapshot lineage.") + if latest_snapshot.id == checkpoint.id: + if self._same_snapshot(checkpoint, latest_snapshot): + return None + return RuntimeError( + "Resumable update checkpoint snapshot was replaced.") + + checkpoint_entries = ( + self.commit_scanner.read_entries_for_row_id_ranges( + checkpoint, self._protected_row_id_ranges)) + latest_entries = self.commit_scanner.read_entries_for_row_id_ranges( + latest_snapshot, self._protected_row_id_ranges) + if (self._row_id_entry_signatures(checkpoint_entries) + != self._row_id_entry_signatures(latest_entries)): + return RuntimeError( + "Concurrent overwrite changed a row-id group already " + "completed by the resumable update.") + return None + + def protect_planned_row_id_files( + self, row_id_ranges, file_signatures): + self._planned_row_id_ranges = Range.sort_and_merge_overlap( + list(row_id_ranges or []), True, True) + self._planned_row_id_signatures = set(file_signatures or []) + + def clear_planned_row_id_files(self): + self._planned_row_id_ranges = [] + self._planned_row_id_signatures = set() + + def check_planned_row_id_files(self, latest_snapshot): + if not self._planned_row_id_ranges: + return None + latest_entries = self.commit_scanner.read_entries_for_row_id_ranges( + latest_snapshot, self._planned_row_id_ranges) + if (self._row_id_entry_signatures(latest_entries) + != self._planned_row_id_signatures): + return RuntimeError( + "Target files changed after the resumable update group " + "was planned.") + return None + + @staticmethod + def _same_snapshot(left, right): + if left is None or right is None: + return left is right + if left.uuid is not None or right.uuid is not None: + return left.id == right.id and left.uuid == right.uuid + return left == right + + @staticmethod + def _row_id_entry_signatures(entries): + return { + row_id_file_signature( + entry.partition, entry.bucket, entry.file) + for entry in entries + if entry.kind == 0 + } + def should_be_overwrite_commit(self, append_file_entries=None, append_index_files=None): for entry in append_file_entries or []: if entry.kind == 1: diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index 73f3d22429f5..4dc2004ff6a6 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -139,10 +139,27 @@ def __init__(self, snapshot_commit: SnapshotCommit, table, commit_user: str, table=table, commit_scanner=self.commit_scanner ) + self.snapshot_properties = None table_rollback = table.catalog_environment.catalog_table_rollback() self.rollback = CommitRollback(table_rollback) if table_rollback is not None else None + def with_snapshot_properties(self, properties): + self.snapshot_properties = ( + dict(properties) if properties is not None else None) + return self + + def protect_row_id_ranges(self, checkpoint_snapshot, row_id_ranges): + self.conflict_detection.protect_row_id_ranges( + checkpoint_snapshot, row_id_ranges) + return self + + def protect_planned_row_id_files( + self, row_id_ranges, file_signatures): + self.conflict_detection.protect_planned_row_id_files( + row_id_ranges, file_signatures) + return self + def commit(self, commit_messages: List[CommitMessage], commit_identifier: int): """Commit the given commit messages in normal append mode.""" if not commit_messages: @@ -411,6 +428,7 @@ def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, self.conflict_detection._row_id_check_from_snapshot = ( latest_snapshot.id ) + self.conflict_detection.clear_planned_row_id_files() # No snapshot commit was attempted for the conflicting files, # so the rewritten attempt is still deterministic. retry_result = None @@ -480,6 +498,14 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str if self._is_duplicate_commit(retry_result, latest_snapshot, commit_identifier, commit_kind): return SuccessResult() + protected_conflict = ( + self.conflict_detection.check_protected_row_id_ranges( + latest_snapshot)) + if protected_conflict is not None: + if retry_result is None: + raise CommitConflictError( + str(protected_conflict)) from protected_conflict + raise protected_conflict latest_snapshot_id = latest_snapshot.id if latest_snapshot else 0 if ( hash_index_base_snapshot is not None @@ -568,6 +594,15 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str # callers do not delete files which a snapshot may reference. raise conflict_exception + planned_conflict = ( + self.conflict_detection.check_planned_row_id_files( + latest_snapshot)) + if planned_conflict is not None: + if retry_result is None: + raise CommitConflictError( + str(planned_conflict)) from planned_conflict + raise planned_conflict + # Apply row tracking logic after conflict detection (matches Java ordering) row_tracking_enabled = self.table.options.row_tracking_enabled() next_row_id = None @@ -646,6 +681,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str time_millis=int(time.time() * 1000), next_row_id=next_row_id, index_manifest=index_manifest, + properties=self.snapshot_properties, ) # Generate partition statistics for the commit statistics = self._generate_partition_statistics(commit_entries) diff --git a/paimon-python/pypaimon/write/table_commit.py b/paimon-python/pypaimon/write/table_commit.py index 2f4d8e60a0ce..b326576ca78e 100644 --- a/paimon-python/pypaimon/write/table_commit.py +++ b/paimon-python/pypaimon/write/table_commit.py @@ -61,6 +61,29 @@ def add_commit_callback(self, callback: CommitCallback) -> None: """Register a callback to be invoked after each successful commit.""" self._commit_callbacks.append(callback) + def with_snapshot_properties(self, properties): + self.file_store_commit.with_snapshot_properties(properties) + return self + + def protect_row_id_ranges(self, checkpoint_snapshot, row_id_ranges): + self.file_store_commit.protect_row_id_ranges( + checkpoint_snapshot, row_id_ranges) + return self + + def protect_planned_row_id_files( + self, row_id_ranges, file_signatures): + self.file_store_commit.protect_planned_row_id_files( + row_id_ranges, file_signatures) + return self + + def validate_protected_row_id_ranges(self): + conflict = ( + self.file_store_commit.conflict_detection + .check_protected_row_id_ranges( + self.table.snapshot_manager().get_latest_snapshot())) + if conflict is not None: + raise CommitConflictError(str(conflict)) from conflict + def _commit(self, commit_messages: List[CommitMessage], commit_identifier: int = BATCH_COMMIT_IDENTIFIER): non_empty_messages = [msg for msg in commit_messages if not msg.is_empty()] From 83496dd5cee3e2a52a137f2145b8e8205c5bf717 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 30 Jul 2026 07:41:27 -0700 Subject: [PATCH 05/11] [python][ray] Compact incremental update checkpoints --- docs/docs/pypaimon/ray-data.md | 7 +- .../pypaimon/ray/data_evolution_merge_join.py | 25 ++-- .../pypaimon/ray/update_by_row_id.py | 108 +++++++++++++++--- .../tests/ray_update_by_row_id_test.py | 61 ++++++++++ 4 files changed, 178 insertions(+), 23 deletions(-) diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md index ddcd79cd9c1f..1d6120352ebd 100644 --- a/docs/docs/pypaimon/ray-data.md +++ b/docs/docs/pypaimon/ray-data.md @@ -596,10 +596,11 @@ metrics = update_by_row_id( retry options in `ray_remote_args`. If retries are exhausted, earlier incremental commits remain visible. Worker loss may also discard successful results still buffered in the failed Ray task. -- Each incremental snapshot stores the cumulative completed row-id ranges, and +- Each incremental snapshot stores compressed cumulative row-id ranges, and internal Paimon tags retain the latest checkpoint across normal snapshot - expiration. A retry with the same `operation_id` filters those ranges before - target file-group rewrite and returns the cumulative `num_updated`. + expiration. Larger commit windows reduce checkpoint metadata at the cost of + redoing more groups after a failure. A retry with the same `operation_id` + filters completed ranges before rewrite and returns cumulative `num_updated`. - The checkpoint avoids repeating Paimon's target file read/rewrite work. It cannot checkpoint arbitrary transformations upstream of `update_by_row_id`. Persist a multi-day or nondeterministic source first (for example to Parquet or diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index 6db1bb20184e..ff91fa4326f1 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -59,6 +59,16 @@ def _map_kwargs( return kwargs +def _sorted_range_membership(values, starts, ends): + import numpy as np + + indices = np.searchsorted(ends, values, side="left") + matches = np.zeros(len(values), dtype=bool) + positions = np.flatnonzero(indices < len(starts)) + matches[positions] = starts[indices[positions]] <= values[positions] + return matches + + def _resolve_source_projection( clauses: List[_NormalizedClause], source_on: Sequence[str], @@ -529,6 +539,10 @@ def distributed_update_apply( range_starts = np.asarray([r.from_ for r in valid_ranges], dtype=np.int64) range_ends = np.asarray([r.to for r in valid_ranges], dtype=np.int64) completed_ranges = tuple(skip_row_id_ranges or ()) + completed_starts = np.asarray( + [start for start, _ in completed_ranges], dtype=np.int64) + completed_ends = np.asarray( + [end for _, end in completed_ranges], dtype=np.int64) def _assign_frid(batch: pa.Table) -> pa.Table: if batch.num_rows == 0: @@ -542,10 +556,8 @@ def _assign_frid(batch: pa.Table) -> pa.Table: "or matched rows come from a different table." ) rids = rid_col.to_numpy(zero_copy_only=False) - # Check each row_id belongs to a valid range (vectorized). - in_range = np.zeros(len(rids), dtype=bool) - for s, e in zip(range_starts, range_ends): - in_range |= (rids >= s) & (rids <= e) + in_range = _sorted_range_membership( + rids, range_starts, range_ends) if not in_range.all(): bad = rids[~in_range][0] raise ValueError( @@ -563,9 +575,8 @@ def _assign_frid(batch: pa.Table) -> pa.Table: ) if not completed_ranges: return with_first_row_id - completed = np.zeros(len(rids), dtype=bool) - for start, end in completed_ranges: - completed |= (rids >= start) & (rids <= end) + completed = _sorted_range_membership( + rids, completed_starts, completed_ends) if not completed.any(): return with_first_row_id return with_first_row_id.filter(pa.array(~completed)) diff --git a/paimon-python/pypaimon/ray/update_by_row_id.py b/paimon-python/pypaimon/ray/update_by_row_id.py index 6caa64adb6a1..cb9b4755d67c 100644 --- a/paimon-python/pypaimon/ray/update_by_row_id.py +++ b/paimon-python/pypaimon/ray/update_by_row_id.py @@ -23,9 +23,11 @@ target). Pairs with ``bucket_join``, which produces the row ids. """ +import base64 import hashlib import json import logging +import zlib from typing import Any, Dict, List, Optional import pyarrow as pa @@ -54,6 +56,7 @@ _CHECKPOINT_PROPERTY = "pypaimon.ray.update-by-row-id.checkpoint" _CHECKPOINT_TAG_PREFIX = "_pypaimon_ray_update_" _CHECKPOINT_VERSION = 1 +_RANGES_CODEC = "zlib-delta-json-v1" def _blob_col_names(table: "FileStoreTable") -> set: @@ -315,11 +318,11 @@ def __init__( if checkpoint is not None: snapshot, state = checkpoint self._checkpoint_snapshot = snapshot - self._completed_row_id_ranges = _decode_ranges( - state["completed_ranges"]) + self._completed_row_id_ranges = ( + state["_completed_row_id_ranges"]) self._completed_num_updated = int( state["num_updated"]) - if _checkpoint_state(snapshot) is not None: + if _has_checkpoint_state(snapshot): self._next_commit_identifier = ( snapshot.commit_identifier + 1) if initial_snapshot is not None: @@ -392,7 +395,7 @@ def _commit_pending(self) -> None: candidate_ranges = self._completed_row_id_ranges if self._operation_id is not None: pending_ranges = _row_id_ranges_from_messages(messages) - candidate_ranges = _merge_row_id_ranges( + candidate_ranges = _merge_canonical_ranges( self._completed_row_id_ranges, pending_ranges, ) @@ -560,10 +563,8 @@ def _checkpoint_properties( "operation_id": operation_id, "schema_id": schema_id, "update_cols": list(update_cols), - "completed_ranges": [ - [row_range.from_, row_range.to] - for row_range in completed_ranges - ], + "completed_ranges": _encode_ranges(completed_ranges), + "ranges_codec": _RANGES_CODEC, "num_updated": num_updated, } return { @@ -596,7 +597,7 @@ def _checkpoint_state(snapshot): if not required.issubset(state): raise RuntimeError( "Incomplete resumable update checkpoint.") - _decode_ranges(state["completed_ranges"]) + state["_completed_row_id_ranges"] = _decode_checkpoint_ranges(state) if (isinstance(state["num_updated"], bool) or not isinstance(state["num_updated"], int) or state["num_updated"] < 0): @@ -609,8 +610,10 @@ def _validate_checkpoint_state( snapshot, operation_id, update_cols, - schema_id): - state = _checkpoint_state(snapshot) + schema_id, + state=None): + if state is None: + state = _checkpoint_state(snapshot) if state is None: raise RuntimeError( "Resumable update checkpoint snapshot has no checkpoint state.") @@ -660,7 +663,7 @@ def _load_operation_checkpoint( snapshot = snapshot_manager.get_snapshot_by_id(snapshot_id) if (snapshot is not None and snapshot.commit_user == commit_user - and _checkpoint_state(snapshot) is not None): + and _has_checkpoint_state(snapshot)): checkpoint_snapshot = snapshot break @@ -681,6 +684,7 @@ def _load_operation_checkpoint( operation_id, update_cols, table.table_schema.id, + state, ) _store_operation_checkpoint( catalog, target, checkpoint_tags, checkpoint_snapshot) @@ -711,6 +715,7 @@ def _initial_checkpoint_state(operation_id, schema_id, update_cols): "schema_id": schema_id, "update_cols": list(update_cols), "completed_ranges": [], + "_completed_row_id_ranges": [], "num_updated": 0, } @@ -725,7 +730,7 @@ def _find_operation_snapshot(table, commit_user, commit_identifier): if (snapshot is not None and snapshot.commit_user == commit_user and snapshot.commit_identifier == commit_identifier - and _checkpoint_state(snapshot) is not None): + and _has_checkpoint_state(snapshot)): return snapshot return None @@ -775,6 +780,61 @@ def _delete_checkpoint_tag( return False +def _has_checkpoint_state(snapshot): + return _CHECKPOINT_PROPERTY in (snapshot.properties or {}) + + +def _encode_ranges(ranges): + previous_end = -1 + deltas = [] + for row_range in ranges: + deltas.append([ + row_range.from_ - previous_end - 1, + row_range.to - row_range.from_, + ]) + previous_end = row_range.to + raw = json.dumps(deltas, separators=(",", ":")).encode("utf-8") + return base64.b64encode(zlib.compress(raw)).decode("ascii") + + +def _decode_checkpoint_ranges(state): + codec = state.get("ranges_codec") + encoded = state["completed_ranges"] + if codec is None: + return _decode_ranges(encoded) + if codec != _RANGES_CODEC or not isinstance(encoded, str): + raise RuntimeError( + "Unsupported resumable update checkpoint ranges codec.") + try: + compressed = base64.b64decode(encoded.encode("ascii"), validate=True) + deltas = json.loads(zlib.decompress(compressed).decode("utf-8")) + except Exception as error: + raise RuntimeError( + "Invalid compressed resumable update checkpoint ranges." + ) from error + + encoded_ranges = [] + previous_end = -1 + if not isinstance(deltas, list): + raise RuntimeError( + "Invalid compressed resumable update checkpoint ranges.") + for delta in deltas: + if (not isinstance(delta, list) + or len(delta) != 2 + or isinstance(delta[0], bool) + or isinstance(delta[1], bool) + or not isinstance(delta[0], int) + or not isinstance(delta[1], int) + or delta[1] < 0): + raise RuntimeError( + "Invalid compressed resumable update checkpoint ranges.") + start = previous_end + delta[0] + 1 + end = start + delta[1] + encoded_ranges.append([start, end]) + previous_end = end + return _decode_ranges(encoded_ranges) + + def _decode_ranges(encoded_ranges): ranges = [] if not isinstance(encoded_ranges, list): @@ -810,6 +870,28 @@ def _merge_row_id_ranges(*range_groups): ) +def _merge_canonical_ranges(left, right): + ordered = [] + left_index = 0 + right_index = 0 + while left_index < len(left) or right_index < len(right): + if (right_index >= len(right) + or (left_index < len(left) + and left[left_index].from_ <= right[right_index].from_)): + current = left[left_index] + left_index += 1 + else: + current = right[right_index] + right_index += 1 + + if ordered and current.from_ <= ordered[-1].to + 1: + ordered[-1] = Range( + ordered[-1].from_, max(ordered[-1].to, current.to)) + else: + ordered.append(current) + return ordered + + def _row_id_ranges_from_messages(messages): ranges = [] for message in messages: diff --git a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py index 43c370fb21ae..f3f9846ce34a 100644 --- a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py +++ b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py @@ -365,6 +365,67 @@ def test_resumes_after_partial_failure_without_rewriting_completed_groups(self): self.assertFalse(delete_update_by_row_id_checkpoint( target, self.catalog_options, operation_id)) + def test_sparse_checkpoint_ranges_stay_compact_and_searchable(self): + import importlib + + import numpy as np + + update_module = importlib.import_module( + "pypaimon.ray.update_by_row_id") + join_module = importlib.import_module( + "pypaimon.ray.data_evolution_merge_join") + + ranges = [] + position = 0 + for index in range(100000): + position += ((index * 2654435761) % 10000) + 2 + ranges.append(update_module.Range( + position, position + index % 7)) + position += index % 7 + + properties = update_module._checkpoint_properties( + "sparse-scale", 1, ["age"], ranges, len(ranges)) + encoded = properties[update_module._CHECKPOINT_PROPERTY] + self.assertLess(len(encoded), 400000) + + state = update_module._checkpoint_state( + types.SimpleNamespace(properties=properties)) + self.assertEqual( + ranges, state["_completed_row_id_ranges"]) + + starts = np.asarray( + [row_range.from_ for row_range in ranges], dtype=np.int64) + ends = np.asarray( + [row_range.to for row_range in ranges], dtype=np.int64) + probes = np.concatenate([starts, ends + 1]) + matches = join_module._sorted_range_membership( + probes, starts, ends) + self.assertTrue(matches[:len(ranges)].all()) + self.assertFalse(matches[len(ranges):].any()) + + def test_legacy_checkpoint_ranges_remain_readable(self): + import importlib + + module = importlib.import_module( + "pypaimon.ray.update_by_row_id") + legacy = { + "version": module._CHECKPOINT_VERSION, + "operation_id": "legacy", + "schema_id": 1, + "update_cols": ["age"], + "completed_ranges": [[1, 2], [5, 8]], + "num_updated": 6, + } + snapshot = types.SimpleNamespace(properties={ + module._CHECKPOINT_PROPERTY: module.json.dumps(legacy), + }) + + state = module._checkpoint_state(snapshot) + self.assertEqual( + [module.Range(1, 2), module.Range(5, 8)], + state["_completed_row_id_ranges"], + ) + def test_resume_rejects_external_change_to_completed_group(self): target = self._create() for row_id in range(1, 4): From b48c3ed62a27e32c413346eeb62f9b050355b601 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 30 Jul 2026 07:52:22 -0700 Subject: [PATCH 06/11] [python][ray] Preserve deterministic incremental failures --- .../pypaimon/ray/data_evolution_merge_join.py | 82 ++++++++++--------- .../pypaimon/tests/file_store_commit_test.py | 40 +++++++++ .../tests/ray_update_by_row_id_test.py | 56 +++++++++++++ .../pypaimon/write/file_store_commit.py | 4 +- 4 files changed, 140 insertions(+), 42 deletions(-) diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index ff91fa4326f1..d12f31d45293 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -586,7 +586,7 @@ def _assign_frid(batch: pa.Table) -> pa.Table: captured_table = table captured_cols = cols - capture_group_validation_errors = on_group_result is not None + capture_group_errors = on_group_result is not None def _group_result( msgs_blob, @@ -613,49 +613,51 @@ def _apply_group(group: pa.Table) -> pa.Table: "error": pa.array([], type=pa.string()), }) - if ( - pc.count_distinct(group.column(row_id_name)).as_py() - != group.num_rows - ): - error = ValueError( - "MERGE matched multiple source rows to the same " - "target _ROW_ID. Deduplicate the source before " - "merging." + try: + if ( + pc.count_distinct(group.column(row_id_name)).as_py() + != group.num_rows + ): + raise ValueError( + "MERGE matched multiple source rows to the same " + "target _ROW_ID. Deduplicate the source before " + "merging." + ) + + for_update = group.drop_columns([frid_col]) + row_ids = ( + for_update.column(row_id_name).to_pylist() + if collect_row_ids else [] ) - if not capture_group_validation_errors: - raise error + files_info = ray.get(precomputed_info_ref) + first_row_id = group.column(frid_col)[0].as_py() + split, target_files = files_info.first_row_id_index[first_row_id] + planned_file_signatures = [ + row_id_file_signature( + split.partition, split.bucket, data_file) + for data_file in target_files + ] + worker = TableUpdateByRowId( + captured_table, + "_merge_into_shard_" + uuid.uuid4().hex[:8], + BATCH_COMMIT_IDENTIFIER, + _precomputed_files_info=files_info, + ) + msgs = worker.update_columns(for_update, list(captured_cols)) + + return _group_result( + pickle.dumps(msgs), + for_update.num_rows, + pickle.dumps(row_ids), + pickle.dumps(planned_file_signatures), + None, + ) + except Exception as error: + if not capture_group_errors: + raise return _group_result( b"", 0, b"", b"", _group_error_text(error)) - for_update = group.drop_columns([frid_col]) - row_ids = ( - for_update.column(row_id_name).to_pylist() - if collect_row_ids else [] - ) - files_info = ray.get(precomputed_info_ref) - first_row_id = group.column(frid_col)[0].as_py() - split, target_files = files_info.first_row_id_index[first_row_id] - planned_file_signatures = [ - row_id_file_signature( - split.partition, split.bucket, data_file) - for data_file in target_files - ] - worker = TableUpdateByRowId( - captured_table, - "_merge_into_shard_" + uuid.uuid4().hex[:8], - BATCH_COMMIT_IDENTIFIER, - _precomputed_files_info=files_info, - ) - msgs = worker.update_columns(for_update, list(captured_cols)) - - return _group_result( - pickle.dumps(msgs), - for_update.num_rows, - pickle.dumps(row_ids), - pickle.dumps(planned_file_signatures), - None, - ) - # One group per target data file; bounded by file count and num_partitions. group_partitions = max( 1, min(len(captured_sorted), num_partitions) diff --git a/paimon-python/pypaimon/tests/file_store_commit_test.py b/paimon-python/pypaimon/tests/file_store_commit_test.py index 772b5f460d4b..3f3d2c843e70 100644 --- a/paimon-python/pypaimon/tests/file_store_commit_test.py +++ b/paimon-python/pypaimon/tests/file_store_commit_test.py @@ -25,10 +25,12 @@ from pypaimon.snapshot.snapshot_commit import PartitionStatistics from pypaimon.table.row.generic_row import GenericRow from pypaimon.utils.range import Range +from pypaimon.write.commit.conflict_detection import CommitConflictError from pypaimon.write.commit.row_id_conflict_rewriter import RowIdRewriteResult from pypaimon.write.commit_message import CommitMessage from pypaimon.write.file_store_commit import ( FileStoreCommit, + RetryResult, RewriteResult, ) @@ -546,6 +548,44 @@ def test_row_id_rewrite_respects_commit_retry_limit( file_store_commit.conflict_detection._planned_row_id_ranges, ) + def test_known_failed_retry_keeps_protected_conflict_deterministic( + self, mock_manifest_list_manager, mock_manifest_file_manager): + file_store_commit = self._create_file_store_commit() + snapshot = Mock(id=7) + conflict = RuntimeError("protected conflict") + file_store_commit.conflict_detection.check_protected_row_id_ranges = ( + Mock(return_value=conflict) + ) + + with self.assertRaises(CommitConflictError): + file_store_commit._try_commit_once( + RetryResult(snapshot), + "APPEND", + [Mock()], + [], + 11, + snapshot, + ) + + def test_known_failed_retry_keeps_planned_conflict_deterministic( + self, mock_manifest_list_manager, mock_manifest_file_manager): + file_store_commit = self._create_file_store_commit() + snapshot = Mock(id=7) + conflict = RuntimeError("planned conflict") + file_store_commit.conflict_detection.check_planned_row_id_files = ( + Mock(return_value=conflict) + ) + + with self.assertRaises(CommitConflictError): + file_store_commit._try_commit_once( + RetryResult(snapshot), + "APPEND", + [Mock()], + [], + 11, + snapshot, + ) + @staticmethod def _to_entries(commit_messages): commit_entries = [] diff --git a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py index f3f9846ce34a..c2fa4135c741 100644 --- a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py +++ b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py @@ -258,6 +258,62 @@ def test_group_failure_preserves_completed_groups(self): self._read(target).sort_by("id")["age"].to_pylist(), ) + def test_worker_failure_preserves_sibling_groups(self): + from pypaimon.write.table_update_by_row_id import TableUpdateByRowId + + target = self._create() + for row_id in range(1, 4): + self._write(target, pa.Table.from_pydict( + {"id": [row_id], "name": ["a"], "age": [0]}, + schema=self.pa_schema, + )) + + table = self.catalog.get_table(target) + base_snapshot_id = table.snapshot_manager().get_latest_snapshot().id + row_ids = self._rowid_by_id(target) + source = pa.table( + { + "_ROW_ID": [row_ids[1], row_ids[2], row_ids[3]], + "age": [100, 200, 300], + }, + schema=pa.schema([ + ("_ROW_ID", pa.int64()), + ("age", pa.int32()), + ]), + ) + + class FailingUpdate(TableUpdateByRowId): + + def update_columns(self, updates, update_cols): + if row_ids[3] in updates["_ROW_ID"].to_pylist(): + raise RuntimeError("injected worker failure") + return super().update_columns(updates, update_cols) + + with mock.patch( + "pypaimon.write.table_update_by_row_id.TableUpdateByRowId", + FailingUpdate): + with self.assertRaisesRegex( + RuntimeError, "injected worker failure"): + update_by_row_id( + target, + ray.data.from_arrow(source), + self.catalog_options, + update_cols=["age"], + num_partitions=1, + commit_mode="incremental", + max_groups_per_commit=5, + operation_id="worker-failure", + ) + + self.assertEqual( + base_snapshot_id + 1, + table.snapshot_manager().get_latest_snapshot().id, + ) + self.assertEqual( + [100, 200, 0], + self._read(target).sort_by("id")["age"].to_pylist(), + ) + def test_resumes_after_partial_failure_without_rewriting_completed_groups(self): target = self._create() for row_id in range(1, 4): diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index 4dc2004ff6a6..23d469b2668b 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -502,7 +502,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str self.conflict_detection.check_protected_row_id_ranges( latest_snapshot)) if protected_conflict is not None: - if retry_result is None: + if retry_result is None or retry_result.exception is None: raise CommitConflictError( str(protected_conflict)) from protected_conflict raise protected_conflict @@ -598,7 +598,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str self.conflict_detection.check_planned_row_id_files( latest_snapshot)) if planned_conflict is not None: - if retry_result is None: + if retry_result is None or retry_result.exception is None: raise CommitConflictError( str(planned_conflict)) from planned_conflict raise planned_conflict From afb352d6ea7853d0d70264a8ef035870675f6c7d Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 30 Jul 2026 07:58:29 -0700 Subject: [PATCH 07/11] [python] Preserve SnapshotCommit lifecycle outcomes --- .../pypaimon/tests/file_store_commit_test.py | 74 +++++++++++++++++++ .../pypaimon/write/file_store_commit.py | 70 +++++++++++++----- 2 files changed, 125 insertions(+), 19 deletions(-) diff --git a/paimon-python/pypaimon/tests/file_store_commit_test.py b/paimon-python/pypaimon/tests/file_store_commit_test.py index 3f3d2c843e70..d98873b55ad2 100644 --- a/paimon-python/pypaimon/tests/file_store_commit_test.py +++ b/paimon-python/pypaimon/tests/file_store_commit_test.py @@ -62,6 +62,26 @@ def _create_file_store_commit(self): commit_user='test_user' ) + def _prepare_atomic_attempt(self): + file_store_commit = self._create_file_store_commit() + self.mock_table.identifier = "default.test_table" + self.mock_table.table_schema = Mock(id=7) + self.mock_table.options.row_tracking_enabled.return_value = False + + snapshot_commit = MagicMock() + snapshot_commit.__enter__.return_value = snapshot_commit + snapshot_commit.__exit__.return_value = False + snapshot_commit.commit.return_value = True + file_store_commit.snapshot_commit = snapshot_commit + file_store_commit._write_manifest_files = Mock(return_value=[Mock()]) + file_store_commit._generate_partition_statistics = Mock( + return_value=[]) + file_store_commit.manifest_list_manager.read_all.return_value = [] + + commit_entry = Mock(kind=0) + commit_entry.file = Mock(row_count=1) + return file_store_commit, snapshot_commit, commit_entry + def test_generate_partition_statistics_single_partition_single_file( self, mock_manifest_list_manager, mock_manifest_file_manager): """Test partition statistics generation with single partition and single file.""" @@ -463,6 +483,60 @@ def test_append_commit_inherits_index_manifest( ) self.assertEqual(str(uuid.UUID(committed_snapshot.uuid)), committed_snapshot.uuid) + def test_snapshot_commit_lifecycle_is_preserved( + self, mock_manifest_list_manager, mock_manifest_file_manager): + file_store_commit, snapshot_commit, commit_entry = ( + self._prepare_atomic_attempt()) + + result = file_store_commit._try_commit_once( + None, "APPEND", [commit_entry], [], 1, None) + + self.assertTrue(result.is_success()) + snapshot_commit.__enter__.assert_called_once() + snapshot_commit.__exit__.assert_called_once_with(None, None, None) + + def test_snapshot_commit_enter_failure_is_classified( + self, mock_manifest_list_manager, mock_manifest_file_manager): + error = RuntimeError("enter failed") + file_store_commit, snapshot_commit, commit_entry = ( + self._prepare_atomic_attempt()) + snapshot_commit.__enter__.side_effect = error + + result = file_store_commit._try_commit_once( + None, "APPEND", [commit_entry], [], 1, None) + + self.assertFalse(result.is_success()) + self.assertIs(error, result.exception) + snapshot_commit.commit.assert_not_called() + snapshot_commit.__exit__.assert_not_called() + + def test_snapshot_commit_exit_failure_preserves_success( + self, mock_manifest_list_manager, mock_manifest_file_manager): + file_store_commit, snapshot_commit, commit_entry = ( + self._prepare_atomic_attempt()) + snapshot_commit.__exit__.side_effect = RuntimeError("exit failed") + + result = file_store_commit._try_commit_once( + None, "APPEND", [commit_entry], [], 1, None) + + self.assertTrue(result.is_success()) + + def test_snapshot_commit_error_is_passed_to_exit( + self, mock_manifest_list_manager, mock_manifest_file_manager): + error = RuntimeError("commit failed") + file_store_commit, snapshot_commit, commit_entry = ( + self._prepare_atomic_attempt()) + snapshot_commit.commit.side_effect = error + + result = file_store_commit._try_commit_once( + None, "APPEND", [commit_entry], [], 1, None) + + self.assertFalse(result.is_success()) + self.assertIs(error, result.exception) + exit_args = snapshot_commit.__exit__.call_args.args + self.assertIs(RuntimeError, exit_args[0]) + self.assertIs(error, exit_args[1]) + def test_null_partition_value( self, mock_manifest_list_manager, mock_manifest_file_manager): from pypaimon.data.timestamp import Timestamp diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index 23d469b2668b..44718f484842 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -697,30 +697,62 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str logger.warning(f"Exception occurs when preparing snapshot: {e}", exc_info=True) raise RuntimeError(f"Failed to prepare snapshot: {e}") - # Use SnapshotCommit for atomic commit + success = None + commit_error = None + entered = False try: - with self.snapshot_commit: - success = self.snapshot_commit.commit( - latest_snapshot.uuid if latest_snapshot else None, - snapshot_data, - statistics, + self.snapshot_commit.__enter__() + entered = True + success = self.snapshot_commit.commit( + latest_snapshot.uuid if latest_snapshot else None, + snapshot_data, + statistics, + ) + except Exception as error: + commit_error = error + finally: + if entered: + error_info = ( + (type(commit_error), commit_error, + commit_error.__traceback__) + if commit_error is not None else (None, None, None) ) - if not success: - commit_time_s = (int(time.time() * 1000) - start_millis) / 1000 + try: + self.snapshot_commit.__exit__(*error_info) + except Exception: logger.warning( - "Atomic commit failed for snapshot #%d by user %s " - "with identifier %s and kind %s after %.0f seconds. Try again.", - new_snapshot_id, - self.commit_user, - commit_identifier, - commit_kind, - commit_time_s, + "Failed to close snapshot commit.", + exc_info=True, ) - return RetryResult(latest_snapshot, None, base_data_files=base_data_files) - except Exception as e: + + if commit_error is not None: # Commit exception, not sure about the situation and should not clean up the files - logger.warning("Retry commit for exception.", exc_info=True) - return RetryResult(latest_snapshot, e, base_data_files=base_data_files) + logger.warning( + "Retry commit for exception.", + exc_info=( + type(commit_error), + commit_error, + commit_error.__traceback__, + ), + ) + return RetryResult( + latest_snapshot, + commit_error, + base_data_files=base_data_files, + ) + if not success: + commit_time_s = (int(time.time() * 1000) - start_millis) / 1000 + logger.warning( + "Atomic commit failed for snapshot #%d by user %s " + "with identifier %s and kind %s after %.0f seconds. Try again.", + new_snapshot_id, + self.commit_user, + commit_identifier, + commit_kind, + commit_time_s, + ) + return RetryResult( + latest_snapshot, None, base_data_files=base_data_files) logger.info( "Successfully commit snapshot %d to table %s by user %s " From 49e2d045daf1c1d70bc94ab158b1327d82587562 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 30 Jul 2026 08:47:49 -0700 Subject: [PATCH 08/11] [python][ray] Retry checkpoint tag advancement --- docs/docs/pypaimon/ray-data.md | 2 ++ .../pypaimon/ray/update_by_row_id.py | 19 +++++++++++++++ .../tests/ray_update_by_row_id_test.py | 23 ++++++++++++++++++- 3 files changed, 43 insertions(+), 1 deletion(-) diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md index 1d6120352ebd..469ceb902674 100644 --- a/docs/docs/pypaimon/ray-data.md +++ b/docs/docs/pypaimon/ray-data.md @@ -601,6 +601,8 @@ metrics = update_by_row_id( expiration. Larger commit windows reduce checkpoint metadata at the cost of redoing more groups after a failure. A retry with the same `operation_id` filters completed ranges before rewrite and returns cumulative `num_updated`. + If a process stops between snapshot commit and tag advancement, recovery may + redo at most the latest commit window. - The checkpoint avoids repeating Paimon's target file read/rewrite work. It cannot checkpoint arbitrary transformations upstream of `update_by_row_id`. Persist a multi-day or nondeterministic source first (for example to Parquet or diff --git a/paimon-python/pypaimon/ray/update_by_row_id.py b/paimon-python/pypaimon/ray/update_by_row_id.py index cb9b4755d67c..8904ada79e96 100644 --- a/paimon-python/pypaimon/ray/update_by_row_id.py +++ b/paimon-python/pypaimon/ray/update_by_row_id.py @@ -27,6 +27,7 @@ import hashlib import json import logging +import time import zlib from typing import Any, Dict, List, Optional @@ -57,6 +58,7 @@ _CHECKPOINT_TAG_PREFIX = "_pypaimon_ray_update_" _CHECKPOINT_VERSION = 1 _RANGES_CODEC = "zlib-delta-json-v1" +_CHECKPOINT_TAG_UPDATE_ATTEMPTS = 3 def _blob_col_names(table: "FileStoreTable") -> set: @@ -737,6 +739,23 @@ def _find_operation_snapshot(table, commit_user, commit_identifier): def _store_operation_checkpoint( catalog, target, checkpoint_tags, snapshot): + for attempt in range(_CHECKPOINT_TAG_UPDATE_ATTEMPTS): + try: + _store_operation_checkpoint_once( + catalog, target, checkpoint_tags, snapshot) + return + except Exception: + if attempt + 1 == _CHECKPOINT_TAG_UPDATE_ATTEMPTS: + raise + logger.warning( + "Retry resumable update checkpoint tag.", + exc_info=True, + ) + time.sleep(0.1 * (2 ** attempt)) + + +def _store_operation_checkpoint_once( + catalog, target, checkpoint_tags, snapshot): slot = snapshot.commit_identifier % 2 checkpoint_tag = checkpoint_tags[slot] previous = _get_checkpoint_tag(catalog, target, checkpoint_tag) diff --git a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py index c2fa4135c741..070cd064a457 100644 --- a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py +++ b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py @@ -640,7 +640,7 @@ def test_resume_recovers_commit_before_checkpoint_tag_update(self): def fail_progress_tag(catalog, *args, **kwargs): nonlocal create_calls create_calls += 1 - if create_calls == 3: + if create_calls >= 3: raise RuntimeError("tag update failed") return create_tag(catalog, *args, **kwargs) @@ -661,6 +661,7 @@ def fail_progress_tag(catalog, *args, **kwargs): operation_id=operation_id, ) + self.assertEqual(5, create_calls) table = self.catalog.get_table(target) partial_snapshot_id = ( table.snapshot_manager().get_latest_snapshot().id) @@ -686,6 +687,26 @@ def fail_progress_tag(catalog, *args, **kwargs): table.snapshot_manager().get_latest_snapshot().id, ) + def test_checkpoint_tag_update_retries_transient_failure(self): + import importlib + + module = importlib.import_module( + "pypaimon.ray.update_by_row_id") + snapshot = types.SimpleNamespace(id=2, commit_identifier=1) + error = RuntimeError("transient tag failure") + + with mock.patch.object( + module, + "_store_operation_checkpoint_once", + side_effect=[error, None]) as store, \ + mock.patch.object(module.time, "sleep") as sleep: + module._store_operation_checkpoint( + mock.Mock(), "default.target", + ("tag-0", "tag-1"), snapshot) + + self.assertEqual(2, store.call_count) + sleep.assert_called_once_with(0.1) + def test_atomic_group_failure_commits_nothing(self): target = self._create() for row_id in range(1, 4): From 3b86f9766a5539acd5f554da4997a44330228daf Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 30 Jul 2026 08:52:40 -0700 Subject: [PATCH 09/11] [python] Clean up deterministic commit failures --- .../pypaimon/ray/data_evolution_merge_join.py | 16 +++++----- .../pypaimon/tests/file_store_commit_test.py | 32 +++++++++++++++++++ .../pypaimon/write/file_store_commit.py | 30 +++++++++++++---- 3 files changed, 64 insertions(+), 14 deletions(-) diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_join.py b/paimon-python/pypaimon/ray/data_evolution_merge_join.py index d12f31d45293..58fd405fd0b0 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_join.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_join.py @@ -644,20 +644,20 @@ def _apply_group(group: pa.Table) -> pa.Table: _precomputed_files_info=files_info, ) msgs = worker.update_columns(for_update, list(captured_cols)) - - return _group_result( - pickle.dumps(msgs), - for_update.num_rows, - pickle.dumps(row_ids), - pickle.dumps(planned_file_signatures), - None, - ) except Exception as error: if not capture_group_errors: raise return _group_result( b"", 0, b"", b"", _group_error_text(error)) + return _group_result( + pickle.dumps(msgs), + for_update.num_rows, + pickle.dumps(row_ids), + pickle.dumps(planned_file_signatures), + None, + ) + # One group per target data file; bounded by file count and num_partitions. group_partitions = max( 1, min(len(captured_sorted), num_partitions) diff --git a/paimon-python/pypaimon/tests/file_store_commit_test.py b/paimon-python/pypaimon/tests/file_store_commit_test.py index d98873b55ad2..e49bfbd1e7ad 100644 --- a/paimon-python/pypaimon/tests/file_store_commit_test.py +++ b/paimon-python/pypaimon/tests/file_store_commit_test.py @@ -537,6 +537,22 @@ def test_snapshot_commit_error_is_passed_to_exit( self.assertIs(RuntimeError, exit_args[0]) self.assertIs(error, exit_args[1]) + def test_rejected_snapshot_commit_cleans_up_manifests( + self, mock_manifest_list_manager, mock_manifest_file_manager): + file_store_commit, snapshot_commit, commit_entry = ( + self._prepare_atomic_attempt()) + snapshot_commit.commit.return_value = False + file_store_commit._clean_up_reuse_tmp_manifests = Mock() + file_store_commit._clean_up_no_reuse_tmp_manifests = Mock() + + result = file_store_commit._try_commit_once( + None, "APPEND", [commit_entry], [], 1, None) + + self.assertFalse(result.is_success()) + self.assertIsNone(result.exception) + file_store_commit._clean_up_reuse_tmp_manifests.assert_called_once() + file_store_commit._clean_up_no_reuse_tmp_manifests.assert_called_once() + def test_null_partition_value( self, mock_manifest_list_manager, mock_manifest_file_manager): from pypaimon.data.timestamp import Timestamp @@ -660,6 +676,22 @@ def test_known_failed_retry_keeps_planned_conflict_deterministic( snapshot, ) + def test_known_failed_retry_exhaustion_is_deterministic( + self, mock_manifest_list_manager, mock_manifest_file_manager): + file_store_commit = self._create_file_store_commit() + file_store_commit.commit_max_retries = 0 + file_store_commit.commit_timeout = 10 ** 9 + snapshot = Mock(id=7) + file_store_commit.snapshot_manager.get_latest_snapshot.return_value = ( + snapshot + ) + file_store_commit._try_commit_once = Mock( + return_value=RetryResult(snapshot)) + + with self.assertRaises(CommitConflictError): + file_store_commit._try_commit( + "APPEND", 11, lambda _snapshot: [Mock()]) + @staticmethod def _to_entries(commit_messages): commit_entries = [] diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index 44718f484842..71122ed9dabc 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -476,10 +476,12 @@ def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, f"after {elapsed_ms} millis with {retry_count} retries, " f"there maybe exist commit conflicts between multiple jobs." ) - if retry_result is not None and retry_result.exception: - raise RuntimeError(error_msg) from retry_result.exception - else: - raise RuntimeError(error_msg) + if retry_result is not None: + if retry_result.exception: + raise RuntimeError( + error_msg) from retry_result.exception + raise CommitConflictError(error_msg) + raise RuntimeError(error_msg) self._commit_retry_wait(retry_count) retry_count += 1 @@ -517,7 +519,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str hash_index_base_snapshot, latest_snapshot_id ) ) - if retry_result is None: + if retry_result is None or retry_result.exception is None: raise CommitConflictError(str(conflict)) from conflict raise conflict @@ -585,7 +587,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str # Rolled back: base/snapshot no longer valid; next attempt # re-scans from scratch (matches Java RollbackRetryResult). return RetryResult(None, conflict_exception) - if retry_result is None: + if retry_result is None or retry_result.exception is None: raise CommitConflictError( str(conflict_exception) ) from conflict_exception @@ -697,6 +699,21 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str logger.warning(f"Exception occurs when preparing snapshot: {e}", exc_info=True) raise RuntimeError(f"Failed to prepare snapshot: {e}") + def clean_up_rejected_commit(): + try: + self._clean_up_reuse_tmp_manifests( + delta_manifest_list, + changelog_manifest_list_name, + new_index_manifest, + ) + self._clean_up_no_reuse_tmp_manifests( + base_manifest_list, merge_new_files) + except Exception: + logger.warning( + "Failed to clean up rejected commit manifests.", + exc_info=True, + ) + success = None commit_error = None entered = False @@ -751,6 +768,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str commit_kind, commit_time_s, ) + clean_up_rejected_commit() return RetryResult( latest_snapshot, None, base_data_files=base_data_files) From 8af3cbd95bd497cad3cd3e143baeede5c0e10697 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 30 Jul 2026 08:59:31 -0700 Subject: [PATCH 10/11] [python][ray] Handle existing empty merge targets --- .../pypaimon/ray/data_evolution_merge_into.py | 6 ++-- .../ray_data_evolution_merge_into_test.py | 33 +++++++++++++++++++ 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/ray/data_evolution_merge_into.py b/paimon-python/pypaimon/ray/data_evolution_merge_into.py index d43ac21f36bc..0114be448c60 100644 --- a/paimon-python/pypaimon/ray/data_evolution_merge_into.py +++ b/paimon-python/pypaimon/ray/data_evolution_merge_into.py @@ -281,6 +281,8 @@ def _build_datasets( # snapshot the caller observed; otherwise concurrent commits in between # would mix data from different snapshots. base_snapshot_id = base_snapshot.id if base_snapshot is not None else None + target_empty = ( + base_snapshot is None or base_snapshot.total_record_count == 0) update_ds = None delete_ds = None @@ -317,7 +319,7 @@ def _build_datasets( # Mirror Spark: matched/not-matched run as two independent joins # (inner / left_anti). One unified left_outer join would force # joined.materialize() to feed both branches, which can OOM on large merges. - if matched_specs and base_snapshot is not None: + if matched_specs and not target_empty: update_cols_union = _union_update_cols(matched_specs) if update_cols_union: update_ds = build_matched_update_ds( @@ -362,7 +364,7 @@ def _build_datasets( catalog_options=ctx.catalog_options, num_partitions=num_partitions, snapshot_id=base_snapshot_id, - target_empty=base_snapshot is None, + target_empty=target_empty, ray_remote_args=ray_remote_args, ) 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 cddb11ac37b1..454cd8549e35 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 @@ -505,6 +505,39 @@ def test_insert_into_empty_target(self): self.assertEqual(out['name'], ['a', 'b', 'c']) self.assertEqual(out['age'], [10, 20, 30]) + def test_update_and_insert_into_existing_empty_snapshot(self): + target = self._create_table() + self._write(target, self._source(ids=(1,))) + write_builder = ( + self.catalog.get_table(target).new_batch_write_builder().overwrite()) + writer = write_builder.new_write() + writer.write_arrow(pa.Table.from_pydict( + { + 'id': pa.array([], type=pa.int32()), + 'name': pa.array([], type=pa.string()), + 'age': pa.array([], type=pa.int32()), + }, + schema=self.pa_schema, + )) + write_builder.new_commit().commit(writer.prepare_commit()) + writer.close() + + metrics = merge_into( + target=target, + source=self._source(ids=(1, 2)), + catalog_options=self.catalog_options, + on=['id'], + when_matched=[WhenMatched.update('*')], + when_not_matched=[WhenNotMatched(insert='*')], + num_partitions=_TEST_NUM_PARTITIONS, + ) + + self.assertEqual([1, 2], self._read_sorted(target)['id']) + self.assertEqual( + {'num_matched': 0, 'num_inserted': 2, 'num_unchanged': 0}, + metrics, + ) + def test_multi_source_match_raises_by_default(self): # One target row matched by several source rows: the winning value is # undefined (Spark DE's checkCardinality=false), so we refuse by default. From b16a26c871a36ec64eb22fe9aefd606173403f66 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 30 Jul 2026 08:59:31 -0700 Subject: [PATCH 11/11] [python] Abort deterministic atomic commit rejections --- .../pypaimon/tests/file_store_commit_test.py | 20 +++++++++ .../tests/ray_update_by_row_id_test.py | 44 +++++++++++++++++++ .../pypaimon/write/file_store_commit.py | 39 ++++++++++++++++ 3 files changed, 103 insertions(+) diff --git a/paimon-python/pypaimon/tests/file_store_commit_test.py b/paimon-python/pypaimon/tests/file_store_commit_test.py index e49bfbd1e7ad..29b1bce40cd5 100644 --- a/paimon-python/pypaimon/tests/file_store_commit_test.py +++ b/paimon-python/pypaimon/tests/file_store_commit_test.py @@ -20,6 +20,7 @@ from datetime import datetime from unittest.mock import MagicMock, Mock, patch +from pypaimon.api.rest_exception import BadRequestException from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.manifest.schema.manifest_entry import ManifestEntry from pypaimon.snapshot.snapshot_commit import PartitionStatistics @@ -553,6 +554,25 @@ def test_rejected_snapshot_commit_cleans_up_manifests( file_store_commit._clean_up_reuse_tmp_manifests.assert_called_once() file_store_commit._clean_up_no_reuse_tmp_manifests.assert_called_once() + def test_deterministic_atomic_error_is_safe_to_abort( + self, mock_manifest_list_manager, mock_manifest_file_manager): + rejection = BadRequestException("rejected") + error = RuntimeError("commit failed") + error.__cause__ = rejection + file_store_commit, snapshot_commit, commit_entry = ( + self._prepare_atomic_attempt()) + snapshot_commit.commit.side_effect = error + file_store_commit._clean_up_reuse_tmp_manifests = Mock() + file_store_commit._clean_up_no_reuse_tmp_manifests = Mock() + + with self.assertRaises(CommitConflictError) as raised: + file_store_commit._try_commit_once( + None, "APPEND", [commit_entry], [], 1, None) + + self.assertIs(error, raised.exception.__cause__) + file_store_commit._clean_up_reuse_tmp_manifests.assert_called_once() + file_store_commit._clean_up_no_reuse_tmp_manifests.assert_called_once() + def test_null_partition_value( self, mock_manifest_list_manager, mock_manifest_file_manager): from pypaimon.data.timestamp import Timestamp diff --git a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py index 070cd064a457..154f6087856b 100644 --- a/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py +++ b/paimon-python/pypaimon/tests/ray_update_by_row_id_test.py @@ -880,6 +880,50 @@ def test_incremental_commit_conflict_aborts_buffered_group_files(self): self._read(target).sort_by("id")["age"].to_pylist(), ) + def test_deterministic_atomic_rejection_aborts_group_files(self): + from pypaimon.catalog.catalog_exception import ( + TableNoPermissionException, + ) + from pypaimon.common.identifier import Identifier + from pypaimon.snapshot.renaming_snapshot_commit import ( + RenamingSnapshotCommit, + ) + + target = self._create() + self._write(target, pa.Table.from_pydict( + {"id": [1], "name": ["a"], "age": [0]}, + schema=self.pa_schema, + )) + table = self.catalog.get_table(target) + snapshot_id = table.snapshot_manager().get_latest_snapshot().id + files_before = self._data_files_under(table) + row_id = self._rowid_by_id(target)[1] + error = TableNoPermissionException( + Identifier.create("default", "target")) + + with mock.patch.object( + RenamingSnapshotCommit, "commit", side_effect=error): + with self.assertRaises(TableNoPermissionException): + update_by_row_id( + target, + pa.table( + {"_ROW_ID": [row_id], "age": [100]}, + schema=pa.schema([ + ("_ROW_ID", pa.int64()), + ("age", pa.int32()), + ]), + ), + self.catalog_options, + update_cols=["age"], + num_partitions=1, + ) + + self.assertEqual(files_before, self._data_files_under(table)) + self.assertEqual( + snapshot_id, + table.snapshot_manager().get_latest_snapshot().id, + ) + def test_pins_base_snapshot_for_conflict_detection(self): # The update pins its base snapshot and threads it to distributed_update_apply, # which uses it for commit-time conflict detection against concurrent writers. diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index 71122ed9dabc..885af4a9389a 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -21,6 +21,17 @@ import uuid from typing import Dict, List, Optional +from pypaimon.api.rest_exception import ( + BadRequestException, + ForbiddenException, + NoSuchResourceException, + NotAuthorizedException, + NotImplementedException, +) +from pypaimon.catalog.catalog_exception import ( + TableNoPermissionException, + TableNotExistException, +) from pypaimon.common.options.core_options import CoreOptions from pypaimon.common.predicate_builder import PredicateBuilder from pypaimon.manifest.manifest_file_manager import ManifestFileManager @@ -55,6 +66,30 @@ logger = logging.getLogger(__name__) +_DETERMINISTIC_ATOMIC_COMMIT_EXCEPTIONS = ( + BadRequestException, + ForbiddenException, + NoSuchResourceException, + NotAuthorizedException, + NotImplementedException, + TableNoPermissionException, + TableNotExistException, + NotImplementedError, +) + + +def _is_deterministic_atomic_commit_failure(error): + current = error + seen = set() + while current is not None and id(current) not in seen: + seen.add(id(current)) + if isinstance( + current, _DETERMINISTIC_ATOMIC_COMMIT_EXCEPTIONS): + return True + current = current.__cause__ + return False + + class CommitResult: """Base class for commit results.""" @@ -743,6 +778,10 @@ def clean_up_rejected_commit(): ) if commit_error is not None: + if _is_deterministic_atomic_commit_failure(commit_error): + clean_up_rejected_commit() + raise CommitConflictError( + str(commit_error)) from commit_error # Commit exception, not sure about the situation and should not clean up the files logger.warning( "Retry commit for exception.",