[python] Support snapshot properties in table commits

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.

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"},
+        )

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)

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,20 +86,15 @@ 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
             logger.info(
                 "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)