Fix duplicate sort order id when adding an existing sort order

pyiceberg/table/update/__init__.py

@@ -268,9 +268,13 @@ class RemoveEncryptionKeyUpdate(IcebergBaseModel):
 
 class _TableMetadataUpdateContext:
     _updates: list[TableUpdate]
+    last_added_sort_order_id: int | None
+    reassigned_sort_order_ids: dict[int, int]
 
     def __init__(self) -> None:
         self._updates = []
+        self.last_added_sort_order_id = None
+        self.reassigned_sort_order_ids = {}
 
     def add_update(self, update: TableUpdate) -> None:
         self._updates.append(update)
@@ -286,11 +290,6 @@ def is_added_schema(self, schema_id: int) -> bool:
     def is_added_partition_spec(self, spec_id: int) -> bool:
         return any(update.spec.spec_id == spec_id for update in self._updates if isinstance(update, AddPartitionSpecUpdate))
 
-    def is_added_sort_order(self, sort_order_id: int) -> bool:
-        return any(
-            update.sort_order.order_id == sort_order_id for update in self._updates if isinstance(update, AddSortOrderUpdate)
-        )
-
     def has_changes(self) -> bool:
         return len(self._updates) > 0
 
@@ -587,15 +586,28 @@ def _(update: RemoveSnapshotRefUpdate, base_metadata: TableMetadata, context: _T
 
 @_apply_table_update.register(AddSortOrderUpdate)
 def _(update: AddSortOrderUpdate, base_metadata: TableMetadata, context: _TableMetadataUpdateContext) -> TableMetadata:
-    for sort in base_metadata.sort_orders:
-        # Only raise in case of a discrepancy
-        if sort.order_id == update.sort_order.order_id and sort != update.sort_order:
-            raise ValueError(f"Sort-order with id {sort.order_id} already exists: {sort}")
+    new_sort_order = update.sort_order
 
-    context.add_update(update)
+    # Reuse the id of an existing sort order with the same fields
+    for sort in base_metadata.sort_orders:
+        if sort.fields == new_sort_order.fields:
+            if sort.order_id != new_sort_order.order_id:
+                # The staged id was assigned against stale metadata; point later updates at the existing order
+                context.reassigned_sort_order_ids[new_sort_order.order_id] = sort.order_id
+            context.last_added_sort_order_id = sort.order_id
+            return base_metadata
+
+    # The id can be taken when the update was built against stale metadata
+    if base_metadata.sort_order_by_id(new_sort_order.order_id) is not None:
+        new_order_id = max(sort.order_id for sort in base_metadata.sort_orders) + 1
+        context.reassigned_sort_order_ids[new_sort_order.order_id] = new_order_id
+        new_sort_order = new_sort_order.model_copy(update={"order_id": new_order_id})
+
+    context.add_update(AddSortOrderUpdate(sort_order=new_sort_order))
+    context.last_added_sort_order_id = new_sort_order.order_id
     return base_metadata.model_copy(
         update={
-            "sort_orders": base_metadata.sort_orders + [update.sort_order],
+            "sort_orders": base_metadata.sort_orders + [new_sort_order],
         }
     )
 
@@ -608,10 +620,12 @@ def _(
 ) -> TableMetadata:
     new_sort_order_id = update.sort_order_id
     if new_sort_order_id == -1:
-        # The last added sort order should be in base_metadata.sort_orders at this point
-        new_sort_order_id = max(sort_order.order_id for sort_order in base_metadata.sort_orders)
-        if not context.is_added_sort_order(new_sort_order_id):
+        if context.last_added_sort_order_id is None:
             raise ValueError("Cannot set current sort order to the last added one when no sort order has been added")
+        new_sort_order_id = context.last_added_sort_order_id
+    else:
+        # Follow ids reassigned while applying this commit, when the update was built against stale metadata
+        new_sort_order_id = context.reassigned_sort_order_ids.get(new_sort_order_id, new_sort_order_id)
 
     if new_sort_order_id == base_metadata.default_sort_order_id:
         return base_metadata

tests/catalog/test_catalog_behaviors.py

@@ -622,6 +622,108 @@ def test_concurrent_commit_table(catalog: Catalog, table_schema_nested: Schema,
             update.add_column(path="c", field_type=IntegerType())
 
 
+def test_stale_update_sort_order_reuses_equal_sort_order(
+    catalog: Catalog, table_schema_nested: Schema, test_table_identifier: Identifier
+) -> None:
+    namespace = Catalog.namespace_from(test_table_identifier)
+    catalog.create_namespace(namespace)
+    table_a = catalog.create_table(test_table_identifier, table_schema_nested)
+    table_b = catalog.load_table(test_table_identifier)
+
+    with table_a.update_sort_order() as update:
+        update.asc("foo", IdentityTransform())
+    with table_a.update_sort_order():
+        pass  # Reset to unsorted
+
+    # Stale handle adds the same sort order that already exists with id 1
+    with table_b.update_sort_order() as update:
+        update.asc("foo", IdentityTransform())
+
+    metadata = catalog.load_table(test_table_identifier).metadata
+    assert [sort_order.order_id for sort_order in metadata.sort_orders] == [0, 1]
+    assert metadata.default_sort_order_id == 1
+
+
+def test_stale_update_sort_order_assigns_new_id(
+    catalog: Catalog, table_schema_nested: Schema, test_table_identifier: Identifier
+) -> None:
+    namespace = Catalog.namespace_from(test_table_identifier)
+    catalog.create_namespace(namespace)
+    table_a = catalog.create_table(test_table_identifier, table_schema_nested)
+    table_b = catalog.load_table(test_table_identifier)
+
+    with table_a.update_sort_order() as update:
+        update.asc("foo", IdentityTransform())
+    with table_a.update_sort_order():
+        pass  # Reset to unsorted
+
+    # Stale handle adds a different sort order with id 1, which is already taken
+    with table_b.update_sort_order() as update:
+        update.asc("bar", IdentityTransform())
+
+    metadata = catalog.load_table(test_table_identifier).metadata
+    assert [sort_order.order_id for sort_order in metadata.sort_orders] == [0, 1, 2]
+    assert metadata.default_sort_order_id == 2
+    assert metadata.sort_order().fields[0].source_id == 2
+
+
+def test_update_sort_order_multiple_in_transaction(
+    catalog: Catalog, table_schema_nested: Schema, test_table_identifier: Identifier
+) -> None:
+    namespace = Catalog.namespace_from(test_table_identifier)
+    catalog.create_namespace(namespace)
+    table = catalog.create_table(test_table_identifier, table_schema_nested)
+
+    with table.transaction() as transaction:
+        with transaction.update_sort_order() as update:
+            update.asc("foo", IdentityTransform())
+        with transaction.update_sort_order() as update:
+            update.asc("bar", IdentityTransform())
+
+    metadata = catalog.load_table(test_table_identifier).metadata
+    assert [sort_order.order_id for sort_order in metadata.sort_orders] == [0, 1, 2]
+    assert metadata.default_sort_order_id == 2
+
+    with table.transaction() as transaction:
+        with transaction.update_sort_order() as update:
+            update.asc("baz", IdentityTransform())
+        with transaction.update_sort_order() as update:
+            update.asc("foo", IdentityTransform())
+
+    metadata = catalog.load_table(test_table_identifier).metadata
+    assert [sort_order.order_id for sort_order in metadata.sort_orders] == [0, 1, 2, 3]
+    assert metadata.default_sort_order_id == 1
+
+
+def test_stale_transaction_keeps_reassigned_sort_order_as_default(
+    catalog: Catalog, table_schema_nested: Schema, test_table_identifier: Identifier
+) -> None:
+    namespace = Catalog.namespace_from(test_table_identifier)
+    catalog.create_namespace(namespace)
+    table_a = catalog.create_table(test_table_identifier, table_schema_nested)
+    table_b = catalog.load_table(test_table_identifier)
+
+    with table_a.update_sort_order() as update:
+        update.asc("foo", IdentityTransform())
+    with table_a.update_sort_order():
+        pass  # Reset to unsorted
+
+    # Stale handle stages bar (id 1), baz (id 2), then bar again, which it sends as an explicit SetDefault(1)
+    with table_b.transaction() as transaction:
+        with transaction.update_sort_order() as update:
+            update.asc("bar", IdentityTransform())
+        with transaction.update_sort_order() as update:
+            update.asc("baz", IdentityTransform())
+        with transaction.update_sort_order() as update:
+            update.asc("bar", IdentityTransform())
+
+    # bar and baz are reassigned to 2 and 3, so the default must follow bar rather than foo's id 1
+    metadata = catalog.load_table(test_table_identifier).metadata
+    assert [sort_order.order_id for sort_order in metadata.sort_orders] == [0, 1, 2, 3]
+    assert metadata.default_sort_order_id == 2
+    assert metadata.sort_order().fields[0].source_id == 2
+
+
 def test_delete_metadata_multiple(catalog: Catalog, table_schema_nested: Schema, random_table_identifier: Identifier) -> None:
     namespace = Catalog.namespace_from(random_table_identifier)
     catalog.create_namespace(namespace)

tests/table/test_init.py

@@ -1215,6 +1215,58 @@ def test_update_metadata_add_update_sort_order(table_v2: Table) -> None:
     assert new_metadata.last_updated_ms > table_v2.metadata.last_updated_ms
 
 
+def test_update_metadata_add_existing_sort_order(table_v2: Table) -> None:
+    existing_sort_order = table_v2.sort_order()
+    for order_id in (existing_sort_order.order_id, existing_sort_order.order_id + 1):
+        new_metadata = update_table_metadata(
+            table_v2.metadata,
+            (
+                AddSortOrderUpdate(sort_order=SortOrder(*existing_sort_order.fields, order_id=order_id)),
+                SetDefaultSortOrderUpdate(sort_order_id=-1),
+            ),
+        )
+        assert new_metadata.sort_orders == table_v2.metadata.sort_orders
+        assert new_metadata.default_sort_order_id == existing_sort_order.order_id
+        assert new_metadata.last_updated_ms == table_v2.metadata.last_updated_ms
+
+
+def test_update_metadata_add_sort_order_with_taken_id(table_v2: Table) -> None:
+    existing_sort_order = table_v2.sort_order()
+    new_sort_order = SortOrder(
+        SortField(source_id=1, transform=IdentityTransform()),
+        order_id=existing_sort_order.order_id,
+    )
+    new_metadata = update_table_metadata(
+        table_v2.metadata,
+        (AddSortOrderUpdate(sort_order=new_sort_order), SetDefaultSortOrderUpdate(sort_order_id=-1)),
+    )
+    assert len(new_metadata.sort_orders) == len(table_v2.metadata.sort_orders) + 1
+    assert new_metadata.sort_orders[-1].fields == new_sort_order.fields
+    assert new_metadata.sort_orders[-1].order_id == existing_sort_order.order_id + 1
+    assert new_metadata.default_sort_order_id == existing_sort_order.order_id + 1
+
+
+def test_update_metadata_add_existing_sort_order_reuses_its_own_id(table_v2: Table) -> None:
+    existing_sort_order = table_v2.sort_order()
+    newer_sort_order = SortOrder(SortField(source_id=1, transform=IdentityTransform()), order_id=existing_sort_order.order_id + 1)
+    metadata = update_table_metadata(
+        table_v2.metadata,
+        (AddSortOrderUpdate(sort_order=newer_sort_order), SetDefaultSortOrderUpdate(sort_order_id=-1)),
+    )
+
+    # A writer that read metadata before newer_sort_order was added picks its own next id, which newer_sort_order took.
+    # Reusing existing_sort_order has to resolve to its own id, which is no longer the highest one.
+    new_metadata = update_table_metadata(
+        metadata,
+        (
+            AddSortOrderUpdate(sort_order=SortOrder(*existing_sort_order.fields, order_id=newer_sort_order.order_id + 1)),
+            SetDefaultSortOrderUpdate(sort_order_id=-1),
+        ),
+    )
+    assert new_metadata.sort_orders == metadata.sort_orders
+    assert new_metadata.default_sort_order_id == existing_sort_order.order_id
+
+
 def test_update_metadata_update_sort_order_invalid(table_v2: Table) -> None:
     with pytest.raises(ValueError, match="Cannot set current sort order to the last added one when no sort order has been added"):
         update_table_metadata(table_v2.metadata, (SetDefaultSortOrderUpdate(sort_order_id=-1),))