[python] Abort failed update-by-row-id writes

paimon-python/pypaimon/tests/table_update_test.py

@@ -15,10 +15,12 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import os
 import random
 import string
 import threading
 import unittest
+from unittest import mock
 
 import pyarrow as pa
 
@@ -1275,6 +1277,58 @@ def _apply_delete_by_row_id(self, table_update, row_ids, cid):
 class TableUpdateBatchTest(_BatchModeMixin, _TableUpdateTestBase, unittest.TestCase):
     """All shared update tests under batch (``BatchWriteBuilder``) semantics."""
 
+    def test_update_by_row_id_aborts_files_after_prepare_commit_failure(self):
+        from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+
+        table_schema = pa.schema([
+            ('id', pa.int32()),
+            ('age', pa.int32()),
+            ('picture', pa.large_binary()),
+        ])
+        table = self._create_table(pa_schema=table_schema)
+        self._write_arrow(table, pa.Table.from_pydict({
+            'id': [1, 2],
+            'age': [10, 20],
+            'picture': [b'blob-1', b'blob-2'],
+        }, schema=table_schema))
+
+        rb = table.new_read_builder().with_projection(['id', '_ROW_ID'])
+        row_ids = rb.new_read().to_arrow(
+            rb.new_scan().plan().splits()).sort_by('id')['_ROW_ID']
+        before_files = self._list_table_files(table)
+
+        def fail_after_prepare_commit(
+                new_files, first_row_id, column_names, blob_columns):
+            raise RuntimeError("forced failure after prepare_commit")
+
+        wb = self._make_write_builder(table)
+        tu = wb.new_update().with_update_type(['age', 'picture'])
+        with mock.patch.object(
+                TableUpdateByRowId,
+                '_assign_update_file_metadata',
+                new=staticmethod(fail_after_prepare_commit)):
+            with self.assertRaisesRegex(
+                    RuntimeError, "forced failure after prepare_commit"):
+                self._apply_update(tu, pa.Table.from_pydict({
+                    '_ROW_ID': [row_ids[0].as_py()],
+                    'age': [99],
+                    'picture': [b'updated-blob'],
+                }, schema=pa.schema([
+                    ('_ROW_ID', pa.int64()),
+                    ('age', pa.int32()),
+                    ('picture', pa.large_binary()),
+                ])), self._next_commit_id())
+
+        self.assertEqual(before_files, self._list_table_files(table))
+
+    @staticmethod
+    def _list_table_files(table):
+        return {
+            os.path.relpath(os.path.join(root, name), table.table_path)
+            for root, _dirs, files in os.walk(table.table_path)
+            for name in files
+        }
+
 
 class TableUpdateStreamTest(_StreamModeMixin, _TableUpdateTestBase, unittest.TestCase):
     """All shared update tests under stream (``StreamWriteBuilder``) semantics,

paimon-python/pypaimon/write/table_update_by_row_id.py

@@ -538,6 +538,7 @@ def _write_group(
         new_files = []
         file_store_write = None
         blob_writers = []
+        success = False
         try:
             if merged_data is not None:
                 file_store_write = FileStoreWrite(self.table, self.commit_user)
@@ -575,11 +576,18 @@ def _write_group(
                         check_from_snapshot=self.snapshot_id,
                     )
                 )
+            success = True
         finally:
-            if file_store_write is not None:
-                file_store_write.close()
-            for blob_writer in blob_writers:
-                blob_writer.close()
+            if success:
+                if file_store_write is not None:
+                    file_store_write.close()
+                for blob_writer in blob_writers:
+                    blob_writer.close()
+            else:
+                if file_store_write is not None:
+                    file_store_write.abort()
+                for blob_writer in blob_writers:
+                    blob_writer.abort()
 
     @staticmethod
     def _assign_update_file_metadata(new_files: List[DataFileMeta], first_row_id: int,