Added ExpireSnapshots Feature

pyiceberg/table/__init__.py

@@ -119,6 +119,7 @@
     ManageSnapshots,
     UpdateSnapshot,
     _FastAppendFiles,
+    ExpireSnapshots
 )
 from pyiceberg.table.update.spec import UpdateSpec
 from pyiceberg.table.update.statistics import UpdateStatistics
@@ -1078,6 +1079,15 @@ def manage_snapshots(self) -> ManageSnapshots:
            ms.create_tag(snapshot_id1, "Tag_A").create_tag(snapshot_id2, "Tag_B")
         """
         return ManageSnapshots(transaction=Transaction(self, autocommit=True))
+
+    def expire_snapshots(self) -> ExpireSnapshots:
+        """
+        Shorthand to run expire snapshots by id or by a timestamp.
+
+        Use table.expire_snapshots().<operation>().commit() to run a specific operation.
+        Use table.expire_snapshots().<operation-one>().<operation-two>().commit() to run multiple operations.
+        """
+        return ExpireSnapshots(transaction=Transaction(self, autocommit=True))
 
     def update_statistics(self) -> UpdateStatistics:
         """

pyiceberg/table/update/snapshot.py

@@ -66,7 +66,9 @@
     AddSnapshotUpdate,
     AssertRefSnapshotId,
     RemoveSnapshotRefUpdate,
+    RemoveSnapshotsUpdate,
     SetSnapshotRefUpdate,
+    TableMetadata,
     TableRequirement,
     TableUpdate,
     U,
@@ -82,7 +84,11 @@
 from pyiceberg.utils.properties import property_as_bool, property_as_int
 
 if TYPE_CHECKING:
-    from pyiceberg.table import Transaction
+    pass
+
+
+from pyiceberg.table.metadata import Snapshot, TableMetadata
+from pyiceberg.table.snapshots import Snapshot
 
 
 def _new_manifest_file_name(num: int, commit_uuid: uuid.UUID) -> str:
@@ -238,7 +244,7 @@ def _summary(self, snapshot_properties: Dict[str, str] = EMPTY_DICT) -> Summary:
             previous_summary=previous_snapshot.summary if previous_snapshot is not None else None,
         )
 
-    def _commit(self) -> UpdatesAndRequirements:
+    def _commit(self, base_metadata: TableMetadata) -> UpdatesAndRequirements:
         new_manifests = self._manifests()
         next_sequence_number = self._transaction.table_metadata.next_sequence_number()
 
@@ -739,11 +745,11 @@ class ManageSnapshots(UpdateTableMetadata["ManageSnapshots"]):
        ms.create_tag(snapshot_id1, "Tag_A").create_tag(snapshot_id2, "Tag_B")
     """
 
+    _snapshot_ids_to_expire = set()
     _updates: Tuple[TableUpdate, ...] = ()
     _requirements: Tuple[TableRequirement, ...] = ()
 
     def _commit(self) -> UpdatesAndRequirements:
-        """Apply the pending changes and commit."""
         return self._updates, self._requirements
 
     def _remove_ref_snapshot(self, ref_name: str) -> ManageSnapshots:
@@ -843,3 +849,64 @@ def remove_branch(self, branch_name: str) -> ManageSnapshots:
             This for method chaining
         """
         return self._remove_ref_snapshot(ref_name=branch_name)
+
+class ExpireSnapshots(UpdateTableMetadata["ExpireSnapshots"]):
+    """
+    Expire snapshots by ID or by timestamp.
+    Use table.expire_snapshots().<operation>().commit() to run a specific operation.
+    Use table.expire_snapshots().<operation-one>().<operation-two>().commit() to run multiple operations.
+    Pending changes are applied on commit.
+    """
+
+    _snapshot_ids_to_expire = set()
+    _updates: Tuple[TableUpdate, ...] = ()
+    _requirements: Tuple[TableRequirement, ...] = ()
+
+    def _commit(self) -> UpdatesAndRequirements:
+        """
+        Commit the staged updates and requirements.
+        This will remove the snapshots with the given IDs.
+
+        Returns:
+            Tuple of updates and requirements to be committed,
+            as required by the calling parent apply functions.
+        """
+        update = RemoveSnapshotsUpdate(snapshot_ids=self._snapshot_ids_to_expire)
+        self._updates += (update,)
+        return self._updates, self._requirements
+
+    def expire_snapshot_by_id(self, snapshot_id: int) -> ExpireSnapshots:
+        """
+        Expire a snapshot by its ID.
+
+        Args:
+            snapshot_id (int): The ID of the snapshot to expire.
+
+        Returns:
+            This for method chaining.
+        """
+        if self._transaction.table_metadata.snapshot_by_id(snapshot_id) is None:
+            raise ValueError(f"Snapshot with ID {snapshot_id} does not exist.")
+        self._snapshot_ids_to_expire.add(snapshot_id)
+        return self
+
+    def expire_snapshots_older_than(self, timestamp_ms: int) -> ExpireSnapshots:
+        """
+        Expire snapshots older than the given timestamp.
+
+        Args:
+            timestamp_ms (int): The timestamp in milliseconds. Snapshots older than this will be expired.
+
+        Returns:
+            This for method chaining.
+        """
+        # Collect IDs of snapshots to be expired
+        snapshots_to_remove = [
+            snapshot.snapshot_id
+            for snapshot in self._transaction.table_metadata.snapshots
+            if snapshot.timestamp_ms < timestamp_ms
+        ]
+        if snapshots_to_remove:
+            for snapshot_id in snapshots_to_remove:
+                self._snapshot_ids_to_expire.add(snapshot_id)
+        return self

tests/expressions/test_literals.py

@@ -744,7 +744,7 @@ def test_invalid_decimal_conversions() -> None:
 def test_invalid_string_conversions() -> None:
     assert_invalid_conversions(
         literal("abc"),
-        [FloatType(), DoubleType(), FixedType(1), BinaryType()],
+        [FixedType(1), BinaryType()],
     )
 
 

tests/integration/test_partition_evolution.py

@@ -140,6 +140,14 @@ def test_add_hour(catalog: Catalog) -> None:
     _validate_new_partition_fields(table, 1000, 1, 1000, PartitionField(2, 1000, HourTransform(), "hour_transform"))
 
 
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
+def test_add_hour_string_transform(catalog: Catalog) -> None:
+    table = _table(catalog)
+    table.update_spec().add_field("event_ts", "hour", "str_hour_transform").commit()
+    _validate_new_partition_fields(table, 1000, 1, 1000, PartitionField(2, 1000, HourTransform(), "str_hour_transform"))
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
 def test_add_hour_generates_default_name(catalog: Catalog) -> None:

tests/integration/test_writes/test_writes.py

@@ -1776,3 +1776,37 @@ def test_write_optional_list(session_catalog: Catalog) -> None:
     session_catalog.load_table(identifier).append(df_2)
 
     assert len(session_catalog.load_table(identifier).scan().to_arrow()) == 4
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("format_version", [1, 2])
+def test_evolve_and_write(
+    spark: SparkSession, session_catalog: Catalog, arrow_table_with_null: pa.Table, format_version: int
+) -> None:
+    identifier = "default.test_evolve_and_write"
+    tbl = _create_table(session_catalog, identifier, properties={"format-version": format_version}, schema=Schema())
+    other_table = session_catalog.load_table(identifier)
+
+    numbers = pa.array([1, 2, 3, 4], type=pa.int32())
+
+    with tbl.update_schema() as upd:
+        # This is not known by other_table
+        upd.add_column("id", IntegerType())
+
+    with other_table.transaction() as tx:
+        # Refreshes the underlying metadata, and the schema
+        other_table.refresh()
+        tx.append(
+            pa.Table.from_arrays(
+                [
+                    numbers,
+                ],
+                schema=pa.schema(
+                    [
+                        pa.field("id", pa.int32(), nullable=True),
+                    ]
+                ),
+            )
+        )
+
+    assert session_catalog.load_table(identifier).scan().to_arrow().column(0).combine_chunks() == numbers

tests/table/test_expire_snapshots.py

@@ -0,0 +1,43 @@
+from unittest.mock import MagicMock
+from uuid import uuid4
+
+from pyiceberg.table import CommitTableResponse, Table
+
+
+def test_expire_snapshot(table_v2: Table) -> None:
+    EXPIRE_SNAPSHOT = 3051729675574597004
+    KEEP_SNAPSHOT = 3055729675574597004
+    # Mock the catalog's commit_table method
+    mock_response = CommitTableResponse(
+        # Use the table's current metadata but keep only the snapshot not to be expired
+        metadata=table_v2.metadata.model_copy(update={"snapshots": [KEEP_SNAPSHOT]}),
+        metadata_location="mock://metadata/location",
+        uuid=uuid4(),
+    )
+
+    # Mock the commit_table method to return the mock response
+    table_v2.catalog.commit_table = MagicMock(return_value=mock_response)
+
+    # Print snapshot IDs for debugging
+    print(f"Snapshot IDs before expiration: {[snapshot.snapshot_id for snapshot in table_v2.metadata.snapshots]}")
+
+    # Assert fixture data to validate test assumptions
+    assert len(table_v2.metadata.snapshots) == 2
+    assert len(table_v2.metadata.snapshot_log) == 2
+    assert len(table_v2.metadata.refs) == 2
+
+    # Expire the snapshot directly without using a transaction
+    try:
+        table_v2.expire_snapshots().expire_snapshot_by_id(EXPIRE_SNAPSHOT).commit()
+    except Exception as e:
+        assert False, f"Commit failed with error: {e}"
+
+    # Assert that commit_table was called once
+    table_v2.catalog.commit_table.assert_called_once()
+
+    # Assert the expired snapshot ID is no longer present
+    remaining_snapshots = table_v2.metadata.snapshots
+    assert EXPIRE_SNAPSHOT not in remaining_snapshots
+
+    # Assert the length of snapshots after expiration
+    assert len(table_v2.metadata.snapshots) == 1