add UpdateSortOrder

Author: KazuhitoT

pyiceberg/table/__init__.py

@@ -282,6 +282,9 @@ def __init__(self, updates: tuple[TableUpdate, ...], requirements: tuple[TableRe
         self._updates = updates
         self._requirements = requirements
 
+    def _reset_state(self) -> None:
+        """No-op for static updates that don't cache metadata-derived state."""
+
     def _commit(self) -> UpdatesAndRequirements:
         """Return the stored updates and requirements."""
         return self._updates, self._requirements
@@ -1090,14 +1093,8 @@ def _reapply_updates(self) -> None:
         This is called on retry to regenerate all pending updates
         based on the latest table metadata, similar to Java's BaseTransaction.applyUpdates().
 
-        All updates are rebuilt from _pending_updates to ensure consistency.
-        This includes:
-        - Static updates (properties, format version) via _StaticUpdate
-        - Snapshot operations via snapshot producers with _reset_state()
-
-        NOTE: Every operation that should survive retry must be tracked in _pending_updates.
-        Simple operations use _StaticUpdate wrapper, complex operations (like snapshot
-        producers) implement _reset_state() and _commit() directly.
+        NOTE: When adding new cached properties to UpdateTableMetadata subclasses,
+        ensure they are cleared in _reset_state() to avoid stale data on retry.
         """
         self._table.refresh()
 
@@ -1106,11 +1103,7 @@ def _reapply_updates(self) -> None:
         self._working_metadata = self._table.metadata
 
         for pending_update in self._pending_updates:
-            # NOTE: When adding new cached properties to snapshot producers,
-            # ensure they are cleared in _reset_state() to avoid stale data on retry
-            if hasattr(pending_update, "_reset_state"):
-                pending_update._reset_state()
-
+            pending_update._reset_state()
             updates, requirements = pending_update._commit()
             self._updates += updates
             self._working_metadata = update_table_metadata(self._working_metadata, updates)

pyiceberg/table/update/__init__.py

@@ -66,7 +66,19 @@ def __init__(self, transaction: Transaction) -> None:
         self._transaction = transaction
 
     @abstractmethod
-    def _commit(self) -> UpdatesAndRequirements: ...
+    def _commit(self) -> UpdatesAndRequirements:
+        """Generate the table updates and requirements for this operation."""
+        ...
+
+    @abstractmethod
+    def _reset_state(self) -> None:
+        """Reset internal state for retry after table metadata refresh.
+
+        This is called by Transaction._reapply_updates() when retrying after a
+        CommitFailedException. Implementations should rebuild any cached state
+        from self._transaction.table_metadata.
+        """
+        ...
 
     def commit(self) -> None:
         updates, requirements = self._commit()

pyiceberg/table/update/snapshot.py

@@ -972,6 +972,9 @@ def __init__(self, transaction: Transaction) -> None:
         self._updates = ()
         self._requirements = ()
 
+    def _reset_state(self) -> None:
+        """No-op: updates contain user-provided snapshot IDs that don't need refresh."""
+
     def _commit(self) -> UpdatesAndRequirements:
         """Apply the pending changes and commit."""
         return self._updates, self._requirements
@@ -1093,6 +1096,15 @@ def __init__(self, transaction: Transaction) -> None:
         self._requirements = ()
         self._snapshot_ids_to_expire = set()
 
+    def _reset_state(self) -> None:
+        """Clear accumulated updates for retry.
+
+        The _snapshot_ids_to_expire are user-provided and preserved.
+        The _updates are cleared so _commit() can rebuild them with
+        refreshed protected snapshot IDs.
+        """
+        self._updates = ()
+
     def _commit(self) -> UpdatesAndRequirements:
         """
         Commit the staged updates and requirements.

pyiceberg/table/update/sorting.py

@@ -39,12 +39,22 @@ class UpdateSortOrder(UpdateTableMetadata["UpdateSortOrder"]):
     _last_assigned_order_id: int | None
     _case_sensitive: bool
     _fields: list[SortField]
+    # Store user intent for retry support: (column_name, transform, direction, null_order)
+    _field_additions: list[tuple[str, Transform[Any, Any], SortDirection, NullOrder]]
 
     def __init__(self, transaction: Transaction, case_sensitive: bool = True) -> None:
         super().__init__(transaction)
-        self._fields: list[SortField] = []
-        self._case_sensitive: bool = case_sensitive
-        self._last_assigned_order_id: int | None = None
+        self._fields = []
+        self._case_sensitive = case_sensitive
+        self._last_assigned_order_id = None
+        self._field_additions = []
+
+    def _reset_state(self) -> None:
+        """Reset state for retry, re-resolving column names from refreshed metadata."""
+        self._fields = []
+        self._last_assigned_order_id = None
+        for column_name, transform, direction, null_order in self._field_additions:
+            self._do_add_sort_field(column_name, transform, direction, null_order)
 
     def _column_name_to_id(self, column_name: str) -> int:
         """Map the column name to the column field id."""
@@ -57,23 +67,22 @@ def _column_name_to_id(self, column_name: str) -> int:
             .field_id
         )
 
-    def _add_sort_field(
+    def _do_add_sort_field(
         self,
-        source_id: int,
+        source_column_name: str,
         transform: Transform[Any, Any],
         direction: SortDirection,
         null_order: NullOrder,
-    ) -> UpdateSortOrder:
-        """Add a sort field to the sort order list."""
+    ) -> None:
+        """Add a sort field to the sort order list (internal implementation for retry support)."""
         self._fields.append(
             SortField(
-                source_id=source_id,
+                source_id=self._column_name_to_id(source_column_name),
                 transform=transform,
                 direction=direction,
                 null_order=null_order,
             )
         )
-        return self
 
     def _reuse_or_create_sort_order_id(self) -> int:
         """Return the last assigned sort order id or create a new one."""
@@ -90,23 +99,17 @@ def asc(
         self, source_column_name: str, transform: Transform[Any, Any], null_order: NullOrder = NullOrder.NULLS_LAST
     ) -> UpdateSortOrder:
         """Add a sort field with ascending order."""
-        return self._add_sort_field(
-            source_id=self._column_name_to_id(source_column_name),
-            transform=transform,
-            direction=SortDirection.ASC,
-            null_order=null_order,
-        )
+        self._field_additions.append((source_column_name, transform, SortDirection.ASC, null_order))
+        self._do_add_sort_field(source_column_name, transform, SortDirection.ASC, null_order)
+        return self
 
     def desc(
         self, source_column_name: str, transform: Transform[Any, Any], null_order: NullOrder = NullOrder.NULLS_LAST
     ) -> UpdateSortOrder:
         """Add a sort field with descending order."""
-        return self._add_sort_field(
-            source_id=self._column_name_to_id(source_column_name),
-            transform=transform,
-            direction=SortDirection.DESC,
-            null_order=null_order,
-        )
+        self._field_additions.append((source_column_name, transform, SortDirection.DESC, null_order))
+        self._do_add_sort_field(source_column_name, transform, SortDirection.DESC, null_order)
+        return self
 
     def _apply(self) -> SortOrder:
         """Return the sort order."""

pyiceberg/table/update/statistics.py

@@ -52,6 +52,9 @@ class UpdateStatistics(UpdateTableMetadata["UpdateStatistics"]):
     def __init__(self, transaction: "Transaction") -> None:
         super().__init__(transaction)
 
+    def _reset_state(self) -> None:
+        """No-op: updates contain user-provided data that doesn't need refresh."""
+
     def set_statistics(self, statistics_file: StatisticsFile) -> "UpdateStatistics":
         self._updates += (
             SetStatisticsUpdate(

tests/table/test_commit_retry.py

@@ -1126,8 +1126,6 @@ def mock_commit(
 
 
 class TestUpdateSchemaRetry:
-    """Tests for UpdateSchema retry behavior (Java-like behavior)."""
-
     def test_update_schema_retried_on_conflict(self, catalog: SqlCatalog, schema: Schema) -> None:
         """Test that UpdateSchema operations are retried on CommitFailedException."""
         from pyiceberg.types import StringType
@@ -1273,3 +1271,196 @@ def mock_commit(
         assert table.schema().schema_id == 1
         assert len(table.schema().fields) == 2
         assert table.schema().find_field("new_col").field_type == StringType()
+
+
+class TestUpdateSortOrderRetry:
+    def test_update_sort_order_retried_on_conflict(self, catalog: SqlCatalog, schema: Schema) -> None:
+        """Test that UpdateSortOrder operations are retried on CommitFailedException."""
+        from pyiceberg.transforms import IdentityTransform
+
+        table = catalog.create_table(
+            "default.test_sort_order_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 sort order conflict")
+            return original_commit(tbl, requirements, updates)
+
+        with patch.object(catalog, "commit_table", side_effect=mock_commit):
+            with table.update_sort_order() as update_sort_order:
+                update_sort_order.asc("id", IdentityTransform())
+
+        assert commit_count == 2
+
+        # Verify sort order was updated
+        table.refresh()
+        sort_order = table.sort_order()
+        assert sort_order.order_id == 1
+        assert len(sort_order.fields) == 1
+        assert sort_order.fields[0].source_id == 1  # "id" column
+
+    def test_update_sort_order_resolves_conflict_on_retry(self, catalog: SqlCatalog, schema: Schema) -> None:
+        """Test that sort order update can resolve conflicts via retry."""
+        from pyiceberg.table.sorting import SortDirection
+        from pyiceberg.transforms import IdentityTransform
+
+        table = catalog.create_table(
+            "default.test_sort_order_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_sort_order() as update_sort_order:
+            update_sort_order.asc("id", IdentityTransform())
+
+        table2 = catalog.load_table("default.test_sort_order_conflict_resolved")
+        with table2.update_sort_order() as update_sort_order2:
+            update_sort_order2.desc("id", IdentityTransform())
+
+        assert table.sort_order().order_id == 1
+        assert table2.sort_order().order_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):
+            with table.update_sort_order() as update_sort_order:
+                update_sort_order.asc("id", IdentityTransform())
+
+        assert commit_count == 2
+
+        table.refresh()
+        sort_order = table.sort_order()
+        assert sort_order.order_id == 1  # Reused existing order with same fields
+        assert len(sort_order.fields) == 1
+        assert sort_order.fields[0].direction == SortDirection.ASC
+
+    def test_transaction_with_sort_order_change_and_append_retries(
+        self, catalog: SqlCatalog, schema: Schema, arrow_table: pa.Table
+    ) -> None:
+        """Test that a transaction with sort order change and append handles retry correctly."""
+        from pyiceberg.transforms import IdentityTransform
+
+        table = catalog.create_table(
+            "default.test_transaction_sort_order_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_sort_order() as update_sort_order:
+                    update_sort_order.asc("id", IdentityTransform())
+                txn.append(arrow_table)
+
+        assert commit_count == 2
+
+        first_attempt_update_types = [type(u).__name__ for u in captured_updates[0]]
+        assert "AddSortOrderUpdate" 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 "AddSortOrderUpdate" in retry_attempt_update_types
+        assert "AddSnapshotUpdate" in retry_attempt_update_types
+
+        assert len(table.scan().to_arrow()) == 3
+
+        sort_order = table.sort_order()
+        assert sort_order.order_id == 1
+        assert len(sort_order.fields) == 1
+        assert sort_order.fields[0].source_id == 1  # "id" column
+
+    def test_sort_order_column_name_re_resolved_on_retry(self, catalog: SqlCatalog, schema: Schema) -> None:
+        """Test that column names are re-resolved from refreshed schema on retry.
+
+        This ensures that if the schema changes between retries (e.g., column ID changes),
+        the sort order will use the correct field ID from the refreshed schema.
+        """
+        from pyiceberg.transforms import IdentityTransform
+
+        table = catalog.create_table(
+            "default.test_sort_order_column_re_resolved",
+            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_sort_fields: list[list[int]] = []
+
+        def mock_commit(
+            tbl: Table, requirements: tuple[TableRequirement, ...], updates: tuple[TableUpdate, ...]
+        ) -> CommitTableResponse:
+            nonlocal commit_count
+            commit_count += 1
+
+            # Extract sort field source IDs from updates
+            from pyiceberg.table.update import AddSortOrderUpdate
+
+            for update in updates:
+                if isinstance(update, AddSortOrderUpdate):
+                    source_ids = [f.source_id for f in update.sort_order.fields]
+                    captured_sort_fields.append(source_ids)
+
+            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.update_sort_order() as update_sort_order:
+                update_sort_order.asc("id", IdentityTransform())
+
+        assert commit_count == 2
+        assert len(captured_sort_fields) == 2
+
+        # Both attempts should resolve "id" to the same source_id (1)
+        # This verifies the column name is being re-resolved correctly
+        assert captured_sort_fields[0] == [1]  # First attempt
+        assert captured_sort_fields[1] == [1]  # Retry attempt