From fa65a4c054e515b7eeb6fd12e4fc1ef1fd3e6559 Mon Sep 17 00:00:00 2001 From: Yann Date: Tue, 8 Sep 2026 21:55:57 +0800 Subject: [PATCH] feat(python): support snapshot properties in table commits Allow batch and stream commits to atomically attach application state to the generated Paimon snapshot, matching the capability available through Java ManifestCommittable. Co-Authored-By: Codex AI-Model: gpt-5 Co-Authored-By: Codex Co-Authored-By: Codex AI-Contributed/Feature: 76/76 AI-Contributed/UT: 101/101 --- .../pypaimon/tests/table/simple_table_test.py | 51 +++++++++++++++++++ .../pypaimon/tests/table_commit_test.py | 50 +++++++++++++++++- .../pypaimon/write/file_store_commit.py | 28 ++++++++-- paimon-python/pypaimon/write/table_commit.py | 48 ++++++++++++----- 4 files changed, 156 insertions(+), 21 deletions(-) diff --git a/paimon-python/pypaimon/tests/table/simple_table_test.py b/paimon-python/pypaimon/tests/table/simple_table_test.py index 5e13579c3349..402db8fb155c 100644 --- a/paimon-python/pypaimon/tests/table/simple_table_test.py +++ b/paimon-python/pypaimon/tests/table/simple_table_test.py @@ -52,6 +52,57 @@ def setUpClass(cls): def tearDownClass(cls): shutil.rmtree(cls.tempdir, ignore_errors=True) + def test_commit_snapshot_properties(self): + schema = Schema.from_pyarrow_schema(self.pa_schema) + self.catalog.create_table( + 'default.test_commit_snapshot_properties', schema, False) + table = self.catalog.get_table( + 'default.test_commit_snapshot_properties') + write_builder = table.new_batch_write_builder() + table_write = write_builder.new_write() + table_commit = write_builder.new_commit() + table_write.write_arrow(pa.Table.from_pydict({ + 'pt': [1], + 'k': [2], + 'v': [3], + }, schema=self.pa_schema)) + + table_commit.commit( + table_write.prepare_commit(), + snapshot_properties={'source': 'capture'}, + ) + table_write.close() + table_commit.close() + + snapshot = table.snapshot_manager().get_latest_snapshot() + self.assertEqual({'source': 'capture'}, snapshot.properties) + + def test_stream_commit_snapshot_properties(self): + schema = Schema.from_pyarrow_schema(self.pa_schema) + self.catalog.create_table( + 'default.test_stream_commit_snapshot_properties', schema, False) + table = self.catalog.get_table( + 'default.test_stream_commit_snapshot_properties') + write_builder = table.new_stream_write_builder() + table_write = write_builder.new_write() + table_commit = write_builder.new_commit() + table_write.write_arrow(pa.Table.from_pydict({ + 'pt': [1], + 'k': [2], + 'v': [3], + }, schema=self.pa_schema)) + + table_commit.commit( + table_write.prepare_commit(42), + 42, + snapshot_properties={'checkpoint': '42'}, + ) + table_write.close() + table_commit.close() + + snapshot = table.snapshot_manager().get_latest_snapshot() + self.assertEqual({'checkpoint': '42'}, snapshot.properties) + def test_tag_scan(self): """ Test reading from a specific tag. diff --git a/paimon-python/pypaimon/tests/table_commit_test.py b/paimon-python/pypaimon/tests/table_commit_test.py index d0b9d62762df..d607e73921cf 100644 --- a/paimon-python/pypaimon/tests/table_commit_test.py +++ b/paimon-python/pypaimon/tests/table_commit_test.py @@ -25,8 +25,7 @@ from pypaimon.write.table_commit import BatchTableCommit, StreamTableCommit -class TestTableCommitEmptyOverwrite(unittest.TestCase): - """Tests for TableCommit._commit handling of empty commit messages in overwrite mode.""" +class TestTableCommit(unittest.TestCase): def _create_commit(self, cls, overwrite_partition=None): commit = cls.__new__(cls) @@ -88,6 +87,35 @@ def test_append_forwards_non_empty_messages(self, name, msg_flags): mock_fsc.commit.assert_not_called() mock_fsc.overwrite.assert_not_called() + def test_batch_commit_forwards_snapshot_properties(self): + commit, mock_fsc = self._create_commit( + BatchTableCommit, overwrite_partition=None) + message = CommitMessage( + partition=(), bucket=0, new_files=[Mock()]) + + commit.commit([message], snapshot_properties={"source": "capture"}) + + mock_fsc.commit.assert_called_once_with( + commit_messages=[message], + commit_identifier=BATCH_COMMIT_IDENTIFIER, + snapshot_properties={"source": "capture"}, + ) + + def test_overwrite_forwards_snapshot_properties(self): + commit, mock_fsc = self._create_commit( + BatchTableCommit, overwrite_partition={"dt": "2024-01-15"}) + message = CommitMessage( + partition=("2024-01-15",), bucket=0, new_files=[Mock()]) + + commit.commit([message], snapshot_properties={"source": "capture"}) + + mock_fsc.overwrite.assert_called_once_with( + overwrite_partition={"dt": "2024-01-15"}, + commit_messages=[message], + commit_identifier=BATCH_COMMIT_IDENTIFIER, + snapshot_properties={"source": "capture"}, + ) + # -- StreamTableCommit overwrite should also reach overwrite() with empty messages -- def test_stream_commit_overwrite_empty_messages(self): @@ -100,3 +128,21 @@ def test_stream_commit_overwrite_empty_messages(self): commit_messages=[], commit_identifier=42, ) + + def test_stream_commit_forwards_snapshot_properties(self): + commit, mock_fsc = self._create_commit( + StreamTableCommit, overwrite_partition=None) + message = CommitMessage( + partition=(), bucket=0, new_files=[Mock()]) + + commit.commit( + [message], + commit_identifier=42, + snapshot_properties={"checkpoint": "42"}, + ) + + mock_fsc.commit.assert_called_once_with( + commit_messages=[message], + commit_identifier=42, + snapshot_properties={"checkpoint": "42"}, + ) diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index fbcda7fffe84..917f4286cd73 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -250,7 +250,11 @@ def __init__(self, snapshot_commit: SnapshotCommit, table, commit_user: str, table_rollback = table.catalog_environment.catalog_table_rollback() self.rollback = CommitRollback(table_rollback) if table_rollback is not None else None - def commit(self, commit_messages: List[CommitMessage], commit_identifier: int): + def commit( + self, + commit_messages: List[CommitMessage], + commit_identifier: int, + snapshot_properties: Optional[Dict[str, str]] = None): """Commit the given commit messages in normal append mode.""" if not commit_messages: return @@ -334,9 +338,15 @@ def commit(self, commit_messages: List[CommitMessage], commit_identifier: int): allow_rollback=allow_rollback, index_deletes=index_deletes, index_adds=index_adds, - hash_index_base_snapshot=hash_index_base_snapshot) + hash_index_base_snapshot=hash_index_base_snapshot, + snapshot_properties=snapshot_properties) - def overwrite(self, overwrite_partition, commit_messages: List[CommitMessage], commit_identifier: int): + def overwrite( + self, + overwrite_partition, + commit_messages: List[CommitMessage], + commit_identifier: int, + snapshot_properties: Optional[Dict[str, str]] = None): """Commit the given commit messages in overwrite mode.""" logger.info( "Ready to overwrite to table %s, number of commit messages: %d", @@ -382,6 +392,7 @@ def overwrite(self, overwrite_partition, commit_messages: List[CommitMessage], c index_deletes=index_deletes, index_adds=index_adds, hash_index_base_snapshot=hash_index_base_snapshot, + snapshot_properties=snapshot_properties, ) @staticmethod @@ -487,7 +498,8 @@ def truncate_table(self, commit_identifier: int) -> None: def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, detect_conflicts=False, allow_rollback=False, index_deletes=None, index_adds=None, changelog_entries=None, - hash_index_base_snapshot=None): + hash_index_base_snapshot=None, + snapshot_properties: Optional[Dict[str, str]] = None): retry_count = 0 retry_result = None @@ -528,6 +540,7 @@ def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, index_adds=index_adds, hash_index_base_snapshot=hash_index_base_snapshot, commit_result_may_be_uncertain=commit_result_may_be_uncertain, + snapshot_properties=snapshot_properties, ) if isinstance(result, RewriteResult): @@ -606,7 +619,9 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str index_deletes=None, index_adds=None, hash_index_base_snapshot=None, - commit_result_may_be_uncertain: bool = False) -> CommitResult: + commit_result_may_be_uncertain: bool = False, + snapshot_properties: Optional[Dict[str, str]] = None + ) -> CommitResult: start_millis = int(time.time() * 1000) if self._is_duplicate_commit( retry_result, @@ -807,6 +822,9 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str latest_snapshot.watermark if latest_snapshot else None), next_row_id=next_row_id, index_manifest=index_manifest, + properties=( + dict(snapshot_properties) + if snapshot_properties else None), ) # 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 3215f7eea7d6..f6eb6748cd0a 100644 --- a/paimon-python/pypaimon/write/table_commit.py +++ b/paimon-python/pypaimon/write/table_commit.py @@ -19,12 +19,12 @@ from typing import Dict, List, Optional from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER - -logger = logging.getLogger(__name__) from pypaimon.write.commit_callback import CommitCallback from pypaimon.write.commit_message import CommitMessage from pypaimon.write.file_store_commit import FileStoreCommit +logger = logging.getLogger(__name__) + class TableCommit: """Common base for batch and stream table commits. @@ -60,8 +60,18 @@ def add_commit_callback(self, callback: CommitCallback) -> None: """Register a callback to be invoked after each successful commit.""" self._commit_callbacks.append(callback) - def _commit(self, commit_messages: List[CommitMessage], commit_identifier: int = BATCH_COMMIT_IDENTIFIER): + def _commit( + self, + commit_messages: List[CommitMessage], + commit_identifier: int = BATCH_COMMIT_IDENTIFIER, + snapshot_properties: Optional[Dict[str, str]] = None): non_empty_messages = [msg for msg in commit_messages if not msg.is_empty()] + commit_kwargs = { + "commit_messages": non_empty_messages, + "commit_identifier": commit_identifier, + } + if snapshot_properties is not None: + commit_kwargs["snapshot_properties"] = snapshot_properties # Never abort files in response to a commit exception. Preserving # possible orphan files is safer than deleting files which another @@ -76,9 +86,7 @@ def _commit(self, commit_messages: List[CommitMessage], commit_identifier: int = ) self.file_store_commit.overwrite( overwrite_partition=self.overwrite_partition, - commit_messages=non_empty_messages, - commit_identifier=commit_identifier - ) + **commit_kwargs) else: if not non_empty_messages: return @@ -86,10 +94,7 @@ def _commit(self, commit_messages: List[CommitMessage], commit_identifier: int = "Committing table %s, %d non-empty messages", self.table.identifier, len(non_empty_messages) ) - self.file_store_commit.commit( - commit_messages=non_empty_messages, - commit_identifier=commit_identifier - ) + self.file_store_commit.commit(**commit_kwargs) def abort(self, commit_messages: List[CommitMessage]): self.file_store_commit.abort(commit_messages) @@ -105,9 +110,16 @@ def __init__(self, table, commit_user: str, static_partition: Optional[dict]): super().__init__(table, commit_user, static_partition) self.batch_committed = False - def commit(self, commit_messages: List[CommitMessage]): + def commit( + self, + commit_messages: List[CommitMessage], + snapshot_properties: Optional[Dict[str, str]] = None): + """Commit once, attaching optional properties to the snapshot.""" self._check_committed() - self._commit(commit_messages, BATCH_COMMIT_IDENTIFIER) + self._commit( + commit_messages, + BATCH_COMMIT_IDENTIFIER, + snapshot_properties=snapshot_properties) def truncate_table(self) -> None: """Truncate the entire table, deleting all data.""" @@ -132,5 +144,13 @@ class StreamTableCommit(TableCommit): :meth:`StreamTableWrite.prepare_commit`. """ - def commit(self, commit_messages: List[CommitMessage], commit_identifier: int): - self._commit(commit_messages, commit_identifier) + def commit( + self, + commit_messages: List[CommitMessage], + commit_identifier: int, + snapshot_properties: Optional[Dict[str, str]] = None): + """Commit a stream checkpoint with optional snapshot properties.""" + self._commit( + commit_messages, + commit_identifier, + snapshot_properties=snapshot_properties)