for UpdateSpec

Author: KazuhitoT

pyiceberg/table/update/spec.py

@@ -58,34 +58,48 @@ class UpdateSpec(UpdateTableMetadata["UpdateSpec"]):
     _adds: list[PartitionField]
     _deletes: set[int]
     _last_assigned_partition_id: int
+    # Store (source_column_name, transform, partition_field_name) for retry support
+    _field_additions: list[tuple[str, Transform[Any, Any], str | None]]
 
     def __init__(self, transaction: Transaction, case_sensitive: bool = True) -> None:
         super().__init__(transaction)
-        self._name_to_field = {field.name: field for field in transaction.table_metadata.spec().fields}
-        self._name_to_added_field = {}
-        self._transform_to_field = {
-            (field.source_id, repr(field.transform)): field for field in transaction.table_metadata.spec().fields
-        }
-        self._transform_to_added_field = {}
-        self._adds = []
-        self._deletes = set()
-        self._last_assigned_partition_id = transaction.table_metadata.last_partition_id or PARTITION_FIELD_ID_START - 1
-        self._renames = {}
         self._transaction = transaction
         self._case_sensitive = case_sensitive
+        self._field_additions = []
+        self._deletes = set()
+        self._renames = {}
+        # Initialize state from current metadata
+        self._init_state_from_metadata()
+
+    def _init_state_from_metadata(self) -> None:
+        """Initialize or reinitialize state from current transaction metadata."""
+        spec = self._transaction.table_metadata.spec()
+        self._name_to_field = {field.name: field for field in spec.fields}
+        self._transform_to_field = {(field.source_id, repr(field.transform)): field for field in spec.fields}
+        self._last_assigned_partition_id = self._transaction.table_metadata.last_partition_id or PARTITION_FIELD_ID_START - 1
+        # Clear intermediate state
+        self._name_to_added_field = {}
+        self._transform_to_added_field = {}
         self._added_time_fields = {}
+        self._adds = []
 
-    def add_field(
+    def _reset_state(self) -> None:
+        """Reset state for retry, rebuilding from refreshed metadata.
+
+        This is called on transaction retry to reapply the spec changes on top of the refreshed table metadata."""
+        self._init_state_from_metadata()
+        for source_column_name, transform, partition_field_name in self._field_additions:
+            self._do_add_field(source_column_name, transform, partition_field_name)
+
+    def _do_add_field(
         self,
         source_column_name: str,
-        transform: str | Transform[Any, Any],
-        partition_field_name: str | None = None,
-    ) -> UpdateSpec:
+        transform: Transform[Any, Any],
+        partition_field_name: str | None,
+    ) -> None:
         ref = Reference(source_column_name)
         bound_ref = ref.bind(self._transaction.table_metadata.schema(), self._case_sensitive)
-        if isinstance(transform, str):
-            transform = parse_transform(transform)
-        # verify transform can actually bind it
+
         output_type = bound_ref.field.field_type
         if not transform.can_transform(output_type):
             raise ValueError(f"{transform} cannot transform {output_type} values from {bound_ref.field.name}")
@@ -121,6 +135,16 @@ def add_field(
 
         self._name_to_added_field[new_field.name] = new_field
         self._adds.append(new_field)
+
+    def add_field(
+        self,
+        source_column_name: str,
+        transform: str | Transform[Any, Any],
+        partition_field_name: str | None = None,
+    ) -> UpdateSpec:
+        transform = parse_transform(transform)
+        self._field_additions.append((source_column_name, transform, partition_field_name))
+        self._do_add_field(source_column_name, transform, partition_field_name)
         return self
 
     def add_identity(self, source_column_name: str) -> UpdateSpec:
@@ -178,6 +202,10 @@ def _commit(self) -> UpdatesAndRequirements:
 
         return updates, requirements
 
+    def commit(self) -> None:
+        updates, requirements = self._commit()
+        self._transaction._apply(updates, requirements, pending_update=self)
+
     def _apply(self) -> PartitionSpec:
         def _check_and_add_partition_name(
             schema: Schema, name: str, source_id: int, transform: Transform[Any, Any], partition_names: set[str]

tests/table/test_commit_retry.py

@@ -997,3 +997,135 @@ def mock_commit(
         result = table.scan().to_arrow()
         assert len(result) == 3
         assert result["id"].to_pylist() == [4, 5, 6]
+
+
+class TestUpdateSpecRetry:
+    def test_update_spec_retried_on_conflict(self, catalog: SqlCatalog, schema: Schema) -> None:
+        """Test that UpdateSpec operations are retried on CommitFailedException."""
+        from pyiceberg.transforms import BucketTransform
+
+        table = catalog.create_table(
+            "default.test_spec_retry",
+            schema=schema,
+            properties={
+                TableProperties.COMMIT_NUM_RETRIES: "3",
+                TableProperties.COMMIT_MIN_RETRY_WAIT_MS: "1",
+                TableProperties.COMMIT_MAX_RETRY_WAIT_MS: "10",
+            },
+        )
+
+        original_commit = catalog.commit_table
+        commit_count = 0
+
+        def mock_commit(
+            tbl: Table, requirements: tuple[TableRequirement, ...], updates: tuple[TableUpdate, ...]
+        ) -> CommitTableResponse:
+            nonlocal commit_count
+            commit_count += 1
+            if commit_count == 1:
+                raise CommitFailedException("Simulated spec conflict")
+            return original_commit(tbl, requirements, updates)
+
+        with patch.object(catalog, "commit_table", side_effect=mock_commit):
+            with table.update_spec() as update_spec:
+                update_spec.add_field(
+                    source_column_name="id", transform=BucketTransform(16), partition_field_name="id_bucket"
+                )
+
+        assert commit_count == 2
+
+    def test_update_spec_resolves_conflict_on_retry(self, catalog: SqlCatalog, schema: Schema) -> None:
+        """Test that spec update can resolve conflicts via retry"""
+        from pyiceberg.transforms import BucketTransform
+
+        table = catalog.create_table(
+            "default.test_spec_conflict_resolved",
+            schema=schema,
+            properties={
+                TableProperties.COMMIT_NUM_RETRIES: "5",
+                TableProperties.COMMIT_MIN_RETRY_WAIT_MS: "1",
+                TableProperties.COMMIT_MAX_RETRY_WAIT_MS: "10",
+            },
+        )
+
+        with table.update_spec() as update_spec:
+            update_spec.add_field(source_column_name="id", transform=BucketTransform(16), partition_field_name="id_bucket")
+
+        table2 = catalog.load_table("default.test_spec_conflict_resolved")
+        with table2.update_spec() as update_spec2:
+            update_spec2.add_identity("id")
+
+        assert table.spec().spec_id == 1
+        assert table2.spec().spec_id == 2
+
+        original_commit = catalog.commit_table
+        commit_count = 0
+
+        def mock_commit(
+            tbl: Table, requirements: tuple[TableRequirement, ...], updates: tuple[TableUpdate, ...]
+        ) -> CommitTableResponse:
+            nonlocal commit_count
+            commit_count += 1
+            return original_commit(tbl, requirements, updates)
+
+        with patch.object(catalog, "commit_table", side_effect=mock_commit):
+            # Retry resolves conflicts caused by mismatch spec_id
+            with table.update_spec() as update_spec:
+                update_spec.add_field(source_column_name="id", transform=BucketTransform(8), partition_field_name="id_bucket_new")
+
+        assert commit_count == 2
+
+    def test_transaction_with_spec_change_and_append_retries(
+        self, catalog: SqlCatalog, schema: Schema, arrow_table: pa.Table
+    ) -> None:
+        """Test that a transaction with spec change and append handles retry correctly."""
+        table = catalog.create_table(
+            "default.test_transaction_spec_and_append",
+            schema=schema,
+            properties={
+                TableProperties.COMMIT_NUM_RETRIES: "3",
+                TableProperties.COMMIT_MIN_RETRY_WAIT_MS: "1",
+                TableProperties.COMMIT_MAX_RETRY_WAIT_MS: "10",
+            },
+        )
+
+        original_commit = catalog.commit_table
+        commit_count = 0
+        captured_updates: list[tuple[TableUpdate, ...]] = []
+
+        def mock_commit(
+            tbl: Table, requirements: tuple[TableRequirement, ...], updates: tuple[TableUpdate, ...]
+        ) -> CommitTableResponse:
+            nonlocal commit_count
+            commit_count += 1
+            captured_updates.append(updates)
+            if commit_count == 1:
+                raise CommitFailedException("Simulated conflict")
+            return original_commit(tbl, requirements, updates)
+
+        with patch.object(catalog, "commit_table", side_effect=mock_commit):
+            with table.transaction() as txn:
+                with txn.update_spec() as update_spec:
+                    update_spec.add_identity("id")
+                txn.append(arrow_table)
+
+        assert commit_count == 2
+
+        first_attempt_update_types = [type(u).__name__ for u in captured_updates[0]]
+        assert "AddPartitionSpecUpdate" in first_attempt_update_types
+        assert "AddSnapshotUpdate" in first_attempt_update_types
+
+        retry_attempt_update_types = [type(u).__name__ for u in captured_updates[1]]
+        assert "AddPartitionSpecUpdate" in retry_attempt_update_types
+        assert "AddSnapshotUpdate" in retry_attempt_update_types
+
+        assert len(table.scan().to_arrow()) == 3
+
+        from pyiceberg.transforms import IdentityTransform
+
+        assert table.spec().spec_id == 1
+        assert len(table.spec().fields) == 1
+        partition_field = table.spec().fields[0]
+        assert partition_field.name == "id"
+        assert partition_field.source_id == 1  # "id" column's field_id
+        assert isinstance(partition_field.transform, IdentityTransform)