diff --git a/pyiceberg/table/update/__init__.py b/pyiceberg/table/update/__init__.py index aa1a09ebb2..3cdf25e5f7 100644 --- a/pyiceberg/table/update/__init__.py +++ b/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 diff --git a/tests/catalog/test_catalog_behaviors.py b/tests/catalog/test_catalog_behaviors.py index 1148ca8527..6c7347b6f3 100644 --- a/tests/catalog/test_catalog_behaviors.py +++ b/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) diff --git a/tests/table/test_init.py b/tests/table/test_init.py index f6b984f1a3..0cf55bd294 100644 --- a/tests/table/test_init.py +++ b/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),))