feat: update maintenance features with deduplication and retention strategies

ForeverAngry · ForeverAngry · commit 311c44246360 · 2025-07-05T01:08:59.000-04:00
diff --git a/.gitignore b/.gitignore
@@ -50,3 +50,4 @@ htmlcov
 pyiceberg/avro/decoder_fast.c
 pyiceberg/avro/*.html
 pyiceberg/avro/*.so
+ruff.toml
diff --git a/mkdocs/docs/api.md b/mkdocs/docs/api.md
@@ -1164,7 +1164,7 @@ maintenance.expire_snapshots_with_retention_policy(
 )
 ```
 
-#### Use Cases
+#### Deduplication Use Cases
 
 - **Operational Resilience**: Always keep recent snapshots for rollback.
 - **Space Reclamation**: Remove old, unneeded snapshots.
diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py
@@ -907,6 +907,7 @@ def inspect(self) -> InspectTable:
     @property
     def maintenance(self) -> MaintenanceTable:
         """Return the MaintenanceTable object for maintenance.
+
         Returns:
             MaintenanceTable object based on this Table.
         """
diff --git a/pyiceberg/table/maintenance.py b/pyiceberg/table/maintenance.py
@@ -17,10 +17,10 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, List, Optional, Set, Union
+from typing import TYPE_CHECKING, Any, List, Optional, Set, Union
 
-from pyiceberg.manifest import DataFile
-from pyiceberg.utils.concurrent import ThreadPoolExecutor
+from pyiceberg.manifest import DataFile, ManifestFile
+from pyiceberg.utils.concurrent import ThreadPoolExecutor  # type: ignore[attr-defined]
 
 logger = logging.getLogger(__name__)
 
@@ -52,7 +52,7 @@ def expire_snapshot_by_id(self, snapshot_id: int) -> None:
         """
         with self.tbl.transaction() as txn:
             # Check if snapshot exists
-            if txn.table_metadata.snapshot_by_id(snapshot_id) is None:
+            if not any(snapshot.snapshot_id == snapshot_id for snapshot in txn.table_metadata.snapshots):
                 raise ValueError(f"Snapshot with ID {snapshot_id} does not exist.")
 
             # Check if snapshot is protected
@@ -97,7 +97,7 @@ def expire_snapshots_older_than(self, timestamp_ms: int) -> None:
         """
         # First check if there are any snapshots to expire to avoid unnecessary transactions
         protected_ids = self._get_protected_snapshot_ids(self.tbl.metadata)
-        snapshots_to_expire = []
+        snapshots_to_expire: List[int] = []
 
         for snapshot in self.tbl.metadata.snapshots:
             if snapshot.timestamp_ms < timestamp_ms and snapshot.snapshot_id not in protected_ids:
@@ -110,10 +110,7 @@ def expire_snapshots_older_than(self, timestamp_ms: int) -> None:
                 txn._apply((RemoveSnapshotsUpdate(snapshot_ids=snapshots_to_expire),))
 
     def expire_snapshots_older_than_with_retention(
-        self, 
-        timestamp_ms: int, 
-        retain_last_n: Optional[int] = None,
-        min_snapshots_to_keep: Optional[int] = None
+        self, timestamp_ms: int, retain_last_n: Optional[int] = None, min_snapshots_to_keep: Optional[int] = None
     ) -> None:
         """Expire all unprotected snapshots with a timestamp older than a given value, with retention strategies.
 
@@ -123,9 +120,7 @@ def expire_snapshots_older_than_with_retention(
             min_snapshots_to_keep: Minimum number of snapshots to keep in total.
         """
         snapshots_to_expire = self._get_snapshots_to_expire_with_retention(
-            timestamp_ms=timestamp_ms,
-            retain_last_n=retain_last_n,
-            min_snapshots_to_keep=min_snapshots_to_keep
+            timestamp_ms=timestamp_ms, retain_last_n=retain_last_n, min_snapshots_to_keep=min_snapshots_to_keep
         )
 
         if snapshots_to_expire:
@@ -147,25 +142,21 @@ def retain_last_n_snapshots(self, n: int) -> None:
             raise ValueError("Number of snapshots to retain must be at least 1")
 
         protected_ids = self._get_protected_snapshot_ids(self.tbl.metadata)
-        
+
         # Sort snapshots by timestamp (most recent first)
-        sorted_snapshots = sorted(
-            self.tbl.metadata.snapshots, 
-            key=lambda s: s.timestamp_ms, 
-            reverse=True
-        )
-        
+        sorted_snapshots = sorted(self.tbl.metadata.snapshots, key=lambda s: s.timestamp_ms, reverse=True)
+
         # Keep the last N snapshots and all protected ones
         snapshots_to_keep = set()
         snapshots_to_keep.update(protected_ids)
-        
+
         # Add the N most recent snapshots
         for i, snapshot in enumerate(sorted_snapshots):
             if i < n:
                 snapshots_to_keep.add(snapshot.snapshot_id)
-        
+
         # Find snapshots to expire
-        snapshots_to_expire = []
+        snapshots_to_expire: List[int] = []
         for snapshot in self.tbl.metadata.snapshots:
             if snapshot.snapshot_id not in snapshots_to_keep:
                 snapshots_to_expire.append(snapshot.snapshot_id)
@@ -177,10 +168,7 @@ def retain_last_n_snapshots(self, n: int) -> None:
                 txn._apply((RemoveSnapshotsUpdate(snapshot_ids=snapshots_to_expire),))
 
     def _get_snapshots_to_expire_with_retention(
-        self,
-        timestamp_ms: Optional[int] = None,
-        retain_last_n: Optional[int] = None,
-        min_snapshots_to_keep: Optional[int] = None
+        self, timestamp_ms: Optional[int] = None, retain_last_n: Optional[int] = None, min_snapshots_to_keep: Optional[int] = None
     ) -> List[int]:
         """Get snapshots to expire considering retention strategies.
 
@@ -193,54 +181,46 @@ def _get_snapshots_to_expire_with_retention(
             List of snapshot IDs to expire.
         """
         protected_ids = self._get_protected_snapshot_ids(self.tbl.metadata)
-        
+
         # Sort snapshots by timestamp (most recent first)
-        sorted_snapshots = sorted(
-            self.tbl.metadata.snapshots, 
-            key=lambda s: s.timestamp_ms, 
-            reverse=True
-        )
-        
+        sorted_snapshots = sorted(self.tbl.metadata.snapshots, key=lambda s: s.timestamp_ms, reverse=True)
+
         # Start with all snapshots that could be expired
         candidates_for_expiration = []
         snapshots_to_keep = set(protected_ids)
-        
+
         # Apply retain_last_n constraint
         if retain_last_n is not None:
             for i, snapshot in enumerate(sorted_snapshots):
                 if i < retain_last_n:
                     snapshots_to_keep.add(snapshot.snapshot_id)
-        
+
         # Apply timestamp constraint
         for snapshot in self.tbl.metadata.snapshots:
-            if (snapshot.snapshot_id not in snapshots_to_keep and 
-                (timestamp_ms is None or snapshot.timestamp_ms < timestamp_ms)):
+            if snapshot.snapshot_id not in snapshots_to_keep and (timestamp_ms is None or snapshot.timestamp_ms < timestamp_ms):
                 candidates_for_expiration.append(snapshot)
-        
+
         # Sort candidates by timestamp (oldest first) for potential expiration
         candidates_for_expiration.sort(key=lambda s: s.timestamp_ms)
-        
+
         # Apply min_snapshots_to_keep constraint
         total_snapshots = len(self.tbl.metadata.snapshots)
-        snapshots_to_expire = []
-        
+        snapshots_to_expire: List[int] = []
+
         for candidate in candidates_for_expiration:
             # Check if expiring this snapshot would violate min_snapshots_to_keep
             remaining_after_expiration = total_snapshots - len(snapshots_to_expire) - 1
-            
+
             if min_snapshots_to_keep is None or remaining_after_expiration >= min_snapshots_to_keep:
                 snapshots_to_expire.append(candidate.snapshot_id)
             else:
                 # Stop expiring to maintain minimum count
                 break
-        
+
         return snapshots_to_expire
 
     def expire_snapshots_with_retention_policy(
-        self,
-        timestamp_ms: Optional[int] = None,
-        retain_last_n: Optional[int] = None,
-        min_snapshots_to_keep: Optional[int] = None
+        self, timestamp_ms: Optional[int] = None, retain_last_n: Optional[int] = None, min_snapshots_to_keep: Optional[int] = None
     ) -> List[int]:
         """Comprehensive snapshot expiration with multiple retention strategies.
 
@@ -266,13 +246,13 @@ def expire_snapshots_with_retention_policy(
         Examples:
             # Keep last 5 snapshots regardless of age
             maintenance.expire_snapshots_with_retention_policy(retain_last_n=5)
-            
+
             # Expire snapshots older than timestamp but keep at least 3 total
             maintenance.expire_snapshots_with_retention_policy(
                 timestamp_ms=1234567890000,
                 min_snapshots_to_keep=3
             )
-            
+
             # Combined policy: expire old snapshots but keep last 10 and at least 5 total
             maintenance.expire_snapshots_with_retention_policy(
                 timestamp_ms=1234567890000,
@@ -282,14 +262,12 @@ def expire_snapshots_with_retention_policy(
         """
         if retain_last_n is not None and retain_last_n < 1:
             raise ValueError("retain_last_n must be at least 1")
-        
+
         if min_snapshots_to_keep is not None and min_snapshots_to_keep < 1:
             raise ValueError("min_snapshots_to_keep must be at least 1")
 
         snapshots_to_expire = self._get_snapshots_to_expire_with_retention(
-            timestamp_ms=timestamp_ms,
-            retain_last_n=retain_last_n,
-            min_snapshots_to_keep=min_snapshots_to_keep
+            timestamp_ms=timestamp_ms, retain_last_n=retain_last_n, min_snapshots_to_keep=min_snapshots_to_keep
         )
 
         if snapshots_to_expire:
@@ -326,12 +304,10 @@ def _get_all_datafiles(
         target_file_path: Optional[str] = None,
         parallel: bool = True,
     ) -> List[DataFile]:
-        """
-        Collect all DataFiles in the table, optionally filtering by file path.
-        """
+        """Collect all DataFiles in the table, optionally filtering by file path."""
         datafiles: List[DataFile] = []
 
-        def process_manifest(manifest) -> list[DataFile]:
+        def process_manifest(manifest: ManifestFile) -> list[DataFile]:
             found: list[DataFile] = []
             for entry in manifest.fetch_manifest_entry(io=self.tbl.io):
                 if hasattr(entry, "data_file"):
@@ -356,7 +332,7 @@ def process_manifest(manifest) -> list[DataFile]:
             # Only current snapshot
             for chunk in self.tbl.inspect.data_files().to_pylist():
                 file_path = chunk.get("file_path")
-                partition = chunk.get("partition", {})
+                partition: dict[str, Any] = dict(chunk.get("partition", {}) or {})
                 if target_file_path is None or file_path == target_file_path:
                     datafiles.append(DataFile(file_path=file_path, partition=partition))
         return datafiles
@@ -389,16 +365,16 @@ def deduplicate_data_files(
             seen = {}
             duplicates = []
             for df in all_datafiles:
-                partition = dict(df.partition) if hasattr(df.partition, "items") else df.partition
+                partition: dict[str, Any] = df.partition.to_dict() if hasattr(df.partition, "to_dict") else {}
                 if scan_all_partitions:
-                    key = (df.file_path, tuple(sorted(partition.items())) if partition else None)
+                    key = (df.file_path, tuple(sorted(partition.items())) if partition else ())
                 else:
-                    key = df.file_path
+                    key = (df.file_path, ())  # Add an empty tuple for partition when scan_all_partitions is False
                 if key in seen:
                     duplicates.append(df)
                 else:
                     seen[key] = df
-            to_remove = duplicates
+            to_remove = duplicates  # type: ignore[assignment]
 
         # Normalize to DataFile objects
         normalized_to_remove: List[DataFile] = []
diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py
@@ -82,10 +82,7 @@
 from pyiceberg.utils.properties import property_as_bool, property_as_int
 
 if TYPE_CHECKING:
-    pass
-
-
-from pyiceberg.table.metadata import Snapshot
+    from pyiceberg.table import Transaction
 
 
 def _new_manifest_file_name(num: int, commit_uuid: uuid.UUID) -> str:
diff --git a/ruff.toml b/ruff.toml
@@ -34,13 +34,7 @@ exclude = [
     ".svn",
     ".tox",
     ".venv",
-    "__pypackages__",
-    "_build",
-    "buck-out",
-    "build",
-    "dist",
-    "node_modules",
-    "venv",
+    "vendor",
 ]
 
 # Ignore _all_ violations.
diff --git a/tests/expressions/test_literals.py b/tests/expressions/test_literals.py
@@ -744,6 +744,21 @@ def test_invalid_decimal_conversions() -> None:
 def test_invalid_string_conversions() -> None:
     assert_invalid_conversions(
         literal("abc"),
+        [
+            BooleanType(),
+            IntegerType(),
+            LongType(),
+            FloatType(),
+            DoubleType(),
+            DateType(),
+            TimeType(),
+            TimestampType(),
+            TimestamptzType(),
+            DecimalType(9, 2),
+            UUIDType(),
+            FixedType(1),
+            BinaryType(),
+        ],
     )
 
 
diff --git a/tests/table/test_dedup_data_file_filepaths.py b/tests/table/test_dedup_data_file_filepaths.py
diff --git a/tests/table/test_retention_strategies.py b/tests/table/test_retention_strategies.py

Original file line number	Diff line number	Diff line change
`@@ -1164,7 +1164,7 @@ maintenance.expire_snapshots_with_retention_policy(`
`1164`	`1164`	`)`
`1165`	`1165`	```
`1166`	`1166`
`1167`		`-#### Use Cases`
	`1167`	`+#### Deduplication Use Cases`
`1168`	`1168`
`1169`	`1169`	`- Operational Resilience: Always keep recent snapshots for rollback.`
`1170`	`1170`	`- Space Reclamation: Remove old, unneeded snapshots.`