diff --git a/pyiceberg/table/update/__init__.py b/pyiceberg/table/update/__init__.py index 30315b0cc1..038b952bb3 100644 --- a/pyiceberg/table/update/__init__.py +++ b/pyiceberg/table/update/__init__.py @@ -193,6 +193,11 @@ class RemoveStatisticsUpdate(IcebergBaseModel): snapshot_id: int = Field(alias="snapshot-id") +class RemovePartitionSpecsUpdate(IcebergBaseModel): + action: Literal["remove-partition-specs"] = Field(default="remove-partition-specs") + spec_ids: List[int] = Field(alias="spec-ids") + + class RemoveSchemasUpdate(IcebergBaseModel): action: Literal["remove-schemas"] = Field(default="remove-schemas") schema_ids: List[int] = Field(alias="schema-ids") @@ -227,6 +232,7 @@ class RemovePartitionStatisticsUpdate(IcebergBaseModel): RemovePropertiesUpdate, SetStatisticsUpdate, RemoveStatisticsUpdate, + RemovePartitionSpecsUpdate, RemoveSchemasUpdate, SetPartitionStatisticsUpdate, RemovePartitionStatisticsUpdate, @@ -595,6 +601,21 @@ def _(update: RemoveStatisticsUpdate, base_metadata: TableMetadata, context: _Ta return base_metadata.model_copy(update={"statistics": statistics}) +@_apply_table_update.register(RemovePartitionSpecsUpdate) +def _(update: RemovePartitionSpecsUpdate, base_metadata: TableMetadata, context: _TableMetadataUpdateContext) -> TableMetadata: + for remove_spec_id in update.spec_ids: + if not any(spec.spec_id == remove_spec_id for spec in base_metadata.partition_specs): + raise ValueError(f"Partition spec with id {remove_spec_id} does not exist") + + if base_metadata.default_spec_id in update.spec_ids: + raise ValueError(f"Cannot remove default partition spec: {base_metadata.default_spec_id}") + + partition_specs = [spec for spec in base_metadata.partition_specs if spec.spec_id not in update.spec_ids] + + context.add_update(update) + return base_metadata.model_copy(update={"partition_specs": partition_specs}) + + @_apply_table_update.register(RemoveSchemasUpdate) def _(update: RemoveSchemasUpdate, base_metadata: TableMetadata, context: _TableMetadataUpdateContext) -> TableMetadata: # This method should error if any schemas do not exist. diff --git a/tests/table/test_init.py b/tests/table/test_init.py index cd81df4d97..5f64738d9b 100644 --- a/tests/table/test_init.py +++ b/tests/table/test_init.py @@ -66,6 +66,7 @@ ) from pyiceberg.table.statistics import BlobMetadata, PartitionStatisticsFile, StatisticsFile from pyiceberg.table.update import ( + AddPartitionSpecUpdate, AddSnapshotUpdate, AddSortOrderUpdate, AssertCreate, @@ -76,6 +77,7 @@ AssertLastAssignedPartitionId, AssertRefSnapshotId, AssertTableUUID, + RemovePartitionSpecsUpdate, RemovePartitionStatisticsUpdate, RemovePropertiesUpdate, RemoveSchemasUpdate, @@ -1294,6 +1296,38 @@ def test_update_metadata_log_overflow(table_v2: Table) -> None: assert len(new_metadata.metadata_log) == 1 +def test_remove_partition_spec_update(table_v2: Table) -> None: + base_metadata = table_v2.metadata + new_spec = PartitionSpec(PartitionField(source_id=2, field_id=1001, transform=IdentityTransform(), name="y"), spec_id=1) + metadata_with_new_spec = update_table_metadata(base_metadata, (AddPartitionSpecUpdate(spec=new_spec),)) + + assert len(metadata_with_new_spec.partition_specs) == 2 + + update = RemovePartitionSpecsUpdate(spec_ids=[1]) + updated_metadata = update_table_metadata( + metadata_with_new_spec, + (update,), + ) + + assert len(updated_metadata.partition_specs) == 1 + + +def test_remove_partition_spec_update_spec_does_not_exist(table_v2: Table) -> None: + update = RemovePartitionSpecsUpdate( + spec_ids=[123], + ) + with pytest.raises(ValueError, match="Partition spec with id 123 does not exist"): + update_table_metadata(table_v2.metadata, (update,)) + + +def test_remove_partition_spec_update_default_spec(table_v2: Table) -> None: + update = RemovePartitionSpecsUpdate( + spec_ids=[0], + ) + with pytest.raises(ValueError, match="Cannot remove default partition spec: 0"): + update_table_metadata(table_v2.metadata, (update,)) + + def test_remove_schemas_update(table_v2: Table) -> None: base_metadata = table_v2.metadata assert len(base_metadata.schemas) == 2