[python][ray] Incrementally commit update_by_row_id file groups

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": <rows>}`.
 
@@ -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
 

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,46 +562,59 @@ 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:
             return pa.Table.from_pydict({
                 "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
 
 

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": <rows>}``.
     """
     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