diff --git a/paimon-python/pypaimon/tests/table_update_test.py b/paimon-python/pypaimon/tests/table_update_test.py index bb1feebd2d82..adb08dd995a4 100644 --- a/paimon-python/pypaimon/tests/table_update_test.py +++ b/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, diff --git a/paimon-python/pypaimon/write/table_update_by_row_id.py b/paimon-python/pypaimon/write/table_update_by_row_id.py index 623668ad1192..014240799318 100644 --- a/paimon-python/pypaimon/write/table_update_by_row_id.py +++ b/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,