refactor: streamline data file retrieval in MaintenanceTable and enhance deduplication tests

Author: ForeverAngry

pyiceberg/table/maintenance.py

@@ -294,27 +294,9 @@ def _get_protected_snapshot_ids(self, table_metadata: TableMetadata) -> Set[int]
 
     def _get_all_datafiles(self) -> List[DataFile]:
         """Collect all DataFiles in the current snapshot only."""
-        datafiles: List[DataFile] = []
-
-        current_snapshot = self.tbl.current_snapshot()
-        if not current_snapshot:
-            return datafiles
-
-        def process_manifest(manifest: ManifestFile) -> list[DataFile]:
-            found: list[DataFile] = []
-            for entry in manifest.fetch_manifest_entry(io=self.tbl.io, discard_deleted=True):
-                if hasattr(entry, "data_file"):
-                    found.append(entry.data_file)
-            return found
-
-        # Scan only the current snapshot's manifests
-        manifests = current_snapshot.manifests(io=self.tbl.io)
-        with ThreadPoolExecutor() as executor:
-            results = executor.map(process_manifest, manifests)
-            for res in results:
-                datafiles.extend(res)
-
-        return datafiles
+        data_file_structs = self.tbl.inspect.data_files()
+        data_files = [DataFile(df) for df in data_file_structs]
+        return data_files
 
     def deduplicate_data_files(self) -> List[DataFile]:
         """

pyiceberg/table/update/snapshot.py

@@ -903,104 +903,4 @@ def remove_branch(self, branch_name: str) -> ManageSnapshots:
         Returns:
             This for method chaining
         """
-        return self._remove_ref_snapshot(ref_name=branch_name)
-
-
-class ExpireSnapshots(UpdateTableMetadata["ExpireSnapshots"]):
-    """
-    Expire snapshots by ID.
-
-    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[int] = 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, but will always skip protected snapshots (branch/tag heads).
-
-        Returns:
-            Tuple of updates and requirements to be committed,
-            as required by the calling parent apply functions.
-        """
-        # Remove any protected snapshot IDs from the set to expire, just in case
-        protected_ids = self._get_protected_snapshot_ids()
-        self._snapshot_ids_to_expire -= protected_ids
-        update = RemoveSnapshotsUpdate(snapshot_ids=self._snapshot_ids_to_expire)
-        self._updates += (update,)
-        return self._updates, self._requirements
-
-    def _get_protected_snapshot_ids(self) -> Set[int]:
-        """
-        Get the IDs of protected snapshots.
-
-        These are the HEAD snapshots of all branches and all tagged snapshots.  These ids are to be excluded from expiration.
-
-        Returns:
-            Set of protected snapshot IDs to exclude from expiration.
-        """
-        protected_ids: Set[int] = set()
-
-        for ref in self._transaction.table_metadata.refs.values():
-            if ref.snapshot_ref_type in [SnapshotRefType.TAG, SnapshotRefType.BRANCH]:
-                protected_ids.add(ref.snapshot_id)
-
-        return protected_ids
-
-    def expire_snapshot_by_id(self, snapshot_id: int) -> ExpireSnapshots:
-        """
-        Expire a snapshot by its ID.
-
-        This will mark the snapshot for expiration.
-
-        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.")
-
-        if snapshot_id in self._get_protected_snapshot_ids():
-            raise ValueError(f"Snapshot with ID {snapshot_id} is protected and cannot be expired.")
-
-        self._snapshot_ids_to_expire.add(snapshot_id)
-
-        return self
-
-    def expire_snapshots_by_ids(self, snapshot_ids: List[int]) -> "ExpireSnapshots":
-        """
-        Expire multiple snapshots by their IDs.
-
-        This will mark the snapshots for expiration.
-
-        Args:
-            snapshot_ids (List[int]): List of snapshot IDs to expire.
-        Returns:
-            This for method chaining.
-        """
-        for snapshot_id in snapshot_ids:
-            self.expire_snapshot_by_id(snapshot_id)
-        return self
-
-    def expire_snapshots_older_than(self, timestamp_ms: int) -> "ExpireSnapshots":
-        """
-        Expire all unprotected snapshots with a timestamp older than a given value.
-
-        Args:
-            timestamp_ms (int): Only snapshots with timestamp_ms < this value will be expired.
-
-        Returns:
-            This for method chaining.
-        """
-        protected_ids = self._get_protected_snapshot_ids()
-        for snapshot in self._transaction.table_metadata.snapshots:
-            if snapshot.timestamp_ms < timestamp_ms and snapshot.snapshot_id not in protected_ids:
-                self._snapshot_ids_to_expire.add(snapshot.snapshot_id)
-        return self
+        return self._remove_ref_snapshot(ref_name=branch_name)

tests/table/test_dedup_data_file_filepaths.py

@@ -124,6 +124,9 @@ def test_deduplicate_data_files_removes_duplicates_in_current_snapshot(
 ) -> None:
     mt = MaintenanceTable(tbl=prepopulated_table)
 
+    print("=== Before deduplication ===")
+    check_data_files(prepopulated_table)
+
     all_datafiles: List[DataFile] = mt._get_all_datafiles()
     file_names: List[str] = [os.path.basename(df.file_path) for df in all_datafiles]
     # There should be more than one reference before deduplication
@@ -132,8 +135,11 @@ def test_deduplicate_data_files_removes_duplicates_in_current_snapshot(
     )
     removed: List[DataFile] = mt.deduplicate_data_files()
 
+    print("=== After deduplication ===")
+    check_data_files(prepopulated_table)
+
     all_datafiles_after: List[DataFile] = mt._get_all_datafiles()
     file_names_after: List[str] = [os.path.basename(df.file_path) for df in all_datafiles_after]
     # Only one reference should remain after deduplication
     assert file_names_after.count(dupe_data_file_path.name) == 1
-    assert all(isinstance(df, DataFile) for df in removed)
+    assert all(isinstance(df, DataFile) for df in removed)