fix bug and add tests

Author: KazuhitoT

pyiceberg/table/__init__.py

@@ -1032,38 +1032,25 @@ def commit_transaction(self) -> Table:
         return self._table
 
     def _commit_with_retry(self) -> None:
-        """Commit transaction with retry logic for snapshot operations.
+        """Commit transaction with retry logic.
 
         On retry, refreshes table metadata and regenerates snapshots from
         the pending snapshot producers.
         """
         properties = self._table.metadata.properties
 
+        max_attempts = property_as_int(properties, TableProperties.COMMIT_NUM_RETRIES)
+        min_wait_ms = property_as_int(properties, TableProperties.COMMIT_MIN_RETRY_WAIT_MS)
+        max_wait_ms = property_as_int(properties, TableProperties.COMMIT_MAX_RETRY_WAIT_MS)
+        total_timeout_ms = property_as_int(properties, TableProperties.COMMIT_TOTAL_RETRY_TIME_MS)
+
         retry_config = RetryConfig(
-            max_attempts=property_as_int(
-                properties,
-                TableProperties.COMMIT_NUM_RETRIES,
-                TableProperties.COMMIT_NUM_RETRIES_DEFAULT,
-            )
-            or TableProperties.COMMIT_NUM_RETRIES_DEFAULT,
-            min_wait_ms=property_as_int(
-                properties,
-                TableProperties.COMMIT_MIN_RETRY_WAIT_MS,
-                TableProperties.COMMIT_MIN_RETRY_WAIT_MS_DEFAULT,
-            )
-            or TableProperties.COMMIT_MIN_RETRY_WAIT_MS_DEFAULT,
-            max_wait_ms=property_as_int(
-                properties,
-                TableProperties.COMMIT_MAX_RETRY_WAIT_MS,
-                TableProperties.COMMIT_MAX_RETRY_WAIT_MS_DEFAULT,
-            )
-            or TableProperties.COMMIT_MAX_RETRY_WAIT_MS_DEFAULT,
-            total_timeout_ms=property_as_int(
-                properties,
-                TableProperties.COMMIT_TOTAL_RETRY_TIME_MS,
-                TableProperties.COMMIT_TOTAL_RETRY_TIME_MS_DEFAULT,
-            )
-            or TableProperties.COMMIT_TOTAL_RETRY_TIME_MS_DEFAULT,
+            max_attempts=max_attempts if max_attempts is not None else TableProperties.COMMIT_NUM_RETRIES_DEFAULT,
+            min_wait_ms=min_wait_ms if min_wait_ms is not None else TableProperties.COMMIT_MIN_RETRY_WAIT_MS_DEFAULT,
+            max_wait_ms=max_wait_ms if max_wait_ms is not None else TableProperties.COMMIT_MAX_RETRY_WAIT_MS_DEFAULT,
+            total_timeout_ms=total_timeout_ms
+            if total_timeout_ms is not None
+            else TableProperties.COMMIT_TOTAL_RETRY_TIME_MS_DEFAULT,
         )
 
         first_attempt = True
@@ -1073,7 +1060,6 @@ def do_commit() -> None:
             if first_attempt:
                 first_attempt = False
             else:
-                # On retry, reapply all updates with refreshed metadata
                 self._reapply_updates()
 
             self._table._do_commit(  # pylint: disable=W0212

pyiceberg/table/update/schema.py

@@ -90,14 +90,9 @@ class UpdateSchema(UpdateTableMetadata["UpdateSchema"]):
     _allow_incompatible_changes: bool
     _case_sensitive: bool
 
-    # Store user intent for retry support
-    _column_additions: list[tuple[str | tuple[str, ...], IcebergType, str | None, bool, Any]]
-    _column_updates: list[tuple[str | tuple[str, ...], IcebergType | None, bool | None, str | None]]
-    _column_deletions: list[str | tuple[str, ...]]
-    _column_renames: list[tuple[str | tuple[str, ...], str]]
-    _move_operations: list[tuple[str, str | tuple[str, ...], str | tuple[str, ...] | None]]
-    _optional_columns: list[str | tuple[str, ...]]
-    _default_value_updates: list[tuple[str | tuple[str, ...], Any]]
+    # Store all operations and order for retry support.
+    # This _operations overlaps _adds, _updates and other intermediate variables.
+    _operations: list[tuple[Any, ...]]
 
     def __init__(
         self,
@@ -113,18 +108,8 @@ def __init__(
         self._case_sensitive = case_sensitive
         self._name_mapping = name_mapping
         self._provided_schema = schema  # Store for _reset_state
+        self._operations = []  # for retry support
 
-        # Initialize user intent storage
-        self._column_additions = []
-        self._column_updates = []
-        self._column_deletions = []
-        self._column_renames = []
-        self._move_operations = []
-        self._optional_columns = []
-        self._default_value_updates = []
-        self._identifier_field_updates: set[str] | None = None
-
-        # Initialize state from metadata
         self._init_state_from_metadata(schema)
 
     def _init_state_from_metadata(self, schema: Schema | None = None) -> None:
@@ -155,38 +140,48 @@ def get_column_name(field_id: int) -> str:
         }
 
     def _reset_state(self) -> None:
-        """Reset state for retry, rebuilding from refreshed metadata."""
-        self._init_state_from_metadata(self._provided_schema)
-
-        for path, field_type, doc, required, default_value in self._column_additions:
-            self._do_add_column(path, field_type, doc, required, default_value)
-
-        for path in self._column_deletions:
-            self._do_delete_column(path)
-
-        for path_from, new_name in self._column_renames:
-            self._do_rename_column(path_from, new_name)
-
-        for upd_path, upd_field_type, upd_required, upd_doc in self._column_updates:
-            self._do_update_column(upd_path, upd_field_type, upd_required, upd_doc)
+        """Reset state for retry, rebuilding from refreshed metadata.
 
-        for path in self._optional_columns:
-            self._set_column_requirement(path, required=False)
-
-        for path, default_value in self._default_value_updates:
-            self._set_column_default_value(path, default_value)
+        This is called on transaction retry to reapply the schema changes on top of the refreshed table metadata.
+        """
+        self._init_state_from_metadata(self._provided_schema)
 
-        for op, path, other_path in self._move_operations:
-            if op == "first":
+        # Refresh name mapping from the latest table metadata to avoid overwriting concurrent changes
+        if self._name_mapping is not None:
+            self._name_mapping = self._transaction.table_metadata.name_mapping()
+
+        for operation in self._operations:
+            op_type = operation[0]
+            if op_type == "add":
+                _, path, field_type, doc, required, default_value = operation
+                self._do_add_column(path, field_type, doc, required, default_value)
+            elif op_type == "delete":
+                _, path = operation
+                self._do_delete_column(path)
+            elif op_type == "rename":
+                _, path_from, new_name = operation
+                self._do_rename_column(path_from, new_name)
+            elif op_type == "update":
+                _, path, field_type, required, doc = operation
+                self._do_update_column(path, field_type, required, doc)
+            elif op_type == "optional":
+                _, path = operation
+                self._set_column_requirement(path, required=False)
+            elif op_type == "default_value":
+                _, path, default_value = operation
+                self._set_column_default_value(path, default_value)
+            elif op_type == "move_first":
+                _, path = operation
                 self._do_move_first(path)
-            elif op == "before":
-                self._do_move_before(path, other_path)  # type: ignore
-            elif op == "after":
-                self._do_move_after(path, other_path)  # type: ignore
-
-        # Restore identifier fields if they were explicitly set
-        if self._identifier_field_updates is not None:
-            self._identifier_field_names = self._identifier_field_updates.copy()
+            elif op_type == "move_before":
+                _, path, before_path = operation
+                self._do_move_before(path, before_path)
+            elif op_type == "move_after":
+                _, path, after_name = operation
+                self._do_move_after(path, after_name)
+            elif op_type == "set_identifier_fields":
+                _, fields = operation
+                self._identifier_field_names = set(fields)
 
     def case_sensitive(self, case_sensitive: bool) -> UpdateSchema:
         """Determine if the case of schema needs to be considered when comparing column names.
@@ -243,7 +238,7 @@ def add_column(
         Returns:
             This for method chaining.
         """
-        self._column_additions.append((path, field_type, doc, required, default_value))
+        self._operations.append(("add", path, field_type, doc, required, default_value))
         self._do_add_column(path, field_type, doc, required, default_value)
         return self
 
@@ -335,7 +330,7 @@ def delete_column(self, path: str | tuple[str, ...]) -> UpdateSchema:
         Returns:
             The UpdateSchema with the delete operation staged.
         """
-        self._column_deletions.append(path)
+        self._operations.append(("delete", path))
         self._do_delete_column(path)
         return self
 
@@ -362,7 +357,7 @@ def set_default_value(self, path: str | tuple[str, ...], default_value: L | None
         Returns:
             The UpdateSchema with the delete operation staged.
         """
-        self._default_value_updates.append((path, default_value))
+        self._operations.append(("default_value", path, default_value))
         self._set_column_default_value(path, default_value)
         return self
 
@@ -376,7 +371,7 @@ def rename_column(self, path_from: str | tuple[str, ...], new_name: str) -> Upda
         Returns:
             The UpdateSchema with the rename operation staged.
         """
-        self._column_renames.append((path_from, new_name))
+        self._operations.append(("rename", path_from, new_name))
         self._do_rename_column(path_from, new_name)
         return self
 
@@ -425,12 +420,12 @@ def make_column_optional(self, path: str | tuple[str, ...]) -> UpdateSchema:
         Returns:
             The UpdateSchema with the requirement change staged.
         """
-        self._optional_columns.append(path)
+        self._operations.append(("optional", path))
         self._set_column_requirement(path, required=False)
         return self
 
     def set_identifier_fields(self, *fields: str) -> None:
-        self._identifier_field_updates = set(fields)
+        self._operations.append(("set_identifier_fields", fields))
         self._identifier_field_names = set(fields)
 
     def _set_column_requirement(self, path: str | tuple[str, ...], required: bool) -> None:
@@ -535,8 +530,7 @@ def update_column(
         if field_type is None and required is None and doc is None:
             return self
 
-        # Store intent for retry support
-        self._column_updates.append((path, field_type, required, doc))
+        self._operations.append(("update", path, field_type, required, doc))
         self._do_update_column(path, field_type, required, doc)
         return self
 
@@ -633,7 +627,7 @@ def move_first(self, path: str | tuple[str, ...]) -> UpdateSchema:
         Returns:
             The UpdateSchema with the move operation staged.
         """
-        self._move_operations.append(("first", path, None))
+        self._operations.append(("move_first", path))
         self._do_move_first(path)
         return self
 
@@ -657,7 +651,7 @@ def move_before(self, path: str | tuple[str, ...], before_path: str | tuple[str,
         Returns:
             The UpdateSchema with the move operation staged.
         """
-        self._move_operations.append(("before", path, before_path))
+        self._operations.append(("move_before", path, before_path))
         self._do_move_before(path, before_path)
         return self
 
@@ -695,7 +689,7 @@ def move_after(self, path: str | tuple[str, ...], after_name: str | tuple[str, .
         Returns:
             The UpdateSchema with the move operation staged.
         """
-        self._move_operations.append(("after", path, after_name))
+        self._operations.append(("move_after", path, after_name))
         self._do_move_after(path, after_name)
         return self
 

pyiceberg/table/update/spec.py

@@ -58,17 +58,15 @@ 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]]
+    # Store all operations and order for retry support.
+    _operations: list[tuple[Any, ...]]
 
     def __init__(self, transaction: Transaction, case_sensitive: bool = True) -> None:
         super().__init__(transaction)
         self._transaction = transaction
         self._case_sensitive = case_sensitive
-        self._field_additions = []
-        self._deletes = set()
-        self._renames = {}
-        # Initialize state from current metadata
+        self._operations = []
+
         self._init_state_from_metadata()
 
     def _init_state_from_metadata(self) -> None:
@@ -78,6 +76,8 @@ def _init_state_from_metadata(self) -> None:
         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._deletes = set()
+        self._renames = {}
         self._name_to_added_field = {}
         self._transform_to_added_field = {}
         self._added_time_fields = {}
@@ -89,8 +89,18 @@ def _reset_state(self) -> None:
         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)
+
+        for operation in self._operations:
+            op_type = operation[0]
+            if op_type == "add":
+                _, source_column_name, transform, partition_field_name = operation
+                self._do_add_field(source_column_name, transform, partition_field_name)
+            elif op_type == "remove":
+                _, name = operation
+                self._do_remove_field(name)
+            elif op_type == "rename":
+                _, name, new_name = operation
+                self._do_rename_field(name, new_name)
 
     def _do_add_field(
         self,
@@ -128,7 +138,7 @@ def _do_add_field(
         existing_partition_field = self._name_to_field.get(new_field.name)
         if existing_partition_field and new_field.field_id not in self._deletes:
             if isinstance(existing_partition_field.transform, VoidTransform):
-                self.rename_field(
+                self._do_rename_field(
                     existing_partition_field.name, existing_partition_field.name + "_" + str(existing_partition_field.field_id)
                 )
             else:
@@ -144,14 +154,15 @@ def add_field(
         partition_field_name: str | None = None,
     ) -> UpdateSpec:
         transform = parse_transform(transform)
-        self._field_additions.append((source_column_name, transform, partition_field_name))
+        self._operations.append(("add", 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:
         return self.add_field(source_column_name, IdentityTransform(), None)
 
-    def remove_field(self, name: str) -> UpdateSpec:
+    def _do_remove_field(self, name: str) -> None:
+        """Remove a partition field (internal implementation for retry support)."""
         added = self._name_to_added_field.get(name)
         if added:
             raise ValueError(f"Cannot delete newly added field {name}")
@@ -163,12 +174,18 @@ def remove_field(self, name: str) -> UpdateSpec:
             raise ValueError(f"No such partition field: {name}")
 
         self._deletes.add(field.field_id)
+
+    def remove_field(self, name: str) -> UpdateSpec:
+        self._operations.append(("remove", name))
+        self._do_remove_field(name)
         return self
 
-    def rename_field(self, name: str, new_name: str) -> UpdateSpec:
+    def _do_rename_field(self, name: str, new_name: str) -> None:
+        """Rename a partition field (internal implementation for retry support)."""
         existing_field = self._name_to_field.get(new_name)
         if existing_field and isinstance(existing_field.transform, VoidTransform):
-            return self.rename_field(name, name + "_" + str(existing_field.field_id))
+            self._do_rename_field(name, name + "_" + str(existing_field.field_id))
+            return
         added = self._name_to_added_field.get(name)
         if added:
             raise ValueError("Cannot rename recently added partitions")
@@ -178,6 +195,10 @@ def rename_field(self, name: str, new_name: str) -> UpdateSpec:
         if field.field_id in self._deletes:
             raise ValueError(f"Cannot delete and rename partition field {name}")
         self._renames[name] = new_name
+
+    def rename_field(self, name: str, new_name: str) -> UpdateSpec:
+        self._operations.append(("rename", name, new_name))
+        self._do_rename_field(name, new_name)
         return self
 
     def _commit(self) -> UpdatesAndRequirements:

pyiceberg/utils/retry.py

@@ -15,8 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-"""Retry utilities for Iceberg operations."""
-
 from __future__ import annotations
 
 import logging